Jump to…
snowfix: 500 error because of database is patchedyuzpxzopsouq1mo
1//! Rate limiting and per-user concurrency caps (spec §9, M6).
2//!
3//! > Cap: pack size on receive […] and **request concurrency per user**.
4//!
5//! Two different protections, because they stop two different things:
6//!
7//! * A **token bucket** bounds the *rate* of requests. It is what stops a script
8//! walking every change in a repository, or grinding at the login endpoint.
9//! * A **concurrency cap** bounds how many requests one identity may have in
10//! flight. A rate limit alone does not stop ten simultaneous diffs of the
11//! Linux kernel, and those are the requests that actually consume the box.
12//!
13//! Both are keyed by user id when there is one and by IP otherwise. Keying
14//! authenticated traffic by user rather than IP matters in both directions: a
15//! team behind one NAT is not one attacker, and one attacker on a hundred
16//! addresses is still one account.
17//!
18//! "By IP" means [`client_ip`], not the TCP peer. Behind a reverse proxy the
19//! peer is the proxy for every request in the world, so keying on it does not
20//! produce a strict limit — it produces a *global* one, and a single client can
21//! spend the whole budget and lock every signed-out visitor out of the site.
22//! `TRUSTED_PROXIES` is what closes that, and an empty setting is loud about it.
23//!
24//! There are two token buckets in the request path, at different depths.
25//! Resolving a session is a database round trip and it happens *before* the
26//! limiter that knows who you are, so a coarse per-address bucket
27//! ([`edge_layer`]) runs ahead of it, and the real one ([`layer`]) runs after.
28//!
29//! In-memory, deliberately. A shared limiter would mean Redis, and this is a
30//! single-instance product (spec §1). The state is small, bounded, and reset by
31//! a restart — which is the correct behaviour for a limiter whose only job is to
32//! keep one process healthy.
33
34use std::collections::HashMap;
35use std::net::IpAddr;
36use std::sync::{Arc, Mutex};
37use std::time::{Duration, Instant};
38
39use axum::extract::{ConnectInfo, Request, State};
40use axum::http::{HeaderMap, StatusCode};
41use axum::middleware::Next;
42use axum::response::{IntoResponse, Response};
43use sqlx::types::ipnetwork::IpNetwork;
44
45use crate::state::{AppState, CurrentUser};
46
47/// Sustained requests per second, per identity.
48const REFILL_PER_SEC: f64 = 8.0;
49
50/// Burst above the sustained rate. A page load is one document plus its assets,
51/// and a reviewer clicking through a stack fires several in a second.
52const BURST: f64 = 40.0;
53
54/// The much tighter bucket for endpoints that are worth grinding at: the login
55/// redirect, the OIDC callback, and the setup-token claim.
56const AUTH_REFILL_PER_SEC: f64 = 0.5;
57const AUTH_BURST: f64 = 10.0;
58
59/// Simultaneous in-flight requests per identity.
60///
61/// A diff or a highlight can occupy a thread for a while; this is what stops one
62/// identity holding all of them.
63const MAX_CONCURRENT: u32 = 12;
64
65/// The coarse bucket applied before the session is resolved.
66///
67/// Deliberately far above [`BURST`]: it is not a second rate limit, it is a
68/// bound on how much work an unauthenticated flood can force *ahead of* the
69/// real limiter. Session resolution is a database round trip and it runs first,
70/// so without this a client can spend a query per request no matter what the
71/// bucket below decides.
72const EDGE_REFILL_PER_SEC: f64 = 50.0;
73const EDGE_BURST: f64 = 200.0;
74
75/// Entries idle longer than this are dropped, so the map does not grow with
76/// every address that has ever connected.
77const IDLE_EVICT: Duration = Duration::from_secs(600);
78
79/// The bucket every request without a resolvable peer address shares.
80///
81/// `UNSPECIFIED` is not a routable address, so it cannot collide with a real
82/// client, and sharing one bucket is the conservative choice: unattributable
83/// traffic is limited together rather than not at all.
84const UNKNOWN_PEER: IpAddr = IpAddr::V4(std::net::Ipv4Addr::UNSPECIFIED);
85
86/// How many identities to track before evicting aggressively. Reached only
87/// under a distributed flood, where the correct behaviour is to keep working
88/// rather than to allocate.
89const MAX_TRACKED: usize = 50_000;
90
91#[derive(Clone)]
92pub struct Limiter(Arc<Mutex<Inner>>);
93
94struct Inner {
95 buckets: HashMap<Key, Bucket>,
96 last_sweep: Instant,
97}
98
99#[derive(Clone, PartialEq, Eq, Hash, Debug)]
100enum Key {
101 User(uuid::Uuid),
102 Addr(IpAddr),
103}
104
105struct Bucket {
106 tokens: f64,
107 auth_tokens: f64,
108 edge_tokens: f64,
109 in_flight: u32,
110 last: Instant,
111}
112
113impl Bucket {
114 fn full(now: Instant) -> Bucket {
115 Bucket {
116 tokens: BURST,
117 auth_tokens: AUTH_BURST,
118 edge_tokens: EDGE_BURST,
119 in_flight: 0,
120 last: now,
121 }
122 }
123
124 /// Refill for the time that passed.
125 ///
126 /// `saturating_duration_since` because a clock that went backwards must not
127 /// mint tokens.
128 fn refill(&mut self, now: Instant) {
129 let elapsed = now.saturating_duration_since(self.last).as_secs_f64();
130 self.tokens = (self.tokens + elapsed * REFILL_PER_SEC).min(BURST);
131 self.auth_tokens = (self.auth_tokens + elapsed * AUTH_REFILL_PER_SEC).min(AUTH_BURST);
132 self.edge_tokens = (self.edge_tokens + elapsed * EDGE_REFILL_PER_SEC).min(EDGE_BURST);
133 self.last = now;
134 }
135}
136
137impl Default for Limiter {
138 fn default() -> Self {
139 Self::new()
140 }
141}
142
143impl Limiter {
144 pub fn new() -> Self {
145 Limiter(Arc::new(Mutex::new(Inner {
146 buckets: HashMap::new(),
147 last_sweep: Instant::now(),
148 })))
149 }
150
151 /// Take one token, and a concurrency slot.
152 ///
153 /// Returns `None` when the identity is over a limit. The returned guard
154 /// releases the concurrency slot when dropped — including when the handler
155 /// panics, which is why it is a guard and not a pair of calls.
156 fn acquire(&self, key: Key, auth: bool, now: Instant) -> Option<Guard> {
157 let mut inner = self.0.lock().expect("rate limiter poisoned");
158 inner.sweep(now);
159
160 // A flood of distinct keys must not be able to grow the map without
161 // bound. Past the cap, unknown keys are refused rather than admitted —
162 // the alternative is admitting everything precisely when under attack.
163 if inner.buckets.len() >= MAX_TRACKED && !inner.buckets.contains_key(&key) {
164 return None;
165 }
166
167 let bucket = inner.buckets.entry(key.clone()).or_insert(Bucket::full(now));
168 bucket.refill(now);
169
170 if bucket.in_flight >= MAX_CONCURRENT {
171 return None;
172 }
173 if bucket.tokens < 1.0 {
174 return None;
175 }
176 if auth && bucket.auth_tokens < 1.0 {
177 return None;
178 }
179
180 bucket.tokens -= 1.0;
181 if auth {
182 bucket.auth_tokens -= 1.0;
183 }
184 bucket.in_flight += 1;
185
186 Some(Guard { limiter: self.clone(), key })
187 }
188
189 /// Take one token from the coarse pre-session bucket.
190 ///
191 /// No concurrency slot: this runs before the handler is chosen and releases
192 /// nothing, so it bounds arrival rate only.
193 fn acquire_edge(&self, ip: IpAddr, now: Instant) -> bool {
194 let key = Key::Addr(ip);
195 let mut inner = self.0.lock().expect("rate limiter poisoned");
196 inner.sweep(now);
197
198 if inner.buckets.len() >= MAX_TRACKED && !inner.buckets.contains_key(&key) {
199 return false;
200 }
201
202 let bucket = inner.buckets.entry(key).or_insert(Bucket::full(now));
203 bucket.refill(now);
204
205 if bucket.edge_tokens < 1.0 {
206 return false;
207 }
208 bucket.edge_tokens -= 1.0;
209 true
210 }
211
212 fn release(&self, key: &Key) {
213 if let Ok(mut inner) = self.0.lock() {
214 if let Some(b) = inner.buckets.get_mut(key) {
215 b.in_flight = b.in_flight.saturating_sub(1);
216 }
217 }
218 }
219
220 #[cfg(test)]
221 fn tracked(&self) -> usize {
222 self.0.lock().unwrap().buckets.len()
223 }
224}
225
226impl Inner {
227 fn sweep(&mut self, now: Instant) {
228 if now.saturating_duration_since(self.last_sweep) < Duration::from_secs(60) {
229 return;
230 }
231 self.last_sweep = now;
232 // An entry with a request in flight is never evicted, however idle its
233 // bucket looks — dropping it would lose the concurrency count and let
234 // the cap be bypassed by a slow request.
235 self.buckets.retain(|_, b| {
236 b.in_flight > 0 || now.saturating_duration_since(b.last) < IDLE_EVICT
237 });
238 }
239}
240
241/// Holds a concurrency slot for the life of a request.
242pub struct Guard {
243 limiter: Limiter,
244 key: Key,
245}
246
247impl Drop for Guard {
248 fn drop(&mut self) {
249 self.limiter.release(&self.key);
250 }
251}
252
253/// Whether a path gets the strict auth bucket.
254fn is_auth_path(path: &str) -> bool {
255 matches!(path, "/login" | "/auth/callback" | "/auth/handle" | "/setup")
256}
257
258/// Whether a path is exempt.
259///
260/// Health checks come from the orchestrator on a fixed interval and must never
261/// be throttled — a rate-limited `/healthz` restarts the container. Static
262/// assets are served from memory and are not worth a bucket.
263fn is_exempt(path: &str) -> bool {
264 matches!(path, "/healthz" | "/readyz") || path.starts_with("/assets/")
265}
266
267/// Not in `http::header`, which only defines registered headers.
268const X_FORWARDED_FOR: &str = "x-forwarded-for";
269
270/// One `X-Forwarded-For` entry as an address.
271///
272/// Proxies vary: bare addresses, `addr:port`, and bracketed IPv6 all appear.
273/// Anything that does not parse is discarded rather than guessed at.
274fn parse_forwarded(entry: &str) -> Option<IpAddr> {
275 let s = entry.trim();
276 if s.is_empty() {
277 return None;
278 }
279 // `[::1]` or `[::1]:8080`
280 if let Some(rest) = s.strip_prefix('[') {
281 let (inner, _) = rest.split_once(']')?;
282 return inner.parse().ok();
283 }
284 if let Ok(ip) = s.parse::<IpAddr>() {
285 return Some(ip);
286 }
287 // `1.2.3.4:5678`. Only IPv4 — a bare IPv6 has colons of its own and was
288 // handled by the parse above.
289 s.rsplit_once(':').and_then(|(host, _)| host.parse().ok())
290}
291
292/// The address to attribute a request to.
293///
294/// The TCP peer is the truth unless it is a proxy we were told to trust, in
295/// which case the client is the rightmost `X-Forwarded-For` entry that is not
296/// itself trusted — walking from the right because the entries an attacker can
297/// forge are on the left, appended before ours.
298///
299/// With no trusted proxies configured the header is ignored entirely. That is
300/// the only safe default: `XFF` is client-controlled, so honouring it from an
301/// arbitrary peer would let anyone claim a fresh bucket per request.
302pub fn client_ip(
303 headers: &HeaderMap,
304 peer: Option<IpAddr>,
305 trusted: &[IpNetwork],
306) -> Option<IpAddr> {
307 let peer = peer?;
308
309 let is_trusted = |ip: IpAddr| trusted.iter().any(|n| n.contains(ip));
310 if !is_trusted(peer) {
311 return Some(peer);
312 }
313
314 headers
315 .get_all(X_FORWARDED_FOR)
316 .iter()
317 .filter_map(|v| v.to_str().ok())
318 .flat_map(|v| v.split(','))
319 .filter_map(parse_forwarded)
320 .collect::<Vec<_>>()
321 .into_iter()
322 .rev()
323 .find(|ip| !is_trusted(*ip))
324 // Every hop was a trusted proxy, or the header was absent: the peer is
325 // the closest thing to a client we can honestly name.
326 .or(Some(peer))
327}
328
329/// The client address for this request, or the shared unknown-peer bucket.
330fn request_ip(state: &AppState, req: &Request) -> IpAddr {
331 let peer = req
332 .extensions()
333 .get::<ConnectInfo<std::net::SocketAddr>>()
334 .map(|c| c.0.ip());
335 client_ip(req.headers(), peer, &state.config.trusted_proxies).unwrap_or(UNKNOWN_PEER)
336}
337
338/// The coarse limiter, which runs *before* the session is resolved.
339///
340/// Its only job is to stop an unauthenticated flood buying a database round trip
341/// per request: session resolution sits between this layer and [`layer`].
342pub async fn edge_layer(State(state): State<AppState>, req: Request, next: Next) -> Response {
343 let path = req.uri().path().to_owned();
344 if is_exempt(&path) {
345 return next.run(req).await;
346 }
347
348 let ip = request_ip(&state, &req);
349 if !state.limiter.acquire_edge(ip, Instant::now()) {
350 tracing::warn!(%path, %ip, "rate limited at the edge");
351 return too_many();
352 }
353
354 next.run(req).await
355}
356
357/// The rate-limiting middleware.
358///
359/// `ConnectInfo` is optional so a missing peer address cannot turn every
360/// request into a 500. It is always present in production — `main` serves with
361/// `into_make_service_with_connect_info` — and its absence falls back to a
362/// single shared bucket, which is stricter than per-address, not looser.
363pub async fn layer(State(state): State<AppState>, req: Request, next: Next) -> Response {
364 let path = req.uri().path().to_owned();
365 if is_exempt(&path) {
366 return next.run(req).await;
367 }
368
369 // The session layer runs before this one, so an authenticated request is
370 // already resolved and gets its own bucket rather than sharing its
371 // neighbours' address.
372 let key = match req.extensions().get::<CurrentUser>().and_then(|u| u.0.as_ref()) {
373 Some(user) => Key::User(user.id),
374 None => Key::Addr(request_ip(&state, &req)),
375 };
376
377 let Some(_guard) = state.limiter.acquire(key.clone(), is_auth_path(&path), Instant::now())
378 else {
379 tracing::warn!(%path, ?key, "rate limited");
380 return too_many();
381 };
382
383 next.run(req).await
384}
385
386fn too_many() -> Response {
387 (
388 StatusCode::TOO_MANY_REQUESTS,
389 [(axum::http::header::RETRY_AFTER, "5")],
390 "Too many requests. Try again in a moment.",
391 )
392 .into_response()
393}
394
395#[cfg(test)]
396mod tests {
397 use super::*;
398
399 fn key() -> Key {
400 Key::Addr("10.0.0.1".parse().unwrap())
401 }
402
403 // ─── attributing a request to a client ───────────────────────────────────
404
405 fn ip(s: &str) -> IpAddr {
406 s.parse().unwrap()
407 }
408
409 fn nets(v: &[&str]) -> Vec<IpNetwork> {
410 v.iter().map(|s| s.parse().unwrap()).collect()
411 }
412
413 fn xff(value: &str) -> HeaderMap {
414 let mut h = HeaderMap::new();
415 h.insert(X_FORWARDED_FOR, value.parse().unwrap());
416 h
417 }
418
419 #[test]
420 fn without_trusted_proxies_the_header_is_ignored() {
421 // The spoofing case: believing this header from an arbitrary peer lets
422 // any client mint a fresh bucket per request.
423 let h = xff("1.2.3.4");
424 assert_eq!(
425 client_ip(&h, Some(ip("203.0.113.9")), &[]),
426 Some(ip("203.0.113.9"))
427 );
428 }
429
430 #[test]
431 fn a_forged_header_from_an_untrusted_peer_is_ignored() {
432 let h = xff("1.2.3.4");
433 let trusted = nets(&["172.23.0.0/16"]);
434 assert_eq!(
435 client_ip(&h, Some(ip("198.51.100.7")), &trusted),
436 Some(ip("198.51.100.7")),
437 "only the configured proxy may speak for a client"
438 );
439 }
440
441 #[test]
442 fn behind_the_proxy_the_client_is_taken_from_the_header() {
443 // The bug this exists to fix: without it every request looks like the
444 // proxy and all anonymous traffic shares one bucket.
445 let h = xff("203.0.113.9");
446 let trusted = nets(&["172.23.0.0/16"]);
447 assert_eq!(
448 client_ip(&h, Some(ip("172.23.0.2")), &trusted),
449 Some(ip("203.0.113.9"))
450 );
451 }
452
453 #[test]
454 fn a_client_cannot_prepend_its_way_to_a_fresh_bucket() {
455 // A client that sends its own XFF has it *prepended* to by the proxy,
456 // so the entries it controls are on the left. Reading from the right is
457 // what makes them inert.
458 let h = xff("9.9.9.9, 8.8.8.8, 203.0.113.9");
459 let trusted = nets(&["172.23.0.0/16"]);
460 assert_eq!(
461 client_ip(&h, Some(ip("172.23.0.2")), &trusted),
462 Some(ip("203.0.113.9")),
463 "the rightmost untrusted entry is the only honest one"
464 );
465 }
466
467 #[test]
468 fn trusted_hops_are_skipped_from_the_right() {
469 let h = xff("203.0.113.9, 172.23.0.5, 172.23.0.9");
470 let trusted = nets(&["172.23.0.0/16"]);
471 assert_eq!(
472 client_ip(&h, Some(ip("172.23.0.2")), &trusted),
473 Some(ip("203.0.113.9"))
474 );
475 }
476
477 #[test]
478 fn an_all_trusted_chain_falls_back_to_the_peer() {
479 let h = xff("172.23.0.5");
480 let trusted = nets(&["172.23.0.0/16"]);
481 assert_eq!(
482 client_ip(&h, Some(ip("172.23.0.2")), &trusted),
483 Some(ip("172.23.0.2"))
484 );
485 }
486
487 #[test]
488 fn a_proxy_that_sends_no_header_falls_back_to_the_peer() {
489 let trusted = nets(&["172.23.0.0/16"]);
490 assert_eq!(
491 client_ip(&HeaderMap::new(), Some(ip("172.23.0.2")), &trusted),
492 Some(ip("172.23.0.2"))
493 );
494 }
495
496 #[test]
497 fn forwarded_entries_parse_in_the_shapes_proxies_actually_send() {
498 assert_eq!(parse_forwarded("1.2.3.4"), Some(ip("1.2.3.4")));
499 assert_eq!(parse_forwarded(" 1.2.3.4 "), Some(ip("1.2.3.4")));
500 assert_eq!(parse_forwarded("1.2.3.4:5678"), Some(ip("1.2.3.4")));
501 assert_eq!(parse_forwarded("::1"), Some(ip("::1")));
502 assert_eq!(parse_forwarded("[::1]"), Some(ip("::1")));
503 assert_eq!(parse_forwarded("[2001:db8::1]:443"), Some(ip("2001:db8::1")));
504 // Junk is discarded, never guessed at.
505 assert_eq!(parse_forwarded(""), None);
506 assert_eq!(parse_forwarded("unknown"), None);
507 assert_eq!(parse_forwarded("_secret"), None);
508 }
509
510 #[test]
511 fn a_missing_peer_yields_no_address() {
512 assert_eq!(client_ip(&xff("1.2.3.4"), None, &nets(&["0.0.0.0/0"])), None);
513 }
514
515 #[test]
516 fn distinct_clients_behind_one_proxy_get_distinct_buckets() {
517 // The property the whole fix is for: two visitors must not be able to
518 // spend each other's budget.
519 let trusted = nets(&["172.23.0.0/16"]);
520 let peer = Some(ip("172.23.0.2"));
521 let a = client_ip(&xff("203.0.113.1"), peer, &trusted).unwrap();
522 let b = client_ip(&xff("203.0.113.2"), peer, &trusted).unwrap();
523 assert_ne!(a, b);
524
525 let l = Limiter::new();
526 let now = Instant::now();
527 let mut held = Vec::new();
528 for _ in 0..MAX_CONCURRENT {
529 held.push(l.acquire(Key::Addr(a), false, now).expect("under the cap"));
530 }
531 assert!(
532 l.acquire(Key::Addr(a), false, now).is_none(),
533 "the first client has spent its own budget"
534 );
535 assert!(
536 l.acquire(Key::Addr(b), false, now).is_some(),
537 "one client must not be able to lock everyone else out"
538 );
539 }
540
541 // ─── the coarse pre-session bucket ───────────────────────────────────────
542
543 // Asserting on constants is the point: this pins a relationship between
544 // them that a later edit could quietly break.
545 #[test]
546 #[allow(clippy::assertions_on_constants)]
547 fn the_edge_bucket_is_far_looser_than_the_real_one() {
548 // It must never be what stops ordinary traffic; the limiter after the
549 // session is where policy lives.
550 assert!(EDGE_BURST > BURST * 4.0);
551 assert!(EDGE_REFILL_PER_SEC > REFILL_PER_SEC * 4.0);
552 }
553
554 #[test]
555 fn the_edge_bucket_bites_eventually() {
556 let l = Limiter::new();
557 let now = Instant::now();
558 for i in 0..EDGE_BURST as usize {
559 assert!(l.acquire_edge(ip("10.0.0.1"), now), "request {i} of the burst");
560 }
561 assert!(!l.acquire_edge(ip("10.0.0.1"), now), "the edge burst is spent");
562 assert!(
563 l.acquire_edge(ip("10.0.0.2"), now),
564 "and it is per-address, not global"
565 );
566 }
567
568 #[test]
569 fn the_edge_bucket_takes_no_concurrency_slot() {
570 // It runs before the handler is chosen and releases nothing, so it must
571 // not consume the cap the real limiter enforces.
572 let l = Limiter::new();
573 let now = Instant::now();
574 for _ in 0..50 {
575 assert!(l.acquire_edge(ip("10.0.0.1"), now));
576 }
577 let mut held = Vec::new();
578 for _ in 0..MAX_CONCURRENT {
579 held.push(
580 l.acquire(Key::Addr(ip("10.0.0.1")), false, now)
581 .expect("the concurrency cap is untouched by the edge bucket"),
582 );
583 }
584 }
585
586 #[test]
587 fn a_burst_is_allowed_then_refused() {
588 let l = Limiter::new();
589 let now = Instant::now();
590
591 // Guards are held, so this also exercises the concurrency cap — which
592 // bites first, and should.
593 let mut held = Vec::new();
594 for _ in 0..MAX_CONCURRENT {
595 held.push(l.acquire(key(), false, now).expect("under the cap"));
596 }
597 assert!(
598 l.acquire(key(), false, now).is_none(),
599 "the concurrency cap must refuse the next request"
600 );
601 }
602
603 #[test]
604 fn a_released_slot_is_reusable() {
605 let l = Limiter::new();
606 let now = Instant::now();
607 {
608 let _g = l.acquire(key(), false, now).unwrap();
609 }
610 assert!(l.acquire(key(), false, now).is_some(), "dropping a guard frees the slot");
611 }
612
613 #[test]
614 fn the_rate_limit_bites_once_the_burst_is_spent() {
615 let l = Limiter::new();
616 let now = Instant::now();
617
618 // Drop each guard immediately so only the token bucket is in play.
619 for i in 0..BURST as usize {
620 assert!(l.acquire(key(), false, now).is_some(), "request {i} of the burst");
621 }
622 assert!(l.acquire(key(), false, now).is_none(), "the burst is spent");
623 }
624
625 #[test]
626 fn tokens_refill_over_time() {
627 let l = Limiter::new();
628 let start = Instant::now();
629 for _ in 0..BURST as usize {
630 let _ = l.acquire(key(), false, start);
631 }
632 assert!(l.acquire(key(), false, start).is_none());
633
634 let later = start + Duration::from_secs(2);
635 assert!(
636 l.acquire(key(), false, later).is_some(),
637 "two seconds should refill {REFILL_PER_SEC} tokens per second"
638 );
639 }
640
641 /// The login endpoint is worth grinding at, so it gets its own much smaller
642 /// bucket — and spending it must not spend the ordinary one.
643 #[test]
644 fn the_auth_bucket_is_separate_and_tighter() {
645 let l = Limiter::new();
646 let now = Instant::now();
647
648 for _ in 0..AUTH_BURST as usize {
649 assert!(l.acquire(key(), true, now).is_some());
650 }
651 assert!(l.acquire(key(), true, now).is_none(), "the auth bucket is spent");
652 assert!(
653 l.acquire(key(), false, now).is_some(),
654 "ordinary requests must still be served"
655 );
656 }
657
658 #[test]
659 fn identities_do_not_share_a_bucket() {
660 let l = Limiter::new();
661 let now = Instant::now();
662 let other = Key::Addr("10.0.0.2".parse().unwrap());
663
664 for _ in 0..BURST as usize {
665 let _ = l.acquire(key(), false, now);
666 }
667 assert!(l.acquire(key(), false, now).is_none());
668 assert!(
669 l.acquire(other, false, now).is_some(),
670 "one address must not exhaust another's budget"
671 );
672 }
673
674 #[test]
675 fn a_user_key_is_distinct_from_an_address_key() {
676 let l = Limiter::new();
677 let now = Instant::now();
678 let user = Key::User(uuid::Uuid::from_u128(1));
679
680 for _ in 0..BURST as usize {
681 let _ = l.acquire(key(), false, now);
682 }
683 assert!(l.acquire(user, false, now).is_some());
684 }
685
686 #[test]
687 fn idle_entries_are_evicted() {
688 let l = Limiter::new();
689 let start = Instant::now();
690 let _ = l.acquire(key(), false, start);
691 assert_eq!(l.tracked(), 1);
692
693 // Past the idle window, and past the sweep interval.
694 let later = start + IDLE_EVICT + Duration::from_secs(1);
695 let _ = l.acquire(Key::Addr("10.0.0.9".parse().unwrap()), false, later);
696 assert_eq!(l.tracked(), 1, "the idle entry should have been swept");
697 }
698
699 #[test]
700 fn a_backwards_clock_does_not_mint_tokens() {
701 let l = Limiter::new();
702 let now = Instant::now();
703 for _ in 0..BURST as usize {
704 let _ = l.acquire(key(), false, now);
705 }
706 // An earlier instant must not refill the bucket.
707 let earlier = now.checked_sub(Duration::from_secs(60)).unwrap_or(now);
708 assert!(l.acquire(key(), false, earlier).is_none());
709 }
710
711 #[test]
712 fn health_checks_and_assets_are_exempt() {
713 assert!(is_exempt("/healthz"));
714 assert!(is_exempt("/readyz"));
715 assert!(is_exempt("/assets/app.css"));
716 assert!(!is_exempt("/"), "ordinary pages are limited");
717 assert!(!is_exempt("/alice/repo/git-upload-pack"));
718 }
719
720 #[test]
721 fn auth_paths_are_recognised() {
722 for p in ["/login", "/auth/callback", "/auth/handle", "/setup"] {
723 assert!(is_auth_path(p), "{p}");
724 }
725 assert!(!is_auth_path("/"), "the dashboard is not an auth endpoint");
726 // Not prefix-matched: a repository called `login` is a page, not an
727 // auth endpoint, and must not inherit the tighter bucket.
728 assert!(!is_auth_path("/login/something"));
729 }
730}

730 lines · Rust