eric / rustwarden

//! Generic in-process rate limiter.
//!
//! Tracks request counts per arbitrary key over a sliding window. Used
//! internally by authentication, public share, and password flow rate
//! limiters, and exposed publicly for apps that need to rate-limit their
//! own endpoints (imports, webhooks, etc.).
//!
//! # Example
//!
//! ```no_run
//! use std::sync::Arc;
//! use rustwarden_core::services::rate_limit_service::RateLimitService;
//!
//! let limiter = Arc::new(RateLimitService::new());
//! // Reject after 10 calls / minute per IP:
//! limiter.check(&format!("import:{}", "1.2.3.4"), 10).unwrap();
//! ```
//!
//! # Limitations
//!
//! - In-process only. Multi-instance deployments need a shared backend
//!   (Redis) — out of scope for v1.
//! - Memory-bounded via `MAX_ENTRIES`. When the map is at cap and the
//!   incoming key is unknown, the least-recently-used key (oldest
//!   most-recent timestamp) is evicted so a legitimate new caller is still
//!   admitted. Unknown keys are never rejected solely because the map is
//!   full.

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

use crate::error::AppError;

/// Maximum distinct keys tracked before LRU-evicting the coldest entry (SEC-M2: memory DoS).
const MAX_ENTRIES: usize = 100_000;

/// Per-key counter plus the instant at which the current window opened.
#[derive(Debug, Clone, Copy)]
struct WindowCounter {
    count: u32,
    window_start: Instant,
}

/// Thread-safe rate limiter keyed by arbitrary strings.
///
/// Supports two semantics:
///
/// - [`check`]/[`check_window`]: fixed-window counter (resets at each window
///   boundary). Cheap, single integer per key.
/// - [`check_sliding_window`]: true sliding window (stores each attempt
///   timestamp). Stricter across boundaries; preferred for security-sensitive
///   flows like password reset throttling.
///
/// Uses a single `Mutex<HashMap>`; contention is acceptable because each
/// check is a handful of microseconds. A sharded `DashMap` backend can be
/// swapped in later without API change.
///
/// [`check`]: Self::check
/// [`check_window`]: Self::check_window
/// [`check_sliding_window`]: Self::check_sliding_window
#[derive(Clone, Default)]
pub struct RateLimitService {
    buckets: Arc<Mutex<HashMap<String, WindowCounter>>>,
    sliding: Arc<Mutex<HashMap<String, Vec<Instant>>>>,
}

impl RateLimitService {
    /// Create a new `RateLimitService` with empty tracking maps.
    pub fn new() -> Self {
        Self::default()
    }

    /// Check and increment the counter for `key`. Returns `Err(TooManyRequests)`
    /// with a `Retry-After` hint when `max_per_minute` is exceeded inside the
    /// rolling 60-second window.
    pub fn check(&self, key: &str, max_per_minute: u32) -> Result<(), AppError> {
        self.check_window(key, max_per_minute, Duration::from_secs(60))
    }

    /// Check and increment the counter for `key` with a custom window size.
    pub fn check_window(&self, key: &str, max: u32, window: Duration) -> Result<(), AppError> {
        self.check_window_capped(key, max, window, MAX_ENTRIES)
    }

    fn check_window_capped(
        &self,
        key: &str,
        max: u32,
        window: Duration,
        cap: usize,
    ) -> Result<(), AppError> {
        let now = Instant::now();
        let mut buckets = self
            .buckets
            .lock()
            .map_err(|_| AppError::Internal("rate limiter mutex poisoned".into()))?;

        // Drop expired entries so the map stays bounded on low-traffic keys.
        buckets.retain(|_, w| now.duration_since(w.window_start) < window);

        if !buckets.contains_key(key) {
            if buckets.len() >= cap {
                tracing::warn!(
                    entries = buckets.len(),
                    cap,
                    "rate limiter map full, evicting oldest window"
                );
            }
            evict_fixed_if_full(&mut buckets, cap);
        }

        let w = buckets.entry(key.to_string()).or_insert(WindowCounter {
            count: 0,
            window_start: now,
        });

        if now.duration_since(w.window_start) >= window {
            w.count = 0;
            w.window_start = now;
        }

        w.count += 1;
        if w.count > max {
            let elapsed = now.duration_since(w.window_start);
            let remaining = window.saturating_sub(elapsed).as_secs().max(1);
            tracing::warn!(
                key = %key,
                count = w.count,
                max,
                retry_after_secs = remaining,
                "rate limit exceeded"
            );
            return Err(AppError::TooManyRequests {
                retry_after_secs: remaining,
            });
        }

        Ok(())
    }

