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::collections::HashMap;
use std::net::IpAddr;
use std::sync::Mutex;
use std::time::Instant;

/// Upper bound on distinct client IPs tracked at once, so a spray from many
/// spoofed/rotating source IPs can't grow the map without limit. When full,
/// idle (refilled-to-burst) buckets are pruned first; if none are idle a new
/// IP falls back to the shared global bucket for that request.
const MAX_TRACKED_IPS: usize = 10_000;

/// 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,
}

impl Bucket {
    fn new(burst: u32) -> Self {
        Self {
            tokens: f64::from(burst),
            last_refill: Instant::now(),
        }
    }

    /// Refill by elapsed time then try to consume one token.
    fn try_take(&mut self, cfg: &RateLimitConfig, now: Instant) -> AcquireResult {
        let elapsed = now.duration_since(self.last_refill).as_secs_f64();
        self.tokens =
            (self.tokens + elapsed * f64::from(cfg.per_sec)).min(f64::from(cfg.burst));
        self.last_refill = now;
        if self.tokens >= 1.0 {
            self.tokens -= 1.0;
            AcquireResult::Allowed
        } else {
            let deficit = 1.0 - self.tokens;
            let wait_secs = (deficit / f64::from(cfg.per_sec)).ceil() as u64;
            AcquireResult::Denied {
                retry_after_secs: wait_secs.max(1),
            }
        }
    }

    /// Whether this bucket has fully refilled (i.e. the client has been idle).
    /// Such buckets carry no state worth keeping — dropping and lazily
    /// recreating one yields an identical full bucket — so they're the safe
    /// eviction target under memory pressure.
    fn is_idle(&self, cfg: &RateLimitConfig, now: Instant) -> bool {
        let elapsed = now.duration_since(self.last_refill).as_secs_f64();
        (self.tokens + elapsed * f64::from(cfg.per_sec)) >= f64::from(cfg.burst)
    }
}

/// 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 },
}

/// Token-bucket limiter keyed by client IP, with a shared global bucket for
/// requests whose IP is unknown (no trusted proxy).
///
/// The old design was a single global bucket. That let one attacker at the
/// sustained rate deny auth to *every* user (the bucket stayed drained). With
/// the SvelteKit proxy now forwarding the peer IP (`X-Forwarded-For`, honoured
/// only when `AUTH_TRUSTED_PROXY` is set), each source IP gets its own bucket,
/// so an attacker only throttles themselves. When no trustworthy IP is
/// available the caller passes `None` and the shared `global` bucket applies —
/// exactly the previous behaviour.
pub struct AuthRateLimiter {
    cfg: RateLimitConfig,
    global: Mutex<Bucket>,
    per_ip: Mutex<HashMap<IpAddr, Bucket>>,
}

impl AuthRateLimiter {
    pub fn new(cfg: RateLimitConfig) -> Self {
        Self {
            cfg,
            global: Mutex::new(Bucket::new(cfg.burst)),
            per_ip: Mutex::new(HashMap::new()),
        }
    }

    /// Consume one token from the bucket for `key` (per-IP when `Some`, the
    /// shared global bucket when `None`). Returns `Denied` with a rounded-up
    /// seconds-until-refill so the caller can emit a `Retry-After` header.
    pub fn try_acquire(&self, key: Option<IpAddr>) -> AcquireResult {
        if self.cfg.per_sec == 0 {
            return AcquireResult::Allowed;
        }
        let now = Instant::now();
        let Some(ip) = key else {
            return self
                .global
                .lock()
                .expect("rate limiter mutex poisoned")
                .try_take(&self.cfg, now);
        };
        let mut map = self.per_ip.lock().expect("rate limiter mutex poisoned");
        if map.len() >= MAX_TRACKED_IPS && !map.contains_key(&ip) {
            map.retain(|_, b| !b.is_idle(&self.cfg, now));
            if map.len() >= MAX_TRACKED_IPS {
                // Still saturated with active attackers — degrade to the shared
                // bucket rather than grow unbounded. Worst case is the old
                // global-bucket behaviour under an extreme distributed flood.
                drop(map);
                return self
                    .global
                    .lock()
                    .expect("rate limiter mutex poisoned")
                    .try_take(&self.cfg, now);
            }
        }
        map.entry(ip)
            .or_insert_with(|| Bucket::new(self.cfg.burst))
            .try_take(&self.cfg, now)
    }
}

#[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(None), 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(None), AcquireResult::Allowed);
        assert_eq!(rl.try_acquire(None), AcquireResult::Allowed);
        assert_eq!(rl.try_acquire(None), AcquireResult::Allowed);
        match rl.try_acquire(None) {
            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(None), AcquireResult::Allowed);
        assert!(matches!(rl.try_acquire(None), AcquireResult::Denied { .. }));
        std::thread::sleep(std::time::Duration::from_millis(150));
        assert_eq!(
            rl.try_acquire(None),
            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(None);
        match slow.try_acquire(None) {
            AcquireResult::Denied { retry_after_secs } => assert_eq!(retry_after_secs, 1),
            _ => panic!("expected Denied"),
        }
    }

    fn ip(s: &str) -> Option<IpAddr> {
        Some(s.parse().unwrap())
    }

    #[test]
    fn per_ip_buckets_are_independent() {
        // The whole point of the fix: one IP draining its bucket must not deny
        // a different IP.
        let rl = AuthRateLimiter::new(RateLimitConfig {
            per_sec: 1,
            burst: 2,
        });
        let a = ip("203.0.113.7");
        let b = ip("198.51.100.9");
        assert_eq!(rl.try_acquire(a), AcquireResult::Allowed);
        assert_eq!(rl.try_acquire(a), AcquireResult::Allowed);
        assert!(matches!(rl.try_acquire(a), AcquireResult::Denied { .. }));
        // b is untouched.
        assert_eq!(rl.try_acquire(b), AcquireResult::Allowed);
        assert_eq!(rl.try_acquire(b), AcquireResult::Allowed);
        assert!(matches!(rl.try_acquire(b), AcquireResult::Denied { .. }));
    }

    #[test]
    fn none_key_shares_the_global_bucket_independent_of_per_ip() {
        let rl = AuthRateLimiter::new(RateLimitConfig {
            per_sec: 1,
            burst: 2,
        });
        // Drain the global (None) bucket.
        assert_eq!(rl.try_acquire(None), AcquireResult::Allowed);
        assert_eq!(rl.try_acquire(None), AcquireResult::Allowed);
        assert!(matches!(rl.try_acquire(None), AcquireResult::Denied { .. }));
        // A real IP is on its own bucket, unaffected by the drained global one.
        assert_eq!(rl.try_acquire(ip("203.0.113.1")), AcquireResult::Allowed);
    }

    #[test]
    fn tracked_ip_map_is_bounded() {
        let rl = AuthRateLimiter::new(RateLimitConfig {
            per_sec: 1,
            burst: 1,
        });
        // Distinct IPs beyond the cap must not grow the map without bound —
        // excess requests fall back to the shared bucket instead.
        for i in 0..(MAX_TRACKED_IPS as u64 + 50) {
            let octet_a = (i >> 8) as u8;
            let octet_b = (i & 0xff) as u8;
            let addr = format!("10.20.{octet_a}.{octet_b}");
            let _ = rl.try_acquire(Some(addr.parse().unwrap()));
        }
        assert!(rl.per_ip.lock().unwrap().len() <= MAX_TRACKED_IPS);
    }
}