From 2b96b4f13d70c2697f99a0190ce3d83d55453202 Mon Sep 17 00:00:00 2001
From: Zsolt Felfoldi <zsfelfoldi@gmail.com>
Date: Fri, 15 Nov 2019 21:08:18 +0100
Subject: [PATCH] les: add talk request handler

---
 les/server.go         |  7 +++++++
 les/server_handler.go |  6 ++++++
 p2p/discv5/net.go     | 10 ++++++++--
 3 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/les/server.go b/les/server.go
index 429e866cb8..ccae237c68 100644
--- a/les/server.go
+++ b/les/server.go
@@ -42,6 +42,7 @@ type LesServer struct {
 	handler     *serverHandler
 	lesTopics   []discv5.Topic
 	privateKey  *ecdsa.PrivateKey
+	srvr        *p2p.Server
 
 	// Flow control and capacity management
 	fcManager    *flowcontrol.ClientManager
@@ -165,6 +166,7 @@ func (s *LesServer) Protocols() []p2p.Protocol {
 
 // Start starts the LES server
 func (s *LesServer) Start(srvr *p2p.Server) {
+	s.srvr = srvr
 	s.privateKey = srvr.PrivateKey
 	s.handler.start()
 
@@ -172,6 +174,7 @@ func (s *LesServer) Start(srvr *p2p.Server) {
 	go s.capacityManagement()
 
 	if srvr.DiscV5 != nil {
+		srvr.DiscV5.RegisterTalkHandler("les", s.handler.talkRequestHandler)
 		for _, topic := range s.lesTopics {
 			topic := topic
 			go func() {
@@ -189,6 +192,10 @@ func (s *LesServer) Start(srvr *p2p.Server) {
 func (s *LesServer) Stop() {
 	close(s.closeCh)
 
+	if s.srvr.DiscV5 != nil {
+		s.srvr.DiscV5.RemoveTalkHandler("les")
+	}
+
 	// Disconnect existing sessions.
 	// This also closes the gate for any new registrations on the peer set.
 	// sessions which are already established but not added to pm.peers yet
diff --git a/les/server_handler.go b/les/server_handler.go
index 4b505c2bc4..5c80b9ffb1 100644
--- a/les/server_handler.go
+++ b/les/server_handler.go
@@ -20,6 +20,7 @@ import (
 	"encoding/binary"
 	"encoding/json"
 	"errors"
+	"net"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -35,6 +36,7 @@ import (
 	"github.com/ethereum/go-ethereum/log"
 	"github.com/ethereum/go-ethereum/metrics"
 	"github.com/ethereum/go-ethereum/p2p"
+	"github.com/ethereum/go-ethereum/p2p/enode"
 	"github.com/ethereum/go-ethereum/rlp"
 	"github.com/ethereum/go-ethereum/trie"
 )
@@ -953,3 +955,7 @@ func (h *serverHandler) broadcastHeaders() {
 		}
 	}
 }
+
+func (h *serverHandler) talkRequestHandler(id enode.ID, addr *net.UDPAddr, payload rlp.RawValue) (rlp.RawValue, bool) {
+	return payload, true
+}
diff --git a/p2p/discv5/net.go b/p2p/discv5/net.go
index adbc45bb8b..1c407d1d87 100644
--- a/p2p/discv5/net.go
+++ b/p2p/discv5/net.go
@@ -87,7 +87,7 @@ type Network struct {
 }
 
 type (
-	TalkRequestHandler  func(enode.ID, rlp.RawValue) (rlp.RawValue, bool)
+	TalkRequestHandler  func(enode.ID, *net.UDPAddr, rlp.RawValue) (rlp.RawValue, bool)
 	TalkResponseHandler func(rlp.RawValue) bool
 )
 
@@ -1220,7 +1220,7 @@ func (net *Network) handleQueryEvent(n *Node, ev nodeEvent, pkt *ingressPacket)
 		subFn := net.talkRequestSubs[string(p.TalkID)]
 		net.talkRequestSubLock.RUnlock()
 		if subFn != nil {
-			resp, ok := subFn(enode.ID(n.sha), p.Payload)
+			resp, ok := subFn(enode.ID(n.sha), n.addr(), p.Payload)
 			if ok {
 				net.conn.send(n, talkResponsePacket, talkResponse{ReplyTok: pkt.hash, Payload: resp})
 			} else {
@@ -1329,6 +1329,12 @@ func (net *Network) RegisterTalkHandler(talkID string, handler TalkRequestHandle
 	net.talkRequestSubLock.Unlock()
 }
 
+func (net *Network) RemoveTalkHandler(talkID string) {
+	net.talkRequestSubLock.Lock()
+	delete(net.talkRequestSubs, talkID)
+	net.talkRequestSubLock.Unlock()
+}
+
 func (net *Network) SendTalkRequest(to *enode.Node, talkID string, payload rlp.RawValue, handler TalkResponseHandler) func() bool {
 	var nodeID NodeID
 	copy(nodeID[:], crypto.FromECDSAPub(to.Pubkey())[1:])