[hotfix 0.107.x] fix: fix identify unregister

network/src/peer_store/addr_manager.rs

@@ -22,15 +22,8 @@ impl AddrManager {
                 .get(&addr_info.addr)
                 .map(|addr| addr.last_connected_at_ms)
             {
-                // replace exists addr if has later last_connected_at_ms
-                if addr_info.last_connected_at_ms > exists_last_connected_at_ms {
-                    if let Some(old) = self.remove(&addr_info.addr) {
-                        // init from `add_outbound_addr`
-                        if addr_info.flags == 0 {
-                            addr_info.flags = old.flags;
-                        }
-                    }
-                } else {
+                // Get time earlier than record time, return directly
+                if addr_info.last_connected_at_ms < exists_last_connected_at_ms {
                     return;
                 }
             }

network/src/peer_store/peer_store_impl.rs

@@ -85,6 +85,16 @@ impl PeerStore {
         ));
     }
 
+    /// Update outbound peer last connected ms
+    pub fn update_outbound_addr_last_connected_ms(&mut self, addr: Multiaddr) {
+        if self.ban_list.is_addr_banned(&addr) {
+            return;
+        }
+        if let Some(info) = self.addr_manager.get_mut(&addr) {
+            info.last_connected_at_ms = faketime::unix_time_as_millis()
+        }
+    }
+
     /// Get address manager
     pub fn addr_manager(&self) -> &AddrManager {
         &self.addr_manager

network/src/protocols/discovery/addr.rs

@@ -46,7 +46,7 @@ pub trait AddressManager {
     fn misbehave(&mut self, session: &SessionContext, kind: &Misbehavior) -> MisbehaveResult;
     fn get_random(&mut self, n: usize, target: Flags) -> Vec<(Multiaddr, Flags)>;
     fn required_flags(&self) -> Flags;
-    fn node_flags(&self, id: SessionId) -> Flags;
+    fn node_flags(&self, id: SessionId) -> Option<Flags>;
 }
 
 // bitcoin: bloom.h, bloom.cpp => CRollingBloomFilter

network/src/protocols/discovery/mod.rs

@@ -129,8 +129,10 @@ impl<M: AddressManager + Send + Sync> ServiceProtocol for DiscoveryProtocol<M> {
                                 // add client listen address to manager
                                 if let RemoteAddress::Listen(ref addr) = state.remote_addr {
                                     let flags = self.addr_mgr.node_flags(session.id);
-                                    self.addr_mgr
-                                        .add_new_addr(session.id, (addr.clone(), flags));
+                                    self.addr_mgr.add_new_addr(
+                                        session.id,
+                                        (addr.clone(), flags.unwrap_or(Flags::COMPATIBILITY)),
+                                    );
                                 }
                             }
                             if version >= state::REUSE_PORT_VERSION {
@@ -232,7 +234,9 @@ impl<M: AddressManager + Send + Sync> ServiceProtocol for DiscoveryProtocol<M> {
                 .check_timer(now, ANNOUNCE_INTERVAL)
                 .filter(|addr| self.addr_mgr.is_valid_addr(addr))
             {
-                announce_list.push((addr.clone(), self.addr_mgr.node_flags(*id)));
+                if let Some(flags) = self.addr_mgr.node_flags(*id) {
+                    announce_list.push((addr.clone(), flags));
+                }
             }
         }
 
@@ -384,16 +388,10 @@ impl AddressManager for DiscoveryAddressManager {
         self.network_state.required_flags
     }
 
-    fn node_flags(&self, id: SessionId) -> Flags {
+    fn node_flags(&self, id: SessionId) -> Option<Flags> {
         self.network_state.with_peer_registry(|reg| {
-            if let Some(peer) = reg.get_peer(id) {
-                peer.identify_info
-                    .as_ref()
-                    .map(|a| a.flags)
-                    .unwrap_or(Flags::COMPATIBILITY)
-            } else {
-                Flags::COMPATIBILITY
-            }
+            reg.get_peer(id)
+                .and_then(|peer| peer.identify_info.as_ref().map(|a| a.flags))
         })
     }
 }

network/src/protocols/identify/mod.rs

@@ -374,19 +374,9 @@ impl Callback for IdentifyCallback {
             // Due to the filtering strategy of the peer store, if the node is
             // disconnected after a long connection is maintained for more than seven days,
             // it is possible that the node will be accidentally evicted, so it is necessary
-            // to reset the information of the node when disconnected.
-            let flags = self.network_state.with_peer_registry(|reg| {
-                if let Some(p) = reg.get_peer(context.session.id) {
-                    p.identify_info
-                        .as_ref()
-                        .map(|i| i.flags)
-                        .unwrap_or(Flags::COMPATIBILITY)
-                } else {
-                    Flags::COMPATIBILITY
-                }
-            });
+            // to reset the last_connected_time of the node when disconnected.
             self.network_state.with_peer_store_mut(|peer_store| {
-                peer_store.add_outbound_addr(context.session.address.clone(), flags);
+                peer_store.update_outbound_addr_last_connected_ms(context.session.address.clone());
             });
         }
     }

network/src/protocols/tests/mod.rs

@@ -48,6 +48,12 @@ impl Node {
         self.control.dial(addr, protocol).unwrap();
     }
 
+    fn disconnect_all(&self) {
+        for id in self.connected_sessions() {
+            self.control.disconnect(id).unwrap();
+        }
+    }
+
     fn session_num(&self) -> usize {
         self.connected_sessions().len()
     }
@@ -680,3 +686,57 @@ fn test_bootnode_mode_inbound_eviction() {
     // Normal connection, 2 + 1
     wait_connect_state(&node1, 3);
 }
+
+#[test]
+fn test_dont_reset_peer_flags_on_disconnect() {
+    let node1 = net_service_start(
+        "/test/1".to_string(),
+        true,
+        Flags::COMPATIBILITY,
+        Flags::all(),
+    );
+    let node2 = net_service_start(
+        "/test/1".to_string(),
+        true,
+        Flags::COMPATIBILITY,
+        Flags::all(),
+    );
+
+    node1.dial(
+        &node2,
+        TargetProtocol::Single(SupportProtocols::Identify.protocol_id()),
+    );
+
+    wait_connect_state(&node1, 1);
+    wait_connect_state(&node1, 1);
+
+    let check_flags = |node: &Node| {
+        for info in node
+            .network_state
+            .peer_store
+            .lock()
+            .addr_manager()
+            .addrs_iter()
+        {
+            assert_eq!(info.flags, Flags::all().bits())
+        }
+    };
+
+    check_flags(&node1);
+    check_flags(&node2);
+    node1.disconnect_all();
+
+    check_flags(&node1);
+    check_flags(&node2);
+
+    node1.dial(
+        &node2,
+        TargetProtocol::Single(SupportProtocols::Identify.protocol_id()),
+    );
+
+    wait_connect_state(&node1, 1);
+    wait_connect_state(&node1, 1);
+
+    check_flags(&node1);
+    check_flags(&node2);
+}