fix(kad): potentially incorrect return value of Addresses::remove

Adding a check when there is only one address left in the list of `Addresses` to verify that the remaining one does match the provided one when calling `remove`.

Fixes: #4815.

Pull-Request: #4816.

protocols/kad/CHANGELOG.md

@@ -1,5 +1,7 @@
 ## 0.45.1 - unreleased
 
+- Fix a bug where calling `Behaviour::remove_address` with an address not in the peer's bucket would remove the peer from the routing table if the bucket has only one address left.
+  See [PR 4816](https://github.com/libp2p/rust-libp2p/pull/4816)
 - Add `std::fmt::Display` implementation on `QueryId`.
   See [PR 4814](https://github.com/libp2p/rust-libp2p/pull/4814).
 

protocols/kad/src/addresses.rs

@@ -67,7 +67,7 @@ impl Addresses {
     /// otherwise unreachable.
     #[allow(clippy::result_unit_err)]
     pub fn remove(&mut self, addr: &Multiaddr) -> Result<(), ()> {
-        if self.addrs.len() == 1 {
+        if self.addrs.len() == 1 && self.addrs[0] == *addr {
             return Err(());
         }
 
@@ -113,3 +113,76 @@ impl fmt::Debug for Addresses {
         f.debug_list().entries(self.addrs.iter()).finish()
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn given_one_address_when_removing_different_one_returns_ok() {
+        let mut addresses = make_addresses([tcp_addr(1234)]);
+
+        let result = addresses.remove(&tcp_addr(4321));
+
+        assert!(result.is_ok());
+        assert_eq!(
+            addresses.into_vec(),
+            vec![tcp_addr(1234)],
+            "`Addresses` to not change because we tried to remove a non-present address"
+        );
+    }
+
+    #[test]
+    fn given_one_address_when_removing_correct_one_returns_err() {
+        let mut addresses = make_addresses([tcp_addr(1234)]);
+
+        let result = addresses.remove(&tcp_addr(1234));
+
+        assert!(result.is_err());
+        assert_eq!(
+            addresses.into_vec(),
+            vec![tcp_addr(1234)],
+            "`Addresses` to not be empty because it would have been the last address to be removed"
+        );
+    }
+
+    #[test]
+    fn given_many_addresses_when_removing_different_one_does_not_remove_and_returns_ok() {
+        let mut addresses = make_addresses([tcp_addr(1234), tcp_addr(4321)]);
+
+        let result = addresses.remove(&tcp_addr(5678));
+
+        assert!(result.is_ok());
+        assert_eq!(
+            addresses.into_vec(),
+            vec![tcp_addr(1234), tcp_addr(4321)],
+            "`Addresses` to not change because we tried to remove a non-present address"
+        );
+    }
+
+    #[test]
+    fn given_many_addresses_when_removing_correct_one_removes_and_returns_ok() {
+        let mut addresses = make_addresses([tcp_addr(1234), tcp_addr(4321)]);
+
+        let result = addresses.remove(&tcp_addr(1234));
+
+        assert!(result.is_ok());
+        assert_eq!(
+            addresses.into_vec(),
+            vec![tcp_addr(4321)],
+            "`Addresses to no longer contain address was present and then removed`"
+        );
+    }
+
+    /// Helper function to easily initialize Addresses struct with multiple addresses.
+    fn make_addresses(addresses: impl IntoIterator<Item = Multiaddr>) -> Addresses {
+        Addresses {
+            addrs: SmallVec::from_iter(addresses),
+        }
+    }
+
+    /// Helper function to create a tcp Multiaddr with a specific port
+    fn tcp_addr(port: u16) -> Multiaddr {
+        format!("/ip4/127.0.0.1/tcp/{port}").parse().unwrap()
+    }
+}