diff --git a/easytier-core/src/foundation/token_bucket.rs b/easytier-core/src/foundation/token_bucket.rs index 785be405..e8bffae2 100644 --- a/easytier-core/src/foundation/token_bucket.rs +++ b/easytier-core/src/foundation/token_bucket.rs @@ -1,13 +1,14 @@ -use atomic_shim::AtomicU64; use dashmap::DashMap; -use std::sync::atomic::{AtomicBool, Ordering}; -use std::sync::{Arc, Mutex}; +use parking_lot::Mutex; +use std::sync::{ + Arc, Mutex as StdMutex, + atomic::{AtomicBool, Ordering}, +}; use std::time::{Duration, Instant}; use tokio::sync::Notify; use tokio_util::task::AbortOnDropHandle; use crate::foundation::time; -use easytier_proto::common::LimiterConfig; #[async_trait::async_trait] pub(crate) trait ByteLimiter: Send + Sync { @@ -27,49 +28,46 @@ impl ByteLimiter for () { pub(crate) type ArcByteLimiter = Arc; -/// Token Bucket rate limiter using atomic operations -pub struct TokenBucket { - available_tokens: AtomicU64, // Current token count (atomic) - last_refill_time: AtomicU64, // Last refill time as micros since epoch - config: BucketConfig, // Immutable configuration - refill_task: Mutex>>, // Background refill task - start_time: Instant, // Bucket creation time +const MIN_FILL_RATE: u64 = 8_196; +const NANOS_PER_SECOND: u128 = 1_000_000_000; - refill_notifier: Arc, +/// Token Bucket rate limiter with on-demand refill. +pub struct TokenBucket { + state: Mutex, + config: BucketConfig, + stop_notifier: Notify, stopped: AtomicBool, } +struct BucketState { + available_tokens: u64, + last_refill: Instant, + refill_remainder: u64, +} + #[derive(Clone, Copy)] pub struct BucketConfig { - capacity: u64, // Maximum token capacity - fill_rate: u64, // Tokens added per second - refill_interval: Duration, // Time between refill operations + capacity: u64, + fill_rate: u64, } -impl From for BucketConfig { - fn from(cfg: LimiterConfig) -> Self { - let burst_rate = 1.max(cfg.burst_rate.unwrap_or(1)); - let fill_rate = 8196.max(cfg.bps.unwrap_or(u64::MAX / burst_rate)); - let refill_interval = cfg - .fill_duration_ms - .map(|x| Duration::from_millis(1.max(x))) - .unwrap_or(Duration::from_millis(10)); - BucketConfig { - capacity: burst_rate * fill_rate, +impl BucketConfig { + pub fn new(capacity: u64, fill_rate: u64) -> Self { + Self { + capacity, fill_rate, - refill_interval, } } + + pub fn with_default_capacity(fill_rate: u64) -> Self { + let fill_rate = fill_rate.max(MIN_FILL_RATE); + Self::new(fill_rate, fill_rate) + } } impl TokenBucket { - pub fn new(capacity: u64, bps: u64, refill_interval: Duration) -> Arc { - let config = BucketConfig { - capacity, - fill_rate: bps, - refill_interval, - }; - Self::new_from_cfg(config) + pub fn new(capacity: u64, bps: u64) -> Arc { + Self::new_from_cfg(BucketConfig::new(capacity, bps)) } /// Creates a new Token Bucket rate limiter @@ -77,79 +75,50 @@ impl TokenBucket { /// # Arguments /// * `capacity` - Bucket capacity in bytes /// * `bps` - Bandwidth limit in bytes per second - /// * `refill_interval` - Refill interval (recommended 10-50ms) - pub fn new_from_cfg(config: BucketConfig) -> Arc { - // Create Arc instance with placeholder task - let arc_self = Arc::new(Self { - available_tokens: AtomicU64::new(config.capacity), - last_refill_time: AtomicU64::new(0), + pub fn new_from_cfg(mut config: BucketConfig) -> Arc { + config.capacity = config.capacity.max(1); + config.fill_rate = config.fill_rate.max(1); + + Arc::new(Self { + state: Mutex::new(BucketState { + available_tokens: config.capacity, + last_refill: Instant::now(), + refill_remainder: 0, + }), config, - refill_task: Mutex::new(None), - start_time: std::time::Instant::now(), - refill_notifier: Arc::new(Notify::new()), + stop_notifier: Notify::new(), stopped: AtomicBool::new(false), - }); - - // Start background refill task - let weak_bucket = Arc::downgrade(&arc_self); - let refill_interval = arc_self.config.refill_interval; - let refill_notifer = arc_self.refill_notifier.clone(); - let refill_task = tokio::spawn(async move { - let mut interval = time::interval(refill_interval); - loop { - interval.tick().await; - let Some(bucket) = weak_bucket.upgrade() else { - break; - }; - bucket.refill(); - refill_notifer.notify_waiters(); - } - }); - - // Replace placeholder task with actual one - arc_self - .refill_task - .lock() - .unwrap() - .replace(AbortOnDropHandle::new(refill_task)); - arc_self + }) } - /// Internal refill method (called only by background task) - fn refill(&self) { - let now_micros = self.elapsed_micros(); - let prev_time = self.last_refill_time.swap(now_micros, Ordering::Acquire); - - // Calculate elapsed time in seconds - let elapsed_secs = (now_micros.saturating_sub(prev_time)) as f64 / 1_000_000.0; - - // Calculate tokens to add - let tokens_to_add = (self.config.fill_rate as f64 * elapsed_secs) as u64; - if tokens_to_add == 0 { + /// Refill tokens based on elapsed time since last refill. + /// Called while holding the bucket state lock. + fn refill(&self, state: &mut BucketState, now: Instant) { + let elapsed_nanos = now.saturating_duration_since(state.last_refill).as_nanos(); + if elapsed_nanos == 0 { return; } - // Add tokens without exceeding capacity - let mut current = self.available_tokens.load(Ordering::Relaxed); - loop { - let new = current - .saturating_add(tokens_to_add) - .min(self.config.capacity); - match self.available_tokens.compare_exchange_weak( - current, - new, - Ordering::Release, - Ordering::Relaxed, - ) { - Ok(_) => break, - Err(actual) => current = actual, - } + state.last_refill = now; + if state.available_tokens == self.config.capacity { + state.refill_remainder = 0; + return; } - } - /// Calculate microseconds since bucket creation - fn elapsed_micros(&self) -> u64 { - self.start_time.elapsed().as_micros() as u64 + let generated = (self.config.fill_rate as u128) + .saturating_mul(elapsed_nanos) + .saturating_add(state.refill_remainder as u128); + let tokens_to_add = generated / NANOS_PER_SECOND; + let refill_remainder = generated % NANOS_PER_SECOND; + let available_capacity = self.config.capacity - state.available_tokens; + + if tokens_to_add >= available_capacity as u128 { + state.available_tokens = self.config.capacity; + state.refill_remainder = 0; + } else { + state.available_tokens += tokens_to_add as u64; + state.refill_remainder = refill_remainder as u64; + } } /// Attempt to consume tokens without blocking @@ -165,44 +134,63 @@ impl TokenBucket { return false; } - let mut current = self.available_tokens.load(Ordering::Relaxed); - loop { - if current < tokens { - return false; - } + let mut state = self.state.lock(); + self.refill(&mut state, Instant::now()); + if state.available_tokens < tokens { + return false; + } - let new = current - tokens; - match self.available_tokens.compare_exchange_weak( - current, - new, - Ordering::AcqRel, - Ordering::Relaxed, - ) { - Ok(_) => return true, - Err(actual) => current = actual, + state.available_tokens -= tokens; + true + } + + /// Consume tokens, sleeping until they become available. + pub async fn consume(&self, tokens: u64) { + let mut remaining = tokens; + while remaining > 0 { + if self.stopped.load(Ordering::Acquire) { + return; } + let chunk = remaining.min(self.config.capacity); + self.consume_chunk(chunk).await; + remaining -= chunk; } } - /// Consume tokens, blocking if not available - pub async fn consume(&self, tokens: u64) { + async fn consume_chunk(&self, tokens: u64) { loop { - let notified = self.refill_notifier.notified(); - if self.try_consume(tokens) { + let stopped = self.stop_notifier.notified(); + if self.stopped.load(Ordering::Acquire) { return; } - notified.await; + + let sleep_dur = { + let mut state = self.state.lock(); + self.refill(&mut state, Instant::now()); + if state.available_tokens >= tokens { + state.available_tokens -= tokens; + return; + } + + let deficit = tokens - state.available_tokens; + let required = deficit as u128 * NANOS_PER_SECOND; + let remaining = required - state.refill_remainder as u128; + let sleep_nanos = remaining + .div_ceil(self.config.fill_rate as u128) + .min(u64::MAX as u128) as u64; + Duration::from_nanos(sleep_nanos.max(1_000_000)) + }; + + tokio::select! { + _ = time::sleep(sleep_dur) => {} + _ = stopped => {} + } } } async fn stop(&self) { self.stopped.store(true, Ordering::Release); - self.refill_notifier.notify_waiters(); - let task = self.refill_task.lock().unwrap().take(); - if let Some(task) = task { - task.abort(); - let _ = task.await; - } + self.stop_notifier.notify_waiters(); } } @@ -219,7 +207,7 @@ impl ByteLimiter for TokenBucket { pub struct TokenBucketManager { buckets: Arc>>, - retain_task: Mutex>>, + retain_task: StdMutex>>, } impl Default for TokenBucketManager { @@ -252,7 +240,7 @@ impl TokenBucketManager { Self { buckets, - retain_task: Mutex::new(Some(AbortOnDropHandle::new(retain_task))), + retain_task: StdMutex::new(Some(AbortOnDropHandle::new(retain_task))), } } @@ -285,12 +273,36 @@ impl TokenBucketManager { #[cfg(test)] mod tests { use super::*; - use tokio::time::{Duration, sleep}; + use tokio::time::{Duration, sleep, timeout}; + + #[test] + fn bucket_config_uses_one_second_default_capacity() { + let config = BucketConfig::with_default_capacity(100_000); + + assert_eq!(config.capacity, 100_000); + assert_eq!(config.fill_rate, 100_000); + } + + #[test] + fn bucket_config_default_capacity_fits_regular_packets() { + let config = BucketConfig::with_default_capacity(8_196); + + assert_eq!(config.capacity, 8_196); + assert_eq!(config.fill_rate, 8_196); + } + + #[test] + fn bucket_config_preserves_explicit_capacity() { + let config = BucketConfig::new(200_000, 100_000); + + assert_eq!(config.capacity, 200_000); + assert_eq!(config.fill_rate, 100_000); + } /// Test initial state after creation #[tokio::test] async fn test_initial_state() { - let bucket = TokenBucket::new(1000, 1000, Duration::from_millis(10)); + let bucket = TokenBucket::new(1000, 1000); // Should have full capacity initially assert!(bucket.try_consume(1000)); @@ -300,7 +312,7 @@ mod tests { /// Test token consumption behavior #[tokio::test] async fn test_consumption() { - let bucket = TokenBucket::new(1500, 1000, Duration::from_millis(10)); + let bucket = TokenBucket::new(1500, 1000); // First packet should succeed assert!(bucket.try_consume(1000)); @@ -314,7 +326,7 @@ mod tests { #[tokio::test] async fn stop_releases_waiting_consumers() { - let bucket = TokenBucket::new(1, 1, Duration::from_secs(60)); + let bucket = TokenBucket::new(1, 1); assert!(bucket.try_consume(1)); let waiting = tokio::spawn({ let bucket = bucket.clone(); @@ -332,27 +344,85 @@ mod tests { assert!(bucket.try_consume(u64::MAX)); } - /// Test background refill functionality + /// Test lazy refill functionality #[tokio::test] async fn test_refill() { - let bucket = TokenBucket::new(1000, 1000, Duration::from_millis(10)); + let bucket = TokenBucket::new(1_000_000, 10_000); // Drain the bucket - assert!(bucket.try_consume(1000)); - assert!(!bucket.try_consume(1)); + assert!(bucket.try_consume(1_000_000)); - // Wait for refill (1 refill interval + buffer) + // Wait for time to pass (tokens accumulate lazily on next consume) sleep(Duration::from_millis(25)).await; + let tokens = { + let mut state = bucket.state.lock(); + bucket.refill(&mut state, Instant::now()); + state.available_tokens + }; + assert!(tokens > 0, "Expected some refilled tokens"); + assert!( + tokens < bucket.config.capacity, + "Bucket unexpectedly refilled to capacity: {}", + tokens + ); + } - // Should have approximately 20 tokens (1000 tokens/s * 0.02s) - assert!(bucket.try_consume(15)); - assert!(!bucket.try_consume(10)); // But not full capacity + #[test] + fn test_refill_preserves_fractional_tokens() { + let bucket = TokenBucket::new(100, 100); + let start = Instant::now(); + let mut state = bucket.state.lock(); + state.available_tokens = 0; + state.last_refill = start; + state.refill_remainder = 0; + + for step in 1..=10 { + bucket.refill(&mut state, start + Duration::from_millis(step * 15)); + } + + assert_eq!(state.available_tokens, 15); + assert_eq!(state.refill_remainder, 0); + } + + #[test] + fn test_refill_preserves_submicrosecond_time() { + let bucket = TokenBucket::new(100, 1_000_000); + let start = Instant::now(); + let mut state = bucket.state.lock(); + state.available_tokens = 0; + state.last_refill = start; + state.refill_remainder = 0; + + for step in 1..=10 { + bucket.refill(&mut state, start + Duration::from_nanos(step * 1_500)); + } + + assert_eq!(state.available_tokens, 15); + assert_eq!(state.refill_remainder, 0); + } + + #[test] + fn test_refill_discards_excess_credit_at_capacity() { + let bucket = TokenBucket::new(10, 10); + let start = Instant::now(); + let mut state = bucket.state.lock(); + state.available_tokens = 0; + state.last_refill = start; + + let refill_time = start + Duration::from_secs(2); + bucket.refill(&mut state, refill_time); + assert_eq!(state.available_tokens, 10); + assert_eq!(state.refill_remainder, 0); + + state.available_tokens = 0; + bucket.refill(&mut state, refill_time); + assert_eq!(state.available_tokens, 0); } /// Test capacity enforcement #[tokio::test] async fn test_capacity_limit() { - let bucket = TokenBucket::new(500, 1000, Duration::from_millis(10)); + let bucket = TokenBucket::new(500, 1000); // Wait longer than refill interval sleep(Duration::from_millis(50)).await; @@ -363,39 +433,39 @@ mod tests { } /// Test high load with concurrent access - #[tokio::test] + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn test_concurrent_access() { - let bucket = TokenBucket::new(10_000, 1_000_000, Duration::from_millis(10)); + let bucket = TokenBucket::new(10_000, 1); let mut handles = vec![]; // Spawn 100 tasks to consume tokens concurrently for _ in 0..100 { let bucket = bucket.clone(); handles.push(tokio::spawn(async move { + let mut consumed = 0; for _ in 0..100 { - let _ = bucket.try_consume(10); + if bucket.try_consume(10) { + consumed += 10; + } } + consumed })); } // Wait for all tasks to complete + let mut consumed = 0; for handle in handles { - handle.await.unwrap(); + consumed += handle.await.unwrap(); } - // Verify we didn't exceed capacity - let tokens_left = bucket.available_tokens.load(Ordering::Relaxed); - assert!( - tokens_left <= 10_000, - "Tokens exceeded capacity: {}", - tokens_left - ); + assert_eq!(consumed, 10_000); + assert_eq!(bucket.state.lock().available_tokens, 0); } /// Test behavior when packet size exceeds capacity #[tokio::test] async fn test_oversized_packet() { - let bucket = TokenBucket::new(1500, 1000, Duration::from_millis(10)); + let bucket = TokenBucket::new(1500, 1000); // Packet larger than capacity should be rejected assert!(!bucket.try_consume(1600)); @@ -404,23 +474,35 @@ mod tests { assert!(bucket.try_consume(1000)); } - /// Test refill precision with small intervals #[tokio::test] - async fn test_refill_precision() { - let bucket = TokenBucket::new(10_000, 10_000, Duration::from_micros(100)); // 100μs interval + async fn test_zero_fill_rate_is_normalized() { + let bucket = TokenBucket::new(1000, 0); - // Drain most tokens - assert!(bucket.try_consume(9900)); + assert_eq!(bucket.config.fill_rate, 1); + } - // Wait for multiple refills - sleep(Duration::from_millis(1)).await; + #[tokio::test] + async fn test_consume_oversized_packet_in_chunks() { + let bucket = TokenBucket::new(10, 1_000_000); - // Should have accumulated about 100 tokens (10,000 tokens/s * 0.001s) - let tokens = bucket.available_tokens.load(Ordering::Relaxed); - assert!( - (100..=200).contains(&tokens), - "Unexpected token count: {}", - tokens - ); + timeout(Duration::from_millis(100), bucket.consume(25)) + .await + .expect("oversized consume should be split into capacity-sized chunks"); + } + + /// Test refill precision after elapsed time. + #[test] + fn test_refill_precision() { + let bucket = TokenBucket::new(10_000, 10_000); + let start = Instant::now(); + let mut state = bucket.state.lock(); + state.available_tokens = 100; + state.last_refill = start; + state.refill_remainder = 0; + + bucket.refill(&mut state, start + Duration::from_micros(1_234)); + + assert_eq!(state.available_tokens, 112); + assert_eq!(state.refill_remainder, 340_000_000); } } diff --git a/easytier-core/src/packet/mod.rs b/easytier-core/src/packet/mod.rs index 739222e2..87340314 100644 --- a/easytier-core/src/packet/mod.rs +++ b/easytier-core/src/packet/mod.rs @@ -569,11 +569,7 @@ impl ZCPacket { ret.mut_payload()[foreign_network_hdr.get_header_len()..] .copy_from_slice(foreign_zc_packet.tunnel_payload()); - let hdr = ret.mut_peer_manager_header().unwrap(); - hdr.from_peer_id = 0.into(); - hdr.to_peer_id = 0.into(); - hdr.packet_type = PacketType::ForeignNetworkPacket as u8; - hdr.len.set(total_payload_len as u32); + ret.fill_peer_manager_hdr(0, 0, PacketType::ForeignNetworkPacket as u8); ret } @@ -689,6 +685,7 @@ impl ZCPacket { hdr.packet_type = packet_type; hdr.flags = 0; hdr.forward_counter = 1; + hdr.reserved = 0; hdr.len.set(payload_len as u32); } @@ -783,6 +780,11 @@ impl ZCPacket { } pub fn foreign_network_inner_packet_type(&self) -> Option { + self.foreign_network_inner_packet_info() + .map(|(hdr, _)| hdr.packet_type) + } + + pub fn foreign_network_inner_packet_info(&self) -> Option<(&PeerManagerHeader, usize)> { if self.peer_manager_header()?.packet_type != PacketType::ForeignNetworkPacket as u8 { return None; } @@ -790,7 +792,9 @@ impl ZCPacket { let payload = self.payload(); let hdr = ForeignNetworkPacketHeader::ref_from_prefix(payload)?; let inner_packet = payload.get(hdr.get_header_len()..)?; - PeerManagerHeader::ref_from_prefix(inner_packet).map(|hdr| hdr.packet_type) + let peer_manager_header = PeerManagerHeader::ref_from_prefix(inner_packet)?; + let payload_len = inner_packet.len() - PEER_MANAGER_HEADER_SIZE; + Some((peer_manager_header, payload_len)) } pub fn foreign_network_packet(mut self) -> Self { diff --git a/easytier-core/src/peers/acl/processor.rs b/easytier-core/src/peers/acl/processor.rs index 649f112a..b1913fe0 100644 --- a/easytier-core/src/peers/acl/processor.rs +++ b/easytier-core/src/peers/acl/processor.rs @@ -786,11 +786,7 @@ impl AclProcessor { panic!("Rate limit bucket not found"); } RateLimitValue { - token_bucket: TokenBucket::new( - burst as u64, - rate as u64, - Duration::from_millis(10), - ), + token_bucket: TokenBucket::new(burst as u64, rate as u64), last_update: Instant::now(), } }); diff --git a/easytier-core/src/peers/conn/peer_conn.rs b/easytier-core/src/peers/conn/peer_conn.rs index 377abe97..15f3f8fb 100644 --- a/easytier-core/src/peers/conn/peer_conn.rs +++ b/easytier-core/src/peers/conn/peer_conn.rs @@ -36,6 +36,7 @@ use crate::peers::{ PacketRecvChan, context::{ArcPeerContext, NetworkIdentity, NetworkSecretDigest}, send_packet_to_chan, + traffic_metrics::data_packet_payload_len, }; use crate::{ config::PeerId, @@ -1290,6 +1291,7 @@ 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); let Some(peer_mgr_hdr) = zc_packet.mut_peer_manager_header() else { tracing::error!( "unexpected packet: {:?}, cannot decode peer manager hdr", @@ -1315,8 +1317,10 @@ impl PeerConn { break; } - if let Some(limiter) = recv_limiter.as_ref() { - limiter.consume(buf_len).await; + if let Some(payload_len) = limited_payload_len + && let Some(limiter) = recv_limiter.as_ref() + { + limiter.consume(payload_len).await; } } diff --git a/easytier-core/src/peers/context.rs b/easytier-core/src/peers/context.rs index 149a7128..e790facb 100644 --- a/easytier-core/src/peers/context.rs +++ b/easytier-core/src/peers/context.rs @@ -13,9 +13,7 @@ use cidr::{Ipv4Cidr, Ipv4Inet, Ipv6Cidr, Ipv6Inet}; use dashmap::DashMap; use easytier_proto::{ acl::Acl, - common::{ - FlagsInConfig, LimiterConfig, PeerFeatureFlag, SecureModeConfig, StunInfo, TunnelInfo, - }, + common::{FlagsInConfig, PeerFeatureFlag, SecureModeConfig, StunInfo, TunnelInfo}, peer_rpc::{PeerGroupInfo, TrustedCredentialPubkeyProof}, }; use hmac::{Hmac, Mac}; @@ -31,7 +29,7 @@ use crate::{ }, events::{CoreEvent, CoreEventSink}, foundation::stats::{LabelSet, LabelType, MetricName, StatsManager}, - foundation::token_bucket::{ArcByteLimiter, TokenBucketManager}, + foundation::token_bucket::{ArcByteLimiter, BucketConfig, TokenBucketManager}, peers::{ credential_manager::{CredentialManager, CredentialStorage}, util::shrink_dashmap, @@ -362,17 +360,7 @@ impl CorePeerContext { return None; } let manager = state.manager.get_or_insert_with(TokenBucketManager::new); - Some( - manager.get_or_create( - key, - LimiterConfig { - burst_rate: None, - bps: Some(bps), - fill_duration_ms: None, - } - .into(), - ), - ) + Some(manager.get_or_create(key, BucketConfig::with_default_capacity(bps))) } pub(crate) async fn stop(&self) { diff --git a/easytier-core/src/peers/foreign_network/mod.rs b/easytier-core/src/peers/foreign_network/mod.rs index 8983371c..46c440c4 100644 --- a/easytier-core/src/peers/foreign_network/mod.rs +++ b/easytier-core/src/peers/foreign_network/mod.rs @@ -51,7 +51,8 @@ use super::{ relay_peer_map::RelayPeerMap, route::{NextHopPolicy, Route, RouteInterface, peer_ospf_route::PeerRoute}, traffic_metrics::{ - TrafficKind, TrafficMetricRecorder, is_relay_data_packet_type, traffic_kind, + TrafficKind, TrafficMetricRecorder, data_packet_payload_len, is_relay_data_packet_type, + traffic_kind, }, util::shrink_dashmap, whitelist::check_network_in_relay_whitelist, @@ -1347,8 +1348,9 @@ impl ForeignNetworkPacketRouter { ); continue; } - if let Some(bps_limiter) = bps_limiter.as_ref() - && !bps_limiter.try_consume(len.into()) + if let Some(payload_len) = data_packet_payload_len(&zc_packet) + && let Some(bps_limiter) = bps_limiter.as_ref() + && !bps_limiter.try_consume(payload_len) { continue; } diff --git a/easytier-core/src/peers/traffic_metrics.rs b/easytier-core/src/peers/traffic_metrics.rs index be1f3238..724bed6c 100644 --- a/easytier-core/src/peers/traffic_metrics.rs +++ b/easytier-core/src/peers/traffic_metrics.rs @@ -5,7 +5,7 @@ use futures::future::BoxFuture; use crate::config::PeerId; use crate::foundation::stats::{CounterHandle, LabelSet, LabelType, MetricName, StatsManager}; -use crate::packet::PacketType; +use crate::packet::{PacketType, ZCPacket}; use crate::peers::util::shrink_dashmap; use crate::proto::peer_rpc::RoutePeerInfo; @@ -200,6 +200,26 @@ pub fn is_relay_data_packet_type(packet_type: u8) -> bool { || packet_type == PacketType::ForeignNetworkPacket as u8 } +pub(crate) fn data_packet_payload_len(packet: &ZCPacket) -> Option { + let header = packet.peer_manager_header()?; + if header.packet_type == PacketType::ForeignNetworkPacket as u8 { + if header.is_encrypted() { + return Some(packet.payload_len() as u64); + } + return match packet.foreign_network_inner_packet_info() { + Some((inner_header, payload_len)) + if traffic_kind(inner_header.packet_type) == TrafficKind::Data => + { + Some(payload_len as u64) + } + Some(_) => None, + None => Some(packet.payload_len() as u64), + }; + } + + (traffic_kind(header.packet_type) == TrafficKind::Data).then(|| packet.payload_len() as u64) +} + #[derive(Clone)] struct TrafficMetricGroup { data: Arc, @@ -327,6 +347,86 @@ mod tests { .with_label_type(LabelType::ToInstanceId(instance_id.to_string())) } + #[test] + fn relay_limiter_classifies_only_data_packets() { + for packet_type in [ + PacketType::Data, + PacketType::KcpSrc, + PacketType::KcpDst, + PacketType::QuicSrc, + PacketType::QuicDst, + PacketType::DataWithKcpSrcModified, + PacketType::DataWithQuicSrcModified, + PacketType::ForeignNetworkPacket, + ] { + assert!(is_relay_data_packet_type(packet_type as u8)); + } + + for packet_type in [ + PacketType::Ping, + PacketType::Pong, + PacketType::RpcReq, + PacketType::RpcResp, + PacketType::RelayHandshake, + PacketType::RelayHandshakeAck, + ] { + assert!(!is_relay_data_packet_type(packet_type as u8)); + } + } + + fn packet_with_type(packet_type: PacketType, payload_len: usize) -> ZCPacket { + let mut packet = ZCPacket::new_with_payload(&vec![0; payload_len]); + packet.fill_peer_manager_hdr(1, 2, packet_type as u8); + packet + } + + #[test] + fn data_packet_payload_len_uses_easytier_payload() { + let mut data_packet = packet_with_type(PacketType::Data, 1_024); + data_packet.mut_peer_manager_header().unwrap().len.set(0); + assert_eq!(data_packet_payload_len(&data_packet), Some(1_024)); + + let control_packet = packet_with_type(PacketType::RpcReq, 128); + assert_eq!(data_packet_payload_len(&control_packet), None); + } + + #[test] + fn data_packet_payload_len_unwraps_foreign_network_packet() { + let network_name = "foreign".to_string(); + let mut inner_data = packet_with_type(PacketType::Data, 1_024); + inner_data.mut_peer_manager_header().unwrap().len.set(0); + let foreign_data = ZCPacket::new_for_foreign_network(&network_name, 2, &inner_data); + assert!(!foreign_data.peer_manager_header().unwrap().is_encrypted()); + assert_eq!(data_packet_payload_len(&foreign_data), Some(1_024)); + + let inner_control = packet_with_type(PacketType::RpcReq, 128); + let foreign_control = ZCPacket::new_for_foreign_network(&network_name, 2, &inner_control); + assert_eq!(data_packet_payload_len(&foreign_control), None); + + let malformed_foreign = packet_with_type(PacketType::ForeignNetworkPacket, 7); + assert_eq!( + data_packet_payload_len(&malformed_foreign), + Some(malformed_foreign.payload_len() as u64) + ); + } + + #[test] + fn data_packet_payload_len_charges_encrypted_foreign_payload() { + let network_name = "foreign".to_string(); + let inner_control = packet_with_type(PacketType::RpcReq, 128); + let mut encrypted_foreign = + ZCPacket::new_for_foreign_network(&network_name, 2, &inner_control); + encrypted_foreign + .mut_peer_manager_header() + .unwrap() + .set_encrypted(true); + + assert_eq!( + data_packet_payload_len(&encrypted_foreign), + Some(encrypted_foreign.payload_len() as u64) + ); + } + #[tokio::test] async fn logical_traffic_metrics_upgrade_unknown_instance_label() { let stats_mgr = Arc::new(StatsManager::new()); diff --git a/easytier-proto/proto/common.proto b/easytier-proto/proto/common.proto index 7a0bf404..632cab79 100644 --- a/easytier-proto/proto/common.proto +++ b/easytier-proto/proto/common.proto @@ -250,14 +250,6 @@ message PortForwardConfigPb { message ProxyDstInfo { SocketAddr dst_addr = 1; } -message LimiterConfig { - optional uint64 burst_rate = - 1; // default 1 means no burst (capacity is same with bps) - optional uint64 bps = 2; // default 0 means no limit (unit is B/s) - optional uint64 fill_duration_ms = - 3; // default 10ms, the period to fill the bucket -} - message SecureModeConfig { bool enabled = 1; diff --git a/easytier/src/tests/three_node.rs b/easytier/src/tests/three_node.rs index cfc6d275..15ffa902 100644 --- a/easytier/src/tests/three_node.rs +++ b/easytier/src/tests/three_node.rs @@ -2296,13 +2296,7 @@ pub async fn relay_bps_limit_test(#[values(100, 200, 400, 800)] bps_limit: u64) println!("bps: {}", bps); let bps = bps as u64 / 1024; - // allow 50kb jitter - assert!( - bps >= bps_limit - 50 && bps <= bps_limit + 50, - "bps: {}, bps_limit: {}", - bps, - bps_limit - ); + assert_limited_payload_bps(bps, bps_limit); drop_insts(insts).await; } @@ -2339,16 +2333,25 @@ pub async fn instance_recv_bps_limit_test(#[values(100, 800)] bps_limit: u64) { println!("bps: {}", bps); let bps = bps as u64 / 1024; - assert!( - bps >= bps_limit - 50 && bps <= bps_limit + 50, - "bps: {}, bps_limit: {}", - bps, - bps_limit - ); + assert_limited_payload_bps(bps, bps_limit); drop_insts(insts).await; } +fn assert_limited_payload_bps(bps: u64, bps_limit: u64) { + // The benchmark measures TCP application payload while the limiter counts + // EasyTier data payload, including the inner IP and transport headers. + let min_bps = bps_limit.saturating_sub((bps_limit / 10).max(50)); + let max_bps = bps_limit + 50; + assert!( + bps >= min_bps && bps <= max_bps, + "bps: {}, expected: {}..={}", + bps, + min_bps, + max_bps + ); +} + async fn assert_peer_admission_blocked(inst: &Instance, url: url::Url) { let ip = url .host_str()