zesdex/apps/infrastructure/src/middleware/rate_limit.rs

//! Simple in-memory rate limiter for Axum.

use std::collections::HashMap;
use std::sync::Mutex;

/// In-memory sliding-window rate limiter.
#[derive(Debug)]
pub struct RateLimiter {
    windows: Mutex<HashMap<String, Vec<i64>>>,
}

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

    pub fn check_rate_limit(
        &self,
        client_id: &str,
        max_requests: u32,
        window_secs: u64,
    ) -> anyhow::Result<bool> {
        let now = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap_or_default()
            .as_secs() as i64;

        let cutoff = now.saturating_sub(window_secs as i64);
        let mut windows = self
            .windows
            .lock()
            .map_err(|e| anyhow::anyhow!("rate limiter lock poisoned: {e}"))?;

        let timestamps = windows
            .entry(client_id.to_string())
            .or_insert_with(Vec::new);
        timestamps.retain(|&ts| ts >= cutoff);

        if timestamps.len() >= max_requests as usize {
            return Ok(false);
        }

        timestamps.push(now);
        Ok(true)
    }

    pub fn reset(&self) -> anyhow::Result<()> {
        let mut windows = self
            .windows
            .lock()
            .map_err(|e| anyhow::anyhow!("rate limiter lock poisoned: {e}"))?;
        windows.clear();
        Ok(())
    }
}

impl Default for RateLimiter {
    fn default() -> Self {
        Self::new()
    }
}