From 0f3d8e443492477b97eb2ec7e10f04396fe2fbf0 Mon Sep 17 00:00:00 2001 From: KKRainbow <5665404+KKRainbow@users.noreply.github.com> Date: Sat, 19 Sep 2026 08:22:12 +0800 Subject: [PATCH] fix(core): keep peer liveness responsive under receive limits (#2590) --- easytier-core/src/foundation/token_bucket.rs | 36 +++++---- easytier-core/src/peers/conn/peer_conn.rs | 13 ++-- easytier-core/src/peers/tests.rs | 79 +++++++++++++++++++- 3 files changed, 101 insertions(+), 27 deletions(-) diff --git a/easytier-core/src/foundation/token_bucket.rs b/easytier-core/src/foundation/token_bucket.rs index 583aa566..c692cd65 100644 --- a/easytier-core/src/foundation/token_bucket.rs +++ b/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] diff --git a/easytier-core/src/peers/conn/peer_conn.rs b/easytier-core/src/peers/conn/peer_conn.rs index a49c4937..19c1840a 100644 --- a/easytier-core/src/peers/conn/peer_conn.rs +++ b/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"); diff --git a/easytier-core/src/peers/tests.rs b/easytier-core/src/peers/tests.rs index 4da0107d..bc928a98 100644 --- a/easytier-core/src/peers/tests.rs +++ b/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 { + 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 {