feat: mark dns address connected time on peer store

nervosnetwork · Jan 24, 2025 · 5a8dec2 · 5a8dec2
1 parent b44d2d2
commit 5a8dec2
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 5 deletions.
diff --git a/network/src/network.rs b/network/src/network.rs
@@ -74,7 +74,8 @@ pub struct NetworkState {
     pub(crate) peer_store: Mutex<PeerStore>,
     /// Node listened addresses
     pub(crate) listened_addrs: RwLock<Vec<Multiaddr>>,
-    dialing_addrs: RwLock<HashMap<PeerId, Instant>>,
+    dialing_addrs: RwLock<HashMap<PeerId, (Instant, Multiaddr)>>,
+    pub(crate) pending_dns_addrs: RwLock<HashMap<PeerId, Multiaddr>>,
     /// Node public addresses,
     /// includes manually public addrs and remote peer observed addrs
     public_addrs: RwLock<HashSet<Multiaddr>>,
@@ -135,6 +136,7 @@ impl NetworkState {
             bootnodes,
             peer_registry: RwLock::new(peer_registry),
             dialing_addrs: RwLock::new(HashMap::default()),
+            pending_dns_addrs: RwLock::new(HashMap::default()),
             public_addrs: RwLock::new(public_addrs),
             listened_addrs: RwLock::new(Vec::new()),
             pending_observed_addrs: RwLock::new(HashSet::default()),
@@ -183,6 +185,7 @@ impl NetworkState {
             bootnodes,
             peer_registry: RwLock::new(peer_registry),
             dialing_addrs: RwLock::new(HashMap::default()),
+            pending_dns_addrs: RwLock::new(HashMap::default()),
             public_addrs: RwLock::new(public_addrs),
             listened_addrs: RwLock::new(Vec::new()),
             pending_observed_addrs: RwLock::new(HashSet::default()),
@@ -405,7 +408,7 @@ impl NetworkState {
             return false;
         }
 
-        if let Some(dial_started) = self.dialing_addrs.read().get(peer_id) {
+        if let Some((dial_started, _)) = self.dialing_addrs.read().get(peer_id) {
             trace!(
                 "Do not send repeated dial commands to network service: {:?}, {}",
                 peer_id,
@@ -437,7 +440,15 @@ impl NetworkState {
 
     pub(crate) fn dial_success(&self, addr: &Multiaddr) {
         if let Some(peer_id) = extract_peer_id(addr) {
-            self.dialing_addrs.write().remove(&peer_id);
+            if let Some(dial_addr) = self.dialing_addrs.write().remove(&peer_id).map(|a| a.1) {
+                let has_dns = dial_addr
+                    .iter()
+                    .any(|p| matches!(p, Protocol::Dns4(_) | Protocol::Dns6(_)));
+
+                if &dial_addr != addr && has_dns {
+                    self.pending_dns_addrs.write().insert(peer_id, dial_addr);
+                }
+            }
         }
     }
 
@@ -467,7 +478,7 @@ impl NetworkState {
         p2p_control.dial(addr.clone(), target)?;
         self.dialing_addrs.write().insert(
             extract_peer_id(&addr).expect("verified addr"),
-            Instant::now(),
+            (Instant::now(), addr),
         );
         Ok(())
     }
@@ -721,6 +732,7 @@ impl ServiceHandle for EventHandler {
                     "SessionOpen({}, {})",
                     session_context.id, session_context.address,
                 );
+
                 self.network_state.dial_success(&session_context.address);
 
                 let iter = self.inbound_eviction();
@@ -803,6 +815,10 @@ impl ServiceHandle for EventHandler {
                         peer_store.remove_disconnected_peer(&session_context.address);
                     });
                 }
+                self.network_state
+                    .pending_dns_addrs
+                    .write()
+                    .remove(&extract_peer_id(&session_context.address).expect("must have peerid"));
             }
             _ => {
                 info!("p2p service event: {:?}", event);

diff --git a/network/src/peer_store/peer_store_impl.rs b/network/src/peer_store/peer_store_impl.rs
@@ -12,6 +12,7 @@ use crate::{
     Flags, PeerId, SessionType,
 };
 use ipnetwork::IpNetwork;
+use p2p::multiaddr::Protocol;
 use rand::prelude::IteratorRandom;
 use std::collections::{hash_map::Entry, HashMap};
 
@@ -110,7 +111,20 @@ impl PeerStore {
         if self.ban_list.is_addr_banned(&addr) {
             return;
         }
-        if let Some(info) = self.addr_manager.get_mut(&addr) {
+        let base_addr = addr
+            .iter()
+            .filter_map(|p| {
+                if matches!(
+                    p,
+                    Protocol::Ws | Protocol::Wss | Protocol::Memory(_) | Protocol::Tls(_)
+                ) {
+                    None
+                } else {
+                    Some(p)
+                }
+            })
+            .collect();
+        if let Some(info) = self.addr_manager.get_mut(&base_addr) {
             info.last_connected_at_ms = ckb_systemtime::unix_time_as_millis()
         }
     }

diff --git a/network/src/protocols/identify/mod.rs b/network/src/protocols/identify/mod.rs
@@ -457,8 +457,15 @@ impl Callback for IdentifyCallback {
                     // and if set it to peer store, it will be broadcast to the entire network,
                     // but this is an unverified address
                     if renew {
+                        let dns_addr = self.network_state.pending_dns_addrs.write().remove(
+                            &extract_peer_id(&context.session.address).expect("must have peerid"),
+                        );
                         self.network_state.with_peer_store_mut(|peer_store| {
                             peer_store.add_outbound_addr(context.session.address.clone(), flags);
+                            if let Some(dns) = dns_addr {
+                                // mark dns address connected time
+                                peer_store.add_outbound_addr(dns, flags);
+                            }
                         });
                     }