dogfood
Jump to…
snowcorepublic2 bookmarks
snowfix: 500 error because of database is patchedyuzpxzopsouq1mo
mainroot / crates / df-web / src / ratelimit.rsHide BlameRaw
Matt W1//! Rate limiting and per-user concurrency caps (spec §9, M6).
Matt W2//!
Matt W3//! > Cap: pack size on receive […] and **request concurrency per user**.
Matt W4//!
Matt W5//! Two different protections, because they stop two different things:
Matt W6//!
Matt W7//! * A **token bucket** bounds the *rate* of requests. It is what stops a script
Matt W8//! walking every change in a repository, or grinding at the login endpoint.
Matt W9//! * A **concurrency cap** bounds how many requests one identity may have in
Matt W10//! flight. A rate limit alone does not stop ten simultaneous diffs of the
Matt W11//! Linux kernel, and those are the requests that actually consume the box.
Matt W12//!
Matt W13//! Both are keyed by user id when there is one and by IP otherwise. Keying
Matt W14//! authenticated traffic by user rather than IP matters in both directions: a
Matt W15//! team behind one NAT is not one attacker, and one attacker on a hundred
Matt W16//! addresses is still one account.
Matt W17//!
Matt W18//! "By IP" means [`client_ip`], not the TCP peer. Behind a reverse proxy the
Matt W19//! peer is the proxy for every request in the world, so keying on it does not
Matt W20//! produce a strict limit — it produces a *global* one, and a single client can
Matt W21//! spend the whole budget and lock every signed-out visitor out of the site.
Matt W22//! `TRUSTED_PROXIES` is what closes that, and an empty setting is loud about it.
Matt W23//!
Matt W24//! There are two token buckets in the request path, at different depths.
Matt W25//! Resolving a session is a database round trip and it happens *before* the
Matt W26//! limiter that knows who you are, so a coarse per-address bucket
Matt W27//! ([`edge_layer`]) runs ahead of it, and the real one ([`layer`]) runs after.
Matt W28//!
Matt W29//! In-memory, deliberately. A shared limiter would mean Redis, and this is a
Matt W30//! single-instance product (spec §1). The state is small, bounded, and reset by
Matt W31//! a restart — which is the correct behaviour for a limiter whose only job is to
Matt W32//! keep one process healthy.
Matt W33
Matt W34use std::collections::HashMap;
Matt W35use std::net::IpAddr;
Matt W36use std::sync::{Arc, Mutex};
Matt W37use std::time::{Duration, Instant};
Matt W38
Matt W39use axum::extract::{ConnectInfo, Request, State};
Matt W40use axum::http::{HeaderMap, StatusCode};
Matt W41use axum::middleware::Next;
Matt W42use axum::response::{IntoResponse, Response};
Matt W43use sqlx::types::ipnetwork::IpNetwork;
Matt W44
Matt W45use crate::state::{AppState, CurrentUser};
Matt W46
Matt W47/// Sustained requests per second, per identity.
Matt W48const REFILL_PER_SEC: f64 = 8.0;
Matt W49
Matt W50/// Burst above the sustained rate. A page load is one document plus its assets,
Matt W51/// and a reviewer clicking through a stack fires several in a second.
Matt W52const BURST: f64 = 40.0;
Matt W53
Matt W54/// The much tighter bucket for endpoints that are worth grinding at: the login
Matt W55/// redirect, the OIDC callback, and the setup-token claim.
Matt W56const AUTH_REFILL_PER_SEC: f64 = 0.5;
Matt W57const AUTH_BURST: f64 = 10.0;
Matt W58
Matt W59/// Simultaneous in-flight requests per identity.
Matt W60///
Matt W61/// A diff or a highlight can occupy a thread for a while; this is what stops one
Matt W62/// identity holding all of them.
Matt W63const MAX_CONCURRENT: u32 = 12;
Matt W64
Matt W65/// The coarse bucket applied before the session is resolved.
Matt W66///
Matt W67/// Deliberately far above [`BURST`]: it is not a second rate limit, it is a
Matt W68/// bound on how much work an unauthenticated flood can force *ahead of* the
Matt W69/// real limiter. Session resolution is a database round trip and it runs first,
Matt W70/// so without this a client can spend a query per request no matter what the
Matt W71/// bucket below decides.
Matt W72const EDGE_REFILL_PER_SEC: f64 = 50.0;
Matt W73const EDGE_BURST: f64 = 200.0;
Matt W74
Matt W75/// Entries idle longer than this are dropped, so the map does not grow with
Matt W76/// every address that has ever connected.
Matt W77const IDLE_EVICT: Duration = Duration::from_secs(600);
Matt W78
Matt W79/// The bucket every request without a resolvable peer address shares.
Matt W80///
Matt W81/// `UNSPECIFIED` is not a routable address, so it cannot collide with a real
Matt W82/// client, and sharing one bucket is the conservative choice: unattributable
Matt W83/// traffic is limited together rather than not at all.
Matt W84const UNKNOWN_PEER: IpAddr = IpAddr::V4(std::net::Ipv4Addr::UNSPECIFIED);
Matt W85
Matt W86/// How many identities to track before evicting aggressively. Reached only
Matt W87/// under a distributed flood, where the correct behaviour is to keep working
Matt W88/// rather than to allocate.
Matt W89const MAX_TRACKED: usize = 50_000;
Matt W90
Matt W91#[derive(Clone)]
Matt W92pub struct Limiter(Arc<Mutex<Inner>>);
Matt W93
Matt W94struct Inner {
Matt W95 buckets: HashMap<Key, Bucket>,
Matt W96 last_sweep: Instant,
Matt W97}
Matt W98
Matt W99#[derive(Clone, PartialEq, Eq, Hash, Debug)]
Matt W100enum Key {
Matt W101 User(uuid::Uuid),
Matt W102 Addr(IpAddr),
Matt W103}
Matt W104
Matt W105struct Bucket {
Matt W106 tokens: f64,
Matt W107 auth_tokens: f64,
Matt W108 edge_tokens: f64,
Matt W109 in_flight: u32,
Matt W110 last: Instant,
Matt W111}
Matt W112
Matt W113impl Bucket {
Matt W114 fn full(now: Instant) -> Bucket {
Matt W115 Bucket {
Matt W116 tokens: BURST,
Matt W117 auth_tokens: AUTH_BURST,
Matt W118 edge_tokens: EDGE_BURST,
Matt W119 in_flight: 0,
Matt W120 last: now,
Matt W121 }
Matt W122 }
Matt W123
Matt W124 /// Refill for the time that passed.
Matt W125 ///
Matt W126 /// `saturating_duration_since` because a clock that went backwards must not
Matt W127 /// mint tokens.
Matt W128 fn refill(&mut self, now: Instant) {
Matt W129 let elapsed = now.saturating_duration_since(self.last).as_secs_f64();
Matt W130 self.tokens = (self.tokens + elapsed * REFILL_PER_SEC).min(BURST);
Matt W131 self.auth_tokens = (self.auth_tokens + elapsed * AUTH_REFILL_PER_SEC).min(AUTH_BURST);
Matt W132 self.edge_tokens = (self.edge_tokens + elapsed * EDGE_REFILL_PER_SEC).min(EDGE_BURST);
Matt W133 self.last = now;
Matt W134 }
Matt W135}
Matt W136
Matt W137impl Default for Limiter {
Matt W138 fn default() -> Self {
Matt W139 Self::new()
Matt W140 }
Matt W141}
Matt W142
Matt W143impl Limiter {
Matt W144 pub fn new() -> Self {
Matt W145 Limiter(Arc::new(Mutex::new(Inner {
Matt W146 buckets: HashMap::new(),
Matt W147 last_sweep: Instant::now(),
Matt W148 })))
Matt W149 }
Matt W150
Matt W151 /// Take one token, and a concurrency slot.
Matt W152 ///
Matt W153 /// Returns `None` when the identity is over a limit. The returned guard
Matt W154 /// releases the concurrency slot when dropped — including when the handler
Matt W155 /// panics, which is why it is a guard and not a pair of calls.
Matt W156 fn acquire(&self, key: Key, auth: bool, now: Instant) -> Option<Guard> {
Matt W157 let mut inner = self.0.lock().expect("rate limiter poisoned");
Matt W158 inner.sweep(now);
Matt W159
Matt W160 // A flood of distinct keys must not be able to grow the map without
Matt W161 // bound. Past the cap, unknown keys are refused rather than admitted —
Matt W162 // the alternative is admitting everything precisely when under attack.
Matt W163 if inner.buckets.len() >= MAX_TRACKED && !inner.buckets.contains_key(&key) {
Matt W164 return None;
Matt W165 }
Matt W166
Matt W167 let bucket = inner.buckets.entry(key.clone()).or_insert(Bucket::full(now));
Matt W168 bucket.refill(now);
Matt W169
Matt W170 if bucket.in_flight >= MAX_CONCURRENT {
Matt W171 return None;
Matt W172 }
Matt W173 if bucket.tokens < 1.0 {
Matt W174 return None;
Matt W175 }
Matt W176 if auth && bucket.auth_tokens < 1.0 {
Matt W177 return None;
Matt W178 }
Matt W179
Matt W180 bucket.tokens -= 1.0;
Matt W181 if auth {
Matt W182 bucket.auth_tokens -= 1.0;
Matt W183 }
Matt W184 bucket.in_flight += 1;
Matt W185
Matt W186 Some(Guard { limiter: self.clone(), key })
Matt W187 }
Matt W188
Matt W189 /// Take one token from the coarse pre-session bucket.
Matt W190 ///
Matt W191 /// No concurrency slot: this runs before the handler is chosen and releases
Matt W192 /// nothing, so it bounds arrival rate only.
Matt W193 fn acquire_edge(&self, ip: IpAddr, now: Instant) -> bool {
Matt W194 let key = Key::Addr(ip);
Matt W195 let mut inner = self.0.lock().expect("rate limiter poisoned");
Matt W196 inner.sweep(now);
Matt W197
Matt W198 if inner.buckets.len() >= MAX_TRACKED && !inner.buckets.contains_key(&key) {
Matt W199 return false;
Matt W200 }
Matt W201
Matt W202 let bucket = inner.buckets.entry(key).or_insert(Bucket::full(now));
Matt W203 bucket.refill(now);
Matt W204
Matt W205 if bucket.edge_tokens < 1.0 {
Matt W206 return false;
Matt W207 }
Matt W208 bucket.edge_tokens -= 1.0;
Matt W209 true
Matt W210 }
Matt W211
Matt W212 fn release(&self, key: &Key) {
Matt W213 if let Ok(mut inner) = self.0.lock() {
Matt W214 if let Some(b) = inner.buckets.get_mut(key) {
Matt W215 b.in_flight = b.in_flight.saturating_sub(1);
Matt W216 }
Matt W217 }
Matt W218 }
Matt W219
Matt W220 #[cfg(test)]
Matt W221 fn tracked(&self) -> usize {
Matt W222 self.0.lock().unwrap().buckets.len()
Matt W223 }
Matt W224}
Matt W225
Matt W226impl Inner {
Matt W227 fn sweep(&mut self, now: Instant) {
Matt W228 if now.saturating_duration_since(self.last_sweep) < Duration::from_secs(60) {
Matt W229 return;
Matt W230 }
Matt W231 self.last_sweep = now;
Matt W232 // An entry with a request in flight is never evicted, however idle its
Matt W233 // bucket looks — dropping it would lose the concurrency count and let
Matt W234 // the cap be bypassed by a slow request.
Matt W235 self.buckets.retain(|_, b| {
Matt W236 b.in_flight > 0 || now.saturating_duration_since(b.last) < IDLE_EVICT
Matt W237 });
Matt W238 }
Matt W239}
Matt W240
Matt W241/// Holds a concurrency slot for the life of a request.
Matt W242pub struct Guard {
Matt W243 limiter: Limiter,
Matt W244 key: Key,
Matt W245}
Matt W246
Matt W247impl Drop for Guard {
Matt W248 fn drop(&mut self) {
Matt W249 self.limiter.release(&self.key);
Matt W250 }
Matt W251}
Matt W252
Matt W253/// Whether a path gets the strict auth bucket.
Matt W254fn is_auth_path(path: &str) -> bool {
Matt W255 matches!(path, "/login" | "/auth/callback" | "/auth/handle" | "/setup")
Matt W256}
Matt W257
Matt W258/// Whether a path is exempt.
Matt W259///
Matt W260/// Health checks come from the orchestrator on a fixed interval and must never
Matt W261/// be throttled — a rate-limited `/healthz` restarts the container. Static
Matt W262/// assets are served from memory and are not worth a bucket.
Matt W263fn is_exempt(path: &str) -> bool {
Matt W264 matches!(path, "/healthz" | "/readyz") || path.starts_with("/assets/")
Matt W265}
Matt W266
Matt W267/// Not in `http::header`, which only defines registered headers.
Matt W268const X_FORWARDED_FOR: &str = "x-forwarded-for";
Matt W269
Matt W270/// One `X-Forwarded-For` entry as an address.
Matt W271///
Matt W272/// Proxies vary: bare addresses, `addr:port`, and bracketed IPv6 all appear.
Matt W273/// Anything that does not parse is discarded rather than guessed at.
Matt W274fn parse_forwarded(entry: &str) -> Option<IpAddr> {
Matt W275 let s = entry.trim();
Matt W276 if s.is_empty() {
Matt W277 return None;
Matt W278 }
Matt W279 // `[::1]` or `[::1]:8080`
Matt W280 if let Some(rest) = s.strip_prefix('[') {
Matt W281 let (inner, _) = rest.split_once(']')?;
Matt W282 return inner.parse().ok();
Matt W283 }
Matt W284 if let Ok(ip) = s.parse::<IpAddr>() {
Matt W285 return Some(ip);
Matt W286 }
Matt W287 // `1.2.3.4:5678`. Only IPv4 — a bare IPv6 has colons of its own and was
Matt W288 // handled by the parse above.
Matt W289 s.rsplit_once(':').and_then(|(host, _)| host.parse().ok())
Matt W290}
Matt W291
Matt W292/// The address to attribute a request to.
Matt W293///
Matt W294/// The TCP peer is the truth unless it is a proxy we were told to trust, in
Matt W295/// which case the client is the rightmost `X-Forwarded-For` entry that is not
Matt W296/// itself trusted — walking from the right because the entries an attacker can
Matt W297/// forge are on the left, appended before ours.
Matt W298///
Matt W299/// With no trusted proxies configured the header is ignored entirely. That is
Matt W300/// the only safe default: `XFF` is client-controlled, so honouring it from an
Matt W301/// arbitrary peer would let anyone claim a fresh bucket per request.
Matt W302pub fn client_ip(
Matt W303 headers: &HeaderMap,
Matt W304 peer: Option<IpAddr>,
Matt W305 trusted: &[IpNetwork],
Matt W306) -> Option<IpAddr> {
Matt W307 let peer = peer?;
Matt W308
Matt W309 let is_trusted = |ip: IpAddr| trusted.iter().any(|n| n.contains(ip));
Matt W310 if !is_trusted(peer) {
Matt W311 return Some(peer);
Matt W312 }
Matt W313
Matt W314 headers
Matt W315 .get_all(X_FORWARDED_FOR)
Matt W316 .iter()
Matt W317 .filter_map(|v| v.to_str().ok())
Matt W318 .flat_map(|v| v.split(','))
Matt W319 .filter_map(parse_forwarded)
Matt W320 .collect::<Vec<_>>()
Matt W321 .into_iter()
Matt W322 .rev()
Matt W323 .find(|ip| !is_trusted(*ip))
Matt W324 // Every hop was a trusted proxy, or the header was absent: the peer is
Matt W325 // the closest thing to a client we can honestly name.
Matt W326 .or(Some(peer))
Matt W327}
Matt W328
Matt W329/// The client address for this request, or the shared unknown-peer bucket.
Matt W330fn request_ip(state: &AppState, req: &Request) -> IpAddr {
Matt W331 let peer = req
Matt W332 .extensions()
Matt W333 .get::<ConnectInfo<std::net::SocketAddr>>()
Matt W334 .map(|c| c.0.ip());
Matt W335 client_ip(req.headers(), peer, &state.config.trusted_proxies).unwrap_or(UNKNOWN_PEER)
Matt W336}
Matt W337
Matt W338/// The coarse limiter, which runs *before* the session is resolved.
Matt W339///
Matt W340/// Its only job is to stop an unauthenticated flood buying a database round trip
Matt W341/// per request: session resolution sits between this layer and [`layer`].
Matt W342pub async fn edge_layer(State(state): State<AppState>, req: Request, next: Next) -> Response {
Matt W343 let path = req.uri().path().to_owned();
Matt W344 if is_exempt(&path) {
Matt W345 return next.run(req).await;
Matt W346 }
Matt W347
Matt W348 let ip = request_ip(&state, &req);
Matt W349 if !state.limiter.acquire_edge(ip, Instant::now()) {
Matt W350 tracing::warn!(%path, %ip, "rate limited at the edge");
Matt W351 return too_many();
Matt W352 }
Matt W353
Matt W354 next.run(req).await
Matt W355}
Matt W356
Matt W357/// The rate-limiting middleware.
Matt W358///
Matt W359/// `ConnectInfo` is optional so a missing peer address cannot turn every
Matt W360/// request into a 500. It is always present in production — `main` serves with
Matt W361/// `into_make_service_with_connect_info` — and its absence falls back to a
Matt W362/// single shared bucket, which is stricter than per-address, not looser.
Matt W363pub async fn layer(State(state): State<AppState>, req: Request, next: Next) -> Response {
Matt W364 let path = req.uri().path().to_owned();
Matt W365 if is_exempt(&path) {
Matt W366 return next.run(req).await;
Matt W367 }
Matt W368
Matt W369 // The session layer runs before this one, so an authenticated request is
Matt W370 // already resolved and gets its own bucket rather than sharing its
Matt W371 // neighbours' address.
Matt W372 let key = match req.extensions().get::<CurrentUser>().and_then(|u| u.0.as_ref()) {
Matt W373 Some(user) => Key::User(user.id),
Matt W374 None => Key::Addr(request_ip(&state, &req)),
Matt W375 };
Matt W376
Matt W377 let Some(_guard) = state.limiter.acquire(key.clone(), is_auth_path(&path), Instant::now())
Matt W378 else {
Matt W379 tracing::warn!(%path, ?key, "rate limited");
Matt W380 return too_many();
Matt W381 };
Matt W382
Matt W383 next.run(req).await
Matt W384}
Matt W385
Matt W386fn too_many() -> Response {
Matt W387 (
Matt W388 StatusCode::TOO_MANY_REQUESTS,
Matt W389 [(axum::http::header::RETRY_AFTER, "5")],
Matt W390 "Too many requests. Try again in a moment.",
Matt W391 )
Matt W392 .into_response()
Matt W393}
Matt W394
Matt W395#[cfg(test)]
Matt W396mod tests {
Matt W397 use super::*;
Matt W398
Matt W399 fn key() -> Key {
Matt W400 Key::Addr("10.0.0.1".parse().unwrap())
Matt W401 }
Matt W402
Matt W403 // ─── attributing a request to a client ───────────────────────────────────
Matt W404
Matt W405 fn ip(s: &str) -> IpAddr {
Matt W406 s.parse().unwrap()
Matt W407 }
Matt W408
Matt W409 fn nets(v: &[&str]) -> Vec<IpNetwork> {
Matt W410 v.iter().map(|s| s.parse().unwrap()).collect()
Matt W411 }
Matt W412
Matt W413 fn xff(value: &str) -> HeaderMap {
Matt W414 let mut h = HeaderMap::new();
Matt W415 h.insert(X_FORWARDED_FOR, value.parse().unwrap());
Matt W416 h
Matt W417 }
Matt W418
Matt W419 #[test]
Matt W420 fn without_trusted_proxies_the_header_is_ignored() {
Matt W421 // The spoofing case: believing this header from an arbitrary peer lets
Matt W422 // any client mint a fresh bucket per request.
Matt W423 let h = xff("1.2.3.4");
Matt W424 assert_eq!(
Matt W425 client_ip(&h, Some(ip("203.0.113.9")), &[]),
Matt W426 Some(ip("203.0.113.9"))
Matt W427 );
Matt W428 }
Matt W429
Matt W430 #[test]
Matt W431 fn a_forged_header_from_an_untrusted_peer_is_ignored() {
Matt W432 let h = xff("1.2.3.4");
Matt W433 let trusted = nets(&["172.23.0.0/16"]);
Matt W434 assert_eq!(
Matt W435 client_ip(&h, Some(ip("198.51.100.7")), &trusted),
Matt W436 Some(ip("198.51.100.7")),
Matt W437 "only the configured proxy may speak for a client"
Matt W438 );
Matt W439 }
Matt W440
Matt W441 #[test]
Matt W442 fn behind_the_proxy_the_client_is_taken_from_the_header() {
Matt W443 // The bug this exists to fix: without it every request looks like the
Matt W444 // proxy and all anonymous traffic shares one bucket.
Matt W445 let h = xff("203.0.113.9");
Matt W446 let trusted = nets(&["172.23.0.0/16"]);
Matt W447 assert_eq!(
Matt W448 client_ip(&h, Some(ip("172.23.0.2")), &trusted),
Matt W449 Some(ip("203.0.113.9"))
Matt W450 );
Matt W451 }
Matt W452
Matt W453 #[test]
Matt W454 fn a_client_cannot_prepend_its_way_to_a_fresh_bucket() {
Matt W455 // A client that sends its own XFF has it *prepended* to by the proxy,
Matt W456 // so the entries it controls are on the left. Reading from the right is
Matt W457 // what makes them inert.
Matt W458 let h = xff("9.9.9.9, 8.8.8.8, 203.0.113.9");
Matt W459 let trusted = nets(&["172.23.0.0/16"]);
Matt W460 assert_eq!(
Matt W461 client_ip(&h, Some(ip("172.23.0.2")), &trusted),
Matt W462 Some(ip("203.0.113.9")),
Matt W463 "the rightmost untrusted entry is the only honest one"
Matt W464 );
Matt W465 }
Matt W466
Matt W467 #[test]
Matt W468 fn trusted_hops_are_skipped_from_the_right() {
Matt W469 let h = xff("203.0.113.9, 172.23.0.5, 172.23.0.9");
Matt W470 let trusted = nets(&["172.23.0.0/16"]);
Matt W471 assert_eq!(
Matt W472 client_ip(&h, Some(ip("172.23.0.2")), &trusted),
Matt W473 Some(ip("203.0.113.9"))
Matt W474 );
Matt W475 }
Matt W476
Matt W477 #[test]
Matt W478 fn an_all_trusted_chain_falls_back_to_the_peer() {
Matt W479 let h = xff("172.23.0.5");
Matt W480 let trusted = nets(&["172.23.0.0/16"]);
Matt W481 assert_eq!(
Matt W482 client_ip(&h, Some(ip("172.23.0.2")), &trusted),
Matt W483 Some(ip("172.23.0.2"))
Matt W484 );
Matt W485 }
Matt W486
Matt W487 #[test]
Matt W488 fn a_proxy_that_sends_no_header_falls_back_to_the_peer() {
Matt W489 let trusted = nets(&["172.23.0.0/16"]);
Matt W490 assert_eq!(
Matt W491 client_ip(&HeaderMap::new(), Some(ip("172.23.0.2")), &trusted),
Matt W492 Some(ip("172.23.0.2"))
Matt W493 );
Matt W494 }
Matt W495
Matt W496 #[test]
Matt W497 fn forwarded_entries_parse_in_the_shapes_proxies_actually_send() {
Matt W498 assert_eq!(parse_forwarded("1.2.3.4"), Some(ip("1.2.3.4")));
Matt W499 assert_eq!(parse_forwarded(" 1.2.3.4 "), Some(ip("1.2.3.4")));
Matt W500 assert_eq!(parse_forwarded("1.2.3.4:5678"), Some(ip("1.2.3.4")));
Matt W501 assert_eq!(parse_forwarded("::1"), Some(ip("::1")));
Matt W502 assert_eq!(parse_forwarded("[::1]"), Some(ip("::1")));
Matt W503 assert_eq!(parse_forwarded("[2001:db8::1]:443"), Some(ip("2001:db8::1")));
Matt W504 // Junk is discarded, never guessed at.
Matt W505 assert_eq!(parse_forwarded(""), None);
Matt W506 assert_eq!(parse_forwarded("unknown"), None);
Matt W507 assert_eq!(parse_forwarded("_secret"), None);
Matt W508 }
Matt W509
Matt W510 #[test]
Matt W511 fn a_missing_peer_yields_no_address() {
Matt W512 assert_eq!(client_ip(&xff("1.2.3.4"), None, &nets(&["0.0.0.0/0"])), None);
Matt W513 }
Matt W514
Matt W515 #[test]
Matt W516 fn distinct_clients_behind_one_proxy_get_distinct_buckets() {
Matt W517 // The property the whole fix is for: two visitors must not be able to
Matt W518 // spend each other's budget.
Matt W519 let trusted = nets(&["172.23.0.0/16"]);
Matt W520 let peer = Some(ip("172.23.0.2"));
Matt W521 let a = client_ip(&xff("203.0.113.1"), peer, &trusted).unwrap();
Matt W522 let b = client_ip(&xff("203.0.113.2"), peer, &trusted).unwrap();
Matt W523 assert_ne!(a, b);
Matt W524
Matt W525 let l = Limiter::new();
Matt W526 let now = Instant::now();
Matt W527 let mut held = Vec::new();
Matt W528 for _ in 0..MAX_CONCURRENT {
Matt W529 held.push(l.acquire(Key::Addr(a), false, now).expect("under the cap"));
Matt W530 }
Matt W531 assert!(
Matt W532 l.acquire(Key::Addr(a), false, now).is_none(),
Matt W533 "the first client has spent its own budget"
Matt W534 );
Matt W535 assert!(
Matt W536 l.acquire(Key::Addr(b), false, now).is_some(),
Matt W537 "one client must not be able to lock everyone else out"
Matt W538 );
Matt W539 }
Matt W540
Matt W541 // ─── the coarse pre-session bucket ───────────────────────────────────────
Matt W542
Matt W543 // Asserting on constants is the point: this pins a relationship between
Matt W544 // them that a later edit could quietly break.
Matt W545 #[test]
Matt W546 #[allow(clippy::assertions_on_constants)]
Matt W547 fn the_edge_bucket_is_far_looser_than_the_real_one() {
Matt W548 // It must never be what stops ordinary traffic; the limiter after the
Matt W549 // session is where policy lives.
Matt W550 assert!(EDGE_BURST > BURST * 4.0);
Matt W551 assert!(EDGE_REFILL_PER_SEC > REFILL_PER_SEC * 4.0);
Matt W552 }
Matt W553
Matt W554 #[test]
Matt W555 fn the_edge_bucket_bites_eventually() {
Matt W556 let l = Limiter::new();
Matt W557 let now = Instant::now();
Matt W558 for i in 0..EDGE_BURST as usize {
Matt W559 assert!(l.acquire_edge(ip("10.0.0.1"), now), "request {i} of the burst");
Matt W560 }
Matt W561 assert!(!l.acquire_edge(ip("10.0.0.1"), now), "the edge burst is spent");
Matt W562 assert!(
Matt W563 l.acquire_edge(ip("10.0.0.2"), now),
Matt W564 "and it is per-address, not global"
Matt W565 );
Matt W566 }
Matt W567
Matt W568 #[test]
Matt W569 fn the_edge_bucket_takes_no_concurrency_slot() {
Matt W570 // It runs before the handler is chosen and releases nothing, so it must
Matt W571 // not consume the cap the real limiter enforces.
Matt W572 let l = Limiter::new();
Matt W573 let now = Instant::now();
Matt W574 for _ in 0..50 {
Matt W575 assert!(l.acquire_edge(ip("10.0.0.1"), now));
Matt W576 }
Matt W577 let mut held = Vec::new();
Matt W578 for _ in 0..MAX_CONCURRENT {
Matt W579 held.push(
Matt W580 l.acquire(Key::Addr(ip("10.0.0.1")), false, now)
Matt W581 .expect("the concurrency cap is untouched by the edge bucket"),
Matt W582 );
Matt W583 }
Matt W584 }
Matt W585
Matt W586 #[test]
Matt W587 fn a_burst_is_allowed_then_refused() {
Matt W588 let l = Limiter::new();
Matt W589 let now = Instant::now();
Matt W590
Matt W591 // Guards are held, so this also exercises the concurrency cap — which
Matt W592 // bites first, and should.
Matt W593 let mut held = Vec::new();
Matt W594 for _ in 0..MAX_CONCURRENT {
Matt W595 held.push(l.acquire(key(), false, now).expect("under the cap"));
Matt W596 }
Matt W597 assert!(
Matt W598 l.acquire(key(), false, now).is_none(),
Matt W599 "the concurrency cap must refuse the next request"
Matt W600 );
Matt W601 }
Matt W602
Matt W603 #[test]
Matt W604 fn a_released_slot_is_reusable() {
Matt W605 let l = Limiter::new();
Matt W606 let now = Instant::now();
Matt W607 {
Matt W608 let _g = l.acquire(key(), false, now).unwrap();
Matt W609 }
Matt W610 assert!(l.acquire(key(), false, now).is_some(), "dropping a guard frees the slot");
Matt W611 }
Matt W612
Matt W613 #[test]
Matt W614 fn the_rate_limit_bites_once_the_burst_is_spent() {
Matt W615 let l = Limiter::new();
Matt W616 let now = Instant::now();
Matt W617
Matt W618 // Drop each guard immediately so only the token bucket is in play.
Matt W619 for i in 0..BURST as usize {
Matt W620 assert!(l.acquire(key(), false, now).is_some(), "request {i} of the burst");
Matt W621 }
Matt W622 assert!(l.acquire(key(), false, now).is_none(), "the burst is spent");
Matt W623 }
Matt W624
Matt W625 #[test]
Matt W626 fn tokens_refill_over_time() {
Matt W627 let l = Limiter::new();
Matt W628 let start = Instant::now();
Matt W629 for _ in 0..BURST as usize {
Matt W630 let _ = l.acquire(key(), false, start);
Matt W631 }
Matt W632 assert!(l.acquire(key(), false, start).is_none());
Matt W633
Matt W634 let later = start + Duration::from_secs(2);
Matt W635 assert!(
Matt W636 l.acquire(key(), false, later).is_some(),
Matt W637 "two seconds should refill {REFILL_PER_SEC} tokens per second"
Matt W638 );
Matt W639 }
Matt W640
Matt W641 /// The login endpoint is worth grinding at, so it gets its own much smaller
Matt W642 /// bucket — and spending it must not spend the ordinary one.
Matt W643 #[test]
Matt W644 fn the_auth_bucket_is_separate_and_tighter() {
Matt W645 let l = Limiter::new();
Matt W646 let now = Instant::now();
Matt W647
Matt W648 for _ in 0..AUTH_BURST as usize {
Matt W649 assert!(l.acquire(key(), true, now).is_some());
Matt W650 }
Matt W651 assert!(l.acquire(key(), true, now).is_none(), "the auth bucket is spent");
Matt W652 assert!(
Matt W653 l.acquire(key(), false, now).is_some(),
Matt W654 "ordinary requests must still be served"
Matt W655 );
Matt W656 }
Matt W657
Matt W658 #[test]
Matt W659 fn identities_do_not_share_a_bucket() {
Matt W660 let l = Limiter::new();
Matt W661 let now = Instant::now();
Matt W662 let other = Key::Addr("10.0.0.2".parse().unwrap());
Matt W663
Matt W664 for _ in 0..BURST as usize {
Matt W665 let _ = l.acquire(key(), false, now);
Matt W666 }
Matt W667 assert!(l.acquire(key(), false, now).is_none());
Matt W668 assert!(
Matt W669 l.acquire(other, false, now).is_some(),
Matt W670 "one address must not exhaust another's budget"
Matt W671 );
Matt W672 }
Matt W673
Matt W674 #[test]
Matt W675 fn a_user_key_is_distinct_from_an_address_key() {
Matt W676 let l = Limiter::new();
Matt W677 let now = Instant::now();
Matt W678 let user = Key::User(uuid::Uuid::from_u128(1));
Matt W679
Matt W680 for _ in 0..BURST as usize {
Matt W681 let _ = l.acquire(key(), false, now);
Matt W682 }
Matt W683 assert!(l.acquire(user, false, now).is_some());
Matt W684 }
Matt W685
Matt W686 #[test]
Matt W687 fn idle_entries_are_evicted() {
Matt W688 let l = Limiter::new();
Matt W689 let start = Instant::now();
Matt W690 let _ = l.acquire(key(), false, start);
Matt W691 assert_eq!(l.tracked(), 1);
Matt W692
Matt W693 // Past the idle window, and past the sweep interval.
Matt W694 let later = start + IDLE_EVICT + Duration::from_secs(1);
Matt W695 let _ = l.acquire(Key::Addr("10.0.0.9".parse().unwrap()), false, later);
Matt W696 assert_eq!(l.tracked(), 1, "the idle entry should have been swept");
Matt W697 }
Matt W698
Matt W699 #[test]
Matt W700 fn a_backwards_clock_does_not_mint_tokens() {
Matt W701 let l = Limiter::new();
Matt W702 let now = Instant::now();
Matt W703 for _ in 0..BURST as usize {
Matt W704 let _ = l.acquire(key(), false, now);
Matt W705 }
Matt W706 // An earlier instant must not refill the bucket.
Matt W707 let earlier = now.checked_sub(Duration::from_secs(60)).unwrap_or(now);
Matt W708 assert!(l.acquire(key(), false, earlier).is_none());
Matt W709 }
Matt W710
Matt W711 #[test]
Matt W712 fn health_checks_and_assets_are_exempt() {
Matt W713 assert!(is_exempt("/healthz"));
Matt W714 assert!(is_exempt("/readyz"));
Matt W715 assert!(is_exempt("/assets/app.css"));
Matt W716 assert!(!is_exempt("/"), "ordinary pages are limited");
Matt W717 assert!(!is_exempt("/alice/repo/git-upload-pack"));
Matt W718 }
Matt W719
Matt W720 #[test]
Matt W721 fn auth_paths_are_recognised() {
Matt W722 for p in ["/login", "/auth/callback", "/auth/handle", "/setup"] {
Matt W723 assert!(is_auth_path(p), "{p}");
Matt W724 }
Matt W725 assert!(!is_auth_path("/"), "the dashboard is not an auth endpoint");
Matt W726 // Not prefix-matched: a repository called `login` is a page, not an
Matt W727 // auth endpoint, and must not inherit the tighter bucket.
Matt W728 assert!(!is_auth_path("/login/something"));
Matt W729 }
Matt W730}

730 lines · Rust