diff --git a/les/peer.go b/les/peer.go index 72156814ddb2..c529145239c7 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 ec856c977c7c..ecb65150a891 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 70c6a310af07..c474363232af 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