swarm/: Patch reporting on banned peer connections

swarm/Cargo.toml

@@ -19,6 +19,7 @@ smallvec = "1.6.1"
 void = "1"
 futures-timer = "3.0.2"
 instant = "0.1.11"
+env_logger = "0.9"
 
 [dev-dependencies]
 libp2p = { path = "../", default-features = false, features = ["yamux", "plaintext"] }

swarm/src/lib.rs

@@ -275,6 +275,9 @@ where
     /// List of nodes for which we deny any incoming connection.
     banned_peers: HashSet<PeerId>,
 
+    /// Connections for which we withhold any reporting. These belong to banned peers.
+    banned_connections: HashSet<ConnectionId>,
+
     /// Pending event to be delivered to connection handlers
     /// (or dropped if the peer disconnected) before the `behaviour`
     /// can be polled again.
@@ -603,22 +606,32 @@ where
                 Poll::Pending => network_not_ready = true,
                 Poll::Ready(NetworkEvent::ConnectionEvent { connection, event }) => {
                     let peer = connection.peer_id();
-                    let connection = connection.id();
-                    this.behaviour.inject_event(peer, connection, event);
+                    let conn_id = connection.id();
+                    if !this.banned_connections.contains(&conn_id) {
+                        this.behaviour.inject_event(peer, conn_id, event);
+                    } else {
+                        log::debug!(
+                            "Ignoring event from disallowed connection: {:?} {:?}.",
+                            peer,
+                            conn_id,
+                        );
+                    }
                 }
                 Poll::Ready(NetworkEvent::AddressChange {
                     connection,
                     new_endpoint,
                     old_endpoint,
                 }) => {
                     let peer = connection.peer_id();
-                    let connection = connection.id();
-                    this.behaviour.inject_address_change(
-                        &peer,
-                        &connection,
-                        &old_endpoint,
-                        &new_endpoint,
-                    );
+                    let conn_id = connection.id();
+                    if !this.banned_connections.contains(&conn_id) {
+                        this.behaviour.inject_address_change(
+                            &peer,
+                            &conn_id,
+                            &old_endpoint,
+                            &new_endpoint,
+                        );
+                    }
                 }
                 Poll::Ready(NetworkEvent::ConnectionEstablished {
                     connection,
@@ -628,6 +641,8 @@ where
                     let peer_id = connection.peer_id();
                     let endpoint = connection.endpoint().clone();
                     if this.banned_peers.contains(&peer_id) {
+                        // Mark the connection for the banned peer as disallowed.
+                        this.banned_connections.insert(connection.id());
                         this.network
                             .peer(peer_id)
                             .into_connected()
@@ -676,14 +691,16 @@ where
                     }
                     let peer_id = connected.peer_id;
                     let endpoint = connected.endpoint;
-                    this.behaviour.inject_connection_closed(
-                        &peer_id,
-                        &id,
-                        &endpoint,
-                        handler.into_protocols_handler(),
-                    );
-                    if num_established == 0 {
-                        this.behaviour.inject_disconnected(&peer_id);
+                    if !this.banned_connections.remove(&id) {
+                        this.behaviour.inject_connection_closed(
+                            &peer_id,
+                            &id,
+                            &endpoint,
+                            handler.into_protocols_handler(),
+                        );
+                        if num_established == 0 {
+                            this.behaviour.inject_disconnected(&peer_id);
+                        }
                     }
                     return Poll::Ready(SwarmEvent::ConnectionClosed {
                         peer_id,
@@ -1253,6 +1270,7 @@ where
             listened_addrs: SmallVec::new(),
             external_addrs: Addresses::default(),
             banned_peers: HashSet::new(),
+            banned_connections: HashSet::new(),
             pending_event: None,
             substream_upgrade_protocol_override: self.substream_upgrade_protocol_override,
         }
@@ -1450,15 +1468,6 @@ mod tests {
         TBehaviour: NetworkBehaviour,
         <<TBehaviour::ProtocolsHandler as IntoProtocolsHandler>::Handler as ProtocolsHandler>::OutEvent: Clone,
     {
-        for s in &[swarm1, swarm2] {
-            if s.behaviour.inject_connection_established.len() > 0 {
-                assert_eq!(s.behaviour.inject_connected.len(), 1);
-            } else {
-                assert_eq!(s.behaviour.inject_connected.len(), 0);
-            }
-            assert!(s.behaviour.inject_connection_closed.is_empty());
-            assert!(s.behaviour.inject_disconnected.is_empty());
-        }
         [swarm1, swarm2]
             .iter()
             .all(|s| s.behaviour.inject_connection_established.len() == num_connections)
@@ -1473,10 +1482,6 @@ mod tests {
         TBehaviour: NetworkBehaviour,
         <<TBehaviour::ProtocolsHandler as IntoProtocolsHandler>::Handler as ProtocolsHandler>::OutEvent: Clone
     {
-        for s in &[swarm1, swarm2] {
-            assert_eq!(s.behaviour.inject_connection_established.len(), 0);
-            assert_eq!(s.behaviour.inject_connected.len(), 0);
-        }
         [swarm1, swarm2]
             .iter()
             .all(|s| s.behaviour.inject_connection_closed.len() == num_connections)
@@ -1490,7 +1495,12 @@ mod tests {
     ///
     /// The test expects both behaviours to be notified via pairs of
     /// inject_connected / inject_disconnected as well as
-    /// inject_connection_established / inject_connection_closed calls.
+    /// inject_connection_established / inject_connection_closed calls
+    /// while unbanned.
+    ///
+    /// While the ban is in effect, further dials occur. For these connections no
+    /// `inject_connected`, `inject_connection_established`, `inject_disconnected`,
+    /// `inject_connection_closed` calls should be registered.
     #[test]
     fn test_connect_disconnect_ban() {
         // Since the test does not try to open any substreams, we can
@@ -1510,55 +1520,104 @@ mod tests {
 
         let swarm1_id = *swarm1.local_peer_id();
 
-        let mut banned = false;
-        let mut unbanned = false;
+        enum Stage {
+            /// Waiting for the peers to connect. Banning has not occurred.
+            Connecting,
+            /// Ban occurred.
+            Banned,
+            // Ban is in place and a dial is ongoing.
+            BannedDial,
+            // Mid-ban dial was registered and the peer was unbanned.
+            Unbanned,
+            // There are dial attempts ongoing for the no longer banned peers.
+            Reconnecting,
+        }
 
-        let num_connections = 10;
+        let num_connections = 2;
 
         for _ in 0..num_connections {
             swarm1.dial(addr2.clone()).unwrap();
         }
-        let mut state = State::Connecting;
 
-        executor::block_on(future::poll_fn(move |cx| {
-            loop {
-                let poll1 = Swarm::poll_next_event(Pin::new(&mut swarm1), cx);
-                let poll2 = Swarm::poll_next_event(Pin::new(&mut swarm2), cx);
-                match state {
-                    State::Connecting => {
-                        if swarms_connected(&swarm1, &swarm2, num_connections) {
-                            if banned {
-                                return Poll::Ready(());
-                            }
-                            swarm2.ban_peer_id(swarm1_id.clone());
-                            swarm1.behaviour.reset();
-                            swarm2.behaviour.reset();
-                            banned = true;
-                            state = State::Disconnecting;
-                        }
+        let mut swarm1_expected_conns = num_connections;
+        let mut swarm2_expected_conns = num_connections;
+
+        log::info!("Starting test");
+        let mut stage = Stage::Connecting;
+
+        executor::block_on(future::poll_fn(move |cx| loop {
+            let poll1 = Swarm::poll_next_event(Pin::new(&mut swarm1), cx);
+            let poll2 = Swarm::poll_next_event(Pin::new(&mut swarm2), cx);
+            match stage {
+                Stage::Connecting => {
+                    if swarms_connected(&swarm1, &swarm2, num_connections) {
+                        log::info!("1");
+                        // Setup to test that already established connections are correctly closed
+                        // and reported as such after the peer in banned.
+                        swarm2.ban_peer_id(swarm1_id.clone());
+                        stage = Stage::Banned;
                     }
-                    State::Disconnecting => {
-                        if swarms_disconnected(&swarm1, &swarm2, num_connections) {
-                            if unbanned {
-                                return Poll::Ready(());
-                            }
-                            // Unban the first peer and reconnect.
-                            swarm2.unban_peer_id(swarm1_id.clone());
-                            swarm1.behaviour.reset();
-                            swarm2.behaviour.reset();
-                            unbanned = true;
-                            for _ in 0..num_connections {
-                                swarm2.dial(addr1.clone()).unwrap();
-                            }
-                            state = State::Connecting;
+                }
+                Stage::Banned => {
+                    if swarms_disconnected(&swarm1, &swarm2, num_connections) {
+                        log::info!("2");
+                        // Setup to test that new connections of banned peers are not reported.
+                        swarm1.dial(addr2.clone()).unwrap();
+                        swarm1_expected_conns += 1;
+                        stage = Stage::BannedDial;
+                    }
+                }
+                Stage::BannedDial => {
+                    if swarm1.behaviour.inject_connection_established.len() == swarm1_expected_conns
+                        && swarm2.network_info().num_peers() == 1
+                    {
+                        log::info!("3");
+                        // The banned connection was established. Check that then banning swarm was
+                        // not reported about the connection.
+                        assert_eq!(
+                            swarm2.behaviour.inject_connection_established.len(), swarm2_expected_conns,
+                            "No additional closed connections should be reported for the banned peer"
+                        );
+                        // Setup to test that the banned connection is not reported upon closing
+                        // even if the peer is unbanned.
+                        swarm2.unban_peer_id(swarm1_id);
+                        stage = Stage::Unbanned;
+                    }
+                }
+                Stage::Unbanned => {
+                    if swarm2.network_info().num_peers() == 0 {
+                        log::info!("4");
+                        // The banned connection has closed. Check that it was not reported.
+                        assert_eq!(
+                            swarm2.behaviour.inject_connection_closed.len(), swarm2_expected_conns,
+                            "No additional closed connections should be reported for the banned peer"
+                        );
+                        assert!(swarm2.banned_connections.is_empty());
+                        // Setup to test that a ban lifted does not affect future connections.
+                        for _ in 0..num_connections {
+                            swarm1.dial(addr2.clone()).unwrap();
                         }
+                        swarm1_expected_conns += num_connections;
+                        swarm2_expected_conns += num_connections;
+                        stage = Stage::Reconnecting;
                     }
                 }
-
-                if poll1.is_pending() && poll2.is_pending() {
-                    return Poll::Pending;
+                Stage::Reconnecting => {
+                    if swarm1.behaviour.inject_connection_established.len() == swarm1_expected_conns
+                    {
+                        log::info!("5");
+                        assert_eq!(
+                            swarm2.behaviour.inject_connection_established.len(),
+                            swarm2_expected_conns
+                        );
+                        return Poll::Ready(());
+                    }
                 }
             }
+
+            if poll1.is_pending() && poll2.is_pending() {
+                return Poll::Pending;
+            }
         }))
     }
 
@@ -1607,8 +1666,6 @@ mod tests {
                         swarm2
                             .disconnect_peer_id(swarm1_id.clone())
                             .expect("Error disconnecting");
-                        swarm1.behaviour.reset();
-                        swarm2.behaviour.reset();
                         state = State::Disconnecting;
                     }
                 }
@@ -1618,8 +1675,6 @@ mod tests {
                             return Poll::Ready(());
                         }
                         reconnected = true;
-                        swarm1.behaviour.reset();
-                        swarm2.behaviour.reset();
                         for _ in 0..num_connections {
                             swarm2.dial(addr1.clone()).unwrap();
                         }
@@ -1683,8 +1738,6 @@ mod tests {
                                 connection: CloseConnection::All,
                             },
                         );
-                        swarm1.behaviour.reset();
-                        swarm2.behaviour.reset();
                         state = State::Disconnecting;
                     }
                 }
@@ -1694,8 +1747,6 @@ mod tests {
                             return Poll::Ready(());
                         }
                         reconnected = true;
-                        swarm1.behaviour.reset();
-                        swarm2.behaviour.reset();
                         for _ in 0..num_connections {
                             swarm2.dial(addr1.clone()).unwrap();
                         }
@@ -1762,16 +1813,17 @@ mod tests {
                             );
                             Some(conn_id)
                         };
-                        swarm1.behaviour.reset();
-                        swarm2.behaviour.reset();
                         state = State::Disconnecting;
                     }
                 }
                 State::Disconnecting => {
                     for s in &[&swarm1, &swarm2] {
                         assert_eq!(s.behaviour.inject_disconnected.len(), 0);
-                        assert_eq!(s.behaviour.inject_connection_established.len(), 0);
-                        assert_eq!(s.behaviour.inject_connected.len(), 0);
+                        assert_eq!(
+                            s.behaviour.inject_connection_established.len(),
+                            num_connections
+                        );
+                        assert_eq!(s.behaviour.inject_connected.len(), 1);
                     }
                     if [&swarm1, &swarm2]
                         .iter()