add token bucket limiter on peer conn recv (#1842)

We should limit peer conn recv to make sure we don't recv too much from peers.

easytier/src/peers/foreign_network_manager.rs

@@ -221,7 +221,7 @@ impl ForeignNetworkEntry {
 
             async fn recv(&self) -> Result<ZCPacket, Error> {
                 if let Some(o) = self.packet_recv.lock().await.recv().await {
-                    tracing::info!("recv rpc packet in foreign network manager rpc transport");
+                    tracing::trace!("recv rpc packet in foreign network manager rpc transport");
                     Ok(o)
                 } else {
                     Err(Error::Unknown)
@@ -335,7 +335,7 @@ impl ForeignNetworkEntry {
                     tracing::warn!("invalid packet, skip");
                     continue;
                 };
-                tracing::info!(?hdr, "recv packet in foreign network manager");
+                tracing::trace!(?hdr, "recv packet in foreign network manager");
                 let to_peer_id = hdr.to_peer_id.get();
                 if to_peer_id == my_node_id {
                     if hdr.packet_type == PacketType::TaRpc as u8

easytier/src/peers/peer_conn.rs

@@ -40,7 +40,7 @@ use crate::{
     peers::peer_session::{PeerSessionStore, SessionKey, UpsertResponderSessionReturn},
     proto::{
         api::instance::{PeerConnInfo, PeerConnStats},
-        common::{SecureModeConfig, TunnelInfo},
+        common::{LimiterConfig, SecureModeConfig, TunnelInfo},
         peer_rpc::{
             HandshakeRequest, PeerConnNoiseMsg1Pb, PeerConnNoiseMsg2Pb, PeerConnNoiseMsg3Pb,
             PeerConnSessionActionPb, SecureAuthLevel,
@@ -1171,6 +1171,23 @@ impl PeerConn {
         };
         self.counters.store(Some(Arc::new(counters)));
 
+        let is_foreign_network = conn_info_for_instrument.network_name
+            != self.global_ctx.get_network_identity().network_name;
+        let recv_limiter = if is_foreign_network {
+            let relay_network_bps_limit = self.global_ctx.get_flags().foreign_relay_bps_limit;
+            let limiter_config = LimiterConfig {
+                burst_rate: None,
+                bps: Some(relay_network_bps_limit),
+                fill_duration_ms: None,
+            };
+            Some(self.global_ctx.token_bucket_manager().get_or_create(
+                &format!("{}:recv", conn_info_for_instrument.network_name),
+                limiter_config.into(),
+            ))
+        } else {
+            None
+        };
+
         let counters = self.counters.load_full().unwrap();
 
         self.tasks.spawn(
@@ -1185,8 +1202,9 @@ impl PeerConn {
                     }
 
                     let mut zc_packet = ret.unwrap();
+                    let buf_len = zc_packet.buf_len() as u64;
 
-                    counters.traffic_rx_bytes.add(zc_packet.buf_len() as u64);
+                    counters.traffic_rx_bytes.add(buf_len);
                     counters.traffic_rx_packets.inc();
 
                     let Some(peer_mgr_hdr) = zc_packet.mut_peer_manager_header() else {
@@ -1194,7 +1212,7 @@ impl PeerConn {
                             "unexpected packet: {:?}, cannot decode peer manager hdr",
                             zc_packet
                         );
-                        continue;
+                        break;
                     };
 
                     if peer_mgr_hdr.packet_type == PacketType::Ping as u8 {
@@ -1209,6 +1227,10 @@ impl PeerConn {
                     } else if sender.send(zc_packet).await.is_err() {
                         break;
                     }
+
+                    if let Some(limiter) = recv_limiter.as_ref() {
+                        limiter.consume(buf_len).await;
+                    }
                 }
 
                 tracing::info!("end recving peer conn packet");

easytier/src/peers/peer_ospf_route.rs

@@ -530,7 +530,7 @@ impl SyncedRouteInfo {
         for (peer_idx, peer_id_version) in conn_bitmap.peer_ids.iter().enumerate() {
             let connceted_peers = conn_bitmap.get_connected_peers(peer_idx);
             self.fill_empty_peer_info(&connceted_peers);
-            need_inc_version = self.update_conn_info_one_peer(peer_id_version, connceted_peers);
+            need_inc_version |= self.update_conn_info_one_peer(peer_id_version, connceted_peers);
         }
         if need_inc_version {
             self.version.inc();
@@ -548,7 +548,7 @@ impl SyncedRouteInfo {
                 peer_conn_info.connected_peer_ids.iter().copied().collect();
 
             self.fill_empty_peer_info(&connected_peers);
-            need_inc_version = self.update_conn_info_one_peer(&peer_id_version, connected_peers);
+            need_inc_version |= self.update_conn_info_one_peer(&peer_id_version, connected_peers);
         }
         if need_inc_version {
             self.version.inc();