feat(p2p/discover): UDP pre-negotiation of available server capacity #22183

p2p/discover/v5_udp.go

@@ -74,7 +74,7 @@ type UDPv5 struct {
 
 	// talkreq handler registry
 	trlock     sync.Mutex
-	trhandlers map[string]func([]byte) []byte
+	trhandlers map[string]TalkRequestHandler
 
 	// channels into dispatch
 	packetInCh    chan ReadPacket
@@ -96,6 +96,9 @@ type UDPv5 struct {
 	wg             sync.WaitGroup
 }
 
+// TalkRequestHandler callback processes a talk request and optionally returns a reply
+type TalkRequestHandler func(enode.ID, *net.UDPAddr, []byte) []byte
+
 // callV5 represents a remote procedure call against another node.
 type callV5 struct {
 	node         *enode.Node
@@ -145,7 +148,7 @@ func newUDPv5(conn UDPConn, ln *enode.LocalNode, cfg Config) (*UDPv5, error) {
 		log:          cfg.Log,
 		validSchemes: cfg.ValidSchemes,
 		clock:        cfg.Clock,
-		trhandlers:   make(map[string]func([]byte) []byte),
+		trhandlers:   make(map[string]TalkRequestHandler),
 		// channels into dispatch
 		packetInCh:    make(chan ReadPacket, 1),
 		readNextCh:    make(chan struct{}, 1),
@@ -233,7 +236,7 @@ func (t *UDPv5) LocalNode() *enode.LocalNode {
 // RegisterTalkHandler adds a handler for 'talk requests'. The handler function is called
 // whenever a request for the given protocol is received and should return the response
 // data or nil.
-func (t *UDPv5) RegisterTalkHandler(protocol string, handler func([]byte) []byte) {
+func (t *UDPv5) RegisterTalkHandler(protocol string, handler TalkRequestHandler) {
 	t.trlock.Lock()
 	defer t.trlock.Unlock()
 	t.trhandlers[protocol] = handler
@@ -841,7 +844,7 @@ func (t *UDPv5) handleTalkRequest(p *v5wire.TalkRequest, fromID enode.ID, fromAd
 
 	var response []byte
 	if handler != nil {
-		response = handler(p.Message)
+		response = handler(fromID, fromAddr, p.Message)
 	}
 	resp := &v5wire.TalkResponse{ReqID: p.ReqID, Message: response}
 	t.sendResponse(fromID, fromAddr, resp)

p2p/discover/v5_udp_test.go

@@ -442,7 +442,7 @@ func TestUDPv5_talkHandling(t *testing.T) {
 	defer test.close()
 
 	var recvMessage []byte
-	test.udp.RegisterTalkHandler("test", func(message []byte) []byte {
+	test.udp.RegisterTalkHandler("test", func(id enode.ID, addr *net.UDPAddr, message []byte) []byte {
 		recvMessage = message
 		return []byte("test response")
 	})