//! Rate limiting (SPEC §8.6, AC-2.12, §8.5), as a tower layer that `router.rs` attaches to every route. //! //! * Unauthenticated endpoints (E-1 to E-5 and unknown paths): at most 60 requests per minute per client IP. //! * APK downloads (E-6, `/apk/files/{file_name}`, D-119): at most 10 requests per minute per client IP, in a bucket //! of their own that does not count against the 60 (security review L-3: each answer can be a 10 MB file). //! * 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. //! * 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. //! //! Refusals are 429 `rate_limited` with `Retry-After` in whole seconds. Windows are fixed: they start at the first //! request (or failure) counted and last one minute (one hour for accept failures). Everything is in memory and resets //! on restart; there is no table. //! //! The client IP is the TCP peer (axum `ConnectInfo`), or the last `X-Forwarded-For` entry when the peer //! is an address of `KVITT_TRUSTED_PROXY` (Nginx Proxy Manager, SPEC §8.5). The host name is resolved lazily and //! re-resolved at most once per [`PROXY_RESOLVE_INTERVAL`]; a name that does not resolve means no peer is trusted //! (D-022), and so does the default `none` ([`NO_TRUSTED_PROXY`]), which is never looked up. Without //! `ConnectInfo` (a server not started with `into_make_service_with_connect_info`) all clients share one IP bucket. use std::collections::HashMap; use std::convert::Infallible; use std::net::{IpAddr, Ipv4Addr, SocketAddr}; use std::sync::{Arc, Mutex}; use std::task::{Context, Poll}; use std::time::Instant; use axum::extract::{ConnectInfo, MatchedPath, Request}; use axum::http::StatusCode; use axum::response::{IntoResponse, Response}; use chrono::{DateTime, Duration, Utc}; use tower::{Layer, Service}; use crate::audit::BoxFuture; use crate::auth::{bearer_token, is_token_shaped, sha256_hex}; use crate::config::NO_TRUSTED_PROXY; use crate::error::AppError; use crate::time::Clock; use crate::{AppState, lock}; /// Unauthenticated requests per window per client IP (SPEC §8.6). pub const UNAUTHENTICATED_LIMIT: u32 = 60; /// APK downloads (E-6) per window per client IP (SPEC §8.6, D-119). pub const APK_DOWNLOAD_LIMIT: u32 = 10; /// Authenticated requests per window per device (SPEC §8.6). pub const AUTHENTICATED_LIMIT: u32 = 600; /// The request window (SPEC §8.6: per minute). pub const REQUEST_WINDOW: Duration = Duration::minutes(1); /// Failed invite accepts tolerated per IP within [`ACCEPT_FAILURE_WINDOW`] (AC-2.12). pub const ACCEPT_FAILURE_LIMIT: u32 = 10; /// The invite-accept failure window (AC-2.12). pub const ACCEPT_FAILURE_WINDOW: Duration = Duration::minutes(60); /// How long a resolution of `KVITT_TRUSTED_PROXY` is reused. pub const PROXY_RESOLVE_INTERVAL: std::time::Duration = std::time::Duration::from_secs(60); /// How often expired windows are pruned from memory. const PRUNE_INTERVAL: Duration = Duration::minutes(10); /// Route templates of the unauthenticated endpoints (E-1 to E-6). const PUBLIC_ROUTES: [&str; 6] = [ "/health", "/setup/status", "/invites/accept", "/i/{code}", "/apk/latest", "/apk/files/{file_name}", ]; const ACCEPT_ROUTE: &str = "/invites/accept"; /// Route template of E-6, limited by [`APK_DOWNLOAD_LIMIT`] instead of [`UNAUTHENTICATED_LIMIT`]. const APK_ROUTE: &str = "/apk/files/{file_name}"; /// The layer for `router.rs`. One limiter per router: build the router once and reuse it (as `serve` does). pub fn layer(state: &AppState) -> RateLimitLayer { RateLimitLayer { limiter: Arc::new(Limiter::new( state.clock.clone(), state.config.trusted_proxy.clone(), )), } } /// A fixed counting window. #[derive(Debug, Clone, Copy)] struct Window { start: DateTime, count: u32, } impl Window { fn expired(&self, now: DateTime, len: Duration) -> bool { now >= self.start + len } fn retry_after(&self, now: DateTime, len: Duration) -> u64 { let remaining = (self.start + len - now).num_milliseconds().max(0); u64::try_from((remaining + 999) / 1000).unwrap_or(1).max(1) } } /// Counts one request in `slot` unless the window is full; `Err(retry_after_secs)` when it is. fn hit( slot: &mut Option, now: DateTime, len: Duration, limit: u32, ) -> Result<(), u64> { let window = match slot { Some(w) if !w.expired(now, len) => w, _ => slot.insert(Window { start: now, count: 0, }), }; if window.count >= limit { return Err(window.retry_after(now, len)); } window.count += 1; Ok(()) } /// Whether `slot` is full (without counting). fn full(slot: Option<&Window>, now: DateTime, len: Duration, limit: u32) -> Option { slot.filter(|w| !w.expired(now, len) && w.count >= limit) .map(|w| w.retry_after(now, len)) } #[derive(Debug, Default)] struct TokenState { window: Option, /// The token's last authenticated response was 2xx (and none was 401 since): it belongs to a device. accepted: bool, } #[derive(Debug, Default)] struct Buckets { ip: HashMap, apk: HashMap, tokens: HashMap, accept_failures: HashMap, last_prune: Option>, } impl Buckets { fn prune(&mut self, now: DateTime) { if self.last_prune.is_some_and(|t| now - t < PRUNE_INTERVAL) { return; } self.last_prune = Some(now); self.ip.retain(|_, w| !w.expired(now, REQUEST_WINDOW)); self.apk.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.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. enum Key { Ip(IpAddr), /// E-6 from this client IP. Apk(IpAddr), Token { ip: IpAddr, hash: String, }, } /// The shared limiter state. pub struct Limiter { clock: Arc, trusted_proxy: String, proxy: Mutex)>>, buckets: Mutex, } impl Limiter { pub fn new(clock: Arc, trusted_proxy: String) -> Self { Self { clock, trusted_proxy, proxy: Mutex::new(None), buckets: Mutex::new(Buckets::default()), } } /// The addresses of `KVITT_TRUSTED_PROXY`, resolved at most once per [`PROXY_RESOLVE_INTERVAL`]. async fn proxy_addrs(&self) -> Vec { if let Some((at, addrs)) = lock(&self.proxy).as_ref() && at.elapsed() < PROXY_RESOLVE_INTERVAL { return addrs.clone(); } let name = self.trusted_proxy.trim(); let addrs: Vec = if name.is_empty() || name.eq_ignore_ascii_case(NO_TRUSTED_PROXY) { Vec::new() } else if let Ok(ip) = name.parse::() { vec![ip] } else { match tokio::net::lookup_host((name, 0)).await { Ok(found) => found.map(|a| a.ip()).collect(), Err(_) => Vec::new(), } }; *lock(&self.proxy) = Some((Instant::now(), addrs.clone())); addrs } /// The TCP peer and the last `X-Forwarded-For` entry of `request`. fn peer_and_forwarded(request: &Request) -> (Option, Option) { let peer = request .extensions() .get::>() .map(|c| c.0.ip()); let forwarded = request .headers() .get("x-forwarded-for") .and_then(|v| v.to_str().ok()) .and_then(|v| v.rsplit(',').next()) .and_then(|v| v.trim().parse::().ok()); (peer, forwarded) } /// The client IP (SPEC §8.5): the forwarded address when the peer is the trusted proxy, else the peer. async fn client_ip(&self, peer: Option, forwarded: Option) -> IpAddr { let Some(peer) = peer else { return IpAddr::V4(Ipv4Addr::UNSPECIFIED); }; match forwarded { Some(client) if self.proxy_addrs().await.contains(&peer) => client, _ => peer, } } fn key(request: &Request, route: &str, ip: IpAddr) -> Key { if route == APK_ROUTE { return Key::Apk(ip); } if PUBLIC_ROUTES.contains(&route) || route.is_empty() { return Key::Ip(ip); } match bearer_token(request.headers()).filter(|token| is_token_shaped(token)) { Some(token) => Key::Token { ip, hash: sha256_hex(token), }, None => Key::Ip(ip), } } /// Counts the request before it runs; `Err(retry_after_secs)` refuses it. fn admit(&self, key: &Key, accept: bool, now: DateTime) -> Result { let mut buckets = lock(&self.buckets); buckets.prune(now); match key { Key::Ip(ip) => { if accept && let Some(retry) = full( buckets.accept_failures.get(ip), now, ACCEPT_FAILURE_WINDOW, ACCEPT_FAILURE_LIMIT, ) { return Err(retry); } buckets.hit_ip(*ip, now)?; Ok(Admitted::default()) } Key::Apk(ip) => { let mut slot = buckets.apk.get(ip).copied(); let result = hit(&mut slot, now, REQUEST_WINDOW, APK_DOWNLOAD_LIMIT); if let Some(w) = slot { buckets.apk.insert(*ip, w); } result.map(|()| Admitted::default()) } Key::Token { ip, hash } => { let accepted = buckets.tokens.get(hash).is_some_and(|t| t.accepted); 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)?; Ok(Admitted { ip_window }) } } } /// Records what the response says about the request (401s, accept failures, accepted tokens). fn record( &self, key: &Key, accept: bool, admitted: Admitted, status: StatusCode, now: DateTime, ) { let mut buckets = lock(&self.buckets); match key { Key::Ip(ip) => { if accept && status.is_client_error() && status != StatusCode::TOO_MANY_REQUESTS { let mut slot = buckets.accept_failures.get(ip).copied(); // The window is refused once full (admit), so this count never exceeds the limit. let _ = hit(&mut slot, now, ACCEPT_FAILURE_WINDOW, ACCEPT_FAILURE_LIMIT); if let Some(w) = slot { buckets.accept_failures.insert(*ip, w); } } } Key::Apk(_) => {} Key::Token { ip, hash } => { 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; } } } } } async fn handle(self: Arc, request: Request, mut inner: S) -> Response where S: Service, { let route = request .extensions() .get::() .map(|p| p.as_str().to_owned()) .unwrap_or_default(); let accept = route == ACCEPT_ROUTE; let (peer, forwarded) = Self::peer_and_forwarded(&request); let ip = self.client_ip(peer, forwarded).await; let key = Self::key(&request, &route, ip); 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, admitted, response.status(), self.clock.now()); response } } /// The tower layer of [`layer`]. #[derive(Clone)] pub struct RateLimitLayer { limiter: Arc, } impl Layer for RateLimitLayer { type Service = RateLimit; fn layer(&self, inner: S) -> Self::Service { RateLimit { inner, limiter: self.limiter.clone(), } } } /// The rate-limiting service wrapping one route. #[derive(Clone)] pub struct RateLimit { inner: S, limiter: Arc, } impl Service for RateLimit where S: Service + Clone + Send + 'static, S::Future: Send + 'static, { type Response = Response; type Error = Infallible; type Future = BoxFuture<'static, Result>; fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll> { self.inner.poll_ready(cx) } fn call(&mut self, request: Request) -> Self::Future { // The ready clone goes into the future; a fresh clone stays for the next call (tower's usual pattern). let clone = self.inner.clone(); let inner = std::mem::replace(&mut self.inner, clone); let limiter = self.limiter.clone(); Box::pin(async move { Ok(limiter.handle(request, inner).await) }) } } #[cfg(test)] mod tests { use super::*; use crate::time::{FixedClock, parse_ts}; #[test] fn fixed_window_counts_and_reports_retry_after() { let now = parse_ts("2026-10-17T12:00:00.000Z").unwrap(); let mut slot = None; for _ in 0..3 { assert!(hit(&mut slot, now, REQUEST_WINDOW, 3).is_ok()); } assert_eq!( hit(&mut slot, now + Duration::seconds(20), REQUEST_WINDOW, 3), Err(40) ); assert!(hit(&mut slot, now + Duration::seconds(60), REQUEST_WINDOW, 3).is_ok()); } #[test] fn retry_after_rounds_up_and_is_at_least_one() { let now = parse_ts("2026-10-17T12:00:00.000Z").unwrap(); let w = Window { start: now, count: 1, }; assert_eq!( w.retry_after(now + Duration::milliseconds(59_500), REQUEST_WINDOW), 1 ); assert_eq!( w.retry_after(now + Duration::milliseconds(100), REQUEST_WINDOW), 60 ); } fn limiter_trusting(trusted_proxy: &str) -> Limiter { let now = parse_ts("2026-10-17T12:00:00.000Z").unwrap(); Limiter::new(Arc::new(FixedClock::new(now)), trusted_proxy.to_owned()) } #[tokio::test] async fn the_default_none_trusts_no_peer() { let peer = IpAddr::V4(Ipv4Addr::new(192, 168, 1, 20)); let forwarded = IpAddr::V4(Ipv4Addr::new(203, 0, 113, 5)); for value in [crate::config::DEFAULT_TRUSTED_PROXY, " NONE ", ""] { let limiter = limiter_trusting(value); assert_eq!( limiter.proxy_addrs().await, Vec::::new(), "{value:?}" ); assert_eq!(limiter.client_ip(Some(peer), Some(forwarded)).await, peer); } } #[tokio::test] async fn an_unresolvable_name_trusts_no_peer() { // `.invalid` never resolves (RFC 6761). let limiter = limiter_trusting("proxy.kvitt.invalid"); let peer = IpAddr::V4(Ipv4Addr::new(192, 168, 1, 20)); let forwarded = IpAddr::V4(Ipv4Addr::new(203, 0, 113, 5)); assert_eq!(limiter.client_ip(Some(peer), Some(forwarded)).await, peer); } #[tokio::test] async fn an_ip_literal_is_trusted() { let limiter = limiter_trusting("192.168.1.20"); let peer = IpAddr::V4(Ipv4Addr::new(192, 168, 1, 20)); let forwarded = IpAddr::V4(Ipv4Addr::new(203, 0, 113, 5)); assert_eq!( limiter.client_ip(Some(peer), Some(forwarded)).await, forwarded ); } 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 apk_downloads_have_their_own_bucket_and_are_pruned() { let now = parse_ts("2026-10-17T12:00:00.000Z").unwrap(); let limiter = Limiter::new(Arc::new(FixedClock::new(now)), String::new()); let ip = IpAddr::V4(Ipv4Addr::LOCALHOST); for _ in 0..APK_DOWNLOAD_LIMIT { assert!(limiter.admit(&Key::Apk(ip), false, now).is_ok()); } assert_eq!(limiter.admit(&Key::Apk(ip), false, now), Err(60)); assert_eq!( ip_count(&limiter), 0, "downloads do not count against the 60" ); assert!(limiter.admit(&Key::Ip(ip), false, now).is_ok()); let later = now + Duration::minutes(11); assert!( limiter .admit(&Key::Apk(IpAddr::V4(Ipv4Addr::BROADCAST)), false, later) .is_ok() ); assert_eq!(lock(&limiter.buckets).apk.len(), 1, "only the new entry"); } #[test] fn pruning_drops_expired_windows() { let now = parse_ts("2026-10-17T12:00:00.000Z").unwrap(); let limiter = Limiter::new(Arc::new(FixedClock::new(now)), String::new()); let ip = IpAddr::V4(Ipv4Addr::LOCALHOST); assert!(limiter.admit(&Key::Ip(ip), false, now).is_ok()); assert_eq!(lock(&limiter.buckets).ip.len(), 1); assert!( limiter .admit( &Key::Ip(IpAddr::V4(Ipv4Addr::BROADCAST)), false, now + Duration::minutes(11) ) .is_ok() ); assert_eq!(lock(&limiter.buckets).ip.len(), 1); } }