    /// True sliding-window check: rejects if more than `max` attempts occurred
    /// in the last `window` duration, counted attempt-by-attempt.
    ///
    /// Stricter than [`check_window`] because it never allows a burst across a
    /// window boundary. Use for security-sensitive throttling (password reset,
    /// email verification resend, etc.).
    ///
    /// [`check_window`]: Self::check_window
    pub fn check_sliding_window(
        &self,
        key: &str,
        max: u32,
        window: Duration,
    ) -> Result<(), AppError> {
        self.check_sliding_window_capped(key, max, window, MAX_ENTRIES)
    }

    fn check_sliding_window_capped(
        &self,
        key: &str,
        max: u32,
        window: Duration,
        cap: usize,
    ) -> Result<(), AppError> {
        let now = Instant::now();
        let mut sliding = self
            .sliding
            .lock()
            .map_err(|_| AppError::Internal("rate limiter mutex poisoned".into()))?;

        if !sliding.contains_key(key) {
            if sliding.len() >= cap {
                tracing::warn!(
                    entries = sliding.len(),
                    cap,
                    "sliding rate limiter map full, evicting least-recently-used key"
                );
            }
            evict_if_full(&mut sliding, cap, now);
        }

        let attempts = sliding.entry(key.to_string()).or_default();
        attempts.retain(|t| now.duration_since(*t) < window);

        if attempts.len() as u32 >= max {
            let oldest = attempts.first().copied().unwrap_or(now);
            let elapsed = now.duration_since(oldest);
            let retry = window.saturating_sub(elapsed).as_secs().max(1);
            tracing::warn!(
                key = %key,
                attempts = attempts.len(),
                max,
                retry_after_secs = retry,
                "sliding rate limit exceeded"
            );
            return Err(AppError::TooManyRequests {
                retry_after_secs: retry,
            });
        }

        attempts.push(now);
        Ok(())
    }

    /// Background task that periodically drops expired buckets.
    ///
    /// Not strictly required (every `check` already evicts stale entries for
    /// its own window), but useful when most keys stop being accessed yet the
    /// map should not keep growing through bursty traffic.
    pub async fn cleanup_loop(self, interval: Duration, window: Duration) {
        let mut ticker = tokio::time::interval(interval);
        ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
        loop {
            ticker.tick().await;
            let now = Instant::now();
            if let Ok(mut buckets) = self.buckets.lock() {
                buckets.retain(|_, w| now.duration_since(w.window_start) < window);
            }
            if let Ok(mut sliding) = self.sliding.lock() {
                sliding.retain(|_, attempts| {
                    attempts.retain(|t| now.duration_since(*t) < window);
                    !attempts.is_empty()
                });
            }
        }
    }

    /// Number of distinct keys currently tracked in the fixed-window map.
    pub fn len(&self) -> usize {
        self.buckets.lock().map(|b| b.len()).unwrap_or(0)
    }

    /// Return `true` if no keys are currently tracked in the fixed-window map.
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// Number of distinct keys currently tracked in the sliding-window map.
    pub fn sliding_len(&self) -> usize {
        self.sliding.lock().map(|s| s.len()).unwrap_or(0)
    }
}

