From 046793b6fa710e6b913170080e9453701d15aac8 Mon Sep 17 00:00:00 2001 From: thinkAfCod Date: Sat, 29 Mar 2025 15:41:14 +0800 Subject: [PATCH] session cache enr and replace handleTalkRequest parameter type --- p2p/discover/v5_talk.go | 13 ++- p2p/discover/v5_udp.go | 163 +++++++++++++++------------ p2p/discover/v5_udp_test.go | 6 +- p2p/discover/v5wire/encoding.go | 10 +- p2p/discover/v5wire/encoding_test.go | 14 +-- p2p/discover/v5wire/session.go | 16 ++- 6 files changed, 136 insertions(+), 86 deletions(-) diff --git a/p2p/discover/v5_talk.go b/p2p/discover/v5_talk.go index 2246b47141..10c4983d5a 100644 --- a/p2p/discover/v5_talk.go +++ b/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. diff --git a/p2p/discover/v5_udp.go b/p2p/discover/v5_udp.go index 0204792a26..44f7acff90 100644 --- a/p2p/discover/v5_udp.go +++ b/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,22 +64,24 @@ 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. type UDPv5 struct { // static fields - conn UDPConn - tab *Table - cachedAddrNode *lru.Cache[string, *enode.Node] - netrestrict *netutil.Netlist - priv *ecdsa.PrivateKey - localNode *enode.LocalNode - db *enode.DB - log log.Logger - clock mclock.Clock - validSchemes enr.IdentityScheme - respTimeout time.Duration + conn UDPConn + tab *Table + netrestrict *netutil.Netlist + priv *ecdsa.PrivateKey + localNode *enode.LocalNode + db *enode.DB + log log.Logger + clock mclock.Clock + validSchemes enr.IdentityScheme + respTimeout time.Duration // misc buffers used during message handling logcontext []interface{} @@ -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 @@ -161,28 +169,29 @@ func newUDPv5(conn UDPConn, ln *enode.LocalNode, cfg Config) (*UDPv5, error) { cfg = cfg.withDefaults() t := &UDPv5{ // static fields - conn: newMeteredConn(conn), - cachedAddrNode: lru.NewCache[string, *enode.Node](1024), - localNode: ln, - db: ln.Database(), - netrestrict: cfg.NetRestrict, - priv: cfg.PrivateKey, - log: cfg.Log, - validSchemes: cfg.ValidSchemes, - clock: cfg.Clock, - respTimeout: cfg.V5RespTimeout, + conn: newMeteredConn(conn), + localNode: ln, + db: ln.Database(), + netrestrict: cfg.NetRestrict, + priv: cfg.PrivateKey, + log: cfg.Log, + validSchemes: cfg.ValidSchemes, + clock: cfg.Clock, + respTimeout: cfg.V5RespTimeout, // channels into dispatch packetInCh: make(chan ReadPacket, 1), readNextCh: make(chan struct{}, 1), 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) - } + t.send(r.destID, r.destAddr, r.msg, nil) + 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) - c.handshakeCount++ - c.challenge = p - p.Node = c.node - t.sendCall(c) + 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) -} diff --git a/p2p/discover/v5_udp_test.go b/p2p/discover/v5_udp_test.go index 3026dff538..34450d803d 100644 --- a/p2p/discover/v5_udp_test.go +++ b/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) diff --git a/p2p/discover/v5wire/encoding.go b/p2p/discover/v5wire/encoding.go index e50b7cd16d..21c1f5b96e 100644 --- a/p2p/discover/v5wire/encoding.go +++ b/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 } diff --git a/p2p/discover/v5wire/encoding_test.go b/p2p/discover/v5wire/encoding_test.go index df97e40e89..2304d0f132 100644 --- a/p2p/discover/v5wire/encoding_test.go +++ b/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.nodeB.c.sc.storeNewSession(net.nodeA.id(), net.nodeA.addr(), initKeysB) + 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.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.nodeB.c.sc.storeNewSession(idA, addr, session.keysFlipped()) + net.nodeA.c.sc.storeNewSession(idB, addr, session, net.nodeB.n()) + 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.nodeB.c.sc.storeNewSession(net.nodeA.id(), net.nodeA.addr(), session.keysFlipped()) + 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.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) diff --git a/p2p/discover/v5wire/session.go b/p2p/discover/v5wire/session.go index 862c21fcee..53a6947942 100644 --- a/p2p/discover/v5wire/session.go +++ b/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) }