From 80d4d471592f92f6775e03acf3b046263208f932 Mon Sep 17 00:00:00 2001 From: Tryanks Date: Wed, 23 Sep 2026 02:57:58 +0800 Subject: [PATCH] traverse-server: the pkarr rate limiters read an injected clock manifest_and_pkarr_store_over_http asserts the fourth GET and the fourth PUT are refused, with buckets of three tokens refilled at one per second against the wall clock. A PUT commits a redb transaction with fsync; on a slow Windows runner the sequence took over a second, a token refilled and the request passed with 200. The test failed twice today on unrelated branches. RateLimiter now takes a Clock (Instant::now in production) and Server::spawn_with_clock threads it through. The integration test freezes the clock, so its bursts are consumed by request count alone. The unit test keeps driving allow_at with explicit instants. --- crates/traverse-server/src/pkarr.rs | 37 +++++++++++++++++------- crates/traverse-server/src/server.rs | 23 +++++++++++++-- crates/traverse-server/tests/instance.rs | 10 +++++-- 3 files changed, 55 insertions(+), 15 deletions(-) diff --git a/crates/traverse-server/src/pkarr.rs b/crates/traverse-server/src/pkarr.rs index f037bc2d..50daae98 100644 --- a/crates/traverse-server/src/pkarr.rs +++ b/crates/traverse-server/src/pkarr.rs @@ -139,12 +139,22 @@ impl Store { } } +/// The limiter's source of time. Production reads `Instant::now`; a test +/// freezes it so a slow disk between two requests cannot refill a bucket +/// the test expects to be empty. +pub type Clock = Arc Instant + Send + Sync>; + +pub fn system_clock() -> Clock { + Arc::new(Instant::now) +} + /// A token bucket per client address: `burst` tokens, refilled at /// `per_second`. Idle buckets are dropped once the table grows past /// `PRUNE_ABOVE` entries. pub struct RateLimiter { per_second: f64, burst: f64, + clock: Clock, buckets: Mutex>, } @@ -156,15 +166,20 @@ struct Bucket { const PRUNE_ABOVE: usize = 4096; impl RateLimiter { - pub fn new(per_second: u32, burst: u32) -> Self { + pub fn new(per_second: u32, burst: u32, clock: Clock) -> Self { Self { per_second: f64::from(per_second), burst: f64::from(burst), + clock, buckets: Mutex::new(HashMap::new()), } } - pub fn allow(&self, ip: IpAddr, now: Instant) -> bool { + pub fn allow(&self, ip: IpAddr) -> bool { + self.allow_at(ip, (self.clock)()) + } + + pub fn allow_at(&self, ip: IpAddr, now: Instant) -> bool { let mut buckets = self.buckets.lock().expect("rate limiter lock"); if buckets.len() > PRUNE_ABOVE { let (per_second, burst) = (self.per_second, self.burst); @@ -273,7 +288,7 @@ async fn get_packet( headers: HeaderMap, ) -> Response { let ip = client_ip(&headers, peer, service.trust_forwarded_for); - if !service.get_limiter.allow(ip, Instant::now()) { + if !service.get_limiter.allow(ip) { service.metrics.get_rate_limited.inc(); return rate_limited("pkarr GET rate limit"); } @@ -317,7 +332,7 @@ async fn put_packet( body: Bytes, ) -> Response { let ip = client_ip(&headers, peer, service.trust_forwarded_for); - if !service.put_limiter.allow(ip, Instant::now()) { + if !service.put_limiter.allow(ip) { service.metrics.put_rate_limited.inc(); return rate_limited("pkarr PUT rate limit"); } @@ -415,20 +430,20 @@ mod tests { #[test] fn rate_limit_is_a_per_ip_token_bucket() { - let limiter = RateLimiter::new(4, 8); + let limiter = RateLimiter::new(4, 8, system_clock()); let start = Instant::now(); let a: IpAddr = "10.0.0.1".parse().unwrap(); let b: IpAddr = "10.0.0.2".parse().unwrap(); for _ in 0..8 { - assert!(limiter.allow(a, start)); + assert!(limiter.allow_at(a, start)); } - assert!(!limiter.allow(a, start)); - assert!(limiter.allow(b, start)); + assert!(!limiter.allow_at(a, start)); + assert!(limiter.allow_at(b, start)); // Half a second refills two tokens. let later = start + Duration::from_millis(500); - assert!(limiter.allow(a, later)); - assert!(limiter.allow(a, later)); - assert!(!limiter.allow(a, later)); + assert!(limiter.allow_at(a, later)); + assert!(limiter.allow_at(a, later)); + assert!(!limiter.allow_at(a, later)); } #[test] diff --git a/crates/traverse-server/src/server.rs b/crates/traverse-server/src/server.rs index 0ab3556c..3333d1b1 100644 --- a/crates/traverse-server/src/server.rs +++ b/crates/traverse-server/src/server.rs @@ -54,6 +54,17 @@ impl Server { /// Binds every listener. `updated_at` is the manifest's `updatedAt`: /// the config file's modification time, or the process start. pub async fn spawn(config: Config, updated_at: SystemTime) -> io::Result { + Self::spawn_with_clock(config, updated_at, pkarr::system_clock()).await + } + + /// [`Self::spawn`] with the clock the pkarr rate limiters read. Tests + /// freeze it so their request sequence, not the host's disk speed, + /// decides which request is refused. + pub async fn spawn_with_clock( + config: Config, + updated_at: SystemTime, + clock: pkarr::Clock, + ) -> io::Result { let _ = rustls::crypto::ring::default_provider().install_default(); std::fs::create_dir_all(&config.data_dir)?; let tls = Tls::load(&config).await?.map(Arc::new); @@ -105,8 +116,16 @@ impl Server { let pkarr_metrics = Arc::new(PkarrMetrics::default()); let pkarr = Arc::new(PkarrService { store: Store::open(&config.data_dir.join("pkarr.redb"))?, - put_limiter: RateLimiter::new(config.pkarr.put_per_second, config.pkarr.put_burst), - get_limiter: RateLimiter::new(config.pkarr.get_per_second, config.pkarr.get_burst), + put_limiter: RateLimiter::new( + config.pkarr.put_per_second, + config.pkarr.put_burst, + clock.clone(), + ), + get_limiter: RateLimiter::new( + config.pkarr.get_per_second, + config.pkarr.get_burst, + clock, + ), trust_forwarded_for: config.http.trust_forwarded_for, metrics: pkarr_metrics.clone(), }); diff --git a/crates/traverse-server/tests/instance.rs b/crates/traverse-server/tests/instance.rs index cbb9eb54..bb09afad 100644 --- a/crates/traverse-server/tests/instance.rs +++ b/crates/traverse-server/tests/instance.rs @@ -3,7 +3,8 @@ //! pkarr store and talk through the relay. No external network. use std::{ net::{Ipv4Addr, SocketAddr}, - time::{Duration, SystemTime}, + sync::Arc, + time::{Duration, Instant, SystemTime}, }; use iroh::{ @@ -26,9 +27,14 @@ fn temp_dir(name: &str) -> std::path::PathBuf { async fn spawn(name: &str, configure: impl FnOnce(&mut Config)) -> Server { let mut config = Config::dev(temp_dir(name), SocketAddr::from((Ipv4Addr::LOCALHOST, 0))); configure(&mut config); - Server::spawn( + // The rate limiters see a frozen clock: buckets never refill, so the + // bursts below are consumed by request count alone, however long a + // slow CI disk takes between two requests. + let frozen = Instant::now(); + Server::spawn_with_clock( config, SystemTime::UNIX_EPOCH + Duration::from_secs(1_800_000_000), + Arc::new(move || frozen), ) .await .expect("server binds ephemeral ports")