From fa526e7444fc9a54d0699f1b665a0b137afc1450 Mon Sep 17 00:00:00 2001
From: Dmitry Shulyak <dshulyak@mirantis.com>
Date: Mon, 16 Apr 2018 21:15:03 +0300
Subject: [PATCH] Prevent potential call on discv5 which became nil

It could happen when we found multiple relevant peers and they were
processed before discv5 closed. For example we have max limit of 1:
- found peer A and B in the same kademlia cycle
- process peer A
- max limit is reached -> closed discv5 and set it to nil
- process peer B
- panic!!!
---
 geth/peers/peerpool.go      | 23 ++++++++++++++++-------
 geth/peers/peerpool_test.go | 26 ++++++++++++++++++++++++++
 2 files changed, 42 insertions(+), 7 deletions(-)

diff --git a/geth/peers/peerpool.go b/geth/peers/peerpool.go
index 6e3e6e043..5826924ec 100644
--- a/geth/peers/peerpool.go
+++ b/geth/peers/peerpool.go
@@ -152,18 +152,27 @@ func (p *PeerPool) handleServerPeers(server *p2p.Server, events <-chan *p2p.Peer
 					retryDiscv5 = time.After(0)
 				}
 			case p2p.PeerEventTypeAdd:
-				log.Debug("confirm peer added", "ID", event.Peer)
-				if p.stopOnMax && p.handleAddedPeer(server, event.Peer) {
-					log.Debug("closing discv5 connection", "server", server.Self())
-					server.DiscV5.Close()
-					server.DiscV5 = nil
-					p.feed.Send(Discv5Closed)
-				}
+				p.handleAddedEvent(server, event)
 			}
 		}
 	}
 }
 
+// handledAddedEvent maintains logic of stopping discovery server and added mainly
+// for easier testing.
+func (p *PeerPool) handleAddedEvent(server *p2p.Server, event *p2p.PeerEvent) {
+	log.Debug("confirm peer added", "ID", event.Peer)
+	if p.stopOnMax && p.handleAddedPeer(server, event.Peer) {
+		if server.DiscV5 == nil {
+			return
+		}
+		log.Debug("closing discv5 connection", "server", server.Self())
+		server.DiscV5.Close()
+		server.DiscV5 = nil
+		p.feed.Send(Discv5Closed)
+	}
+}
+
 // handleAddedPeer notifies all topics about added peer and return true if all topics has max limit of connections
 func (p *PeerPool) handleAddedPeer(server *p2p.Server, nodeID discover.NodeID) (all bool) {
 	p.mu.Lock()
diff --git a/geth/peers/peerpool_test.go b/geth/peers/peerpool_test.go
index 28c21e757..d32737086 100644
--- a/geth/peers/peerpool_test.go
+++ b/geth/peers/peerpool_test.go
@@ -12,6 +12,7 @@ import (
 	"github.com/ethereum/go-ethereum/p2p"
 	"github.com/ethereum/go-ethereum/p2p/discover"
 	"github.com/ethereum/go-ethereum/p2p/discv5"
+	"github.com/stretchr/testify/require"
 	"github.com/stretchr/testify/suite"
 
 	"github.com/status-im/status-go/geth/params"
@@ -133,3 +134,28 @@ func (s *PeerPoolSimulationSuite) TearDown() {
 		p.Stop()
 	}
 }
+
+// TestMaxPeersOverflow verifies that following scenario will not occur:
+// - found peer A and B in the same kademlia cycle
+// - process peer A
+// - max limit is reached -> closed discv5 and set it to nil
+// - process peer B
+// - panic because discv5 is nil!!!
+func TestMaxPeersOverflow(t *testing.T) {
+	key, err := crypto.GenerateKey()
+	require.NoError(t, err)
+	peer := &p2p.Server{
+		Config: p2p.Config{
+			PrivateKey:  key,
+			DiscoveryV5: true,
+			NoDiscovery: true,
+		},
+	}
+	require.NoError(t, peer.Start())
+	defer peer.Stop()
+	require.NotNil(t, peer.DiscV5)
+	pool := NewPeerPool(nil, DefaultFastSync, DefaultSlowSync, nil, true)
+	pool.handleAddedEvent(peer, &p2p.PeerEvent{})
+	require.Nil(t, peer.DiscV5)
+	pool.handleAddedEvent(peer, &p2p.PeerEvent{})
+}