chore: implement review feedback

Author: driftluo

network/src/network.rs

@@ -508,7 +508,6 @@ impl NetworkState {
     }
 
     // randomly select count addresses from observed_addrs
-    #[cfg(not(target_family = "wasm"))]
     pub(crate) fn observed_addrs(&self, count: usize) -> Vec<Multiaddr> {
         let observed_addrs = self
             .observed_addrs

network/src/peer_store/peer_store_impl.rs

@@ -217,7 +217,7 @@ impl PeerStore {
         self.addr_manager.fetch_random(count, filter)
     }
 
-    /// Return address that we never connected to, used for penetration.
+    /// Return address that we never connected to, used for hole punching.
     pub fn fetch_nat_addrs(&mut self, count: usize, required_flags: Flags) -> Vec<AddrInfo> {
         // Get info:
         // 1. Never connected
@@ -230,7 +230,7 @@ impl PeerStore {
                 && extract_peer_id(&peer_addr.addr)
                     .map(|peer_id| !peers.contains_key(&peer_id))
                     .unwrap_or_default()
-                && peer_addr.connected(|t| t == 0)
+                && peer_addr.last_connected_at_ms == 0
         };
 
         self.addr_manager.fetch_random(count, filter)

network/src/protocols/hole_punching/component/connection_request.rs

@@ -14,18 +14,80 @@ use crate::{
     protocols::{
         SupportProtocols,
         hole_punching::{
-            ADDRS_COUNT_LIMIT, HolePunching, MAX_TTL, PENETRATED_INTERVAL,
+            ADDRS_COUNT_LIMIT, HolePunching, MAX_HOPS, PENETRATED_INTERVAL,
             component::{forward_request, init_delivered},
             status::{Status, StatusCode},
         },
     },
 };
 
+struct RequestContent {
+    from: PeerId,
+    to: PeerId,
+    listen_addrs: Vec<Multiaddr>,
+    route: Vec<PeerId>,
+    max_hops: u8,
+}
+
+impl TryFrom<&packed::ConnectionRequestReader<'_>> for RequestContent {
+    type Error = Status;
+
+    fn try_from(value: &packed::ConnectionRequestReader<'_>) -> Result<Self, Self::Error> {
+        let from = PeerId::from_bytes(value.from().raw_data().to_vec()).map_err(|_| {
+            StatusCode::InvalidFromPeerId.with_context("the from peer id is invalid")
+        })?;
+        let to = PeerId::from_bytes(value.to().raw_data().to_vec())
+            .map_err(|_| StatusCode::InvalidToPeerId.with_context("the to peer id is invalid"))?;
+        let listen_addrs: Vec<Multiaddr> = value
+            .listen_addrs()
+            .iter()
+            .map(
+                |raw| match Multiaddr::try_from(raw.bytes().raw_data().to_vec()) {
+                    Ok(mut addr) => {
+                        if let Some(peer_id) = extract_peer_id(&addr) {
+                            if peer_id != from {
+                                return Err(StatusCode::InvalidListenAddrLen
+                                    .with_context("peer id in listen address is invalid"));
+                            }
+                        } else {
+                            addr.push(Protocol::P2P(Cow::Borrowed(from.as_bytes())));
+                        }
+                        Ok(addr)
+                    }
+                    Err(_) => Err(StatusCode::InvalidListenAddrLen
+                        .with_context("the listen address is invalid")),
+                },
+            )
+            .collect::<Result<Vec<_>, _>>()?;
+
+        let route: Vec<PeerId> = value
+            .route()
+            .iter()
+            .map(|raw| {
+                PeerId::from_bytes(raw.raw_data().to_vec()).map_err(|_| {
+                    StatusCode::InvalidRoute.with_context("the route peer id is invalid")
+                })
+            })
+            .collect::<Result<Vec<_>, _>>()?;
+
+        let max_hops: u8 = value.max_hops().into();
+
+        Ok(Self {
+            from,
+            to,
+            listen_addrs,
+            route,
+            max_hops,
+        })
+    }
+}
+
 pub(crate) struct ConnectionRequestProcess<'a> {
     message: packed::ConnectionRequestReader<'a>,
     protocol: &'a HolePunching,
     peer: PeerIndex,
     p2p_control: &'a ServiceAsyncControl,
+    msg_item_id: u32,
 }
 
 impl<'a> ConnectionRequestProcess<'a> {
@@ -34,72 +96,68 @@ impl<'a> ConnectionRequestProcess<'a> {
         protocol: &'a HolePunching,
         peer: PeerIndex,
         p2p_control: &'a ServiceAsyncControl,
+        msg_item_id: u32,
     ) -> Self {
         Self {
             message,
             protocol,
             peer,
             p2p_control,
+            msg_item_id,
         }
     }
 
     pub(crate) async fn execute(self) -> Status {
-        if self.message.listen_addrs().len() > ADDRS_COUNT_LIMIT
-            || self.message.listen_addrs().is_empty()
-        {
+        let content = match RequestContent::try_from(&self.message) {
+            Ok(content) => content,
+            Err(status) => return status,
+        };
+        if content.listen_addrs.len() > ADDRS_COUNT_LIMIT || content.listen_addrs.is_empty() {
             return StatusCode::InvalidListenAddrLen
                 .with_context("the listen address count is too large or empty");
         }
-        let ttl: u8 = self.message.ttl().into();
-        if ttl > MAX_TTL {
+
+        if content.max_hops > MAX_HOPS {
             return StatusCode::InvalidMaxTTL.into();
         }
-        if self.message.route().len() > 8 {
+        if content.route.len() > 8 {
             return StatusCode::InvalidRoute.with_context("the route length is too long");
         }
 
         let self_peer_id = self.protocol.network_state.local_peer_id();
-        for peer_id_bytes in self.message.route().iter() {
-            match PeerId::from_bytes(peer_id_bytes.raw_data().to_vec()) {
-                Ok(peer_id) => {
-                    if self_peer_id == &peer_id {
-                        return StatusCode::Ignore.with_context("the message is passed, ignore it");
-                    }
-                }
-                Err(_) => {
-                    return StatusCode::InvalidRoute.into();
-                }
-            }
+        if content.route.contains(self_peer_id) {
+            return StatusCode::Ignore.with_context("the message is passed, ignore it");
         }
 
-        let from_peer_id = match PeerId::from_bytes(self.message.from().raw_data().to_vec()) {
-            Ok(peer_id) => {
-                if self_peer_id == &peer_id {
-                    return StatusCode::Ignore.with_context("the message is passed, ignore it");
-                }
-                peer_id
-            }
-            Err(_) => {
-                return StatusCode::InvalidFromPeerId.into();
-            }
-        };
-        let to_peer_id = match PeerId::from_bytes(self.message.to().raw_data().to_vec()) {
-            Ok(peer_id) => peer_id,
-            Err(_) => {
-                return StatusCode::InvalidToPeerId.into();
-            }
-        };
+        if self
+            .protocol
+            .forward_rate_limiter
+            .check_key(&(content.from.clone(), content.to.clone(), self.msg_item_id))
+            .is_err()
+        {
+            debug!(
+                "from: {}, to {}, item_name: {}, rate limit is reached",
+                content.from, content.to, "ConnectionRequest",
+            );
+            return StatusCode::TooManyRequests.with_context("ConnectionRequest");
+        }
 
-        if self_peer_id == &to_peer_id {
-            self.respond_delivered(from_peer_id, &to_peer_id).await
-        } else if ttl == 0u8 {
-            StatusCode::ReachedMaxTTL.into()
+        if self_peer_id == &content.to {
+            self.respond_delivered(content.from, &content.to, content.listen_addrs)
+                .await
+        } else if content.max_hops == 0u8 {
+            StatusCode::ReachedMaxHops.into()
         } else {
-            self.forward_message(self_peer_id, &to_peer_id).await
+            self.forward_message(self_peer_id, &content.to).await
         }
     }
 
-    async fn respond_delivered(&self, from_peer_id: PeerId, to_peer_id: &PeerId) -> Status {
+    async fn respond_delivered(
+        &self,
+        from_peer_id: PeerId,
+        to_peer_id: &PeerId,
+        remote_listens: Vec<Multiaddr>,
+    ) -> Status {
         if let Some((_, t)) = self.protocol.pending_delivered.read().get(&from_peer_id) {
             let now = unix_time_as_millis();
             if now - t < PENETRATED_INTERVAL {
@@ -140,6 +198,22 @@ impl<'a> ConnectionRequestProcess<'a> {
             to_peer_id
         );
 
+        let remote_listens: Vec<Multiaddr> = remote_listens
+            .into_iter()
+            .filter_map(|addr| match find_type(&addr) {
+                TransportType::Memory
+                | TransportType::Onion
+                | TransportType::Ws
+                | TransportType::Wss
+                | TransportType::Tls => None,
+                TransportType::Tcp => Some(addr),
+            })
+            .collect();
+
+        if remote_listens.is_empty() {
+            return StatusCode::Ignore.with_context("remote listen address is empty");
+        }
+
         if let Err(error) = self
             .p2p_control
             .send_message_to(self.peer, proto_id, new_message)
@@ -148,35 +222,6 @@ impl<'a> ConnectionRequestProcess<'a> {
             return StatusCode::ForwardError.with_context(error);
         }
 
-        let remote_listens = self
-            .message
-            .listen_addrs()
-            .iter()
-            .filter_map(
-                |raw| match Multiaddr::try_from(raw.bytes().raw_data().to_vec()) {
-                    Ok(mut addr) => {
-                        if let Some(peer_id) = extract_peer_id(&addr) {
-                            if peer_id != from_peer_id {
-                                return None;
-                            }
-                        } else {
-                            addr.push(Protocol::P2P(Cow::Borrowed(from_peer_id.as_bytes())));
-                        }
-
-                        match find_type(&addr) {
-                            TransportType::Memory
-                            | TransportType::Onion
-                            | TransportType::Ws
-                            | TransportType::Wss
-                            | TransportType::Tls => None,
-                            TransportType::Tcp => Some(addr),
-                        }
-                    }
-                    Err(_) => None,
-                },
-            )
-            .collect();
-
         let mut pending_delivered = self.protocol.pending_delivered.write();
         let now = unix_time_as_millis();
         pending_delivered.insert(from_peer_id, (remote_listens, now));
@@ -220,11 +265,20 @@ impl<'a> ConnectionRequestProcess<'a> {
                     "target peer {} is not found, broadcast the request to more peers",
                     to_peer_id
                 );
+
+                // Broadcast using gossip while removing the source of the message
                 let sid = self.peer;
+                let mut init = 0usize;
                 if let Err(error) = self
                     .p2p_control
                     .filter_broadcast(
-                        TargetSession::Filter(Box::new(move |id| id != &sid)),
+                        TargetSession::Filter(Box::new(move |id| {
+                            if id == &sid {
+                                return false;
+                            }
+                            init += 1;
+                            init % 3 != 0
+                        })),
                         proto_id,
                         new_message,
                     )