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

/// One key's recent hits, plus the window they were recorded under.
///
/// The `window` field is what makes pruning correct. `prune` used a single fixed 24 h ceiling for
/// every key, on the reasoning that 24 h is the longest window in use (export downloads) — but that
/// meant a `join:{ip}:{name}` key whose 60-SECOND window expired 23 hours ago was still retained.
/// Minting one costs a single 409 and no bcrypt, at 60/min per IP across three endpoints: roughly
/// 172,800 keys/day/IP, about 31 MB/day/IP inside a 1 GB container. The limiter became the
/// memory-exhaustion primitive it exists to prevent.
///
/// Storing the window per key makes the sweep drop each bucket as soon as ITS OWN window has
/// elapsed, which is also what the hot path already does on every check.
struct Bucket {
    hits: Vec<Instant>,
    window: Duration,
}

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 bucket = map.entry(key).or_insert_with(|| Bucket {
            hits: Vec::new(),
            window,
        });
        // A key's window can change under it when an admin edits the limit at runtime. Track the
        // current one so `prune` expires the bucket on the window actually in force.
        bucket.window = window;
        let timestamps = &mut bucket.hits;
        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.
            //
            // `first()`, not `[0]`: with `max == 0` the length check above is false even on an
            // empty vec, so indexing would panic — WHILE HOLDING THIS MUTEX. That poisons it
            // process-wide, so every subsequent `.lock().unwrap()` panics too: upload, feed,
            // join, recover, social, export and the hourly maintenance task all die, and only
            // a container restart brings them back. `max == 0` is not reachable through the
            // admin API (every numeric spec has min = 1) but a direct DB edit would do it, and
            // the blast radius does not justify the sharper syntax.
            let Some(&oldest) = timestamps.first() else {
                return Ok(());
            };
            let elapsed = now.duration_since(oldest);
            let remaining = window.saturating_sub(elapsed);
            Err(remaining.as_secs().max(1))
        }
    }

    /// Is `key` already at or above `max`, WITHOUT recording a hit?
    ///
    /// Needed by limiters whose budget is spent by an outcome rather than by the request — the
    /// per-IP failed-PIN ceiling charges only on a wrong PIN, so the gate at the top of the handler
    /// has to be able to ask "is this IP shut out?" without itself consuming the budget it guards.
    /// Using `check_with_retry` for that would charge every *successful* recovery too, and a venue
    /// full of guests legitimately recovering their own devices would lock itself out.
    ///
    /// Returns `Err(retry_after_secs)` when exhausted, mirroring `check_with_retry` so callers can
    /// build the same 429.
    pub fn peek(&self, key: &str, max: usize, window: Duration) -> Result<(), u64> {
        let now = Instant::now();
        let mut map = self.windows.lock().unwrap();
        let Some(bucket) = map.get_mut(key) else {
            return Ok(());
        };
        bucket
            .hits
            .retain(|&t| now.duration_since(t) < bucket.window);
        if bucket.hits.len() < max {
            return Ok(());
        }
        let Some(&oldest) = bucket.hits.first() else {
            return Ok(());
        };
        Err(window
            .saturating_sub(now.duration_since(oldest))
            .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.
    ///
    /// Expires each bucket against ITS OWN window (see [`Bucket`]), not one global ceiling. The
    /// previous fixed 24 h ceiling retained per-minute keys for a full day — ~172,800 keys/day/IP
    /// at 60/min across three endpoints, each mintable with a single 409 and no bcrypt.
    ///
    /// Holds the one global mutex for the length of the sweep, and that mutex is on the hot path of
    /// upload, feed, join, recover, social and export — so the retain does the cheap thing per
    /// bucket and nothing else. Correct pruning also keeps the map small enough that this stays
    /// cheap, which the old ceiling actively undermined.
    pub fn prune(&self) {
        let now = Instant::now();
        let mut map = self.windows.lock().unwrap();
        let before = map.len();
        map.retain(|_, bucket| {
            bucket
                .hits
                .retain(|&t| now.duration_since(t) < bucket.window);
            !bucket.hits.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 its own window. We can't sleep, so backdate
        // the Instant directly. A ONE-MINUTE window here on purpose: the old prune applied a flat
        // 24 h ceiling to every key, so this bucket — expired for over an hour of wall time —
        // survived the sweep. That is the leak (H3), and pinning it needs a short-window key.
        let ancient = Instant::now()
            .checked_sub(Duration::from_secs(90 * 60))
            .expect("backdating an Instant by 90 minutes");
        rl.windows.lock().unwrap().insert(
            "stale".to_string(),
            Bucket {
                hits: vec![ancient],
                window: MIN,
            },
        );

        // ...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");
    }
}