M3: server core API (E-3 to E-30, E-33), landing page, pictures, rate limiting, recurring generation, bootstrap and pair-device
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
@@ -1,8 +1,444 @@
|
||||
//! Rate limiting (SPEC §8.6). M3 owns this file; in M2a the layer lets every request through.
|
||||
//! 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 tower::layer::util::Identity;
|
||||
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;
|
||||
|
||||
/// The rate-limit layer `router.rs` attaches to every route.
|
||||
pub fn layer() -> Identity {
|
||||
Identity::new()
|
||||
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);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user