From 24cf5f0bea7e5a484bc631a8030510964b1b77cf Mon Sep 17 00:00:00 2001
From: driftluo <driftluo@foxmail.com>
Date: Wed, 15 Jan 2025 11:23:17 +0800
Subject: [PATCH] feat: peer store only retains the TCP protocol

---
 network/src/peer_store/addr_manager.rs    |  6 ------
 network/src/peer_store/peer_store_impl.rs | 10 ----------
 network/src/peer_store/types.rs           | 15 ++++++++++++++-
 network/src/tests/peer_store.rs           | 15 +++++++++++++++
 4 files changed, 29 insertions(+), 17 deletions(-)

diff --git a/network/src/peer_store/addr_manager.rs b/network/src/peer_store/addr_manager.rs
index 127b5e75a1..b4b5179ba3 100644
--- a/network/src/peer_store/addr_manager.rs
+++ b/network/src/peer_store/addr_manager.rs
@@ -1,6 +1,4 @@
 //! Address manager
-#[cfg(target_family = "wasm")]
-use crate::network::{find_type, TransportType};
 use crate::peer_store::types::AddrInfo;
 use p2p::{multiaddr::Multiaddr, utils::multiaddr_to_socketaddr};
 use rand::Rng;
@@ -51,10 +49,6 @@ impl AddrManager {
         let mut addr_infos = Vec::with_capacity(count);
         let mut rng = rand::thread_rng();
         let now_ms = ckb_systemtime::unix_time_as_millis();
-        #[cfg(target_family = "wasm")]
-        let filter = |peer_addr: &AddrInfo| {
-            filter(peer_addr) && matches!(find_type(&peer_addr.addr), TransportType::Ws)
-        };
         for i in 0..self.random_ids.len() {
             // reuse the for loop to shuffle random ids
             // https://en.wikipedia.org/wiki/Fisher%E2%80%93Yates_shuffle
diff --git a/network/src/peer_store/peer_store_impl.rs b/network/src/peer_store/peer_store_impl.rs
index 0dd50d6e1e..3169cb45f7 100644
--- a/network/src/peer_store/peer_store_impl.rs
+++ b/network/src/peer_store/peer_store_impl.rs
@@ -65,10 +65,6 @@ impl PeerStore {
         if self.ban_list.is_addr_banned(&addr) {
             return Ok(());
         }
-        #[cfg(target_family = "wasm")]
-        if !matches!(find_type(&addr), TransportType::Ws) {
-            return Ok(());
-        }
         self.check_purge()?;
         let score = self.score_config.default_score;
         self.addr_manager
@@ -180,12 +176,6 @@ impl PeerStore {
                 && required_flags_filter(required_flags, Flags::from_bits_truncate(peer_addr.flags))
         };
 
-        // Any protocol expect websocket
-        #[cfg(not(target_family = "wasm"))]
-        let filter = |peer_addr: &AddrInfo| {
-            filter(peer_addr) && !matches!(find_type(&peer_addr.addr), TransportType::Ws)
-        };
-
         // get addrs that can attempt.
         self.addr_manager.fetch_random(count, filter)
     }
diff --git a/network/src/peer_store/types.rs b/network/src/peer_store/types.rs
index 58712c7bd0..2db34c306b 100644
--- a/network/src/peer_store/types.rs
+++ b/network/src/peer_store/types.rs
@@ -62,7 +62,20 @@ impl AddrInfo {
     /// Init
     pub fn new(addr: Multiaddr, last_connected_at_ms: u64, score: Score, flags: u64) -> Self {
         AddrInfo {
-            addr,
+            // only store tcp protocol
+            addr: addr
+                .iter()
+                .filter_map(|p| {
+                    if matches!(
+                        p,
+                        Protocol::Ws | Protocol::Wss | Protocol::Memory(_) | Protocol::Tls(_)
+                    ) {
+                        None
+                    } else {
+                        Some(p)
+                    }
+                })
+                .collect(),
             score,
             last_connected_at_ms,
             last_tried_at_ms: 0,
diff --git a/network/src/tests/peer_store.rs b/network/src/tests/peer_store.rs
index ee713f5a7f..e23db5c412 100644
--- a/network/src/tests/peer_store.rs
+++ b/network/src/tests/peer_store.rs
@@ -588,3 +588,18 @@ fn test_addr_unique() {
 
     assert_eq!(peer_store.addr_manager().addrs_iter().count(), 2);
 }
+
+#[test]
+fn test_only_tcp_store() {
+    let mut peer_store = PeerStore::default();
+    let mut addr = random_addr();
+    addr.push(p2p::multiaddr::Protocol::Ws);
+    peer_store
+        .add_addr(addr.clone(), Flags::COMPATIBILITY)
+        .unwrap();
+    assert_eq!(peer_store.fetch_addrs_to_feeler(2).len(), 1);
+    assert_eq!(peer_store.fetch_addrs_to_feeler(1)[0].addr, {
+        addr.pop();
+        addr
+    });
+}