fix(core): keep peer liveness responsive under receive limits (#2590)

easytier-core/src/foundation/token_bucket.rs

@@ -10,17 +10,11 @@ use tokio_util::task::AbortOnDropHandle;
 
 use crate::foundation::time;
 
-#[async_trait::async_trait]
 pub(crate) trait ByteLimiter: Send + Sync {
-    async fn consume(&self, bytes: u64);
-
     fn try_consume(&self, bytes: u64) -> bool;
 }
 
-#[async_trait::async_trait]
 impl ByteLimiter for () {
-    async fn consume(&self, _bytes: u64) {}
-
     fn try_consume(&self, _bytes: u64) -> bool {
         true
     }
@@ -123,16 +117,19 @@ impl TokenBucket {
 
     /// Attempt to consume tokens without blocking
     ///
+    /// Requests larger than the bucket capacity are charged one full
+    /// bucket instead of being rejected, so oversized packets can still
+    /// pass whenever the bucket holds a full burst; traffic consisting
+    /// only of oversized packets can therefore sustain up to
+    /// packet-size/capacity times the nominal rate.
+    ///
     /// # Returns
     /// `true` if tokens were consumed, `false` if insufficient tokens
     pub fn try_consume(&self, tokens: u64) -> bool {
         if self.stopped.load(Ordering::Acquire) {
             return true;
         }
-        // Fast path for oversized packets
-        if tokens > self.config.capacity {
-            return false;
-        }
+        let tokens = tokens.min(self.config.capacity);
 
         let mut state = self.state.lock();
         self.refill(&mut state, Instant::now());
@@ -194,12 +191,7 @@ impl TokenBucket {
     }
 }
 
-#[async_trait::async_trait]
 impl ByteLimiter for TokenBucket {
-    async fn consume(&self, bytes: u64) {
-        TokenBucket::consume(self, bytes).await;
-    }
-
     fn try_consume(&self, bytes: u64) -> bool {
         TokenBucket::try_consume(self, bytes)
     }
@@ -469,13 +461,19 @@ mod tests {
     /// Test behavior when packet size exceeds capacity
     #[tokio::test]
     async fn test_oversized_packet() {
-        let bucket = TokenBucket::new(1500, 1000);
+        // Fill rate of 1 token/s keeps the test independent of timing:
+        // even a long scheduling stall cannot refill a meaningful amount.
+        let bucket = TokenBucket::new(1500, 1);
 
-        // Packet larger than capacity should be rejected
-        assert!(!bucket.try_consume(1600));
+        // Oversized packets are charged one full bucket instead of
+        // being rejected, so a full bucket always admits them.
+        assert!(bucket.try_consume(1600));
+        assert_eq!(bucket.state.lock().available_tokens, 0);
 
-        // Regular packets should still work
+        // A partially drained bucket still rejects them.
+        let bucket = TokenBucket::new(1500, 1);
         assert!(bucket.try_consume(1000));
+        assert!(!bucket.try_consume(1600));
     }
 
     #[tokio::test]

easytier-core/src/peers/conn/peer_conn.rs

@@ -1317,7 +1317,12 @@ impl PeerConn {
 
                     let mut zc_packet = ret.unwrap();
                     let buf_len = zc_packet.buf_len() as u64;
-                    let limited_payload_len = data_packet_payload_len(&zc_packet);
+                    if let Some(payload_len) = data_packet_payload_len(&zc_packet)
+                        && let Some(limiter) = recv_limiter.as_ref()
+                        && !limiter.try_consume(payload_len)
+                    {
+                        continue;
+                    }
                     let Some(peer_mgr_hdr) = zc_packet.mut_peer_manager_header() else {
                         tracing::error!(
                             "unexpected packet: {:?}, cannot decode peer manager hdr",
@@ -1345,12 +1350,6 @@ impl PeerConn {
                     {
                         break;
                     }
-
-                    if let Some(payload_len) = limited_payload_len
-                        && let Some(limiter) = recv_limiter.as_ref()
-                    {
-                        limiter.consume(payload_len).await;
-                    }
                 }
 
                 tracing::info!("end recving peer conn packet");

easytier-core/src/peers/tests.rs

@@ -6,6 +6,7 @@ use x25519_dalek::{PublicKey, StaticSecret};
 use crate::foundation::time::{Duration, timeout};
 
 use crate::{
+    foundation::token_bucket::{ArcByteLimiter, ByteLimiter},
     packet::{PacketType, ZCPacket},
     peers::{
         PeerConnectionOrigin, PeerPacketIngress,
@@ -14,7 +15,7 @@ use crate::{
             peer_map::PeerMap,
             peer_session::PeerSessionStore,
         },
-        context::NetworkIdentity,
+        context::{NetworkIdentity, PeerContext},
         create_packet_recv_chan,
         error::Error,
         recv_packet_envelope_from_chan,
@@ -23,6 +24,32 @@ use crate::{
     tunnel::ring::create_ring_tunnel_pair,
 };
 
+struct RejectingRecvLimiter;
+
+impl ByteLimiter for RejectingRecvLimiter {
+    fn try_consume(&self, _bytes: u64) -> bool {
+        false
+    }
+}
+
+struct LimitedPeerContext {
+    limiter: ArcByteLimiter,
+}
+
+impl PeerContext for LimitedPeerContext {
+    fn network_identity(&self) -> NetworkIdentity {
+        NetworkIdentity::default()
+    }
+
+    fn recv_limiter(
+        &self,
+        _network_name: &str,
+        _is_foreign_network: bool,
+    ) -> Option<ArcByteLimiter> {
+        Some(self.limiter.clone())
+    }
+}
+
 impl PeerConn {
     #[tracing::instrument]
     async fn do_handshake_as_server(&mut self) -> Result<(), Error> {
@@ -53,6 +80,56 @@ async fn peer_conn_handshake_over_memory_tunnel() {
     assert_eq!(server.get_conn_info().features, ["liveness-echo-v1"]);
 }
 
+#[tokio::test]
+async fn peer_recv_limit_drops_excess_data_without_blocking_ping() {
+    let peer_session_store = Arc::new(PeerSessionStore::new());
+    let (client_tunnel, server_tunnel) = create_ring_tunnel_pair();
+    let client_ctx = Arc::new(NoopPeerContext::default());
+    let server_ctx = Arc::new(LimitedPeerContext {
+        limiter: Arc::new(RejectingRecvLimiter),
+    });
+
+    let mut client = PeerConn::new(1, client_ctx, client_tunnel, peer_session_store.clone());
+    let mut server = PeerConn::new(2, server_ctx, server_tunnel, peer_session_store);
+
+    let (client_ret, server_ret) = tokio::join!(
+        client.do_handshake_as_client(),
+        server.do_handshake_as_server()
+    );
+    client_ret.unwrap();
+    server_ret.unwrap();
+
+    let (client_tx, _client_rx) = create_packet_recv_chan();
+    let (server_tx, mut server_rx) = create_packet_recv_chan();
+    client.start_recv_loop(client_tx).await;
+    server.start_recv_loop(server_tx).await;
+
+    let client_rx_packets = client.get_stats().rx_packets;
+
+    let mut data = ZCPacket::new_with_payload(b"over limit");
+    data.fill_peer_manager_hdr(1, 2, PacketType::Data as u8);
+    client.send_msg(data).await.unwrap();
+
+    let mut ping = ZCPacket::new_with_payload(&[]);
+    ping.fill_peer_manager_hdr(1, 2, PacketType::Ping as u8);
+    client.send_msg(ping).await.unwrap();
+
+    timeout(Duration::from_secs(1), async {
+        while client.get_stats().rx_packets == client_rx_packets {
+            tokio::task::yield_now().await;
+        }
+    })
+    .await
+    .expect("receive limiting must not block the following pong");
+
+    assert!(
+        timeout(Duration::from_millis(50), server_rx.recv())
+            .await
+            .is_err(),
+        "over-limit data packet should be dropped"
+    );
+}
+
 #[tokio::test]
 async fn peer_conn_noise_handshake_advertises_liveness_echo() {
     fn context(peer_key: u8) -> Arc<NoopPeerContext> {