From 9c26137c7ec18d7413eae7f03487155fe4d581da Mon Sep 17 00:00:00 2001 From: driftluo Date: Fri, 1 Sep 2023 11:11:45 +0800 Subject: [PATCH] fix: keep peer store address unique --- network/src/peer_store/addr_manager.rs | 15 +++++++++------ network/src/peer_store/peer_store_impl.rs | 2 +- network/src/tests/mod.rs | 3 ++- network/src/tests/peer_store.rs | 19 +++++++++++++++++++ 4 files changed, 31 insertions(+), 8 deletions(-) diff --git a/network/src/peer_store/addr_manager.rs b/network/src/peer_store/addr_manager.rs index ab9fbbaa70b..218d28769eb 100644 --- a/network/src/peer_store/addr_manager.rs +++ b/network/src/peer_store/addr_manager.rs @@ -18,14 +18,17 @@ impl AddrManager { /// Add an address information to address manager pub fn add(&mut self, mut addr_info: AddrInfo) { if let Some(key) = multiaddr_to_socketaddr(&addr_info.addr) { - if let Some(exists_last_connected_at_ms) = self - .get(&addr_info.addr) - .map(|addr| addr.last_connected_at_ms) - { + if let Some(&id) = self.addr_to_id.get(&key) { + let (exist_last_connected_at_ms, random_id_pos) = { + let info = self.id_to_info.get(&id).expect("must exists"); + (info.last_connected_at_ms, info.random_id_pos) + }; // Get time earlier than record time, return directly - if addr_info.last_connected_at_ms < exists_last_connected_at_ms { - return; + if addr_info.last_connected_at_ms >= exist_last_connected_at_ms { + addr_info.random_id_pos = random_id_pos; + self.id_to_info.insert(id, addr_info); } + return; } let id = self.next_id; diff --git a/network/src/peer_store/peer_store_impl.rs b/network/src/peer_store/peer_store_impl.rs index c68f831ab26..1f347a58ff5 100644 --- a/network/src/peer_store/peer_store_impl.rs +++ b/network/src/peer_store/peer_store_impl.rs @@ -46,7 +46,7 @@ impl PeerStore { .entry(extract_peer_id(&addr).expect("connected addr should have peer id")) { Entry::Occupied(mut entry) => { - let mut peer = entry.get_mut(); + let peer = entry.get_mut(); peer.connected_addr = addr; peer.last_connected_at_ms = now_ms; peer.session_type = session_type; diff --git a/network/src/tests/mod.rs b/network/src/tests/mod.rs index 954775112e5..dde66545f05 100644 --- a/network/src/tests/mod.rs +++ b/network/src/tests/mod.rs @@ -6,7 +6,8 @@ mod peer_store_db; fn random_addr() -> crate::multiaddr::Multiaddr { format!( - "/ip4/127.0.0.1/tcp/42/p2p/{}", + "/ip4/127.0.0.1/tcp/{}/p2p/{}", + rand::random::(), crate::PeerId::random().to_base58() ) .parse() diff --git a/network/src/tests/peer_store.rs b/network/src/tests/peer_store.rs index eeeba15dede..30bca810203 100644 --- a/network/src/tests/peer_store.rs +++ b/network/src/tests/peer_store.rs @@ -533,3 +533,22 @@ fn test_eviction() { .unwrap(); assert!(peer_store.mut_addr_manager().get(&new_peer_addr).is_some()); } + +#[test] +fn test_addr_unique() { + let mut peer_store = PeerStore::default(); + let addr = random_addr(); + let addr_1 = random_addr(); + + peer_store + .add_addr(addr.clone(), Flags::COMPATIBILITY) + .unwrap(); + peer_store.add_addr(addr_1, Flags::COMPATIBILITY).unwrap(); + assert_eq!(peer_store.addr_manager().addrs_iter().count(), 2); + assert_eq!(peer_store.fetch_addrs_to_feeler(2).len(), 2); + + peer_store.add_addr(addr, Flags::COMPATIBILITY).unwrap(); + assert_eq!(peer_store.fetch_addrs_to_feeler(2).len(), 2); + + assert_eq!(peer_store.addr_manager().addrs_iter().count(), 2); +}