session cache enr and replace handleTalkRequest parameter type

p2p/discover/v5_talk.go

@@ -39,7 +39,7 @@ const talkHandlerLaunchTimeout = 400 * time.Millisecond
 // Note that talk handlers are expected to come up with a response very quickly, within at
 // most 200ms or so. If the handler takes longer than that, the remote end may time out
 // and wont receive the response.
-type TalkRequestHandler func(enode.ID, *net.UDPAddr, []byte) []byte
+type TalkRequestHandler func(*enode.Node, *net.UDPAddr, []byte) []byte
 
 type talkSystem struct {
 	transport *UDPv5
@@ -72,13 +72,18 @@ func (t *talkSystem) register(protocol string, handler TalkRequestHandler) {
 
 // handleRequest handles a talk request.
 func (t *talkSystem) handleRequest(id enode.ID, addr netip.AddrPort, req *v5wire.TalkRequest) {
+	var n *enode.Node
+	if n = t.transport.codec.SessionNode(id, addr.String()); n == nil {
+		log.Error("Got a TALKREQ from a node that has not completed the handshake", "id", id, "addr", addr)
+		return
+	}
 	t.mutex.Lock()
 	handler, ok := t.handlers[req.Protocol]
 	t.mutex.Unlock()
 
 	if !ok {
 		resp := &v5wire.TalkResponse{ReqID: req.ReqID}
-		t.transport.sendResponse(id, addr, resp)
+		t.transport.sendResponse(n.ID(), addr, resp)
 		return
 	}
 
@@ -90,9 +95,9 @@ func (t *talkSystem) handleRequest(id enode.ID, addr netip.AddrPort, req *v5wire
 		go func() {
 			defer func() { t.slots <- struct{}{} }()
 			udpAddr := &net.UDPAddr{IP: addr.Addr().AsSlice(), Port: int(addr.Port())}
-			respMessage := handler(id, udpAddr, req.Message)
+			respMessage := handler(n, udpAddr, req.Message)
 			resp := &v5wire.TalkResponse{ReqID: req.ReqID, Message: respMessage}
-			t.transport.sendFromAnotherThread(id, addr, resp)
+			t.transport.sendFromAnotherThread(n.ID(), addr, resp)
 		}()
 	case <-timeout.C:
 		// Couldn't get it in time, drop the request.

p2p/discover/v5_udp.go

@@ -30,7 +30,6 @@ import (
 	"sync"
 	"time"
 
-	"github.com/ethereum/go-ethereum/common/lru"
 	"github.com/ethereum/go-ethereum/common/mclock"
 	"github.com/ethereum/go-ethereum/log"
 	"github.com/ethereum/go-ethereum/p2p/discover/v5wire"
@@ -65,6 +64,9 @@ type codecV5 interface {
 	// CurrentChallenge returns the most recent WHOAREYOU challenge that was encoded to given node.
 	// This will return a non-nil value if there is an active handshake attempt with the node, and nil otherwise.
 	CurrentChallenge(id enode.ID, addr string) *v5wire.Whoareyou
+
+	// SessionNode returns a node that has completed the handshake.
+	SessionNode(enode.ID, string) *enode.Node
 }
 
 // UDPv5 is the implementation of protocol version 5.
@@ -72,7 +74,6 @@ type UDPv5 struct {
 	// static fields
 	conn         UDPConn
 	tab          *Table
-	cachedAddrNode *lru.Cache[string, *enode.Node]
 	netrestrict  *netutil.Netlist
 	priv         *ecdsa.PrivateKey
 	localNode    *enode.LocalNode
@@ -95,12 +96,14 @@ type UDPv5 struct {
 	callDoneCh    chan *callV5
 	respTimeoutCh chan *callTimeout
 	sendCh        chan sendRequest
+	sendNoRespCh  chan *sendNoRespRequest
 	unhandled     chan<- ReadPacket
 
 	// state of dispatch
 	codec            codecV5
 	activeCallByNode map[enode.ID]*callV5
 	activeCallByAuth map[v5wire.Nonce]*callV5
+	noRespCallByAuth map[v5wire.Nonce]*callV5
 	callQueue        map[enode.ID][]*callV5
 
 	// shutdown stuff
@@ -112,6 +115,11 @@ type UDPv5 struct {
 
 type sendRequest struct {
 	destID   enode.ID
+	destAddr netip.AddrPort
+	msg      v5wire.Packet
+}
+
+type sendNoRespRequest struct {
 	destNode *enode.Node
 	destAddr netip.AddrPort
 	msg      v5wire.Packet
@@ -162,7 +170,6 @@ func newUDPv5(conn UDPConn, ln *enode.LocalNode, cfg Config) (*UDPv5, error) {
 	t := &UDPv5{
 		// static fields
 		conn:         newMeteredConn(conn),
-		cachedAddrNode: lru.NewCache[string, *enode.Node](1024),
 		localNode:    ln,
 		db:           ln.Database(),
 		netrestrict:  cfg.NetRestrict,
@@ -177,12 +184,14 @@ func newUDPv5(conn UDPConn, ln *enode.LocalNode, cfg Config) (*UDPv5, error) {
 		callCh:        make(chan *callV5),
 		callDoneCh:    make(chan *callV5),
 		sendCh:        make(chan sendRequest),
+		sendNoRespCh:  make(chan *sendNoRespRequest),
 		respTimeoutCh: make(chan *callTimeout),
 		unhandled:     cfg.Unhandled,
 		// state of dispatch
 		codec:            v5wire.NewCodec(ln, cfg.PrivateKey, cfg.Clock, cfg.V5ProtocolID),
 		activeCallByNode: make(map[enode.ID]*callV5),
 		activeCallByAuth: make(map[v5wire.Nonce]*callV5),
+		noRespCallByAuth: make(map[v5wire.Nonce]*callV5),
 		callQueue:        make(map[enode.ID][]*callV5),
 		// shutdown
 		closeCtx:       closeCtx,
@@ -582,13 +591,17 @@ func (t *UDPv5) dispatch() {
 		case c := <-t.callCh:
 			t.callQueue[c.id] = append(t.callQueue[c.id], c)
 			t.sendNextCall(c.id)
+		case cnr := <-t.sendNoRespCh:
+			// send a TalkReq call but not waiting for TalkResp
+			// may more used by portal network
+			t.sendCallNotWaitResp(cnr)
 
 		case ct := <-t.respTimeoutCh:
 			active := t.activeCallByNode[ct.c.id]
 			if ct.c == active && ct.timer == active.timeout {
 				ct.c.err <- errTimeout
 			}
-			delete(t.activeCallByAuth, ct.c.nonce)
+			delete(t.noRespCallByAuth, ct.c.nonce)
 			ct.c.timeout.Stop()
 
 		case c := <-t.callDoneCh:
@@ -602,30 +615,8 @@ func (t *UDPv5) dispatch() {
 			t.sendNextCall(c.id)
 
 		case r := <-t.sendCh:
-			if _, ok := r.msg.(*v5wire.TalkResponse); ok {
-				// send a TalkResponse out
 			t.send(r.destID, r.destAddr, r.msg, nil)
-			} else {
+
-				// send out a TalkRequest that is a UTP packet for portal network
-				// request should be cached to handle WHOAREYOU
-				c := &callV5{id: r.destID, addr: r.destAddr}
-				c.node = r.destNode
-				c.packet = r.msg
-				c.reqid = make([]byte, 8)
-				c.ch = make(chan v5wire.Packet, 1)
-				c.err = make(chan error, 1)
-				// Assign request ID.
-				if tq, ok := r.msg.(*v5wire.TalkRequest); ok {
-					if len(tq.ReqID) == 0 {
-						crand.Read(c.reqid)
-						c.packet.SetRequestID(c.reqid)
-					}
-				}
-				nonce, _ := t.send(c.id, c.addr, c.packet, nil)
-				c.nonce = nonce
-				t.activeCallByAuth[nonce] = c
-				t.startResponseTimeout(c)
-			}
 		case p := <-t.packetInCh:
 			t.handlePacket(p.Data, p.Addr)
 			// Arm next read.
@@ -644,6 +635,9 @@ func (t *UDPv5) dispatch() {
 				delete(t.activeCallByNode, id)
 				delete(t.activeCallByAuth, c.nonce)
 			}
+			for nonce := range t.noRespCallByAuth {
+				delete(t.activeCallByAuth, nonce)
+			}
 			return
 		}
 	}
@@ -669,6 +663,41 @@ func (t *UDPv5) startResponseTimeout(c *callV5) {
 	close(done)
 }
 
+// sendCallNotWaitResp send a talk request contains utp packet by call, call will not insert into queue
+func (t *UDPv5) sendCallNotWaitResp(r *sendNoRespRequest) {
+	// send out a TalkRequest that is a UTP packet for portal network
+	// if a handshake has been done, the n will not be nil, and just send the data out
+	if n := t.codec.SessionNode(r.destNode.ID(), r.destAddr.String()); n != nil {
+		t.send(r.destNode.ID(), r.destAddr, r.msg, nil)
+		return
+	}
+
+	// request should be cached to handle WHOAREYOU
+	c := &callV5{id: r.destNode.ID(), addr: r.destAddr}
+	c.node = r.destNode
+	c.packet = r.msg
+	c.reqid = make([]byte, 8)
+	c.ch = make(chan v5wire.Packet, 1)
+	c.err = make(chan error, 1)
+	// Assign request ID.
+	crand.Read(c.reqid)
+	c.packet.SetRequestID(c.reqid)
+
+	nonce, _ := t.send(c.id, c.addr, c.packet, nil)
+	c.nonce = nonce
+	t.noRespCallByAuth[nonce] = c
+	t.startResponseTimeout(c)
+}
+
+// sendNoRespData send a data from a call of no wait resp.
+// The handshake has just been successful, could clear the info in the map
+func (t *UDPv5) sendNoRespData(c *callV5) {
+	// Just resend the data and not use call again
+	delete(t.noRespCallByAuth, c.nonce)
+	c.timeout.Stop()
+	t.send(c.node.ID(), c.addr, c.packet, nil)
+}
+
 // sendNextCall sends the next call in the call queue if there is no active call.
 func (t *UDPv5) sendNextCall(id enode.ID) {
 	queue := t.callQueue[id]
@@ -709,14 +738,15 @@ func (t *UDPv5) sendResponse(toID enode.ID, toAddr netip.AddrPort, packet v5wire
 
 func (t *UDPv5) sendFromAnotherThread(toID enode.ID, toAddr netip.AddrPort, packet v5wire.Packet) {
 	select {
-	case t.sendCh <- sendRequest{toID, nil, toAddr, packet}:
+	case t.sendCh <- sendRequest{destID: toID, destAddr: toAddr, msg: packet}:
 	case <-t.closeCtx.Done():
 	}
 }
 
-func (t *UDPv5) SendFromAnotherThreadWithNode(node *enode.Node, toAddr netip.AddrPort, packet v5wire.Packet) {
+// SendNoResp Send a packet but will not wait for a response
+func (t *UDPv5) SendNoResp(n *enode.Node, toAddr netip.AddrPort, packet v5wire.Packet) {
 	select {
-	case t.sendCh <- sendRequest{node.ID(), node, toAddr, packet}:
+	case t.sendNoRespCh <- &sendNoRespRequest{n, toAddr, packet}:
 	case <-t.closeCtx.Done():
 	}
 }
@@ -733,9 +763,6 @@ func (t *UDPv5) send(toID enode.ID, toAddr netip.AddrPort, packet v5wire.Packet,
 		t.log.Warn(">> "+packet.Name(), t.logcontext...)
 		return nonce, err
 	}
-	if c != nil && c.Node != nil {
-		t.putCache(toAddr.String(), c.Node)
-	}
 
 	_, err = t.conn.WriteToUDPAddrPort(enc, toAddr)
 	t.log.Trace(">> "+packet.Name(), t.logcontext...)
@@ -797,7 +824,6 @@ func (t *UDPv5) handlePacket(rawpacket []byte, fromAddr netip.AddrPort) error {
 	if fromNode != nil {
 		// Handshake succeeded, add to table.
 		t.tab.addInboundNode(fromNode)
-		t.putCache(fromAddr.String(), fromNode)
 	}
 	if packet.Kind() != v5wire.WhoareyouPacket {
 		// WHOAREYOU logged separately to report errors.
@@ -913,17 +939,25 @@ func (t *UDPv5) handleWhoareyou(p *v5wire.Whoareyou, fromID enode.ID, fromAddr n
 		c.err <- errors.New("remote wants handshake, but call has no ENR")
 		return
 	}
-	// Resend the call that was answered by WHOAREYOU.
+
 	t.log.Trace("<< "+p.Name(), "id", c.node.ID(), "addr", fromAddr)
+	if _, ok := t.noRespCallByAuth[p.Nonce]; !ok {
+		// Resend the call that was answered by WHOAREYOU.
 		c.handshakeCount++
 		c.challenge = p
 		p.Node = c.node
 		t.sendCall(c)
+	} else {
+		t.sendNoRespData(c)
+	}
 }
 
 // matchWithCall checks whether a handshake attempt matches the active call.
 func (t *UDPv5) matchWithCall(fromID enode.ID, nonce v5wire.Nonce) (*callV5, error) {
 	c := t.activeCallByAuth[nonce]
+	if c == nil {
+		c = t.noRespCallByAuth[nonce]
+	}
 	if c == nil {
 		return nil, errChallengeNoCall
 	}
@@ -1020,14 +1054,3 @@ func packNodes(reqid []byte, nodes []*enode.Node) []*v5wire.Nodes {
 	}
 	return resp
 }
-
-func (t *UDPv5) putCache(addr string, node *enode.Node) {
-	if n, ok := t.cachedAddrNode.Get(addr); ok {
-		t.log.Debug("Update cached node", "old", n.ID(), "new", node.ID())
-	}
-	t.cachedAddrNode.Add(addr, node)
-}
-
-func (t *UDPv5) GetCachedNode(addr string) (*enode.Node, bool) {
-	return t.cachedAddrNode.Get(addr)
-}

p2p/discover/v5_udp_test.go

@@ -486,7 +486,7 @@ func TestUDPv5_talkHandling(t *testing.T) {
 	defer test.close()
 
 	var recvMessage []byte
-	test.udp.RegisterTalkHandler("test", func(id enode.ID, addr *net.UDPAddr, message []byte) []byte {
+	test.udp.RegisterTalkHandler("test", func(n *enode.Node, addr *net.UDPAddr, message []byte) []byte {
 		recvMessage = message
 		return []byte("test response")
 	})
@@ -795,6 +795,10 @@ func (c *testCodec) Decode(input []byte, addr string) (enode.ID, *enode.Node, v5
 	return frame.NodeID, nil, p, nil
 }
 
+func (c *testCodec) SessionNode(id enode.ID, addr string) *enode.Node {
+	return c.test.nodesByID[id].Node()
+}
+
 func (c *testCodec) decodeFrame(input []byte) (frame testCodecFrame, p v5wire.Packet, err error) {
 	if err = rlp.DecodeBytes(input, &frame); err != nil {
 		return frame, nil, fmt.Errorf("invalid frame: %v", err)

p2p/discover/v5wire/encoding.go

@@ -251,6 +251,12 @@ func (c *Codec) CurrentChallenge(id enode.ID, addr string) *Whoareyou {
 	return c.sc.getHandshake(id, addr)
 }
 
+// SessionNode returns the node associated.
+// This will return nil while a handshake is in progress.
+func (c *Codec) SessionNode(id enode.ID, addr string) *enode.Node {
+	return c.sc.readNode(id, addr)
+}
+
 func (c *Codec) writeHeaders(head *Header) {
 	c.buf.Reset()
 	c.buf.Write(head.IV[:])
@@ -347,7 +353,7 @@ func (c *Codec) encodeHandshakeHeader(toID enode.ID, addr string, challenge *Who
 	}
 
 	// TODO: this should happen when the first authenticated message is received
-	c.sc.storeNewSession(toID, addr, session)
+	c.sc.storeNewSession(toID, addr, session, challenge.Node)
 
 	// Encode the auth header.
 	var (
@@ -522,7 +528,7 @@ func (c *Codec) decodeHandshakeMessage(fromAddr string, head *Header, headerData
 	}
 
 	// Handshake OK, drop the challenge and store the new session keys.
-	c.sc.storeNewSession(auth.h.SrcID, fromAddr, session)
+	c.sc.storeNewSession(auth.h.SrcID, fromAddr, session, node)
 	c.sc.deleteHandshake(auth.h.SrcID, fromAddr)
 	return node, msg, nil
 }

p2p/discover/v5wire/encoding_test.go

@@ -166,7 +166,7 @@ func TestHandshake_rekey(t *testing.T) {
 		readKey:  []byte("BBBBBBBBBBBBBBBB"),
 		writeKey: []byte("AAAAAAAAAAAAAAAA"),
 	}
-	net.nodeA.c.sc.storeNewSession(net.nodeB.id(), net.nodeB.addr(), session)
+	net.nodeA.c.sc.storeNewSession(net.nodeB.id(), net.nodeB.addr(), session, net.nodeB.n())
 
 	// A -> B   FINDNODE (encrypted with zero keys)
 	findnode, authTag := net.nodeA.encode(t, net.nodeB, &Findnode{})
@@ -209,8 +209,8 @@ func TestHandshake_rekey2(t *testing.T) {
 		readKey:  []byte("CCCCCCCCCCCCCCCC"),
 		writeKey: []byte("DDDDDDDDDDDDDDDD"),
 	}
-	net.nodeA.c.sc.storeNewSession(net.nodeB.id(), net.nodeB.addr(), initKeysA)
+	net.nodeA.c.sc.storeNewSession(net.nodeB.id(), net.nodeB.addr(), initKeysA, net.nodeB.n())
-	net.nodeB.c.sc.storeNewSession(net.nodeA.id(), net.nodeA.addr(), initKeysB)
+	net.nodeB.c.sc.storeNewSession(net.nodeA.id(), net.nodeA.addr(), initKeysB, net.nodeA.n())
 
 	// A -> B   FINDNODE encrypted with initKeysA
 	findnode, authTag := net.nodeA.encode(t, net.nodeB, &Findnode{Distances: []uint{3}})
@@ -362,8 +362,8 @@ func TestTestVectorsV5(t *testing.T) {
 				ENRSeq: 2,
 			},
 			prep: func(net *handshakeTest) {
-				net.nodeA.c.sc.storeNewSession(idB, addr, session)
+				net.nodeA.c.sc.storeNewSession(idB, addr, session, net.nodeB.n())
-				net.nodeB.c.sc.storeNewSession(idA, addr, session.keysFlipped())
+				net.nodeB.c.sc.storeNewSession(idA, addr, session.keysFlipped(), net.nodeA.n())
 			},
 		},
 		{
@@ -499,8 +499,8 @@ func BenchmarkV5_DecodePing(b *testing.B) {
 		readKey:  []byte{233, 203, 93, 195, 86, 47, 177, 186, 227, 43, 2, 141, 244, 230, 120, 17},
 		writeKey: []byte{79, 145, 252, 171, 167, 216, 252, 161, 208, 190, 176, 106, 214, 39, 178, 134},
 	}
-	net.nodeA.c.sc.storeNewSession(net.nodeB.id(), net.nodeB.addr(), session)
+	net.nodeA.c.sc.storeNewSession(net.nodeB.id(), net.nodeB.addr(), session, net.nodeB.n())
-	net.nodeB.c.sc.storeNewSession(net.nodeA.id(), net.nodeA.addr(), session.keysFlipped())
+	net.nodeB.c.sc.storeNewSession(net.nodeA.id(), net.nodeA.addr(), session.keysFlipped(), net.nodeA.n())
 	addrB := net.nodeA.addr()
 	ping := &Ping{ReqID: []byte("reqid"), ENRSeq: 5}
 	enc, _, err := net.nodeA.c.Encode(net.nodeB.id(), addrB, ping, nil)

p2p/discover/v5wire/session.go

@@ -54,11 +54,12 @@ type session struct {
 	writeKey     []byte
 	readKey      []byte
 	nonceCounter uint32
+	node         *enode.Node
 }
 
 // keysFlipped returns a copy of s with the read and write keys flipped.
 func (s *session) keysFlipped() *session {
-	return &session{s.readKey, s.writeKey, s.nonceCounter}
+	return &session{s.readKey, s.writeKey, s.nonceCounter, s.node}
 }
 
 func NewSessionCache(maxItems int, clock mclock.Clock) *SessionCache {
@@ -103,8 +104,19 @@ func (sc *SessionCache) readKey(id enode.ID, addr string) []byte {
 	return nil
 }
 
+func (sc *SessionCache) readNode(id enode.ID, addr string) *enode.Node {
+	if s := sc.session(id, addr); s != nil {
+		return s.node
+	}
+	return nil
+}
+
 // storeNewSession stores new encryption keys in the cache.
-func (sc *SessionCache) storeNewSession(id enode.ID, addr string, s *session) {
+func (sc *SessionCache) storeNewSession(id enode.ID, addr string, s *session, n *enode.Node) {
+	if n == nil {
+		panic("session must caches a non-nil node")
+	}
+	s.node = n
 	sc.sessions.Add(sessionID{id, addr}, s)
 }