fix: fix feeler flags fetch fail

nervosnetwork · Jan 16, 2025 · f246db5 · quake · Jan 16, 2025 · f246db5
1 parent c92c849
commit f246db5
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 13 deletions.
diff --git a/network/src/peer_registry.rs b/network/src/peer_registry.rs
@@ -1,6 +1,7 @@
 //! Peer registry
 use crate::network_group::Group;
 use crate::peer_store::PeerStore;
+use crate::Flags;
 use crate::{
     errors::{Error, PeerError},
     extract_peer_id, Peer, PeerId, SessionType,
@@ -24,7 +25,7 @@ pub struct PeerRegistry {
     // Only whitelist peers or allow all peers.
     whitelist_only: bool,
     whitelist_peers: HashSet<PeerId>,
-    feeler_peers: HashSet<PeerId>,
+    feeler_peers: HashMap<PeerId, Flags>,
 }
 
 /// Global network connection status
@@ -63,7 +64,7 @@ impl PeerRegistry {
         PeerRegistry {
             peers: HashMap::with_capacity_and_hasher(20, Default::default()),
             whitelist_peers: whitelist_peers.iter().filter_map(extract_peer_id).collect(),
-            feeler_peers: HashSet::default(),
+            feeler_peers: HashMap::default(),
             max_inbound,
             max_outbound,
             whitelist_only,
@@ -191,10 +192,30 @@ impl PeerRegistry {
     /// Add feeler dail task
     pub fn add_feeler(&mut self, addr: &Multiaddr) {
         if let Some(peer_id) = extract_peer_id(addr) {
-            self.feeler_peers.insert(peer_id);
+            self.feeler_peers.insert(peer_id, Flags::COMPATIBILITY);
         }
     }
 
+    /// Identify change feeler flags
+    pub fn change_feeler_flags(&mut self, addr: &Multiaddr, flags: Flags) -> bool {
+        extract_peer_id(addr)
+            .map(|peer_id| match self.feeler_peers.get_mut(&peer_id) {
+                Some(i) => {
+                    *i = flags;
+                    true
+                }
+                None => false,
+            })
+            .unwrap()
+    }
+
+    /// Get feeler session flags
+    pub fn feeler_flags(&self, addr: &Multiaddr) -> Option<Flags> {
+        extract_peer_id(addr)
+            .map(|peer_id| self.feeler_peers.get(&peer_id).cloned())
+            .flatten()
+    }
+
     /// Remove feeler dail task on session disconnects or fails
     pub fn remove_feeler(&mut self, addr: &Multiaddr) {
         if let Some(peer_id) = extract_peer_id(addr) {
@@ -205,7 +226,7 @@ impl PeerRegistry {
     /// Whether this session is feeler session
     pub fn is_feeler(&self, addr: &Multiaddr) -> bool {
         extract_peer_id(addr)
-            .map(|peer_id| self.feeler_peers.contains(&peer_id))
+            .map(|peer_id| self.feeler_peers.contains_key(&peer_id))
             .unwrap_or_default()
     }
 

diff --git a/network/src/protocols/feeler.rs b/network/src/protocols/feeler.rs
@@ -34,11 +34,8 @@ impl ServiceProtocol for Feeler {
                 .remove(&session.address);
         } else if context.session.ty.is_outbound() {
             let flags = self.network_state.with_peer_registry(|reg| {
-                if let Some(p) = reg.get_peer(session.id) {
-                    p.identify_info
-                        .as_ref()
-                        .map(|i| i.flags)
-                        .unwrap_or(Flags::COMPATIBILITY)
+                if let Some(p) = reg.feeler_flags(&session.address) {
+                    p
                 } else {
                     Flags::COMPATIBILITY
                 }

diff --git a/network/src/protocols/identify/mod.rs b/network/src/protocols/identify/mod.rs
@@ -462,10 +462,9 @@ impl Callback for IdentifyCallback {
                         });
                     }
 
-                    if self
-                        .network_state
-                        .with_peer_registry(|reg| reg.is_feeler(&context.session.address))
-                    {
+                    if self.network_state.with_peer_registry_mut(|reg| {
+                        reg.change_feeler_flags(&context.session.address, flags)
+                    }) {
                         let _ = context
                             .open_protocols(
                                 context.session.id,