From 5179204ee5e096fb4f0ab6d4cc34a8e031337594 Mon Sep 17 00:00:00 2001 From: driftluo Date: Thu, 16 Jan 2025 19:41:41 +0800 Subject: [PATCH] fix: fix feeler flags fetch fail --- network/src/network.rs | 22 ++++------------------ network/src/peer_registry.rs | 25 +++++++++++++++++++++---- network/src/protocols/feeler.rs | 7 ++----- network/src/protocols/identify/mod.rs | 7 +++---- network/src/services/outbound_peer.rs | 17 ++++++++++++----- 5 files changed, 42 insertions(+), 36 deletions(-) diff --git a/network/src/network.rs b/network/src/network.rs index 5a7eb518b9..bc19edf9ac 100644 --- a/network/src/network.rs +++ b/network/src/network.rs @@ -1531,27 +1531,13 @@ pub enum TransportType { Wss, } -impl<'a> From for p2p::multiaddr::Protocol<'a> { - fn from(value: TransportType) -> Self { - match value { - TransportType::Ws => Protocol::Ws, - TransportType::Wss => Protocol::Wss, - _ => unreachable!(), - } - } -} - pub(crate) fn find_type(addr: &Multiaddr) -> TransportType { let mut iter = addr.iter(); - iter.find_map(|proto| { - if let Protocol::Ws = proto { - Some(TransportType::Ws) - } else if let Protocol::Wss = proto { - Some(TransportType::Wss) - } else { - None - } + iter.find_map(|proto| match proto { + Protocol::Ws => Some(TransportType::Ws), + Protocol::Wss => Some(TransportType::Wss), + _ => None, }) .unwrap_or(TransportType::Tcp) } diff --git a/network/src/peer_registry.rs b/network/src/peer_registry.rs index dcc97d160f..c1719c99ad 100644 --- 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, - feeler_peers: HashSet, + feeler_peers: HashMap, } /// 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,26 @@ 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 { + if let Some(peer_id) = extract_peer_id(addr) { + if let Some(i) = self.feeler_peers.get_mut(&peer_id) { + *i = flags; + return true; + } + } + false + } + + /// Get feeler session flags + pub fn feeler_flags(&self, addr: &Multiaddr) -> Option { + extract_peer_id(addr).and_then(|peer_id| self.feeler_peers.get(&peer_id).cloned()) + } + /// 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 +222,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 index 95d68c1bdc..6fe688d9d7 100644 --- 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 index 2fe9d7cef1..84f2f3ffd7 100644 --- 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, diff --git a/network/src/services/outbound_peer.rs b/network/src/services/outbound_peer.rs index 8e1db3bf18..d58cb9af7e 100644 --- a/network/src/services/outbound_peer.rs +++ b/network/src/services/outbound_peer.rs @@ -7,7 +7,10 @@ use ckb_logger::trace; use ckb_systemtime::unix_time_as_millis; use futures::{Future, StreamExt}; use p2p::runtime::{Interval, MissedTickBehavior}; -use p2p::{multiaddr::MultiAddr, service::ServiceControl}; +use p2p::{ + multiaddr::{MultiAddr, Protocol}, + service::ServiceControl, +}; use rand::prelude::IteratorRandom; use std::{ pin::Pin, @@ -71,8 +74,10 @@ impl OutboundPeerService { for mut addr in attempt_peers.into_iter().map(|info| info.addr) { self.network_state.dial_feeler(&self.p2p_control, { - if !matches!(self.transport_type, TransportType::Tcp) { - addr.push(self.transport_type.into()); + match &self.transport_type { + TransportType::Tcp => (), + TransportType::Ws => addr.push(Protocol::Ws), + TransportType::Wss => addr.push(Protocol::Wss), } addr }); @@ -145,8 +150,10 @@ impl OutboundPeerService { for mut addr in peers { self.network_state.dial_identify(&self.p2p_control, { - if !matches!(self.transport_type, TransportType::Tcp) { - addr.push(self.transport_type.into()); + match &self.transport_type { + TransportType::Tcp => (), + TransportType::Ws => addr.push(Protocol::Ws), + TransportType::Wss => addr.push(Protocol::Wss), } addr });