From 3c5e11aaaaa8ae2474603b474c7d66939d517e36 Mon Sep 17 00:00:00 2001
From: driftluo <driftluo@foxmail.com>
Date: Tue, 21 Jan 2025 15:16:10 +0800
Subject: [PATCH] feat: mark dns address connected time on peer store

---
 network/src/network.rs                    | 24 +++++++++++++++++------
 network/src/peer_store/peer_store_impl.rs | 16 ++++++++++++++-
 2 files changed, 33 insertions(+), 7 deletions(-)

diff --git a/network/src/network.rs b/network/src/network.rs
index bc19edf9ac..d73bb90951 100644
--- a/network/src/network.rs
+++ b/network/src/network.rs
@@ -74,7 +74,7 @@ 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)>>,
     /// Node public addresses,
     /// includes manually public addrs and remote peer observed addrs
     public_addrs: RwLock<HashSet<Multiaddr>>,
@@ -405,7 +405,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,
@@ -435,9 +435,11 @@ impl NetworkState {
         true
     }
 
-    pub(crate) fn dial_success(&self, addr: &Multiaddr) {
+    pub(crate) fn dial_success(&self, addr: &Multiaddr) -> Option<Multiaddr> {
         if let Some(peer_id) = extract_peer_id(addr) {
-            self.dialing_addrs.write().remove(&peer_id);
+            self.dialing_addrs.write().remove(&peer_id).map(|a| a.1)
+        } else {
+            None
         }
     }
 
@@ -467,7 +469,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,7 +723,17 @@ impl ServiceHandle for EventHandler {
                     "SessionOpen({}, {})",
                     session_context.id, session_context.address,
                 );
-                self.network_state.dial_success(&session_context.address);
+                // marked dns address connected message
+                if let Some(addr) = self.network_state.dial_success(&session_context.address) {
+                    let has_dns = addr
+                        .iter()
+                        .any(|p| matches!(p, Protocol::Dns4(_) | Protocol::Dns6(_)));
+
+                    if has_dns {
+                        self.network_state
+                            .with_peer_store_mut(|p| p.update_outbound_addr_last_connected_ms(addr))
+                    }
+                }
 
                 let iter = self.inbound_eviction();
 
diff --git a/network/src/peer_store/peer_store_impl.rs b/network/src/peer_store/peer_store_impl.rs
index a8078b8877..06dbaf9187 100644
--- 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()
         }
     }