feat(server): keep hour and day rate limits across restarts

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;

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);

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::<std::net::IpAddr>::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::<uuid::Uuid>::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::<std::net::IpAddr>::new(60, Duration::from_secs(60)));
-    let enrollment_register_limiter =
-        EnrollmentRegisterLimiter(RateLimiter::<uuid::Uuid>::new(500, Duration::from_secs(86400)));
+    let enrollment_register_limiter = EnrollmentRegisterLimiter(
+        RateLimiter::<uuid::Uuid>::new(500, Duration::from_secs(86400)).durable(pool.clone(), "enrollment_register"),
+    );
     let enrollment_policy_limiter =
         EnrollmentPolicyLimiter(RateLimiter::<std::net::IpAddr>::new(300, Duration::from_secs(60)));
-    let invite_limiter = InviteRateLimiter(RateLimiter::<uuid::Uuid>::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::<uuid::Uuid>::new(
         sync_rate,
         Duration::from_secs(3600),
@@ -296,20 +300,13 @@ async fn main() {
         AuditClientRateLimiter(RateLimiter::<uuid::Uuid>::new(100, Duration::from_secs(60)));
     let search_limiter =
         SearchRateLimiter(RateLimiter::<uuid::Uuid>::new(search_rate, Duration::from_secs(60)));
-    let knock_limiter =
-        KnockRateLimiter(RateLimiter::<uuid::Uuid>::new(knock_per_hour, Duration::from_secs(3600)));
-    let session_code_limiter =
-        SessionCodeRateLimiter(RateLimiter::<uuid::Uuid>::new(30, Duration::from_secs(3600)));
-    let redeem_limiter =
-        RedeemRateLimiter(RateLimiter::<uuid::Uuid>::new(20, Duration::from_secs(3600)));
-    let grant_mint_limiter =
-        GrantMintRateLimiter(RateLimiter::<uuid::Uuid>::new(30, Duration::from_secs(3600)));
-    let grant_redeem_limiter =
-        GrantRedeemRateLimiter(RateLimiter::<uuid::Uuid>::new(20, Duration::from_secs(3600)));
-    let enrollment_redeem_limiter =
-        EnrollmentRedeemRateLimiter(RateLimiter::<uuid::Uuid>::new(20, Duration::from_secs(3600)));
-    let resend_verification_limiter =
-        ResendVerificationRateLimiter(RateLimiter::<uuid::Uuid>::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();

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<K = IpAddr> {
     state: Arc<Mutex<Buckets<K>>>,
     max_requests: usize,
     window: Duration,
+    durable: Option<Durable>,
+}
+
+#[derive(Clone)]
+struct Durable {
+    pool: PgPool,
+    bucket: &'static str,
 }
 
 struct Buckets<K> {
@@ -27,11 +39,12 @@ impl<K> Clone for RateLimiter<K> {
             state: Arc::clone(&self.state),
             max_requests: self.max_requests,
             window: self.window,
+            durable: self.durable.clone(),
         }
     }
 }
 
-impl<K: Eq + Hash + Send + 'static> RateLimiter<K> {
+impl<K: Eq + Hash + Display + Send + 'static> RateLimiter<K> {
     pub fn new(max_requests: usize, window: Duration) -> Self {
         Self {
             state: Arc::new(Mutex::new(Buckets {
@@ -40,10 +53,36 @@ impl<K: Eq + Hash + Send + 'static> RateLimiter<K> {
             })),
             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<K: Eq + Hash + Send + 'static> RateLimiter<K> {
     }
 }
 
+async fn admit_durable(durable: &Durable, key: &str, max_requests: usize, window: Duration) -> sqlx::Result<bool> {
+    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<u64> {
+    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<Uuid> {
+        RateLimiter::<Uuid>::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::<IpAddr>::new(3, Duration::from_secs(60));