diff --git a/server/crates/kvitt-server/src/api/rate_limit.rs b/server/crates/kvitt-server/src/api/rate_limit.rs index b33cd6c..4888345 100644 --- a/server/crates/kvitt-server/src/api/rate_limit.rs +++ b/server/crates/kvitt-server/src/api/rate_limit.rs @@ -3,9 +3,13 @@ //! * Unauthenticated endpoints (E-1 to E-6 and unknown paths): at most 60 requests per minute per client IP. //! * Authenticated endpoints: at most 600 requests per minute per device, keyed by the hash of the bearer token //! (the token identifies the device; the limiter never touches the database). A request to an authenticated -//! endpoint without a well-formed bearer token counts against its IP like an unauthenticated one. A response 401 -//! also counts against the IP, and once the IP is over its limit, tokens that have never been accepted are refused -//! too, so random tokens do not bypass the IP limit. +//! endpoint without a well-formed bearer token counts against its IP like an unauthenticated one. +//! * Accepted tokens (D-115): a token becomes accepted only with a 2xx response. A request with a token that is not +//! accepted counts against its IP before it runs (refused once the IP is over its limit); when it is answered 2xx +//! that count is given back, so a newly paired device costs its IP nothing. A 401 for an accepted token counts +//! against the IP and the token is no longer accepted. Random tokens therefore never bypass the IP limit, whatever +//! their response (401, 426 without `X-Kvitt-Protocol`, 405, …), and every 401 counts against the IP exactly once. +//! Token entries, accepted or not, are pruned once their window has expired, so the map stays bounded. //! * Invite accepts: after 10 failed accepts (any 4xx other than 429) from one IP within 60 minutes, every accept //! from that IP is `rate_limited` until the 60 minutes since the first of those failures have passed. //! @@ -121,7 +125,7 @@ fn full(slot: Option<&Window>, now: DateTime, len: Duration, limit: u32) -> #[derive(Debug, Default)] struct TokenState { window: Option, - /// The token has had a response other than 401: it belongs to a device. + /// The token's last authenticated response was 2xx (and none was 401 since): it belongs to a device. accepted: bool, } @@ -140,11 +144,29 @@ impl Buckets { } self.last_prune = Some(now); self.ip.retain(|_, w| !w.expired(now, REQUEST_WINDOW)); + // Accepted tokens too: re-accepting one costs a single request counted against its IP and given back. self.tokens - .retain(|_, t| t.accepted || t.window.is_some_and(|w| !w.expired(now, REQUEST_WINDOW))); + .retain(|_, t| t.window.is_some_and(|w| !w.expired(now, REQUEST_WINDOW))); self.accept_failures .retain(|_, w| !w.expired(now, ACCEPT_FAILURE_WINDOW)); } + + /// Counts one request against `ip`'s unauthenticated window; the window's start, or `Err(retry_after_secs)`. + fn hit_ip(&mut self, ip: IpAddr, now: DateTime) -> Result, u64> { + let mut slot = self.ip.get(&ip).copied(); + let result = hit(&mut slot, now, REQUEST_WINDOW, UNAUTHENTICATED_LIMIT); + if let Some(w) = slot { + self.ip.insert(ip, w); + } + result.map(|()| slot.map_or(now, |w| w.start)) + } +} + +/// What [`Limiter::admit`] counted, for [`Limiter::record`]. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +struct Admitted { + /// The start of the IP window an unaccepted token's request was counted in. + ip_window: Option>, } /// How a request is counted. @@ -233,7 +255,7 @@ impl Limiter { } /// Counts the request before it runs; `Err(retry_after_secs)` refuses it. - fn admit(&self, key: &Key, accept: bool, now: DateTime) -> Result<(), u64> { + fn admit(&self, key: &Key, accept: bool, now: DateTime) -> Result { let mut buckets = lock(&self.buckets); buckets.prune(now); match key { @@ -248,33 +270,32 @@ impl Limiter { { return Err(retry); } - let mut slot = buckets.ip.get(ip).copied(); - let result = hit(&mut slot, now, REQUEST_WINDOW, UNAUTHENTICATED_LIMIT); - if let Some(w) = slot { - buckets.ip.insert(*ip, w); - } - result + buckets.hit_ip(*ip, now)?; + Ok(Admitted::default()) } Key::Token { ip, hash } => { let accepted = buckets.tokens.get(hash).is_some_and(|t| t.accepted); - if !accepted - && let Some(retry) = full( - buckets.ip.get(ip), - now, - REQUEST_WINDOW, - UNAUTHENTICATED_LIMIT, - ) - { - return Err(retry); - } + let ip_window = if accepted { + None + } else { + Some(buckets.hit_ip(*ip, now)?) + }; let state = buckets.tokens.entry(hash.clone()).or_default(); - hit(&mut state.window, now, REQUEST_WINDOW, AUTHENTICATED_LIMIT) + hit(&mut state.window, now, REQUEST_WINDOW, AUTHENTICATED_LIMIT)?; + Ok(Admitted { ip_window }) } } } /// Records what the response says about the request (401s, accept failures, accepted tokens). - fn record(&self, key: &Key, accept: bool, status: StatusCode, now: DateTime) { + fn record( + &self, + key: &Key, + accept: bool, + admitted: Admitted, + status: StatusCode, + now: DateTime, + ) { let mut buckets = lock(&self.buckets); match key { Key::Ip(ip) => { @@ -288,14 +309,22 @@ impl Limiter { } } Key::Token { ip, hash } => { - if status == StatusCode::UNAUTHORIZED { - let mut slot = buckets.ip.get(ip).copied(); - let _ = hit(&mut slot, now, REQUEST_WINDOW, UNAUTHENTICATED_LIMIT); - if let Some(w) = slot { - buckets.ip.insert(*ip, w); - } - } else if status != StatusCode::TOO_MANY_REQUESTS { + if status.is_success() { buckets.tokens.entry(hash.clone()).or_default().accepted = true; + // The request belonged to a device after all: give back what admit counted against the IP. + if let Some(start) = admitted.ip_window + && let Some(w) = buckets.ip.get_mut(ip) + && w.start == start + { + w.count = w.count.saturating_sub(1); + } + } else if status == StatusCode::UNAUTHORIZED && admitted.ip_window.is_none() { + // An accepted token that is now unknown or revoked: count the 401 (an unaccepted token's + // request was counted by admit already) and stop treating the token as a device. + let _ = buckets.hit_ip(*ip, now); + if let Some(state) = buckets.tokens.get_mut(hash) { + state.accepted = false; + } } } } @@ -314,14 +343,15 @@ impl Limiter { let (peer, forwarded) = Self::peer_and_forwarded(&request); let ip = self.client_ip(peer, forwarded).await; let key = Self::key(&request, &route, ip); - if let Err(retry_after) = self.admit(&key, accept, self.clock.now()) { - return AppError::rate_limited(retry_after).into_response(); - } + let admitted = match self.admit(&key, accept, self.clock.now()) { + Ok(admitted) => admitted, + Err(retry_after) => return AppError::rate_limited(retry_after).into_response(), + }; let response = match inner.call(request).await { Ok(response) => response, Err(never) => match never {}, }; - self.record(&key, accept, response.status(), self.clock.now()); + self.record(&key, accept, admitted, response.status(), self.clock.now()); response } } @@ -408,6 +438,94 @@ mod tests { ); } + fn token_key(i: u32) -> Key { + Key::Token { + ip: IpAddr::V4(Ipv4Addr::LOCALHOST), + hash: sha256_hex(&format!("token {i}")), + } + } + + fn ip_count(limiter: &Limiter) -> u32 { + lock(&limiter.buckets) + .ip + .get(&IpAddr::V4(Ipv4Addr::LOCALHOST)) + .map_or(0, |w| w.count) + } + + #[test] + fn unaccepted_tokens_count_against_the_ip_whatever_the_response() { + let now = parse_ts("2026-10-17T12:00:00.000Z").unwrap(); + let limiter = Limiter::new(Arc::new(FixedClock::new(now)), String::new()); + // 426 (no X-Kvitt-Protocol), 405 and 401 alike: none of them accepts the token, each counts once. + let statuses = [ + StatusCode::UPGRADE_REQUIRED, + StatusCode::METHOD_NOT_ALLOWED, + StatusCode::UNAUTHORIZED, + ]; + for i in 0..UNAUTHENTICATED_LIMIT { + let key = token_key(i); + let admitted = limiter.admit(&key, false, now).unwrap(); + let status = statuses[usize::try_from(i).unwrap() % statuses.len()]; + limiter.record(&key, false, admitted, status, now); + } + assert_eq!(ip_count(&limiter), UNAUTHENTICATED_LIMIT); + assert_eq!(limiter.admit(&token_key(999), false, now), Err(60)); + assert!(lock(&limiter.buckets).tokens.values().all(|t| !t.accepted)); + } + + #[test] + fn a_2xx_accepts_the_token_and_gives_its_ip_count_back() { + let now = parse_ts("2026-10-17T12:00:00.000Z").unwrap(); + let limiter = Limiter::new(Arc::new(FixedClock::new(now)), String::new()); + let device = token_key(1); + let admitted = limiter.admit(&device, false, now).unwrap(); + assert_eq!(ip_count(&limiter), 1); + limiter.record(&device, false, admitted, StatusCode::OK, now); + assert_eq!(ip_count(&limiter), 0); + // Accepted: further requests, whatever they are answered, do not count against the IP. + for _ in 0..100 { + let admitted = limiter.admit(&device, false, now).unwrap(); + limiter.record(&device, false, admitted, StatusCode::NOT_FOUND, now); + } + assert_eq!(ip_count(&limiter), 0); + } + + #[test] + fn a_401_for_an_accepted_token_counts_and_withdraws_the_acceptance() { + let now = parse_ts("2026-10-17T12:00:00.000Z").unwrap(); + let limiter = Limiter::new(Arc::new(FixedClock::new(now)), String::new()); + let device = token_key(1); + let admitted = limiter.admit(&device, false, now).unwrap(); + limiter.record(&device, false, admitted, StatusCode::OK, now); + // Revoked meanwhile: the 401 counts once, and from now on every request counts before it runs. + let admitted = limiter.admit(&device, false, now).unwrap(); + limiter.record(&device, false, admitted, StatusCode::UNAUTHORIZED, now); + assert_eq!(ip_count(&limiter), 1); + let admitted = limiter.admit(&device, false, now).unwrap(); + assert!(admitted.ip_window.is_some()); + limiter.record(&device, false, admitted, StatusCode::UNAUTHORIZED, now); + assert_eq!(ip_count(&limiter), 2); + } + + #[test] + fn accepted_tokens_are_pruned_with_their_window() { + let now = parse_ts("2026-10-17T12:00:00.000Z").unwrap(); + let limiter = Limiter::new(Arc::new(FixedClock::new(now)), String::new()); + for i in 0..10 { + let key = token_key(i); + let admitted = limiter.admit(&key, false, now).unwrap(); + limiter.record(&key, false, admitted, StatusCode::OK, now); + } + assert_eq!(lock(&limiter.buckets).tokens.len(), 10); + let later = now + Duration::minutes(11); + let admitted = limiter.admit(&token_key(1), false, later).unwrap(); + assert_eq!(lock(&limiter.buckets).tokens.len(), 1, "only the new entry"); + // Re-accepting costs one request against the IP, given back by the 2xx. + assert!(admitted.ip_window.is_some()); + limiter.record(&token_key(1), false, admitted, StatusCode::OK, later); + assert_eq!(ip_count(&limiter), 0); + } + #[test] fn pruning_drops_expired_windows() { let now = parse_ts("2026-10-17T12:00:00.000Z").unwrap(); diff --git a/server/crates/kvitt-server/tests/core_api/rate_limit.rs b/server/crates/kvitt-server/tests/core_api/rate_limit.rs index 3593fed..591ec8d 100644 --- a/server/crates/kvitt-server/tests/core_api/rate_limit.rs +++ b/server/crates/kvitt-server/tests/core_api/rate_limit.rs @@ -96,6 +96,34 @@ async fn unknown_tokens_count_against_the_ip_but_known_devices_keep_working() { assert_eq!(reply.status, 200); } +#[tokio::test] +async fn random_tokens_without_the_protocol_header_are_limited_by_the_ip() { + // Answered 426 before any token check: such tokens are never accepted, so each request counts against the IP. + let app = TestApp::household().await; + assert_eq!( + app.send(from("192.0.2.1:1", None, "/me", Some(&admin_token()))) + .await + .status, + 200 + ); + let bare = |tok: &str| { + let mut request = from("192.0.2.1:1", None, "/me", Some(tok)); + request.headers_mut().remove("x-kvitt-protocol"); + request + }; + for i in 0..60 { + let reply = app.send(bare(&format!("kvt_{i:0>43}"))).await; + assert_error(&reply, 426, "protocol_unsupported"); + } + let reply = app.send(bare(&token('Z'))).await; + assert_error(&reply, 429, "rate_limited"); + // The paired device keeps working from the same address. + let reply = app + .send(from("192.0.2.1:1", None, "/me", Some(&admin_token()))) + .await; + assert_eq!(reply.status, 200); +} + #[tokio::test] async fn x_forwarded_for_is_trusted_only_from_the_proxy() { let app = TestApp::empty().await;