Make sure connected addresses are inserted when discovering

bin/wasm-node/CHANGELOG.md

@@ -4,6 +4,7 @@
 
 ### Fixed
 
+- Fix panic that occured when connecting to a peer, then discovering it through the background discovery process, then disconnecting from it. ([#2616](https://github.com/paritytech/smoldot/pull/2616))
 - Fix circular dependency between JavaScript modules. ([#2614](https://github.com/paritytech/smoldot/pull/2614))
 
 ## 0.6.29 - 2022-08-09

src/libp2p/peers.rs

@@ -55,17 +55,18 @@ use alloc::{
 };
 use core::{
     hash::Hash,
+    iter,
     num::NonZeroU32,
     ops::{self, Add, Sub},
     time::Duration,
 };
 use rand::{Rng as _, SeedableRng as _};
 
 pub use collection::{
-    ConfigRequestResponse, ConfigRequestResponseIn, ConnectionId, ConnectionToCoordinator,
-    CoordinatorToConnection, InboundError, MultiStreamConnectionTask, NotificationProtocolConfig,
-    NotificationsInClosedErr, NotificationsOutErr, ReadWrite, RequestError, ShutdownCause,
-    SingleStreamConnectionTask, SubstreamId,
+    ConfigRequestResponse, ConfigRequestResponseIn, ConnectionId, ConnectionState,
+    ConnectionToCoordinator, CoordinatorToConnection, InboundError, MultiStreamConnectionTask,
+    NotificationProtocolConfig, NotificationsInClosedErr, NotificationsOutErr, ReadWrite,
+    RequestError, ShutdownCause, SingleStreamConnectionTask, SubstreamId,
 };
 
 /// Configuration for a [`Peers`].
@@ -852,6 +853,65 @@ where
         (connection_id, connection_task)
     }
 
+    /// Returns all the non-handshaking connections that are connected to the given peer. The list
+    /// also includes connections that are shutting down.
+    pub fn established_peer_connections(
+        &'_ self,
+        peer_id: &PeerId,
+    ) -> impl Iterator<Item = ConnectionId> + '_ {
+        let peer_index = match self.peer_indices.get(peer_id) {
+            Some(idx) => *idx,
+            None => return either::Right(iter::empty()),
+        };
+
+        either::Left(
+            self.connections_by_peer
+                .range(
+                    (peer_index, ConnectionId::min_value())
+                        ..=(peer_index, ConnectionId::max_value()),
+                )
+                .map(|(_, connection_id)| *connection_id)
+                .filter(move |connection_id| {
+                    self.inner.connection_state(*connection_id).established
+                }),
+        )
+    }
+
+    /// Returns all the handshaking connections that are expected to reach the given peer. The
+    /// list also includes connections that are shutting down.
+    pub fn handshaking_peer_connections(
+        &'_ self,
+        peer_id: &PeerId,
+    ) -> impl Iterator<Item = ConnectionId> + '_ {
+        let peer_index = match self.peer_indices.get(peer_id) {
+            Some(idx) => *idx,
+            None => return either::Right(iter::empty()),
+        };
+
+        either::Left(
+            self.connections_by_peer
+                .range(
+                    (peer_index, ConnectionId::min_value())
+                        ..=(peer_index, ConnectionId::max_value()),
+                )
+                .map(|(_, connection_id)| *connection_id)
+                .filter(move |connection_id| {
+                    !self.inner.connection_state(*connection_id).established
+                }),
+        )
+    }
+
+    /// Returns the state of the given connection.
+    ///
+    /// # Panic
+    ///
+    /// Panics if the identifier is invalid or corresponds to a connection that has already
+    /// entirely shut down.
+    ///
+    pub fn connection_state(&self, connection_id: ConnectionId) -> ConnectionState {
+        self.inner.connection_state(connection_id)
+    }
+
     /// Returns the list of [`PeerId`]s that have been marked as desired, but that don't have any
     /// associated connection. An associated connection is either a fully established connection
     /// with that peer, or an outgoing connection that is still handshaking but expects to reach

src/network/service.rs

@@ -847,6 +847,7 @@ where
                 kademlia::kbuckets::Entry::LocalKey => return, // TODO: return some diagnostic?
                 kademlia::kbuckets::Entry::Vacant(entry) => {
                     match entry.insert((), now, kademlia::kbuckets::PeerState::Disconnected) {
+                        Err(kademlia::kbuckets::InsertError::Full) => return, // TODO: return some diagnostic?
                         Ok((_, removed_entry)) => {
                             // `removed_entry` is the peer that was removed the k-buckets as the
                             // result of the new insertion. Purge it from `self.kbuckets_peers`
@@ -873,16 +874,52 @@ where
                                         NonZeroUsize::new(e.num_references.get() + 1).unwrap();
                                     e
                                 }
-                                // TODO: is it possible that we're already connected to this peer? since we call set_disconnected() when we disconnect, we would geta panic
-                                hashbrown::hash_map::Entry::Vacant(e) => e.insert(KBucketsPeer {
-                                    num_references: NonZeroUsize::new(1).unwrap(),
-                                    addresses: addresses::Addresses::with_capacity(
+                                hashbrown::hash_map::Entry::Vacant(e) => {
+                                    // The peer was not in the k-buckets, but it is possible that
+                                    // we already have existing connections to it.
+                                    let mut addresses = addresses::Addresses::with_capacity(
                                         self.max_addresses_per_peer.get(),
-                                    ),
-                                }),
+                                    );
+
+                                    for connection_id in
+                                        self.inner.established_peer_connections(&e.key())
+                                    {
+                                        let state = self.inner.connection_state(connection_id);
+                                        debug_assert!(state.established);
+                                        // Because we mark addresses as disconnected when the
+                                        // shutdown process starts, we ignore shutting down
+                                        // connections.
+                                        if state.shutting_down {
+                                            continue;
+                                        }
+                                        addresses
+                                            .insert_discovered(self.inner[connection_id].clone());
+                                        addresses.set_connected(&self.inner[connection_id]);
+                                    }
+
+                                    for connection_id in
+                                        self.inner.handshaking_peer_connections(&e.key())
+                                    {
+                                        let state = self.inner.connection_state(connection_id);
+                                        debug_assert!(!state.established);
+                                        // Because we mark addresses as disconnected when the
+                                        // shutdown process starts, we ignore shutting down
+                                        // connections.
+                                        if state.shutting_down {
+                                            continue;
+                                        }
+                                        addresses
+                                            .insert_discovered(self.inner[connection_id].clone());
+                                        addresses.set_pending(&self.inner[connection_id]);
+                                    }
+
+                                    e.insert(KBucketsPeer {
+                                        num_references: NonZeroUsize::new(1).unwrap(),
+                                        addresses,
+                                    })
+                                }
                             }
                         }
-                        Err(kademlia::kbuckets::InsertError::Full) => return, // TODO: return some diagnostic?
                     }
                 }
                 kademlia::kbuckets::Entry::Occupied(_) => {

src/network/service/addresses.rs

@@ -93,6 +93,19 @@ impl Addresses {
         }
     }
 
+    /// If the given address is in the list, sets its state to "pending".
+    ///
+    /// # Panic
+    ///
+    /// Panics if the state of this address was already pending.
+    ///
+    pub(super) fn set_pending(&mut self, addr: &multiaddr::Multiaddr) {
+        if let Some(index) = self.list.iter().position(|(a, _)| a == addr) {
+            assert!(!matches!(self.list[index].1, State::PendingConnect));
+            self.list[index].1 = State::PendingConnect;
+        }
+    }
+
     /// If the given address is in the list, sets its state to "disconnected".
     ///
     /// # Panic