p2p: throttle peers on remote close too, add throttling tests

p2p/peer_error.go

@@ -73,7 +73,7 @@ const (
 	DiscSelf
 	DiscReadTimeout
 	DiscSubprotocolError
-	DiscFailureCooldown
+	DiscThrottled
 )
 
 var discReasonToString = [...]string{
@@ -90,7 +90,7 @@ var discReasonToString = [...]string{
 	DiscSelf:                "Connected to self",
 	DiscReadTimeout:         "Read timeout",
 	DiscSubprotocolError:    "Subprotocol error",
-	DiscFailureCooldown:     "Failure cooldown",
+	DiscThrottled:           "Node throttled",
 }
 
 func (d DiscReason) String() string {

p2p/server.go

@@ -105,7 +105,7 @@ type Server struct {
 
 	running  bool
 	peers    map[discover.NodeID]*Peer // Currently active peer pool
-	fails   map[discover.NodeID]bool  // List of recently failed connections
+	throttle map[discover.NodeID]bool  // List of recently failed connections
 
 	trusts    map[discover.NodeID]*discover.Node // Map of currently trusted remote nodes
 	trustDial chan *discover.Node                // Dial request channel reserved for the trusted nodes
@@ -204,7 +204,7 @@ func (srv *Server) Start() (err error) {
 	}
 	srv.quit = make(chan struct{})
 	srv.peers = make(map[discover.NodeID]*Peer)
-	srv.fails = make(map[discover.NodeID]bool)
+	srv.throttle = make(map[discover.NodeID]bool)
 
 	// Create the current trust map, and the associated dialing channel
 	srv.trusts = make(map[discover.NodeID]*discover.Node)
@@ -461,24 +461,33 @@ func (srv *Server) dialNode(dest *discover.Node) {
 		// does that when an error occurs.
 		srv.peerWG.Done()
 
-		// Mark the dialing as failed so prevent too fast redials, and set a
-		// timer for re-allowing dialing
-		srv.lock.Lock()
-		srv.fails[dest.ID] = true
-		srv.lock.Unlock()
-		go func() {
-			glog.V(logger.Detail).Infof("Dialing %v failed: %v. Locked until %v\n", dest, err, time.Now().Add(dialFailureCooldown))
-			<-time.After(dialFailureCooldown)
+		// Throttle the node to prevent hammering it with connections
+		glog.V(logger.Detail).Infof("Dialing %v failed: %v, throttling", dest, err)
+		srv.throttleNode(dest.ID)
 
-			srv.lock.Lock()
-			delete(srv.fails, dest.ID)
-			srv.lock.Unlock()
-		}()
 		return
 	}
 	srv.startPeer(conn, dest)
 }
 
+// throttleNode temporarily throttles a node, disabling both egress as well as
+// ingress connections to/from that particular node.
+func (srv *Server) throttleNode(id discover.NodeID) {
+	// Mark the node as throttled
+	srv.lock.Lock()
+	srv.throttle[id] = true
+	srv.lock.Unlock()
+
+	// Wait a while, after which discard the throttling
+	go func() {
+		<-time.After(dialFailureCooldown)
+
+		srv.lock.Lock()
+		delete(srv.throttle, id)
+		srv.lock.Unlock()
+	}()
+}
+
 func (srv *Server) startPeer(fd net.Conn, dest *discover.Node) {
 	// TODO: handle/store session token
 
@@ -525,7 +534,11 @@ func (srv *Server) runPeer(p *Peer) {
 		srv.newPeerHook(p)
 	}
 	discreason := p.run()
+
+	// Drop the node, and throttle it to prevent hammering
 	srv.removePeer(p)
+	srv.throttleNode(p.ID())
+
 	glog.V(logger.Debug).Infof("Removed %v (%v)\n", p, discreason)
 	srvjslog.LogJson(&logger.P2PDisconnected{
 		RemoteId:       p.ID().String(),
@@ -563,8 +576,8 @@ func (srv *Server) checkPeer(id discover.NodeID) (bool, DiscReason) {
 	case id == srv.ntab.Self().ID:
 		return false, DiscSelf
 
-	case srv.fails[id]:
-		return false, DiscFailureCooldown
+	case srv.throttle[id]:
+		return false, DiscThrottled
 
 	default:
 		return true, 0

p2p/server_test.go

@@ -80,7 +80,7 @@ func TestServerDial(t *testing.T) {
 	// run a one-shot TCP server to handle the connection.
 	listener, err := net.Listen("tcp", "127.0.0.1:0")
 	if err != nil {
-		t.Fatalf("could not setup listener: %v")
+		t.Fatalf("could not setup listener: %v", err)
 	}
 	defer listener.Close()
 	accepted := make(chan net.Conn)
@@ -219,6 +219,116 @@ func TestServerDisconnectAtCap(t *testing.T) {
 	}
 }
 
+// Tests that a failed dial will temporarily throttle a peer.
+func TestServerOutboundThrottling(t *testing.T) {
+	defer testlog(t).detach()
+
+	// Start a simple test server
+	server := startTestServer(t, nil)
+	defer server.Stop()
+
+	// Request a failing dial from the server and check throttling
+	node := &discover.Node{
+		ID:      discover.PubkeyID(&newkey().PublicKey),
+		IP:      []byte{127, 0, 0, 1},
+		TCPPort: 65535,
+	}
+	server.trustDial <- node
+
+	time.Sleep(100 * time.Millisecond)
+	if throttle := server.throttle[node.ID]; !throttle {
+		t.Fatalf("throttling mismatch: have %v, want %v", throttle, true)
+	}
+	// Open a one shot TCP server to listen for valid dialing
+	listener, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatalf("failed to setup listener: %v", err)
+	}
+	defer listener.Close()
+
+	connected := make(chan net.Conn)
+	go func() {
+		conn, err := listener.Accept()
+		if err == nil {
+			conn.Close()
+			connected <- conn
+		}
+	}()
+	time.Sleep(100 * time.Millisecond)
+
+	// Request a re-dial from the server, make sure this also fails
+	addr := listener.Addr().(*net.TCPAddr)
+	node.IP, node.TCPPort = addr.IP, addr.Port
+	server.trustDial <- node
+
+	select {
+	case peer := <-connected:
+		t.Fatalf("throttled peer connected: %v", peer)
+
+	case <-time.After(100 * time.Millisecond):
+		// Ok, peer not accepted
+	}
+}
+
+// Tests that a failed peer will get temporarily throttled.
+func TestServerInboundThrottling(t *testing.T) {
+	defer testlog(t).detach()
+
+	// Start a test server and a peer sink for synchronization
+	started := make(chan *Peer)
+	server := &Server{
+		ListenAddr:  "127.0.0.1:0",
+		PrivateKey:  newkey(),
+		MaxPeers:    10,
+		NoDial:      true,
+		newPeerHook: func(p *Peer) { started <- p },
+	}
+	if err := server.Start(); err != nil {
+		t.Fatal("failed to start test server: %v", err)
+	}
+	defer server.Stop()
+
+	// Create a new peer identity to check throttling with
+	dialer := &net.Dialer{Deadline: time.Now().Add(3 * time.Second)}
+	key := newkey()
+
+	// Connect with a new peer, run the handshake and disconnect
+	connection, err := dialer.Dial("tcp", server.ListenAddr)
+	if err != nil {
+		t.Fatalf("failed to dial server: %v", err)
+	}
+	shake := &protoHandshake{Version: baseProtocolVersion, ID: discover.PubkeyID(&key.PublicKey)}
+	if _, err = setupConn(connection, key, shake, server.Self(), false); err != nil {
+		t.Fatalf("failed to run handshake: %v", err)
+	}
+	<-started
+	connection.Close()
+	server.peerWG.Wait()
+
+	// Check that throttling has been enabled
+	if throttle := server.throttle[shake.ID]; !throttle {
+		t.Fatalf("throttling mismatch: have %v, want %v", throttle, true)
+	}
+	// Try to reconnect with the same peer, run the handshake and verify throttling
+	connection, err = dialer.Dial("tcp", server.ListenAddr)
+	if err != nil {
+		t.Fatalf("failed to re-dial server: %v", err)
+	}
+	defer connection.Close()
+
+	shake = &protoHandshake{Version: baseProtocolVersion, ID: discover.PubkeyID(&key.PublicKey)}
+	if _, err = setupConn(connection, key, shake, server.Self(), false); err != nil {
+		t.Fatalf("failed to re-run handshake: %v", err)
+	}
+	select {
+	case peer := <-started:
+		t.Fatalf("throttled peer accepted: %v", peer)
+
+	case <-time.After(100 * time.Millisecond):
+		// Ok, peer not accepted
+	}
+}
+
 func newkey() *ecdsa.PrivateKey {
 	key, err := crypto.GenerateKey()
 	if err != nil {