diff --git a/CLAUDE.md b/CLAUDE.md index 157428d..7a3022f 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -12,7 +12,7 @@ This server cannot run as more than one instance today, and four pieces of state |---|---|---| | `SyncNotifier` | `src/sync_notifier.rs` | a client on A never sees an event caused by B — stale, no error | | `TerminalManager` | `src/terminal_manager.rs` | a shared session only works if every participant lands on the same instance | -| Rate limiters | `src/rate_limit.rs` | N instances means N times every limit, register and auth included | +| Short-window rate limiters | `src/rate_limit.rs` | N instances means N times auth, sync and search; the hour/day limits are already counted in Postgres via `.durable()` | | Presence/usage | `src/routes/presence.rs` | each instance sees only its own users | Every one added between now and that migration is another item on it. Four is a weekend; diff --git a/migrations/071_rate_limit_hits.sql b/migrations/071_rate_limit_hits.sql new file mode 100644 index 0000000..fee0125 --- /dev/null +++ b/migrations/071_rate_limit_hits.sql @@ -0,0 +1,7 @@ +CREATE TABLE rate_limit_hits ( + bucket TEXT NOT NULL, + key TEXT NOT NULL, + at TIMESTAMPTZ NOT NULL DEFAULT now() +); + +CREATE INDEX rate_limit_hits_bucket_key_at ON rate_limit_hits (bucket, key, at); diff --git a/src/main.rs b/src/main.rs index d5d33ec..e68a5c1 100644 --- a/src/main.rs +++ b/src/main.rs @@ -216,6 +216,11 @@ async fn main() { Ok(n) => tracing::info!(deleted = n, "Unattached rule set sweep completed"), Err(e) => tracing::error!(error = %e, "Unattached rule set sweep failed"), } + + match rate_limit::sweep_durable_hits(&retention_pool).await { + Ok(n) => tracing::info!(deleted = n, "Expired rate limit hits swept"), + Err(e) => tracing::error!(error = %e, "Rate limit hit sweep failed"), + } } }); @@ -270,24 +275,23 @@ async fn main() { .unwrap_or(10); let auth_limiter = RateLimiter::::new(auth_per_minute, Duration::from_secs(60)); - let register_limiter = RegisterRateLimiter(RateLimiter::new( - register_per_day, - Duration::from_secs(86400), - )); + let hourly = |max: usize, bucket: &'static str| { + RateLimiter::::new(max, Duration::from_secs(3600)).durable(pool.clone(), bucket) + }; + let register_limiter = RegisterRateLimiter( + RateLimiter::new(register_per_day, Duration::from_secs(86400)).durable(pool.clone(), "register"), + ); let enrollment_resolve_limiter = EnrollmentResolveLimiter(RateLimiter::::new(60, Duration::from_secs(60))); - let enrollment_register_limiter = - EnrollmentRegisterLimiter(RateLimiter::::new(500, Duration::from_secs(86400))); + let enrollment_register_limiter = EnrollmentRegisterLimiter( + RateLimiter::::new(500, Duration::from_secs(86400)).durable(pool.clone(), "enrollment_register"), + ); let enrollment_policy_limiter = EnrollmentPolicyLimiter(RateLimiter::::new(300, Duration::from_secs(60))); - let invite_limiter = InviteRateLimiter(RateLimiter::::new( - invite_per_hour, - Duration::from_secs(3600), - )); - let waitlist_limiter = WaitlistRateLimiter(RateLimiter::new( - waitlist_per_hour, - Duration::from_secs(3600), - )); + let invite_limiter = InviteRateLimiter(hourly(invite_per_hour, "invite")); + let waitlist_limiter = WaitlistRateLimiter( + RateLimiter::new(waitlist_per_hour, Duration::from_secs(3600)).durable(pool.clone(), "waitlist"), + ); let sync_limiter = SyncRateLimiter(RateLimiter::::new( sync_rate, Duration::from_secs(3600), @@ -296,20 +300,13 @@ async fn main() { AuditClientRateLimiter(RateLimiter::::new(100, Duration::from_secs(60))); let search_limiter = SearchRateLimiter(RateLimiter::::new(search_rate, Duration::from_secs(60))); - let knock_limiter = - KnockRateLimiter(RateLimiter::::new(knock_per_hour, Duration::from_secs(3600))); - let session_code_limiter = - SessionCodeRateLimiter(RateLimiter::::new(30, Duration::from_secs(3600))); - let redeem_limiter = - RedeemRateLimiter(RateLimiter::::new(20, Duration::from_secs(3600))); - let grant_mint_limiter = - GrantMintRateLimiter(RateLimiter::::new(30, Duration::from_secs(3600))); - let grant_redeem_limiter = - GrantRedeemRateLimiter(RateLimiter::::new(20, Duration::from_secs(3600))); - let enrollment_redeem_limiter = - EnrollmentRedeemRateLimiter(RateLimiter::::new(20, Duration::from_secs(3600))); - let resend_verification_limiter = - ResendVerificationRateLimiter(RateLimiter::::new(5, Duration::from_secs(3600))); + let knock_limiter = KnockRateLimiter(hourly(knock_per_hour, "knock")); + let session_code_limiter = SessionCodeRateLimiter(hourly(30, "session_code")); + let redeem_limiter = RedeemRateLimiter(hourly(20, "redeem")); + let grant_mint_limiter = GrantMintRateLimiter(hourly(30, "grant_mint")); + let grant_redeem_limiter = GrantRedeemRateLimiter(hourly(20, "grant_redeem")); + let enrollment_redeem_limiter = EnrollmentRedeemRateLimiter(hourly(20, "enrollment_redeem")); + let resend_verification_limiter = ResendVerificationRateLimiter(hourly(5, "resend_verification")); // Lemon Squeezy live metrics cache (background refresh every 5 min). let ls_cache = lemonsqueezy::LsCache::default(); diff --git a/src/rate_limit.rs b/src/rate_limit.rs index 061b7e7..d8760b3 100644 --- a/src/rate_limit.rs +++ b/src/rate_limit.rs @@ -1,18 +1,30 @@ use axum::{extract::Request, http::StatusCode, middleware::Next, response::Response}; +use sqlx::PgPool; use std::collections::HashMap; +use std::fmt::Display; use std::hash::Hash; use std::net::IpAddr; use std::sync::{Arc, OnceLock}; use std::time::{Duration, Instant}; use tokio::sync::Mutex; -use tracing::{warn, Instrument}; +use tracing::{error, warn, Instrument}; use uuid::Uuid; +/// The daily sweep deletes durable hits older than this, so no durable window may exceed it. +pub const DURABLE_WINDOW_MAX: Duration = Duration::from_secs(86_400); + /// Per-key sliding window rate limiter. Key is typically IpAddr or Uuid. pub struct RateLimiter { state: Arc>>, max_requests: usize, window: Duration, + durable: Option, +} + +#[derive(Clone)] +struct Durable { + pool: PgPool, + bucket: &'static str, } struct Buckets { @@ -27,11 +39,12 @@ impl Clone for RateLimiter { state: Arc::clone(&self.state), max_requests: self.max_requests, window: self.window, + durable: self.durable.clone(), } } } -impl RateLimiter { +impl RateLimiter { pub fn new(max_requests: usize, window: Duration) -> Self { Self { state: Arc::new(Mutex::new(Buckets { @@ -40,10 +53,36 @@ impl RateLimiter { })), max_requests, window, + durable: None, } } + /// Also counts hits in Postgres under `bucket`, so the budget survives restarts and is + /// shared by every instance. The in-memory check stays in front to absorb floods. + pub fn durable(mut self, pool: PgPool, bucket: &'static str) -> Self { + assert!(self.window <= DURABLE_WINDOW_MAX, "{bucket}: window outlives the hit sweep"); + self.durable = Some(Durable { pool, bucket }); + self + } + pub async fn check(&self, key: K) -> bool { + let shown = self.durable.as_ref().map(|_| key.to_string()); + if !self.check_memory(key).await { + return false; + } + let (Some(durable), Some(key)) = (&self.durable, shown) else { + return true; + }; + match admit_durable(durable, &key, self.max_requests, self.window).await { + Ok(admitted) => admitted, + Err(e) => { + error!(error = %e, bucket = durable.bucket, "Durable rate limit check failed; using the in-memory count"); + true + } + } + } + + async fn check_memory(&self, key: K) -> bool { let mut state = self.state.lock().await; let now = Instant::now(); let window = self.window; @@ -69,6 +108,47 @@ impl RateLimiter { } } +async fn admit_durable(durable: &Durable, key: &str, max_requests: usize, window: Duration) -> sqlx::Result { + let mut tx = durable.pool.begin().await?; + // The two-key form has its own key space, apart from the single-key instance lock. + sqlx::query("SELECT pg_advisory_xact_lock(hashtext($1), hashtext($2))") + .bind(durable.bucket) + .bind(key) + .execute(&mut *tx) + .await?; + let recent: i64 = sqlx::query_scalar( + "WITH expired AS ( \ + DELETE FROM rate_limit_hits \ + WHERE bucket = $1 AND key = $2 AND at <= now() - make_interval(secs => $3) \ + ) \ + SELECT count(*) FROM rate_limit_hits \ + WHERE bucket = $1 AND key = $2 AND at > now() - make_interval(secs => $3)", + ) + .bind(durable.bucket) + .bind(key) + .bind(window.as_secs_f64()) + .fetch_one(&mut *tx) + .await?; + let admitted = (recent as usize) < max_requests; + if admitted { + sqlx::query("INSERT INTO rate_limit_hits (bucket, key) VALUES ($1, $2)") + .bind(durable.bucket) + .bind(key) + .execute(&mut *tx) + .await?; + } + tx.commit().await?; + Ok(admitted) +} + +pub async fn sweep_durable_hits(pool: &PgPool) -> sqlx::Result { + sqlx::query("DELETE FROM rate_limit_hits WHERE at <= now() - make_interval(secs => $1)") + .bind(DURABLE_WINDOW_MAX.as_secs_f64()) + .execute(pool) + .await + .map(|r| r.rows_affected()) +} + /// A trusted proxy address or subnet. A containerised reverse proxy gets a new /// address on every recreate, so a bare IP cannot be the only accepted form. #[derive(Clone, Copy, Debug)] @@ -429,6 +509,82 @@ mod tests { ); } + fn durable(pool: &PgPool, max: usize) -> RateLimiter { + RateLimiter::::new(max, Duration::from_secs(3600)).durable(pool.clone(), "test") + } + + async fn hits_for(pool: &PgPool, key: Uuid) -> i64 { + sqlx::query_scalar("SELECT count(*) FROM rate_limit_hits WHERE bucket = 'test' AND key = $1") + .bind(key.to_string()) + .fetch_one(pool) + .await + .unwrap() + } + + async fn seed_hit(pool: &PgPool, key: Uuid, age: &str) { + sqlx::query("INSERT INTO rate_limit_hits (bucket, key, at) VALUES ('test', $1, now() - $2::interval)") + .bind(key.to_string()) + .bind(age) + .execute(pool) + .await + .unwrap(); + } + + #[tokio::test] + async fn a_durable_budget_survives_a_restart_and_is_shared_between_instances() { + let pool = crate::test_pool_or_skip!(); + let key = Uuid::new_v4(); + let before = durable(&pool, 2); + assert!(before.check(key).await); + assert!(before.check(key).await); + + let after = durable(&pool, 2); + assert!(!after.check(key).await); + assert!(after.check(Uuid::new_v4()).await, "budgets are per key"); + } + + #[tokio::test] + async fn concurrent_checks_never_overshoot_the_durable_budget() { + let pool = crate::test_pool_or_skip!(); + let key = Uuid::new_v4(); + let checks = (0..12).map(|_| { + let limiter = durable(&pool, 5); + tokio::spawn(async move { limiter.check(key).await }) + }); + let admitted = futures_util::future::join_all(checks) + .await + .into_iter() + .filter(|admitted| *admitted.as_ref().unwrap()) + .count(); + assert_eq!(admitted, 5); + assert_eq!(hits_for(&pool, key).await, 5); + } + + #[tokio::test] + async fn hits_outside_the_window_neither_count_nor_survive() { + let pool = crate::test_pool_or_skip!(); + let key = Uuid::new_v4(); + seed_hit(&pool, key, "2 hours").await; + assert!(durable(&pool, 1).check(key).await); + assert_eq!(hits_for(&pool, key).await, 1, "the expired hit is pruned on check"); + + let quiet = Uuid::new_v4(); + seed_hit(&pool, quiet, "2 days").await; + sweep_durable_hits(&pool).await.unwrap(); + assert_eq!(hits_for(&pool, quiet).await, 0); + assert_eq!(hits_for(&pool, key).await, 1, "the sweep keeps live hits"); + } + + #[tokio::test] + async fn an_unreachable_database_falls_back_to_the_in_memory_count() { + let pool = crate::test_pool_or_skip!(); + let limiter = durable(&pool, 1); + pool.close().await; + let key = Uuid::new_v4(); + assert!(limiter.check(key).await); + assert!(!limiter.check(key).await); + } + #[tokio::test] async fn limiter_admits_up_to_the_configured_maximum() { let limiter = RateLimiter::::new(3, Duration::from_secs(60));