fix(tcp): set conn on peer's endpoint

fumiama · Jul 16, 2024 · 39d8d5b · 39d8d5b
1 parent c7bbcb9
commit 39d8d5b
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/gold/p2p/tcp/tcp.go b/gold/p2p/tcp/tcp.go
@@ -67,6 +67,7 @@ func (ep *EndPoint) Listen() (p2p.Conn, error) {
 
 type connrecv struct {
 	addr *EndPoint // cast from tcpconn.RemoteAddr()
+	conn *net.TCPConn
 	pckt packet
 }
 
@@ -121,6 +122,7 @@ func (conn *Conn) receive(ep *EndPoint) {
 		if tcpconn == nil {
 			return
 		}
+		r.conn = tcpconn
 		_, err := io.Copy(&r.pckt, tcpconn)
 		if err != nil {
 			logrus.Debugln("[tcp] recv from", ep, "err:", err)
@@ -163,6 +165,7 @@ func (conn *Conn) ReadFromPeer(b []byte) (int, p2p.EndPoint, error) {
 		if p == nil {
 			return 0, nil, net.ErrClosed
 		}
+		conn.peers.Set(p.addr.String(), p.conn)
 		if p.pckt.typ == packetTypeNormal {
 			break
 		}