forked from cerc-io/plugeth
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.
This commit is contained in:
parent
23454dcfcb
commit
914e57e49b
@ -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")
|
||||
}
|
||||
|
@ -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)
|
||||
}
|
||||
|
@ -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{
|
||||
|
@ -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
|
||||
}
|
||||
|
Loading…
Reference in New Issue
Block a user