From e16bed30fdf0769f13697d3c8e9662457fac755a Mon Sep 17 00:00:00 2001
From: Pierre Krieger <pierre.krieger1708@gmail.com>
Date: Sat, 18 Nov 2023 12:04:10 +0100
Subject: [PATCH] Ban peer after disconnect

---
 lib/src/network/service.rs        | 15 ++++++++++
 light-base/src/network_service.rs | 47 +++++++++++++++++++++++++++----
 wasm-node/CHANGELOG.md            |  1 +
 3 files changed, 57 insertions(+), 6 deletions(-)

diff --git a/lib/src/network/service.rs b/lib/src/network/service.rs
index 8da2ebafa4..2e7a04a67c 100644
--- a/lib/src/network/service.rs
+++ b/lib/src/network/service.rs
@@ -901,6 +901,21 @@ where
         &self.inner[id].address
     }
 
+    /// Returns the number of connections with the given peer.
+    ///
+    /// Both connections that have and have not finished their handshaking phase are considered.
+    pub fn num_potential_and_established_connections(&self, peer_id: &PeerId) -> usize {
+        let Some(&peer_index) = self.peers_by_peer_id.get(peer_id) else {
+            return 0;
+        };
+
+        self.connections_by_peer_id
+            .range(
+                (peer_index, ConnectionId::min_value())..=(peer_index, ConnectionId::max_value()),
+            )
+            .count()
+    }
+
     /// Pulls a message that must be sent to a connection.
     ///
     /// The message must be passed to [`SingleStreamConnectionTask::inject_coordinator_message`]
diff --git a/light-base/src/network_service.rs b/light-base/src/network_service.rs
index ccf3d44d49..4d4fa88102 100644
--- a/light-base/src/network_service.rs
+++ b/light-base/src/network_service.rs
@@ -1526,6 +1526,7 @@ async fn background_task<TPlat: PlatformRef>(mut task: BackgroundTask<TPlat>) {
 
                     task.peering_strategy
                         .remove_address(expected_peer_id, remote_addr.as_ref());
+                    // TODO: if Bob says that its address is the same as Alice's, and we try to connect to both Alice and Bob, then the Bob connection will reach this path and set Alice's address as connected even though it's already connected; this will later cause a state mismatch when disconnecting
                     let _ = task.peering_strategy.insert_or_set_connected_address(
                         &peer_id,
                         remote_addr.clone().into_vec(),
@@ -1540,12 +1541,46 @@ async fn background_task<TPlat: PlatformRef>(mut task: BackgroundTask<TPlat>) {
                 expected_peer_id,
                 ..
             }) => {
-                if let Some(expected_peer_id) = expected_peer_id {
-                    task.peering_strategy
-                        .disconnect_addr(&expected_peer_id, &address)
-                        .unwrap();
-                    let address = Multiaddr::try_from(address).unwrap();
-                    log::debug!(target: "network", "Connections({}, {}) => Shutdown(handshake_finished=false)", expected_peer_id, address);
+                let Some(expected_peer_id) = expected_peer_id else {
+                    debug_assert!(false);
+                    continue;
+                };
+
+                task.peering_strategy
+                    .disconnect_addr(&expected_peer_id, &address)
+                    .unwrap();
+                let address = Multiaddr::try_from(address).unwrap();
+                log::debug!(target: "network", "Connections({}, {}) => Shutdown(handshake_finished=false)", expected_peer_id, address);
+
+                // Due to race conditions and peerid mismatches, it is possible that there is
+                // another existing connection or connection attempt with that same peer.
+                // If this isn't the case and that this was the last connection attempt, then
+                // we ban the peer.
+                // Banning the peer is necessary in order to avoid trying over and over again the
+                // same address(es).
+                if task
+                    .network
+                    .num_potential_and_established_connections(&expected_peer_id)
+                    == 0
+                {
+                    let ban_duration = Duration::from_secs(5);
+                    for (&chain_id, what_happened) in task.peering_strategy.unassign_slots_and_ban(
+                        &expected_peer_id,
+                        task.platform.now() + ban_duration,
+                    ) {
+                        if matches!(
+                            what_happened,
+                            basic_peering_strategy::UnassignSlotsAndBan::Banned { had_slot: true }
+                        ) {
+                            log::debug!(
+                                target: "network",
+                                "Slots({}) ∌ {} (reason=pre-handshake-disconnect, ban-duration={:?})",
+                                &task.network[chain_id].log_name,
+                                expected_peer_id,
+                                ban_duration
+                            );
+                        }
+                    }
                 }
             }
             WakeUpReason::NetworkEvent(service::Event::Disconnected {
diff --git a/wasm-node/CHANGELOG.md b/wasm-node/CHANGELOG.md
index 56e4ffeda8..9a61cf44d3 100644
--- a/wasm-node/CHANGELOG.md
+++ b/wasm-node/CHANGELOG.md
@@ -5,6 +5,7 @@
 ### Changed
 
 - Addresses that are not supported by the host platform are now ignored during the discovery process. For example, TCP/IP connections are ignored while in a browser. This avoids populating the address book with peers that we know we can't connect to anyway. ([#1359](https://github.com/smol-dot/smoldot/pull/1359))
+- Smoldot will no longer try to connect to the same address over and over again. ([#1358](https://github.com/smol-dot/smoldot/pull/1358))
 
 ## 2.0.10 - 2023-11-17