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 +}