fix: keep peer store address unique

nervosnetwork · Sep 5, 2023 · 416a5f5 · 416a5f5
1 parent b4a71b9
commit 416a5f5
Show file tree

Hide file tree

Showing 4 changed files with 31 additions and 8 deletions.
diff --git 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
@@ -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
@@ -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::<u16>(),
         crate::PeerId::random().to_base58()
     )
     .parse()

diff --git 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);
+}