diff --git a/src/lib.rs b/src/lib.rs index 7189517..01ec89b 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -23,6 +23,7 @@ use axum::extract::DefaultBodyLimit; use axum::http::{header, HeaderValue, Method}; use axum::{routing::get, Router}; use tokio::sync::RwLock; +use tower_governor::{governor::GovernorConfigBuilder, GovernorLayer}; use tower_http::compression::CompressionLayer; use tower_http::cors::{Any, CorsLayer}; use tower_http::timeout::TimeoutLayer; @@ -39,6 +40,15 @@ const MAX_BODY_BYTES: usize = 1024 * 1024; /// so their errors surface first for chain calls; bounds any slow handler). const REQUEST_TIMEOUT: Duration = Duration::from_secs(60); +// Per-IP rate-limit tiers (one token replenished per period; `burst` = bucket). +// STRICT throttles the unauthenticated, abuse-prone surface (auth + website); +// NORMAL covers the authenticated wallet/identity API. Sane defaults; operator- +// tunable config can follow. +const STRICT_PERIOD_MS: u64 = 500; // ~2 req/s sustained +const STRICT_BURST: u32 = 20; +const NORMAL_PERIOD_MS: u64 = 100; // ~10 req/s sustained +const NORMAL_BURST: u32 = 60; + /// Shared application state injected into handlers via `State>`. #[derive(Clone)] pub struct AppState { @@ -59,10 +69,34 @@ pub struct AppState { pub fn build_router(state: Arc) -> Router { let cors = cors_layer(&state.config); + // Per-IP rate limiting (PeerIpKeyExtractor reads ConnectInfo — `main` serves + // with `into_make_service_with_connect_info`; the test harness injects it). + // Tiered: stricter on the unauthenticated auth/website surface (closes the + // website-challenge growth vector), looser on the wallet/identity API. + let mut strict_cfg = GovernorConfigBuilder::default(); + strict_cfg + .per_millisecond(STRICT_PERIOD_MS) + .burst_size(STRICT_BURST); + let strict = GovernorLayer { + config: Arc::new(strict_cfg.finish().expect("valid strict governor config")), + }; + + let mut normal_cfg = GovernorConfigBuilder::default(); + normal_cfg + .per_millisecond(NORMAL_PERIOD_MS) + .burst_size(NORMAL_BURST); + let normal = GovernorLayer { + config: Arc::new(normal_cfg.finish().expect("valid normal governor config")), + }; + let api_v1 = api::auth::routes() .merge(api::website::routes()) - .merge(api::users::routes()) - .merge(api::wallet::routes()); + .layer(strict) + .merge( + api::users::routes() + .merge(api::wallet::routes()) + .layer(normal), + ); Router::new() .route("/health", get(api::health::health)) diff --git a/src/main.rs b/src/main.rs index 368ff0a..9d0088e 100644 --- a/src/main.rs +++ b/src/main.rs @@ -80,6 +80,11 @@ async fn main() -> Result<(), Box> { let listener = tokio::net::TcpListener::bind(&addr).await?; tracing::info!("smirk-backend-core listening on http://{addr}"); - axum::serve(listener, app).await?; + // ConnectInfo carries the peer IP that the per-IP rate limiter keys on. + axum::serve( + listener, + app.into_make_service_with_connect_info::(), + ) + .await?; Ok(()) } diff --git a/tests/common/mod.rs b/tests/common/mod.rs index da1289b..0ccc38f 100644 --- a/tests/common/mod.rs +++ b/tests/common/mod.rs @@ -112,7 +112,18 @@ impl TestApp { token: Option<&str>, body: Option, ) -> (StatusCode, Value) { - let mut builder = Request::builder().method(method).uri(uri); + // The per-IP rate limiter keys on ConnectInfo; in production `main` + // supplies it via into_make_service_with_connect_info. oneshot requests + // don't, so inject a loopback peer (each test builds a fresh router with + // fresh limiter state, so this never trips across tests). + let mut builder = + Request::builder() + .method(method) + .uri(uri) + .extension(axum::extract::ConnectInfo(std::net::SocketAddr::from(( + [127, 0, 0, 1], + 0, + )))); if let Some(t) = token { builder = builder.header("authorization", format!("Bearer {t}")); } diff --git a/tests/rate_limit.rs b/tests/rate_limit.rs new file mode 100644 index 0000000..4fb76fc --- /dev/null +++ b/tests/rate_limit.rs @@ -0,0 +1,27 @@ +//! L1 integration: the per-IP rate limiter actually throttles a burst. +//! +//! The strict tier wraps the unauthenticated auth/website surface. The governor +//! runs before the handler, so an unauthenticated burst to a strict-tier route +//! still counts — past the burst size it must return 429. (All requests share the +//! injected loopback peer, so they hit one bucket.) + +mod common; + +use axum::http::StatusCode; + +#[tokio::test] +async fn strict_tier_throttles_a_burst() { + let app = require_app!(); + let mut saw_429 = false; + for _ in 0..40 { + let (status, _) = app.request("GET", "/api/v1/auth/me", None, None).await; + if status == StatusCode::TOO_MANY_REQUESTS { + saw_429 = true; + break; + } + } + assert!( + saw_429, + "a rapid burst to a strict-tier endpoint should be rate-limited (429)" + ); +}