EventSnap/backend/src/services/rate_limiter.rs

use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};

/// Thread-safe sliding-window rate limiter backed by an in-memory HashMap.
/// Each key (e.g. `"join:{ip}"` or `"upload:{user_id}"`) tracks timestamps
/// of recent requests and rejects new ones once the window is full.
#[derive(Clone)]
pub struct RateLimiter {
    windows: Arc<Mutex<HashMap<String, Vec<Instant>>>>,
}

impl RateLimiter {
    pub fn new() -> Self {
        Self {
            windows: Arc::new(Mutex::new(HashMap::new())),
        }
    }

    /// Returns `Ok(())` if allowed, `Err(retry_after_secs)` if rate-limited.
    /// `retry_after_secs` is how long until the oldest slot in the window expires.
    ///
    /// This is deliberately the ONLY entry point. There used to be a `check()` wrapper
    /// returning a plain bool, and 7 of the 8 call sites used it and then hard-coded
    /// `None` for the response's `Retry-After` — so a throttled client was told to back
    /// off but never for how long. Forcing every caller through the `Result` makes the
    /// retry delay impossible to discard by accident.
    pub fn check_with_retry(
        &self,
        key: impl Into<String>,
        max: usize,
        window: Duration,
    ) -> Result<(), u64> {
        let now = Instant::now();
        let key = key.into();
        let mut map = self.windows.lock().unwrap();
        let timestamps = map.entry(key).or_default();
        timestamps.retain(|&t| now.duration_since(t) < window);
        if timestamps.len() < max {
            timestamps.push(now);
            Ok(())
        } else {
            // The oldest timestamp expires at oldest + window; compute remaining seconds
            let oldest = timestamps[0];
            let elapsed = now.duration_since(oldest);
            let remaining = window.saturating_sub(elapsed);
            Err(remaining.as_secs().max(1))
        }
    }

    /// Wipe every tracked window. Used by the test-mode truncate route so a previous
    /// test's accumulated counters don't bleed into the next test's rate-limit checks.
    pub fn clear(&self) {
        self.windows.lock().unwrap().clear();
    }

    /// Drop keys whose windows are empty after expiring old timestamps. Called from a
    /// background task (see [`crate::services::maintenance`]) so that long-lived
    /// processes don't accumulate one HashMap entry per IP that ever connected.
    ///
    /// Uses a conservative 24h ceiling — anything older than that is gone regardless
    /// of which endpoint's window it was tracked under (the longest window today is
    /// 24h for export downloads). If we ever add longer windows, raise this constant.
    pub fn prune(&self) {
        let now = Instant::now();
        let ceiling = Duration::from_secs(24 * 60 * 60);
        let mut map = self.windows.lock().unwrap();
        let before = map.len();
        map.retain(|_, ts| {
            ts.retain(|&t| now.duration_since(t) < ceiling);
            !ts.is_empty()
        });
        let dropped = before.saturating_sub(map.len());
        if dropped > 0 {
            tracing::debug!("rate limiter pruned {dropped} idle keys");
        }
    }
}

