p2p: fix disconnect at capacity

With the introduction of static/trusted nodes, the peer count
can go above MaxPeers. Update the capacity check to handle this.
While here, decouple the trusted nodes check from the handshake
by passing a function instead.

4 files changed, 46 insertions, 33 deletions

diff --git a/p2p/handshake.go b/p2p/handshake.go
index 8e611cfd5..4cdcee6d4 100644
--- a/p2p/handshake.go
+++ b/p2p/handshake.go
@@ -65,26 +65,26 @@ type protoHandshake struct {
 	ID         discover.NodeID
 }
 
-// setupConn starts a protocol session on the given connection.
-// It runs the encryption handshake and the protocol handshake.
-// If dial is non-nil, the connection the local node is the initiator.
-// If atcap is true, the connection will be disconnected with DiscTooManyPeers
-// after the key exchange.
-func setupConn(fd net.Conn, prv *ecdsa.PrivateKey, our *protoHandshake, dial *discover.Node, atcap bool, trusted map[discover.NodeID]bool) (*conn, error) {
+// setupConn starts a protocol session on the given connection. It
+// runs the encryption handshake and the protocol handshake. If dial
+// is non-nil, the connection the local node is the initiator. If
+// keepconn returns false, the connection will be disconnected with
+// DiscTooManyPeers after the key exchange.
+func setupConn(fd net.Conn, prv *ecdsa.PrivateKey, our *protoHandshake, dial *discover.Node, keepconn func(discover.NodeID) bool) (*conn, error) {
 	if dial == nil {
-		return setupInboundConn(fd, prv, our, atcap, trusted)
+		return setupInboundConn(fd, prv, our, keepconn)
 	} else {
-		return setupOutboundConn(fd, prv, our, dial, atcap, trusted)
+		return setupOutboundConn(fd, prv, our, dial, keepconn)
 	}
 }
 
-func setupInboundConn(fd net.Conn, prv *ecdsa.PrivateKey, our *protoHandshake, atcap bool, trusted map[discover.NodeID]bool) (*conn, error) {
+func setupInboundConn(fd net.Conn, prv *ecdsa.PrivateKey, our *protoHandshake, keepconn func(discover.NodeID) bool) (*conn, error) {
 	secrets, err := receiverEncHandshake(fd, prv, nil)
 	if err != nil {
 		return nil, fmt.Errorf("encryption handshake failed: %v", err)
 	}
 	rw := newRlpxFrameRW(fd, secrets)
-	if atcap && !trusted[secrets.RemoteID] {
+	if !keepconn(secrets.RemoteID) {
 		SendItems(rw, discMsg, DiscTooManyPeers)
 		return nil, errors.New("we have too many peers")
 	}
@@ -99,13 +99,13 @@ func setupInboundConn(fd net.Conn, prv *ecdsa.PrivateKey, our *protoHandshake, a
 	return &conn{rw, rhs}, nil
 }
 
-func setupOutboundConn(fd net.Conn, prv *ecdsa.PrivateKey, our *protoHandshake, dial *discover.Node, atcap bool, trusted map[discover.NodeID]bool) (*conn, error) {
+func setupOutboundConn(fd net.Conn, prv *ecdsa.PrivateKey, our *protoHandshake, dial *discover.Node, keepconn func(discover.NodeID) bool) (*conn, error) {
 	secrets, err := initiatorEncHandshake(fd, prv, dial.ID, nil)
 	if err != nil {
 		return nil, fmt.Errorf("encryption handshake failed: %v", err)
 	}
 	rw := newRlpxFrameRW(fd, secrets)
-	if atcap && !trusted[secrets.RemoteID] {
+	if !keepconn(secrets.RemoteID) {
 		SendItems(rw, discMsg, DiscTooManyPeers)
 		return nil, errors.New("we have too many peers")
 	}
diff --git a/p2p/handshake_test.go b/p2p/handshake_test.go
index 9018e14f2..ab75921a3 100644
--- a/p2p/handshake_test.go
+++ b/p2p/handshake_test.go
@@ -141,9 +141,10 @@ func TestSetupConn(t *testing.T) {
 	fd0, fd1 := net.Pipe()
 
 	done := make(chan struct{})
+	keepalways := func(discover.NodeID) bool { return true }
 	go func() {
 		defer close(done)
-		conn0, err := setupConn(fd0, prv0, hs0, node1, false, nil)
+		conn0, err := setupConn(fd0, prv0, hs0, node1, keepalways)
 		if err != nil {
 			t.Errorf("outbound side error: %v", err)
 			return
@@ -156,7 +157,7 @@ func TestSetupConn(t *testing.T) {
 		}
 	}()
 
-	conn1, err := setupConn(fd1, prv1, hs1, nil, false, nil)
+	conn1, err := setupConn(fd1, prv1, hs1, nil, keepalways)
 	if err != nil {
 		t.Fatalf("inbound side error: %v", err)
 	}
diff --git a/p2p/server.go b/p2p/server.go
index 77f66f167..0598547e4 100644
--- a/p2p/server.go
+++ b/p2p/server.go
@@ -126,7 +126,7 @@ type Server struct {
 	peerWG sync.WaitGroup // active peer goroutines
 }
 
-type setupFunc func(net.Conn, *ecdsa.PrivateKey, *protoHandshake, *discover.Node, bool, map[discover.NodeID]bool) (*conn, error)
+type setupFunc func(net.Conn, *ecdsa.PrivateKey, *protoHandshake, *discover.Node, func(discover.NodeID) bool) (*conn, error)
 type newPeerHook func(*Peer)
 
 // Peers returns all connected peers.
@@ -506,17 +506,7 @@ func (srv *Server) startPeer(fd net.Conn, dest *discover.Node) {
 	// the callers of startPeer added the peer to the wait group already.
 	fd.SetDeadline(time.Now().Add(handshakeTimeout))
 
-	// Check capacity, but override for static nodes
-	srv.lock.RLock()
-	atcap := len(srv.peers) == srv.MaxPeers
-	if dest != nil {
-		if _, ok := srv.staticNodes[dest.ID]; ok {
-			atcap = false
-		}
-	}
-	srv.lock.RUnlock()
-
-	conn, err := srv.setupFunc(fd, srv.PrivateKey, srv.ourHandshake, dest, atcap, srv.trustedNodes)
+	conn, err := srv.setupFunc(fd, srv.PrivateKey, srv.ourHandshake, dest, srv.keepconn)
 	if err != nil {
 		fd.Close()
 		glog.V(logger.Debug).Infof("Handshake with %v failed: %v", fd.RemoteAddr(), err)
@@ -539,6 +529,21 @@ func (srv *Server) startPeer(fd net.Conn, dest *discover.Node) {
 	go srv.runPeer(p)
 }
 
+// preflight checks whether a connection should be kept. it runs
+// after the encryption handshake, as soon as the remote identity is
+// known.
+func (srv *Server) keepconn(id discover.NodeID) bool {
+	srv.lock.RLock()
+	defer srv.lock.RUnlock()
+	if _, ok := srv.staticNodes[id]; ok {
+		return true // static nodes are always allowed
+	}
+	if _, ok := srv.trustedNodes[id]; ok {
+		return true // trusted nodes are always allowed
+	}
+	return len(srv.peers) < srv.MaxPeers
+}
+
 func (srv *Server) runPeer(p *Peer) {
 	glog.V(logger.Debug).Infof("Added %v\n", p)
 	srvjslog.LogJson(&logger.P2PConnected{
diff --git a/p2p/server_test.go b/p2p/server_test.go
index a5e56fa18..bf9df31ab 100644
--- a/p2p/server_test.go
+++ b/p2p/server_test.go
@@ -22,8 +22,11 @@ func startTestServer(t *testing.T, pf newPeerHook) *Server {
 		ListenAddr:  "127.0.0.1:0",
 		PrivateKey:  newkey(),
 		newPeerHook: pf,
-		setupFunc: func(fd net.Conn, prv *ecdsa.PrivateKey, our *protoHandshake, dial *discover.Node, atcap bool, trusted map[discover.NodeID]bool) (*conn, error) {
+		setupFunc: func(fd net.Conn, prv *ecdsa.PrivateKey, our *protoHandshake, dial *discover.Node, keepconn func(discover.NodeID) bool) (*conn, error) {
 			id := randomID()
+			if !keepconn(id) {
+				return nil, DiscAlreadyConnected
+			}
 			rw := newRlpxFrameRW(fd, secrets{
 				MAC:        zero16,
 				AES:        zero16,
@@ -200,7 +203,7 @@ func TestServerDisconnectAtCap(t *testing.T) {
 		// Run the handshakes just like a real peer would.
 		key := newkey()
 		hs := &protoHandshake{Version: baseProtocolVersion, ID: discover.PubkeyID(&key.PublicKey)}
-		_, err = setupConn(conn, key, hs, srv.Self(), false, srv.trustedNodes)
+		_, err = setupConn(conn, key, hs, srv.Self(), keepalways)
 		if i == nconns-1 {
 			// When handling the last connection, the server should
 			// disconnect immediately instead of running the protocol
@@ -250,7 +253,7 @@ func TestServerStaticPeers(t *testing.T) {
 		// Run the handshakes just like a real peer would, and wait for completion
 		key := newkey()
 		shake := &protoHandshake{Version: baseProtocolVersion, ID: discover.PubkeyID(&key.PublicKey)}
-		if _, err = setupConn(conn, key, shake, server.Self(), false, server.trustedNodes); err != nil {
+		if _, err = setupConn(conn, key, shake, server.Self(), keepalways); err != nil {
 			t.Fatalf("conn %d: unexpected error: %v", i, err)
 		}
 		<-started
@@ -344,7 +347,7 @@ func TestServerTrustedPeers(t *testing.T) {
 		// Run the handshakes just like a real peer would, and wait for completion
 		key := newkey()
 		shake := &protoHandshake{Version: baseProtocolVersion, ID: discover.PubkeyID(&key.PublicKey)}
-		if _, err = setupConn(conn, key, shake, server.Self(), false, server.trustedNodes); err != nil {
+		if _, err = setupConn(conn, key, shake, server.Self(), keepalways); err != nil {
 			t.Fatalf("conn %d: unexpected error: %v", i, err)
 		}
 		<-started
@@ -357,7 +360,7 @@ func TestServerTrustedPeers(t *testing.T) {
 	defer conn.Close()
 
 	shake := &protoHandshake{Version: baseProtocolVersion, ID: trusted.ID}
-	if _, err = setupConn(conn, key, shake, server.Self(), false, server.trustedNodes); err != nil {
+	if _, err = setupConn(conn, key, shake, server.Self(), keepalways); err != nil {
 		t.Fatalf("trusted node: unexpected error: %v", err)
 	}
 	select {
@@ -472,7 +475,7 @@ func TestServerMaxPendingAccepts(t *testing.T) {
 	go func() {
 		key := newkey()
 		shake := &protoHandshake{Version: baseProtocolVersion, ID: discover.PubkeyID(&key.PublicKey)}
-		if _, err := setupConn(conns[1], key, shake, server.Self(), false, server.trustedNodes); err != nil {
+		if _, err := setupConn(conns[1], key, shake, server.Self(), keepalways); err != nil {
 			t.Fatalf("failed to run handshake: %v", err)
 		}
 	}()
@@ -486,7 +489,7 @@ func TestServerMaxPendingAccepts(t *testing.T) {
 	go func() {
 		key := newkey()
 		shake := &protoHandshake{Version: baseProtocolVersion, ID: discover.PubkeyID(&key.PublicKey)}
-		if _, err := setupConn(conns[0], key, shake, server.Self(), false, server.trustedNodes); err != nil {
+		if _, err := setupConn(conns[0], key, shake, server.Self(), keepalways); err != nil {
 			t.Fatalf("failed to run handshake: %v", err)
 		}
 	}()
@@ -513,3 +516,7 @@ func randomID() (id discover.NodeID) {
 	}
 	return id
 }
+
+func keepalways(id discover.NodeID) bool {
+	return true
+}