diff --git a/protocols/kad/CHANGELOG.md b/protocols/kad/CHANGELOG.md index c2d117b244c..bac3d952a56 100644 --- a/protocols/kad/CHANGELOG.md +++ b/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). diff --git a/protocols/kad/src/addresses.rs b/protocols/kad/src/addresses.rs index 3c2af4173fd..da0754c9823 100644 --- a/protocols/kad/src/addresses.rs +++ b/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) -> 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() + } +}