/// Extract the client IP from X-Forwarded-For or fall back to a provided socket
/// address string.
///
/// We take the **right-most** entry, not the left-most. Caddy is the sole ingress
/// and the app port is only `expose`d (never published), so the last hop Caddy
/// appends is the real client. A client can prepend arbitrary spoofed values to
/// the left of XFF to dodge throttles — those are ignored here. This assumes
/// exactly one trusted proxy (Caddy); revisit if that changes.
///
/// Pass the peer address as `fallback`, never a constant. Every caller used to pass
/// the literal `"unknown"`, so any request that arrived without XFF — i.e. anything
/// reaching the app directly rather than through Caddy — shared ONE bucket with every
/// other such request, turning the limiter into a self-inflicted global throttle.
pub fn client_ip(headers: &axum::http::HeaderMap, fallback: &str) -> String {
    headers
        .get("x-forwarded-for")
        .and_then(|v| v.to_str().ok())
        .and_then(|s| s.rsplit(',').next())
        .map(|s| s.trim().to_owned())
        .filter(|s| !s.is_empty())
        .unwrap_or_else(|| fallback.to_owned())
}

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

    const MIN: Duration = Duration::from_secs(60);

    #[test]
    fn allows_up_to_max_then_blocks() {
        let rl = RateLimiter::new();
        assert!(rl.check_with_retry("k", 3, MIN).is_ok());
        assert!(rl.check_with_retry("k", 3, MIN).is_ok());
        assert!(rl.check_with_retry("k", 3, MIN).is_ok());
        assert!(
            rl.check_with_retry("k", 3, MIN).is_err(),
            "the 4th request must be blocked"
        );
    }

    #[test]
    fn keys_are_independent() {
        let rl = RateLimiter::new();
        assert!(rl.check_with_retry("a", 1, MIN).is_ok());
        assert!(rl.check_with_retry("a", 1, MIN).is_err());
        assert!(
            rl.check_with_retry("b", 1, MIN).is_ok(),
            "a different key has its own window"
        );
    }

    #[test]
    fn window_slides_and_allows_again_after_expiry() {
        let rl = RateLimiter::new();
        let w = Duration::from_millis(40);
        assert!(rl.check_with_retry("k", 1, w).is_ok());
        assert!(rl.check_with_retry("k", 1, w).is_err());
        std::thread::sleep(Duration::from_millis(55));
        assert!(
            rl.check_with_retry("k", 1, w).is_ok(),
            "the slot should expire once the window passes"
        );
    }

    /// `retry_after` is not a "some number in range" — it is the time until the oldest slot
    /// in the window frees up, and it is surfaced to clients as the backoff they sleep for
    /// (see `upload-queue.ts`). Asserting only `(1..=60)` spans the entire reachable domain
    /// of a 60s window, so a hardcoded `Err(1)` would satisfy it while telling every client
    /// to hammer the server a second later. Pin the actual value.
    #[test]
    fn retry_after_is_the_remaining_window() {
        let rl = RateLimiter::new();

        // The slot was consumed just now, so essentially the whole window remains.
        // `as_secs()` truncates the sub-second remainder, so a 30s window reports 29.
        let w30 = Duration::from_secs(30);
        assert!(rl.check_with_retry("a", 1, w30).is_ok());
        let a = rl.check_with_retry("a", 1, w30).unwrap_err();
        assert_eq!(a, 29, "retry_after must be the remaining window, got {a}");

        // A different window must yield a different retry_after: no single constant can
        // satisfy both this and the assertion above.
        let w10 = Duration::from_secs(10);
        assert!(rl.check_with_retry("b", 1, w10).is_ok());
        let b = rl.check_with_retry("b", 1, w10).unwrap_err();
        assert_eq!(b, 9, "retry_after must scale with the window, got {b}");
    }

    #[test]
    fn retry_after_counts_down_as_the_window_elapses() {
        let rl = RateLimiter::new();
        let w = Duration::from_secs(30);
        assert!(rl.check_with_retry("k", 1, w).is_ok());
        let first = rl.check_with_retry("k", 1, w).unwrap_err();

        std::thread::sleep(Duration::from_millis(1200));
        let second = rl.check_with_retry("k", 1, w).unwrap_err();

        // A client that waits 1.2s must be told to wait ~1.2s less — otherwise the advertised
        // backoff is a constant, not a deadline.
        let shaved = first - second;
        assert!(
            (1..=2).contains(&shaved),
            "1.2s of waiting must shorten the advertised backoff by ~1s (got {first} then {second})"
        );
    }

    #[test]
    fn retry_after_floors_at_one_second() {
        let rl = RateLimiter::new();
        let w = Duration::from_millis(800);
        assert!(rl.check_with_retry("k", 1, w).is_ok());
        let retry = rl.check_with_retry("k", 1, w).unwrap_err();
        // The sub-second remainder truncates to 0; clients must never be told "retry in 0s"
        // (that's a busy-loop). The `.max(1)` floor is what prevents it.
        assert_eq!(
            retry, 1,
            "a sub-second remainder must floor to 1, got {retry}"
        );
    }

    #[test]
    fn clear_resets_every_window() {
        let rl = RateLimiter::new();
        assert!(rl.check_with_retry("k", 1, MIN).is_ok());
        assert!(rl.check_with_retry("k", 1, MIN).is_err());
        rl.clear();
        assert!(
            rl.check_with_retry("k", 1, MIN).is_ok(),
            "clear() must free the window"
        );
    }

    /// `prune()` is a memory-leak guard: without it a long-lived process keeps one HashMap
    /// entry per IP that ever connected. Nothing in the public API observes the map size, so
    /// the only way to catch a no-op body (`fn prune(&self) {}`) is to look at the map — the
    /// tests module can see the private field.
    #[test]
    fn prune_drops_keys_whose_windows_have_fully_expired() {
        let rl = RateLimiter::new();

        // A key whose only timestamp is older than the 24h ceiling. We can't sleep for a day,
        // so backdate the Instant directly.
        let ancient = Instant::now()
            .checked_sub(Duration::from_secs(25 * 60 * 60))
            .expect("backdating an Instant by 25h");
        rl.windows
            .lock()
            .unwrap()
            .insert("stale".to_string(), vec![ancient]);

        // ...alongside a key that is still inside its window.
        assert!(rl.check_with_retry("live", 5, MIN).is_ok());
        assert_eq!(rl.windows.lock().unwrap().len(), 2);

        rl.prune();

        let map = rl.windows.lock().unwrap();
        assert!(
            !map.contains_key("stale"),
            "prune() must drop keys whose timestamps have all expired"
        );
        assert!(
            map.contains_key("live"),
            "prune() must keep keys that still have live timestamps"
        );
        assert_eq!(map.len(), 1, "exactly one key should survive the prune");
    }

    #[test]
    fn prune_does_not_reset_a_live_window() {
        // The counterpart to the test above: pruning must reclaim memory, never quota. If
        // prune() dropped live keys, every background sweep would hand attackers a fresh
        // budget.
        let rl = RateLimiter::new();
        assert!(rl.check_with_retry("k", 1, MIN).is_ok());
        assert!(rl.check_with_retry("k", 1, MIN).is_err());

        rl.prune();

        assert!(
            rl.check_with_retry("k", 1, MIN).is_err(),
            "prune() must not clear a window that is still active"
        );
    }

    #[test]
    fn client_ip_takes_rightmost_forwarded_for_entry() {
        // The right-most entry is the hop our trusted proxy (Caddy) appended.
        let mut h = HeaderMap::new();
        h.insert("x-forwarded-for", "10.0.0.1, 203.0.113.7".parse().unwrap());
        assert_eq!(client_ip(&h, "fallback"), "203.0.113.7");
    }

    #[test]
    fn client_ip_ignores_spoofed_leftmost_entry() {
        // A client prepending a fake IP to dodge throttles must not win.
        let mut h = HeaderMap::new();
        h.insert(
            "x-forwarded-for",
            "1.2.3.4, 9.9.9.9, 203.0.113.7".parse().unwrap(),
        );
        assert_eq!(client_ip(&h, "fallback"), "203.0.113.7");
    }

    #[test]
    fn client_ip_trims_surrounding_whitespace() {
        let mut h = HeaderMap::new();
        h.insert("x-forwarded-for", "  198.51.100.5  ".parse().unwrap());
        assert_eq!(client_ip(&h, "fb"), "198.51.100.5");
    }

    #[test]
    fn client_ip_falls_back_when_header_absent() {
        assert_eq!(client_ip(&HeaderMap::new(), "127.0.0.1"), "127.0.0.1");
    }

    #[test]
    fn client_ip_falls_back_on_trailing_comma_empty_entry() {
        // A trailing comma leaves an empty right-most segment after trimming; the
        // `.filter(!is_empty)` must reject it and fall through to the fallback
        // rather than returning "" (which would collapse callers into one bucket).
        let mut h = HeaderMap::new();
        h.insert("x-forwarded-for", "203.0.113.7, ".parse().unwrap());
        assert_eq!(client_ip(&h, "127.0.0.1"), "127.0.0.1");
    }
}