swarm/: Patch reporting on banned peer connections

core/src/connection.rs

@@ -206,6 +206,8 @@ pub struct Connected {
     pub endpoint: ConnectedPoint,
     /// Information obtained from the transport.
     pub peer_id: PeerId,
+    /// If the connection was allowed by the network.
+    pub is_allowed: bool,
 }
 
 /// Event generated by a [`Connection`].

core/src/connection/pool.rs

@@ -112,6 +112,8 @@ struct EstablishedConnectionInfo<TInEvent> {
     /// [`PeerId`] of the remote peer.
     peer_id: PeerId,
     endpoint: ConnectedPoint,
+    /// Whether the connection was allowed.
+    is_allowed: bool,
     /// Channel endpoint to send commands to the task.
     sender: mpsc::Sender<task::Command<TInEvent>>,
 }
@@ -690,16 +692,23 @@ where
                     .established
                     .get_mut(&peer_id)
                     .expect("`Closed` event for established connection");
-                let EstablishedConnectionInfo { endpoint, .. } =
-                    connections.remove(&id).expect("Connection to be present");
+                let EstablishedConnectionInfo {
+                    endpoint,
+                    is_allowed,
+                    ..
+                } = connections.remove(&id).expect("Connection to be present");
                 self.counters.dec_established(&endpoint);
                 let num_established = u32::try_from(connections.len()).unwrap();
                 if num_established == 0 {
                     self.established.remove(&peer_id);
                 }
                 return Poll::Ready(PoolEvent::ConnectionClosed {
                     id,
-                    connected: Connected { endpoint, peer_id },
+                    connected: Connected {
+                        endpoint,
+                        peer_id,
+                        is_allowed,
+                    },
                     error,
                     num_established,
                     pool: self,
@@ -855,16 +864,23 @@ where
 
                     let (command_sender, command_receiver) =
                         mpsc::channel(self.task_command_buffer_size);
+                    // Any new connection is allowed by default.
+                    let is_allowed = true;
                     conns.insert(
                         id,
                         EstablishedConnectionInfo {
                             peer_id,
                             endpoint: endpoint.clone(),
+                            is_allowed,
                             sender: command_sender,
                         },
                     );
 
-                    let connected = Connected { peer_id, endpoint };
+                    let connected = Connected {
+                        peer_id,
+                        endpoint,
+                        is_allowed,
+                    };
 
                     let connection =
                         super::Connection::new(muxer, handler.into_handler(&connected));
@@ -1010,6 +1026,16 @@ impl<TInEvent> EstablishedConnection<'_, TInEvent> {
         *self.entry.key()
     }
 
+    /// Return whether the connection was allowed from the network.
+    pub fn is_allowed(&self) -> bool {
+        self.entry.get().is_allowed
+    }
+
+    /// Marks the connection as disallowed by the network.
+    pub fn disallow(&mut self) {
+        self.entry.get_mut().is_allowed = false;
+    }
+
     /// (Asynchronously) sends an event to the connection handler.
     ///
     /// If the handler is not ready to receive the event, either because

swarm/src/lib.rs

@@ -603,31 +603,43 @@ 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 connection.is_allowed() {
+                        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 connection.is_allowed() {
+                        this.behaviour.inject_address_change(
+                            &peer,
+                            &conn_id,
+                            &old_endpoint,
+                            &new_endpoint,
+                        );
+                    }
                 }
                 Poll::Ready(NetworkEvent::ConnectionEstablished {
-                    connection,
+                    mut connection,
                     num_established,
                     concurrent_dial_errors,
                 }) => {
                     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.
+                        connection.disallow();
                         this.network
                             .peer(peer_id)
                             .into_connected()
@@ -676,14 +688,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 connected.is_allowed {
+                        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,
@@ -1789,4 +1803,58 @@ mod tests {
             }
         }))
     }
+
+    #[test]
+    fn test_banning_respects_contract() {
+        let handler_proto = DummyProtocolsHandler {
+            keep_alive: KeepAlive::Yes,
+        };
+
+        let mut swarm1 = new_test_swarm::<_, ()>(handler_proto.clone());
+        let mut swarm2 = new_test_swarm::<_, ()>(handler_proto);
+
+        let addr1: Multiaddr = multiaddr::Protocol::Memory(rand::random::<u64>()).into();
+        let addr2: Multiaddr = multiaddr::Protocol::Memory(rand::random::<u64>()).into();
+
+        swarm1.listen_on(addr1.clone().into()).unwrap();
+        swarm2.listen_on(addr2.clone().into()).unwrap();
+
+        let swarm1_id = *swarm1.local_peer_id();
+
+        async fn poll(
+            id: usize,
+            swarm: &mut Swarm<CallTraceBehaviour<MockBehaviour<DummyProtocolsHandler, ()>>>,
+        ) {
+            let ev = swarm.select_next_some().await;
+            log::info!("[{}] {:?}", id, ev);
+        }
+
+        executor::block_on(async {
+            // Wait for the swarms to establish listeners
+            poll(1, &mut swarm1).await;
+            poll(2, &mut swarm2).await;
+
+            swarm1.dial(addr2.clone()).unwrap();
+            log::info!("Swarm 2 banning swarm 1\n\n");
+            swarm2.ban_peer_id(swarm1_id);
+
+            // Incoming connection
+            poll(2, &mut swarm2).await;
+
+            // Connection established
+            poll(2, &mut swarm2).await;
+            // Connection established
+            poll(1, &mut swarm1).await;
+
+            // Both see connections as closed
+            poll(1, &mut swarm1).await;
+            poll(2, &mut swarm2).await;
+
+            // Check what swarm2's behaviour got
+            assert_eq!(
+                swarm2.behaviour.inject_connection_established,
+                swarm2.behaviour.inject_connection_closed
+            );
+        })
+    }
 }