445 lines
16 KiB
Rust
445 lines
16 KiB
Rust
//! 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-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.
|
|
//! * 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<SocketAddr>`), or the last `X-Forwarded-For` entry when the peer
|
|
//! is an address of `KVITT_TRUSTED_PROXY` (Caddy, 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). 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, MutexGuard};
|
|
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::AppState;
|
|
use crate::audit::BoxFuture;
|
|
use crate::auth::{is_token_shaped, sha256_hex};
|
|
use crate::error::AppError;
|
|
use crate::time::Clock;
|
|
|
|
/// Unauthenticated requests per window per client IP (SPEC §8.6).
|
|
pub const UNAUTHENTICATED_LIMIT: u32 = 60;
|
|
/// 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";
|
|
|
|
/// 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<Utc>,
|
|
count: u32,
|
|
}
|
|
|
|
impl Window {
|
|
fn expired(&self, now: DateTime<Utc>, len: Duration) -> bool {
|
|
now >= self.start + len
|
|
}
|
|
|
|
fn retry_after(&self, now: DateTime<Utc>, 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<Window>,
|
|
now: DateTime<Utc>,
|
|
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<Utc>, len: Duration, limit: u32) -> Option<u64> {
|
|
slot.filter(|w| !w.expired(now, len) && w.count >= limit)
|
|
.map(|w| w.retry_after(now, len))
|
|
}
|
|
|
|
#[derive(Debug, Default)]
|
|
struct TokenState {
|
|
window: Option<Window>,
|
|
/// The token has had a response other than 401: it belongs to a device.
|
|
accepted: bool,
|
|
}
|
|
|
|
#[derive(Debug, Default)]
|
|
struct Buckets {
|
|
ip: HashMap<IpAddr, Window>,
|
|
tokens: HashMap<String, TokenState>,
|
|
accept_failures: HashMap<IpAddr, Window>,
|
|
last_prune: Option<DateTime<Utc>>,
|
|
}
|
|
|
|
impl Buckets {
|
|
fn prune(&mut self, now: DateTime<Utc>) {
|
|
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.tokens
|
|
.retain(|_, t| t.accepted || t.window.is_some_and(|w| !w.expired(now, REQUEST_WINDOW)));
|
|
self.accept_failures
|
|
.retain(|_, w| !w.expired(now, ACCEPT_FAILURE_WINDOW));
|
|
}
|
|
}
|
|
|
|
/// How a request is counted.
|
|
enum Key {
|
|
Ip(IpAddr),
|
|
Token { ip: IpAddr, hash: String },
|
|
}
|
|
|
|
/// The shared limiter state.
|
|
pub struct Limiter {
|
|
clock: Arc<dyn Clock>,
|
|
trusted_proxy: String,
|
|
proxy: Mutex<Option<(Instant, Vec<IpAddr>)>>,
|
|
buckets: Mutex<Buckets>,
|
|
}
|
|
|
|
fn lock<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
|
|
match mutex.lock() {
|
|
Ok(guard) => guard,
|
|
Err(poisoned) => poisoned.into_inner(),
|
|
}
|
|
}
|
|
|
|
impl Limiter {
|
|
pub fn new(clock: Arc<dyn Clock>, 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<IpAddr> {
|
|
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<IpAddr> = if name.is_empty() {
|
|
Vec::new()
|
|
} else if let Ok(ip) = name.parse::<IpAddr>() {
|
|
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<IpAddr>, Option<IpAddr>) {
|
|
let peer = request
|
|
.extensions()
|
|
.get::<ConnectInfo<SocketAddr>>()
|
|
.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::<IpAddr>().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<IpAddr>, forwarded: Option<IpAddr>) -> 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 PUBLIC_ROUTES.contains(&route) || route.is_empty() {
|
|
return Key::Ip(ip);
|
|
}
|
|
let token = request
|
|
.headers()
|
|
.get(axum::http::header::AUTHORIZATION)
|
|
.and_then(|v| v.to_str().ok())
|
|
.and_then(|v| v.split_once(' '))
|
|
.filter(|(scheme, _)| scheme.eq_ignore_ascii_case("bearer"))
|
|
.map(|(_, token)| token.trim())
|
|
.filter(|token| is_token_shaped(token));
|
|
match 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<Utc>) -> Result<(), u64> {
|
|
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);
|
|
}
|
|
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
|
|
}
|
|
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 state = buckets.tokens.entry(hash.clone()).or_default();
|
|
hit(&mut state.window, now, REQUEST_WINDOW, AUTHENTICATED_LIMIT)
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Records what the response says about the request (401s, accept failures, accepted tokens).
|
|
fn record(&self, key: &Key, accept: bool, status: StatusCode, now: DateTime<Utc>) {
|
|
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::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 {
|
|
buckets.tokens.entry(hash.clone()).or_default().accepted = true;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
async fn handle<S>(self: Arc<Self>, request: Request, mut inner: S) -> Response
|
|
where
|
|
S: Service<Request, Response = Response, Error = Infallible>,
|
|
{
|
|
let route = request
|
|
.extensions()
|
|
.get::<MatchedPath>()
|
|
.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);
|
|
if let Err(retry_after) = self.admit(&key, accept, self.clock.now()) {
|
|
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());
|
|
response
|
|
}
|
|
}
|
|
|
|
/// The tower layer of [`layer`].
|
|
#[derive(Clone)]
|
|
pub struct RateLimitLayer {
|
|
limiter: Arc<Limiter>,
|
|
}
|
|
|
|
impl<S> Layer<S> for RateLimitLayer {
|
|
type Service = RateLimit<S>;
|
|
|
|
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<S> {
|
|
inner: S,
|
|
limiter: Arc<Limiter>,
|
|
}
|
|
|
|
impl<S> Service<Request> for RateLimit<S>
|
|
where
|
|
S: Service<Request, Response = Response, Error = Infallible> + Clone + Send + 'static,
|
|
S::Future: Send + 'static,
|
|
{
|
|
type Response = Response;
|
|
type Error = Infallible;
|
|
type Future = BoxFuture<'static, Result<Response, Infallible>>;
|
|
|
|
fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
|
|
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
|
|
);
|
|
}
|
|
|
|
#[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);
|
|
}
|
|
}
|