crates/gunnar-ui-server/src/throttle/tests.rs
//! The throttle's own rules, asked of the router's layer and of [`Throttle`] itself.
use std::cell::RefCell;
use std::net::SocketAddr;
use axum::routing::get;
use axum::Router;
use tower::ServiceExt;
use super::*;
thread_local! {
/// The limits the next [`Throttle::for_router`] on this thread takes (tests only).
pub(crate) static NEXT: RefCell<Option<Limits>> = const { RefCell::new(None) };
}
/// The next router built on this thread throttles with `limits`.
pub(crate) fn next_router(limits: Limits) {
NEXT.with(|n| *n.borrow_mut() = Some(limits));
}
/// Limits where nothing is refused but what a test tightens.
pub(crate) fn wide() -> Limits {
let tier = |burst| Tier { budget: Budget { burst, per_second: 0.0 }, in_flight: 1000 };
Limits {
guest: tier(1000),
credentialed: tier(1000),
account: tier(1000),
token: tier(1000),
global: tier(100_000),
render: Budget { burst: 1000, per_second: 0.0 },
render_all: Budget { burst: 1000, per_second: 0.0 },
block_after: 0,
block_window: Duration::from_secs(60),
block_for: Duration::from_secs(300),
max_keys: 1000,
}
}
/// A request from `peer`, with `headers`.
pub(crate) fn from(peer: &str, uri: &str, headers: &[(&str, &str)]) -> Request {
let mut b = http::Request::builder().uri(uri);
for (k, v) in headers {
b = b.header(*k, *v);
}
let mut r = b.body(axum::body::Body::empty()).expect("a request");
let addr: SocketAddr = peer.parse().expect("a socket address");
r.extensions_mut().insert(axum::extract::ConnectInfo(addr));
r
}
/// A router of one route behind the throttle's layer, as `serve::router` mounts it.
fn behind(throttle: &Arc<Throttle>) -> Router {
Router::new()
.route("/x", get(|| async { "ok" }))
.layer(axum::middleware::from_fn_with_state(throttle.clone(), layer))
}
async fn status(app: &Router, r: Request) -> StatusCode {
app.clone().oneshot(r).await.expect("an answer").status()
}
fn frozen(limits: Limits) -> Arc<Throttle> {
let t = Throttle::new(limits);
*t.frozen.lock().unwrap() = Some(Instant::now());
t
}
#[tokio::test]
async fn a_guest_past_its_budget_is_told_to_wait_and_another_address_is_not() {
let mut l = wide();
l.guest.budget.burst = 3;
let t = frozen(l);
let app = behind(&t);
for _ in 0..3 {
assert_eq!(status(&app, from("198.51.100.1:4000", "/x", &[])).await, StatusCode::OK);
}
let r = app.clone().oneshot(from("198.51.100.1:4001", "/x", &[])).await.unwrap();
assert_eq!(r.status(), StatusCode::TOO_MANY_REQUESTS);
assert!(r.headers().get(header::RETRY_AFTER).is_some(), "a 429 says when to come back");
let words = axum::body::to_bytes(r.into_body(), 4096).await.unwrap();
assert!(String::from_utf8_lossy(&words).contains("Too many requests"), "in words");
assert_eq!(status(&app, from("198.51.100.2:4000", "/x", &[])).await, StatusCode::OK, "another address has its own budget");
assert!(t.throttled() >= 1, "the refusal is counted");
}
#[tokio::test]
async fn an_ipv6_slash_64_is_one_client() {
let mut l = wide();
l.guest.budget.burst = 2;
let t = frozen(l);
let app = behind(&t);
assert_eq!(status(&app, from("[2001:db8:1:2::1]:443", "/x", &[])).await, StatusCode::OK);
assert_eq!(status(&app, from("[2001:db8:1:2:ffff::9]:443", "/x", &[])).await, StatusCode::OK);
assert_eq!(status(&app, from("[2001:db8:1:2:abcd::77]:443", "/x", &[])).await, StatusCode::TOO_MANY_REQUESTS, "the same /64");
assert_eq!(status(&app, from("[2001:db8:1:3::1]:443", "/x", &[])).await, StatusCode::OK, "another /64");
}
#[tokio::test]
async fn a_forwarded_for_from_an_untrusted_peer_does_not_change_the_key() {
let mut l = wide();
l.guest.budget.burst = 2;
let t = frozen(l);
let app = behind(&t);
for i in 0..2 {
let spoof = format!("203.0.113.{i}");
assert_eq!(status(&app, from("198.51.100.7:1", "/x", &[("x-forwarded-for", &spoof)])).await, StatusCode::OK);
}
assert_eq!(
status(&app, from("198.51.100.7:1", "/x", &[("x-forwarded-for", "203.0.113.99")])).await,
StatusCode::TOO_MANY_REQUESTS,
"a new x-forwarded-for on each request is still the one peer"
);
// From the front on loopback, the address the front saw is the client.
assert_eq!(status(&app, from("127.0.0.1:1", "/x", &[("x-forwarded-for", "203.0.113.50")])).await, StatusCode::OK);
}
#[tokio::test]
async fn a_credential_gets_the_higher_budget_and_a_token_its_own() {
let mut l = wide();
l.guest.budget.burst = 1;
l.credentialed.budget.burst = 5;
l.token.budget.burst = 2;
let t = frozen(l);
let app = behind(&t);
assert_eq!(status(&app, from("198.51.100.3:1", "/x", &[])).await, StatusCode::OK);
assert_eq!(status(&app, from("198.51.100.3:1", "/x", &[])).await, StatusCode::TOO_MANY_REQUESTS, "the guest budget");
let cookie = format!("{}=whatever", crate::auth::SESSION_COOKIE);
for _ in 0..3 {
assert_eq!(status(&app, from("198.51.100.4:1", "/x", &[("cookie", &cookie)])).await, StatusCode::OK, "a session's address budget is higher");
}
// A token from many addresses is one token.
assert_eq!(status(&app, from("198.51.100.5:1", "/x", &[("authorization", "Bearer machine-one")])).await, StatusCode::OK);
assert_eq!(status(&app, from("198.51.100.6:1", "/x", &[("authorization", "Bearer machine-one")])).await, StatusCode::OK);
assert_eq!(status(&app, from("198.51.100.8:1", "/x", &[("authorization", "Bearer machine-one")])).await, StatusCode::TOO_MANY_REQUESTS, "per token");
assert_eq!(status(&app, from("198.51.100.8:1", "/x", &[("authorization", "Bearer machine-two")])).await, StatusCode::OK, "another token");
// No token is kept as itself.
assert!(t.lock().token.len() >= 2);
}
#[tokio::test]
async fn an_account_is_limited_by_account_from_any_address() {
let mut l = wide();
l.account.budget.burst = 3;
let t = frozen(l);
for i in 0..3 {
assert!(t.admit_account("vera", &format!("198.51.100.{i}")).is_ok());
}
assert!(t.admit_account("vera", "198.51.100.200").is_err(), "a fourth address is the same account");
assert!(t.admit_account("olle", "198.51.100.200").is_ok(), "another account");
}
#[tokio::test]
async fn the_on_demand_renders_have_their_own_lower_budget() {
let mut l = wide();
l.render = Budget { burst: 2, per_second: 0.0 };
l.render_all = Budget { burst: 3, per_second: 0.0 };
let t = frozen(l);
assert!(t.admit_render("198.51.100.1").is_ok());
assert!(t.admit_render("198.51.100.1").is_ok());
assert!(t.admit_render("198.51.100.1").is_err(), "the client's render budget");
// Its page budget is untouched by the renders.
let app = behind(&t);
assert_eq!(status(&app, from("198.51.100.1:1", "/x", &[])).await, StatusCode::OK);
assert!(t.admit_render("198.51.100.2").is_ok());
assert!(t.admit_render("198.51.100.3").is_err(), "every client together");
assert!(Limits::default().render.burst < Limits::default().guest.budget.burst, "lower than a page's");
}
#[tokio::test]
async fn requests_in_flight_are_capped_per_client_and_in_all() {
let mut l = wide();
l.guest.in_flight = 2;
l.global.in_flight = 3;
let t = frozen(l);
let a = t.admit("198.51.100.1", false, None).expect("one");
let _b = t.admit("198.51.100.1", false, None).expect("two");
assert!(t.admit("198.51.100.1", false, None).is_err(), "a third at once");
let _c = t.admit("198.51.100.2", false, None).expect("another client");
assert!(t.admit("198.51.100.3", false, None).is_err(), "the whole box's cap");
drop(a);
assert!(t.admit("198.51.100.1", false, None).is_ok(), "room comes back with the answer");
}
#[tokio::test]
async fn many_addresses_together_meet_the_global_cap() {
let mut l = wide();
l.global.budget.burst = 10;
let t = frozen(l);
let app = behind(&t);
let mut ok = 0;
for i in 0..30 {
if status(&app, from(&format!("198.51.{}.{}:1", i / 200, i % 200 + 1), "/x", &[])).await == StatusCode::OK {
ok += 1;
}
}
assert_eq!(ok, 10);
}
#[tokio::test]
async fn a_repeat_offender_is_blocked_for_a_while_and_said_in_the_audit_log() {
let mut l = wide();
l.guest.budget.burst = 1;
l.block_after = 3;
l.block_for = Duration::from_secs(120);
let t = frozen(l);
crate::captured::install();
assert!(t.admit("198.51.100.9", false, None).is_ok());
for _ in 0..2 {
assert_eq!(t.admit("198.51.100.9", false, None).err().map(|r| r.blocked), Some(false));
}
let third = t.admit("198.51.100.9", false, None).err().expect("refused");
assert!(third.blocked && third.retry_after == 120, "{third:?}");
// Blocked: even with room again, until the time is up.
*t.frozen.lock().unwrap() = Some(t.now() + Duration::from_secs(60));
assert!(t.admit("198.51.100.9", false, None).err().is_some_and(|r| r.blocked));
*t.frozen.lock().unwrap() = Some(t.now() + Duration::from_secs(61));
assert_eq!(t.blocks(), 1);
let log = crate::captured::written();
assert!(log.lines().any(|l| l.contains("throttle.blocked") && l.contains("198.51.100.9") && l.contains(crate::auth::AUDIT_TARGET)), "the block is an audit line: {log}");
let after = t.admit("198.51.100.9", false, None).err().expect("still no room in its bucket");
assert!(!after.blocked, "the block has ended: {after:?}");
}
#[test]
fn the_setting_reads_and_a_bad_entry_keeps_the_default() {
let l = Limits::parse("guest=10/1/2, account=50/5, render=3/0.1, block=5/30/60, nonsense=1, token=x/1");
assert_eq!(l.guest, Tier { budget: Budget { burst: 10, per_second: 1.0 }, in_flight: 2 });
assert_eq!(l.account.budget, Budget { burst: 50, per_second: 5.0 });
assert_eq!(l.account.in_flight, Limits::default().account.in_flight);
assert_eq!(l.render, Budget { burst: 3, per_second: 0.1 });
assert_eq!((l.block_after, l.block_window, l.block_for), (5, Duration::from_secs(30), Duration::from_secs(60)));
assert_eq!(l.token, Limits::default().token);
assert_eq!(Limits::parse(""), Limits::default());
}
#[test]
fn the_defaults_carry_a_page_load_for_one_person() {
let d = Limits::default();
// A document, ~8 assets, ~12 API calls, the feed: well inside a guest's burst.
assert!(d.guest.budget.burst >= 60 && d.guest.in_flight >= 8);
assert!(d.account.budget.burst > d.guest.budget.burst && d.credentialed.budget.burst > d.guest.budget.burst);
assert!(d.global.budget.burst > d.credentialed.budget.burst);
}