From 47b518fb205911a44b460e6afcd990e146e41015 Mon Sep 17 00:00:00 2001
From: Eval EXEC <execvy@gmail.com>
Date: Tue, 31 Dec 2024 16:04:49 +0800
Subject: [PATCH] Use `NetworkAddresses` to represent regular_addrs and
 onion_addrs

Signed-off-by: Eval EXEC <execvy@gmail.com>
---
 network/src/address.rs | 71 ++++++++++++++++++++++++++++++++++++++++++
 network/src/lib.rs     |  1 +
 network/src/network.rs |  3 +-
 network/src/peer.rs    |  3 +-
 4 files changed, 76 insertions(+), 2 deletions(-)
 create mode 100644 network/src/address.rs

diff --git a/network/src/address.rs b/network/src/address.rs
new file mode 100644
index 0000000000..a4dc0f226f
--- /dev/null
+++ b/network/src/address.rs
@@ -0,0 +1,71 @@
+use p2p::multiaddr::{MultiAddr, Protocol};
+
+#[derive(Default)]
+pub struct NetworkAddresses {
+    pub regular_addresses: Vec<MultiAddr>,
+
+    // onion addresses can't be solved by multiaddr_to_socketaddr or socketaddr_to_multiaddr
+    pub onion_addresses: Vec<MultiAddr>,
+}
+
+impl NetworkAddresses {
+    pub fn push(&mut self, address: MultiAddr) {
+        if address
+            .iter()
+            .any(|proto| matches!(proto, Protocol::Onion3(_)))
+        {
+            self.onion_addresses.push(address);
+        } else {
+            self.regular_addresses.push(address);
+        }
+    }
+
+    // contains
+    pub fn contains(&self, address: &MultiAddr) -> bool {
+        self.regular_addresses.contains(address) || self.onion_addresses.contains(address)
+    }
+}
+
+impl IntoIterator for NetworkAddresses {
+    type Item = MultiAddr;
+    type IntoIter = std::vec::IntoIter<MultiAddr>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.regular_addresses
+            .into_iter()
+            .chain(self.onion_addresses.into_iter())
+            .collect::<Vec<_>>()
+            .into_iter()
+    }
+}
+
+// convert Vec<MultiAddr> to NetworkAddresses
+impl From<Vec<MultiAddr>> for NetworkAddresses {
+    fn from(addresses: Vec<MultiAddr>) -> Self {
+        let mut regular_addresses = Vec::new();
+        let mut onion_addresses = Vec::new();
+        for address in addresses {
+            if address
+                .iter()
+                .any(|proto| matches!(proto, Protocol::Onion3(_)))
+            {
+                onion_addresses.push(address);
+            } else {
+                regular_addresses.push(address);
+            }
+        }
+        NetworkAddresses {
+            regular_addresses,
+            onion_addresses,
+        }
+    }
+}
+
+// convert NetworkAddresses to Vec<MultiAddr>
+impl From<NetworkAddresses> for Vec<MultiAddr> {
+    fn from(addresses: NetworkAddresses) -> Self {
+        let mut result = addresses.regular_addresses;
+        result.extend(addresses.onion_addresses);
+        result
+    }
+}
diff --git a/network/src/lib.rs b/network/src/lib.rs
index c672dab65a..6ec1290947 100644
--- a/network/src/lib.rs
+++ b/network/src/lib.rs
@@ -6,6 +6,7 @@
 //! And implemented several basic protocols: identify, discovery, ping, feeler, disconnect_message
 //!
 
+mod address;
 mod behaviour;
 /// compress module
 pub mod compress;
diff --git a/network/src/network.rs b/network/src/network.rs
index a8003aa3d3..4d46d6912c 100644
--- a/network/src/network.rs
+++ b/network/src/network.rs
@@ -1,4 +1,5 @@
 //! Global state struct and start function
+use crate::address::NetworkAddresses;
 use crate::errors::Error;
 #[cfg(not(target_family = "wasm"))]
 use crate::errors::P2PError;
@@ -73,7 +74,7 @@ pub struct NetworkState {
     pub(crate) peer_registry: RwLock<PeerRegistry>,
     pub(crate) peer_store: Mutex<PeerStore>,
     /// Node listened addresses
-    pub(crate) listened_addrs: RwLock<Vec<Multiaddr>>,
+    pub(crate) listened_addrs: RwLock<NetworkAddresses>,
     dialing_addrs: RwLock<HashMap<PeerId, Instant>>,
     /// Node public addresses,
     /// includes manually public addrs and remote peer observed addrs
diff --git a/network/src/peer.rs b/network/src/peer.rs
index fe7ee3dfb3..3c649c7ef3 100644
--- a/network/src/peer.rs
+++ b/network/src/peer.rs
@@ -1,3 +1,4 @@
+use crate::address::NetworkAddresses;
 use crate::network_group::Group;
 use crate::{
     multiaddr::Multiaddr, protocols::identify::Flags, ProtocolId, ProtocolVersion, SessionType,
@@ -21,7 +22,7 @@ pub struct Peer {
     /// Peer address
     pub connected_addr: Multiaddr,
     /// Peer listen addresses
-    pub listened_addrs: Vec<Multiaddr>,
+    pub listened_addrs: NetworkAddresses,
     /// Peer info from identify protocol message
     pub identify_info: Option<PeerIdentifyInfo>,
     /// Ping/Pong message last received time