Mangalord/backend/src/auth/rate_limit.rs

//! Per-process token-bucket rate limiter for the auth endpoints.
//!
//! Protects `/auth/login`, `/auth/register`, and `/auth/me/password`
//! from credential stuffing / password spraying / username probing.
//!
//! The current deploy puts SvelteKit's hooks.server.ts proxy in front
//! of axum without forwarding the original client IP (no
//! `X-Forwarded-For`), so per-IP buckets would all collapse to the
//! proxy container's address. Until the proxy learns to forward the
//! peer address, a single global bucket gives equivalent protection
//! against mass-attack patterns and trades a small DoS surface
//! (legitimate users sharing the limit) for simplicity.
//!
//! Each `AppState` carries its own [`AuthRateLimiter`] instance, so
//! tests run in isolated buckets and won't bleed across `#[sqlx::test]`
//! cases that share a process.

use std::sync::Mutex;
use std::time::Instant;

/// Tunable limits. `per_sec == 0` disables the limiter — used by the
/// test harness and by anyone who wants to opt out via env config.
#[derive(Clone, Copy, Debug)]
pub struct RateLimitConfig {
    pub per_sec: u32,
    pub burst: u32,
}

impl Default for RateLimitConfig {
    /// Disabled by default. The production `AuthConfig::from_env`
    /// overrides to a real limit; the test harness keeps the default
    /// so existing tests don't flake against shared buckets.
    fn default() -> Self {
        Self {
            per_sec: 0,
            burst: 0,
        }
    }
}

/// Production defaults: 5 requests/sec sustained, 10-request burst.
/// Tight enough to make brute force impractical, loose enough that a
/// real user mistyping their password three times in a row doesn't
/// hit it.
pub const PRODUCTION_PER_SEC: u32 = 5;
pub const PRODUCTION_BURST: u32 = 10;

struct Bucket {
    tokens: f64,
    last_refill: Instant,
}

/// Outcome of [`AuthRateLimiter::try_acquire`]. When `Denied`, the
/// caller can use `retry_after_secs` for a `Retry-After: N` header
/// (RFC 6585 §4) so well-behaved clients back off correctly rather
/// than retrying in a tight loop.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AcquireResult {
    Allowed,
    Denied { retry_after_secs: u64 },
}

/// Single-bucket token-bucket limiter. `try_acquire` is cheap (one
/// mutex acquire, no allocations) so the auth path doesn't pay a real
/// cost for the check.
pub struct AuthRateLimiter {
    cfg: RateLimitConfig,
    bucket: Mutex<Bucket>,
}

impl AuthRateLimiter {
    pub fn new(cfg: RateLimitConfig) -> Self {
        Self {
            cfg,
            bucket: Mutex::new(Bucket {
                tokens: cfg.burst as f64,
                last_refill: Instant::now(),
            }),
        }
    }

    /// Consume one token if available. Returns `Denied` with a
    /// rounded-up seconds-until-refill so the caller can emit a
    /// `Retry-After` header.
    pub fn try_acquire(&self) -> AcquireResult {
        if self.cfg.per_sec == 0 {
            return AcquireResult::Allowed;
        }
        let now = Instant::now();
        let mut bucket = self.bucket.lock().expect("rate limiter mutex poisoned");
        let elapsed = now.duration_since(bucket.last_refill).as_secs_f64();
        bucket.tokens =
            (bucket.tokens + elapsed * f64::from(self.cfg.per_sec)).min(f64::from(self.cfg.burst));
        bucket.last_refill = now;
        if bucket.tokens >= 1.0 {
            bucket.tokens -= 1.0;
            AcquireResult::Allowed
        } else {
            // ceil((1 - tokens) / per_sec), minimum 1 — a `Retry-After: 0`
            // would tell clients to retry immediately, which is what we're
            // actively trying to discourage.
            let deficit = 1.0 - bucket.tokens;
            let wait_secs = (deficit / f64::from(self.cfg.per_sec)).ceil() as u64;
            AcquireResult::Denied {
                retry_after_secs: wait_secs.max(1),
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn disabled_limiter_always_allows() {
        let rl = AuthRateLimiter::new(RateLimitConfig {
            per_sec: 0,
            burst: 0,
        });
        for _ in 0..1000 {
            assert_eq!(rl.try_acquire(), AcquireResult::Allowed);
        }
    }

    #[test]
    fn burst_lets_through_initial_window_then_blocks() {
        // 0 refill, burst 3 → first three pass, fourth blocks.
        let rl = AuthRateLimiter::new(RateLimitConfig {
            per_sec: 1,
            burst: 3,
        });
        assert_eq!(rl.try_acquire(), AcquireResult::Allowed);
        assert_eq!(rl.try_acquire(), AcquireResult::Allowed);
        assert_eq!(rl.try_acquire(), AcquireResult::Allowed);
        match rl.try_acquire() {
            AcquireResult::Denied { retry_after_secs } => {
                // Bucket is at ~0 tokens, refill rate 1/sec → ~1s wait.
                assert!(
                    retry_after_secs >= 1,
                    "retry_after must be at least 1s, got {retry_after_secs}"
                );
            }
            AcquireResult::Allowed => panic!("fourth request must be denied"),
        }
    }

    #[test]
    fn tokens_refill_over_time() {
        // 10/sec → after ~120ms we should have at least one token back.
        let rl = AuthRateLimiter::new(RateLimitConfig {
            per_sec: 10,
            burst: 1,
        });
        assert_eq!(rl.try_acquire(), AcquireResult::Allowed);
        assert!(matches!(rl.try_acquire(), AcquireResult::Denied { .. }));
        std::thread::sleep(std::time::Duration::from_millis(150));
        assert_eq!(
            rl.try_acquire(),
            AcquireResult::Allowed,
            "token should have refilled"
        );
    }

    #[test]
    fn retry_after_scales_inversely_with_refill_rate() {
        // 1/sec → wait ~1s after burst exhausted.
        // 10/sec → wait <1s, but we clamp to a minimum of 1s.
        let slow = AuthRateLimiter::new(RateLimitConfig {
            per_sec: 1,
            burst: 1,
        });
        slow.try_acquire();
        match slow.try_acquire() {
            AcquireResult::Denied { retry_after_secs } => assert_eq!(retry_after_secs, 1),
            _ => panic!("expected Denied"),
        }
    }
}