/// Evict the sliding-window key whose most recent timestamp is the oldest
/// (approximate LRU) when `sliding` is at `cap`.
///
/// No-op when the map is under capacity. Empty timestamp lists are treated as
/// "just seen" (`now`) so they are not preferred for eviction.
fn evict_if_full(sliding: &mut HashMap<String, Vec<Instant>>, cap: usize, now: Instant) {
    if sliding.len() < cap {
        return;
    }
    if let Some(oldest) = sliding
        .iter()
        .min_by_key(|(_, v)| v.last().copied().unwrap_or(now))
        .map(|(k, _)| k.clone())
    {
        sliding.remove(&oldest);
    }
}

/// Evict the fixed-window key with the oldest `window_start` when `buckets`
/// is at `cap`. No-op when under capacity.
fn evict_fixed_if_full(buckets: &mut HashMap<String, WindowCounter>, cap: usize) {
    if buckets.len() < cap {
        return;
    }
    if let Some(oldest) = buckets
        .iter()
        .min_by_key(|(_, w)| w.window_start)
        .map(|(k, _)| k.clone())
    {
        buckets.remove(&oldest);
    }
}

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

    #[test]
    fn under_limit_always_allowed() {
        let rl = RateLimitService::new();
        for _ in 0..5 {
            rl.check("k", 10).unwrap();
        }
    }

    #[test]
    fn exceeding_limit_returns_too_many_requests() {
        let rl = RateLimitService::new();
        for _ in 0..3 {
            rl.check("k", 3).unwrap();
        }
        let err = rl.check("k", 3).unwrap_err();
        match err {
            AppError::TooManyRequests { retry_after_secs } => {
                assert!((1..=60).contains(&retry_after_secs));
            }
            other => panic!("expected TooManyRequests, got {other:?}"),
        }
    }

    #[test]
    fn distinct_keys_have_independent_budgets() {
        let rl = RateLimitService::new();
        for _ in 0..3 {
            rl.check("a", 3).unwrap();
            rl.check("b", 3).unwrap();
        }
        assert!(rl.check("a", 3).is_err());
        assert!(rl.check("b", 3).is_err());
    }

    #[test]
    fn window_expiry_releases_budget() {
        let rl = RateLimitService::new();
        let window = Duration::from_millis(50);
        rl.check_window("k", 1, window).unwrap();
        assert!(rl.check_window("k", 1, window).is_err());

        std::thread::sleep(Duration::from_millis(80));
        rl.check_window("k", 1, window)
            .expect("window should reset after expiry");
    }

    #[test]
    fn sliding_window_rejects_burst_across_boundary() {
        let rl = RateLimitService::new();
        let window = Duration::from_millis(100);

        // 3 attempts inside the window — OK.
        rl.check_sliding_window("k", 3, window).unwrap();
        rl.check_sliding_window("k", 3, window).unwrap();
        rl.check_sliding_window("k", 3, window).unwrap();

        // 4th attempt inside the same window — rejected.
        assert!(rl.check_sliding_window("k", 3, window).is_err());

        // After the oldest attempt expires, one slot is released.
        std::thread::sleep(Duration::from_millis(120));
        rl.check_sliding_window("k", 3, window)
            .expect("slot should free after oldest attempt ages out");
    }

    #[test]
    fn sliding_and_fixed_window_are_independent() {
        let rl = RateLimitService::new();
        rl.check("k", 1).unwrap();
        // Fixed window exhausted.
        assert!(rl.check("k", 1).is_err());
        // Sliding window is a separate keyspace.
        rl.check_sliding_window("k", 1, Duration::from_secs(60))
            .expect("sliding map is independent of fixed map");
    }

    #[test]
    fn check_window_respects_max_per_minute_equivalence() {
        // check(k, N) must behave identically to check_window(k, N, 60s).
        let a = RateLimitService::new();
        let b = RateLimitService::new();
        for _ in 0..4 {
            a.check("k", 4).unwrap();
            b.check_window("k", 4, Duration::from_secs(60)).unwrap();
        }
        assert!(a.check("k", 4).is_err());
        assert!(b.check_window("k", 4, Duration::from_secs(60)).is_err());
    }

    #[test]
    fn cloned_limiters_share_buckets() {
        let a = RateLimitService::new();
        let b = a.clone();
        for _ in 0..3 {
            a.check_sliding_window("pwd_reset:t@x.test", 3, Duration::from_secs(3600))
                .unwrap();
        }
        assert!(
            b.check_sliding_window("pwd_reset:t@x.test", 3, Duration::from_secs(3600))
                .is_err()
        );
    }

    #[test]
    fn independent_new_does_not_share() {
        let a = RateLimitService::new();
        let b = RateLimitService::new();
        for _ in 0..3 {
            a.check_sliding_window("k", 3, Duration::from_secs(3600))
                .unwrap();
        }
        assert!(
            b.check_sliding_window("k", 3, Duration::from_secs(3600))
                .is_ok()
        );
    }

    #[test]
    fn evict_if_full_is_noop_under_cap() {
        let now = Instant::now();
        let mut sliding = HashMap::new();
        sliding.insert("a".into(), vec![now]);
        evict_if_full(&mut sliding, 8, now);
        assert_eq!(sliding.len(), 1);
        assert!(sliding.contains_key("a"));
    }

    #[test]
    fn evict_if_full_drops_key_with_oldest_last_timestamp() {
        let now = Instant::now();
        let mut sliding = HashMap::new();
        for i in 0..8u64 {
            sliding.insert(format!("k{i}"), vec![now + Duration::from_millis(i)]);
        }
        assert_eq!(sliding.len(), 8);

        evict_if_full(&mut sliding, 8, now + Duration::from_secs(1));
        assert_eq!(sliding.len(), 7);
        assert!(!sliding.contains_key("k0"));

        sliding.entry("k8".into()).or_default().push(now);
        assert!(sliding.contains_key("k8"));
        assert!(sliding.len() <= 8);
    }

    #[test]
    fn sliding_window_evicts_lru_instead_of_rejecting_new_key() {
        let rl = RateLimitService::new();
        let window = Duration::from_secs(3600);
        for i in 0..8 {
            rl.check_sliding_window_capped(&format!("k{i}"), 10, window, 8)
                .unwrap();
        }
        assert_eq!(rl.sliding_len(), 8);
        rl.check_sliding_window_capped("k8", 10, window, 8)
            .expect("new key must be admitted via LRU eviction");
        assert!(rl.sliding_len() <= 8);
    }

    #[test]
    fn sliding_window_does_not_evict_when_key_already_tracked() {
        let rl = RateLimitService::new();
        let window = Duration::from_secs(3600);
        for i in 0..8 {
            rl.check_sliding_window_capped(&format!("k{i}"), 2, window, 8)
                .unwrap();
        }
        rl.check_sliding_window_capped("k0", 2, window, 8)
            .expect("existing key at cap must keep its bucket");
        assert_eq!(rl.sliding_len(), 8);
    }

    #[test]
    fn evict_fixed_if_full_drops_oldest_window_start() {
        let now = Instant::now();
        let mut buckets = HashMap::new();
        for i in 0..8u64 {
            buckets.insert(
                format!("k{i}"),
                WindowCounter {
                    count: 1,
                    window_start: now + Duration::from_millis(i),
                },
            );
        }
        evict_fixed_if_full(&mut buckets, 8);
        assert_eq!(buckets.len(), 7);
        assert!(!buckets.contains_key("k0"));
        buckets.insert(
            "k8".into(),
            WindowCounter {
                count: 1,
                window_start: now,
            },
        );
        assert!(buckets.contains_key("k8"));
        assert!(buckets.len() <= 8);
    }

    #[test]
    fn fixed_window_evicts_instead_of_rejecting_new_key() {
        let rl = RateLimitService::new();
        let window = Duration::from_secs(60);
        for i in 0..8 {
            rl.check_window_capped(&format!("k{i}"), 10, window, 8)
                .unwrap();
        }
        assert_eq!(rl.len(), 8);
        rl.check_window_capped("k8", 10, window, 8)
            .expect("new key must be admitted via eviction");
        assert!(rl.len() <= 8);
    }
}