From e401f5ff10d7815f77d9fb587291f4fcf69cbdc8 Mon Sep 17 00:00:00 2001
From: gary rong <garyrong0905@gmail.com>
Date: Fri, 7 Aug 2020 21:33:00 +0800
Subject: [PATCH] les: close all connected les-server when shutdown (#21426)

* les: close all connected les-server when shutdown

* les: linter nitpick

Co-authored-by: Martin Holst Swende <martin@swende.se>
---
 les/peer.go           | 39 +++++++++++++++++++++++++++++++++++++++
 les/server.go         |  5 +++++
 les/server_handler.go |  3 +++
 3 files changed, 47 insertions(+)

diff --git a/les/peer.go b/les/peer.go
index 72156814dd..c529145239 100644
--- a/les/peer.go
+++ b/les/peer.go
@@ -1288,3 +1288,42 @@ func (ps *serverPeerSet) close() {
 	}
 	ps.closed = true
 }
+
+// serverSet is a special set which contains all connected les servers.
+// Les servers will also be discovered by discovery protocol because they
+// also run the LES protocol. We can't drop them although they are useless
+// for us(server) but for other protocols(e.g. ETH) upon the devp2p they
+// may be useful.
+type serverSet struct {
+	lock   sync.Mutex
+	set    map[string]*clientPeer
+	closed bool
+}
+
+func newServerSet() *serverSet {
+	return &serverSet{set: make(map[string]*clientPeer)}
+}
+
+func (s *serverSet) register(peer *clientPeer) error {
+	s.lock.Lock()
+	defer s.lock.Unlock()
+
+	if s.closed {
+		return errClosed
+	}
+	if _, exist := s.set[peer.id]; exist {
+		return errAlreadyRegistered
+	}
+	s.set[peer.id] = peer
+	return nil
+}
+
+func (s *serverSet) close() {
+	s.lock.Lock()
+	defer s.lock.Unlock()
+
+	for _, p := range s.set {
+		p.Disconnect(p2p.DiscQuitting)
+	}
+	s.closed = true
+}
diff --git a/les/server.go b/les/server.go
index ec856c977c..ecb65150a8 100644
--- a/les/server.go
+++ b/les/server.go
@@ -39,6 +39,7 @@ type LesServer struct {
 
 	archiveMode bool // Flag whether the ethereum node runs in archive mode.
 	peers       *clientPeerSet
+	serverset   *serverSet
 	handler     *serverHandler
 	lesTopics   []discv5.Topic
 	privateKey  *ecdsa.PrivateKey
@@ -83,6 +84,7 @@ func NewLesServer(node *node.Node, e *eth.Ethereum, config *eth.Config) (*LesSer
 		},
 		archiveMode:  e.ArchiveMode(),
 		peers:        newClientPeerSet(),
+		serverset:    newServerSet(),
 		lesTopics:    lesTopics,
 		fcManager:    flowcontrol.NewClientManager(nil, &mclock.System{}),
 		servingQueue: newServingQueue(int64(time.Millisecond*10), float64(config.LightServ)/100),
@@ -196,6 +198,9 @@ func (s *LesServer) Start() error {
 func (s *LesServer) Stop() error {
 	close(s.closeCh)
 
+	// Disconnect existing connections with other LES servers.
+	s.serverset.close()
+
 	// 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 70c6a310af..c474363232 100644
--- a/les/server_handler.go
+++ b/les/server_handler.go
@@ -123,6 +123,9 @@ func (h *serverHandler) handle(p *clientPeer) error {
 		return err
 	}
 	if p.server {
+		if err := h.server.serverset.register(p); err != nil {
+			return err
+		}
 		// connected to another server, no messages expected, just wait for disconnection
 		_, err := p.rw.ReadMsg()
 		return err