zesdex/crates/zesdex-middleware/src/rate_limit.rs

#![allow(
    clippy::cast_possible_truncation,
    clippy::cast_sign_loss,
    clippy::cast_precision_loss,
    clippy::cast_possible_wrap
)]
//! Simple in-memory rate limiter for Axum.
//!
//! Uses a sliding-window approach: each client has a rolling list of
//! timestamps.  Requests arriving within the window that exceed the
//! configured max are rejected.

use std::collections::HashMap;
use std::future::Future;
use std::pin::Pin;
use std::sync::Mutex;
use std::task::{Context, Poll};
use std::time::{SystemTime, UNIX_EPOCH};

use axum::http::{Request, StatusCode};
use axum::response::{IntoResponse, Response};
use tower::{Layer, Service};

/// In-memory sliding-window rate limiter.
///
/// Thread-safe via interior mutability (`Mutex`).  Each client (identified
/// by a string key, e.g. IP address or session id) has a Vec of entry
/// timestamps (in seconds).  Old entries are cleaned on every check.
#[derive(Debug)]
pub struct RateLimiter {
    windows: Mutex<HashMap<String, Vec<i64>>>,
}

impl RateLimiter {
    /// Create an empty rate limiter.
    pub fn new() -> Self {
        Self {
            windows: Mutex::new(HashMap::new()),
        }
    }

    /// Check whether a request from `client_id` should be allowed.
    ///
    /// * `max_requests` — max number of requests permitted within the
    ///   window.
    /// * `window_secs`  — width of the sliding window in seconds.
    ///
    /// Returns `Ok(true)` if the request is allowed (and records it),
    /// or `Ok(false)` if the client has exceeded the limit.
    ///
    /// The window is **sliding**: only timestamps falling within
    /// `[now - window_secs, now]` are counted.
    pub fn check_rate_limit(
        &self,
        client_id: &str,
        max_requests: u32,
        window_secs: u64,
    ) -> anyhow::Result<bool> {
        let now = SystemTime::now()
            .duration_since(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);

        // Discard entries older than the window.
        timestamps.retain(|&ts| ts >= cutoff);

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

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

    /// Convenience wrapper that returns an Axum [`Response`] on rejection
    /// so it can be used directly in middleware.
    pub fn check_or_429(
        &self,
        client_id: &str,
        max_requests: u32,
        window_secs: u64,
    ) -> Result<(), Response> {
        match self.check_rate_limit(client_id, max_requests, window_secs) {
            Ok(true) => Ok(()),
            Ok(false) => Err((
                StatusCode::TOO_MANY_REQUESTS,
                "rate limit exceeded, try again later",
            )
                .into_response()),
            Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e.to_string()).into_response()),
        }
    }

    /// Remove all stored windows (for testing / reset).
    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()
    }
}

// ---------------------------------------------------------------------------
// Tower Layer / Service
// ---------------------------------------------------------------------------

/// Configuration for the rate-limit middleware layer.
#[derive(Debug, Clone)]
pub struct RateLimitLayer {
    limiter: std::sync::Arc<RateLimiter>,
    max_requests: u32,
    window_secs: u64,
}

impl RateLimitLayer {
    /// Create a new layer with the given limits.
    ///
    /// * `max_requests` — max requests per window per client.
    /// * `window_secs`  — sliding-window width in seconds.
    pub fn new(max_requests: u32, window_secs: u64) -> Self {
        Self {
            limiter: std::sync::Arc::new(RateLimiter::new()),
            max_requests,
            window_secs,
        }
    }

    /// Return a reference to the shared [`RateLimiter`] so callers can
    /// reset it or perform manual checks.
    pub fn limiter(&self) -> &std::sync::Arc<RateLimiter> {
        &self.limiter
    }
}

impl<S> Layer<S> for RateLimitLayer {
    type Service = RateLimitMiddleware<S>;

    fn layer(&self, inner: S) -> Self::Service {
        RateLimitMiddleware {
            inner,
            limiter: std::sync::Arc::clone(&self.limiter),
            max_requests: self.max_requests,
            window_secs: self.window_secs,
        }
    }
}

/// Tower [`Service`] wrapping each request with a rate-limit check.
///
/// Client identity is extracted from the `X-Forwarded-For` header first,
/// falling back to the remote address, then to `"unknown"`.
#[derive(Debug, Clone)]
pub struct RateLimitMiddleware<S> {
    inner: S,
    limiter: std::sync::Arc<RateLimiter>,
    max_requests: u32,
    window_secs: u64,
}

impl<S, ReqBody> Service<Request<ReqBody>> for RateLimitMiddleware<S>
where
    S: Service<Request<ReqBody>, Response = Response> + Send + 'static,
    S::Future: Send + 'static,
    ReqBody: Send + 'static,
{
    type Response = S::Response;
    type Error = S::Error;
    type Future =
        Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send + 'static>>;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.inner.poll_ready(cx)
    }

    fn call(&mut self, req: Request<ReqBody>) -> Self::Future {
        let client_id = req
            .headers()
            .get("X-Forwarded-For")
            .and_then(|v| v.to_str().ok())
            .map(|s| s.split(',').next().unwrap_or(s).trim().to_string())
            .or_else(|| {
                req.headers()
                    .get("X-Real-IP")
                    .and_then(|v| v.to_str().ok())
                    .map(|s| s.to_string())
            })
            .or_else(|| {
                req.extensions()
                    .get::<axum::extract::ConnectInfo<std::net::SocketAddr>>()
                    .map(|ci| ci.0.ip().to_string())
            })
            .unwrap_or_else(|| "unknown".to_string());

        let limiter = std::sync::Arc::clone(&self.limiter);
        let max_requests = self.max_requests;
        let window_secs = self.window_secs;

        match limiter.check_or_429(&client_id, max_requests, window_secs) {
            Ok(()) => {}
            Err(resp) => {
                return Box::pin(async move { Ok(resp) });
            }
        }

        let fut = self.inner.call(req);
        Box::pin(async move { fut.await })
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

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

    #[test]
    fn test_rate_limiter_allows_within_limit() {
        let limiter = RateLimiter::new();
        assert!(limiter.check_rate_limit("client-1", 5, 60).unwrap());
        assert!(limiter.check_rate_limit("client-1", 5, 60).unwrap());
        assert!(limiter.check_rate_limit("client-1", 5, 60).unwrap());
    }

    #[test]
    fn test_rate_limiter_rejects_excess() {
        let limiter = RateLimiter::new();
        assert!(limiter.check_rate_limit("client-2", 3, 60).unwrap());
        assert!(limiter.check_rate_limit("client-2", 3, 60).unwrap());
        assert!(limiter.check_rate_limit("client-2", 3, 60).unwrap());
        assert!(!limiter.check_rate_limit("client-2", 3, 60).unwrap());
    }

    #[test]
    fn test_rate_limiter_independent_clients() {
        let limiter = RateLimiter::new();
        assert!(limiter.check_rate_limit("alice", 2, 60).unwrap());
        assert!(limiter.check_rate_limit("alice", 2, 60).unwrap());
        assert!(!limiter.check_rate_limit("alice", 2, 60).unwrap());
        assert!(limiter.check_rate_limit("bob", 2, 60).unwrap());
    }

    #[test]
    fn test_rate_limiter_reset() {
        let limiter = RateLimiter::new();
        assert!(limiter.check_rate_limit("client-3", 1, 60).unwrap());
        assert!(!limiter.check_rate_limit("client-3", 1, 60).unwrap());
        limiter.reset().unwrap();
        assert!(limiter.check_rate_limit("client-3", 1, 60).unwrap());
    }
}