forks/go-ethereum
1
0
Fork 1
mirror of https://github.com/ethereum/go-ethereum.git synced 2026-07-27 15:16:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

session cache enr and replace handleTalkRequest parameter type

Browse source
This commit is contained in:
thinkAfCod 2025-03-29 15:41:14 +08:00
parent 8261fd8160
commit 046793b6fa
6 changed files with 136 additions and 86 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

13
p2p/discover/v5_talk.go
View file

@ -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.

115
p2p/discover/v5_udp.go
View file

@ -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)
}

6
p2p/discover/v5_udp_test.go
View file

@ -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)

10
p2p/discover/v5wire/encoding.go
View file

@ -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
}

14
p2p/discover/v5wire/encoding_test.go
View file

@ -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)

16
p2p/discover/v5wire/session.go
View file

@ -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)
}