rune_axum_ratelimit/

layer.rs

use http::{HeaderName, Request, Response, StatusCode};
use std::collections::HashMap;
use std::future::Future;
use std::pin::Pin;
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll};
use std::time::{Duration, Instant};
use tower::{Layer, Service};

/// Determines how a per-client key is extracted from the request.
///
/// The key is used to track and limit each client independently. If the configured
/// header is absent the request passes through without being counted.
///
/// # Examples
///
/// ```rust
/// use http::HeaderName;
/// use rune_axum_ratelimit::KeyExtractor;
///
/// // Rate-limit by the first IP in X-Forwarded-For (set by most reverse proxies)
/// let by_ip = KeyExtractor::XForwardedFor;
///
/// // Rate-limit by an API key sent in a custom header
/// let by_key = KeyExtractor::Header(HeaderName::from_static("x-api-key"));
/// ```
#[derive(Clone, Debug)]
pub enum KeyExtractor {
    /// Uses the first (leftmost) IP in the `X-Forwarded-For` header.
    ///
    /// > [!WARNING]
    /// > Clients can spoof this header unless your reverse proxy strips and re-adds it.
    /// > Configure your proxy to overwrite `X-Forwarded-For` with the true client IP.
    XForwardedFor,
    /// Uses the value of a named request header as the client key.
    ///
    /// Suitable for API key–based rate limiting: pass the header that carries the key,
    /// e.g. `HeaderName::from_static("x-api-key")`.
    Header(HeaderName),
}

#[derive(Debug)]
struct WindowState {
    count: u64,
    window_start: Instant,
}

/// Configuration for the rate limiter middleware.
///
/// Build with [`RateLimitConfig::new()`] and chain methods to set the request limit,
/// window duration, rejection status, and key extraction strategy. Then pass to
/// [`RateLimitLayer::new()`].
///
/// # Examples
///
/// ```rust
/// use std::time::Duration;
/// use http::StatusCode;
/// use rune_axum_ratelimit::RateLimitConfig;
///
/// let config = RateLimitConfig::new()
///     .requests(50)
///     .window(Duration::from_secs(30))
///     .status(StatusCode::TOO_MANY_REQUESTS);
/// ```
#[derive(Clone, Debug)]
pub struct RateLimitConfig {
    requests: u64,
    window: Duration,
    status: StatusCode,
    key: KeyExtractor,
}

impl Default for RateLimitConfig {
    fn default() -> Self {
        Self {
            requests: 100,
            window: Duration::from_secs(60),
            status: StatusCode::TOO_MANY_REQUESTS,
            key: KeyExtractor::XForwardedFor,
        }
    }
}

impl RateLimitConfig {
    /// Creates a `RateLimitConfig` with defaults: 100 req/60 s, `X-Forwarded-For` keying,
    /// `429 Too Many Requests` rejection.
    pub fn new() -> Self {
        Self::default()
    }

    /// Sets the maximum number of requests allowed per window per client.
    ///
    /// Defaults to `100`.
    pub fn requests(mut self, n: u64) -> Self {
        self.requests = n;
        self
    }

    /// Sets the duration of each fixed window.
    ///
    /// Defaults to `60` seconds.
    pub fn window(mut self, d: Duration) -> Self {
        self.window = d;
        self
    }

    /// Sets the HTTP status code returned when the limit is exceeded.
    ///
    /// Defaults to `429 Too Many Requests`.
    pub fn status(mut self, s: StatusCode) -> Self {
        self.status = s;
        self
    }

    /// Sets the strategy used to identify each client.
    ///
    /// Defaults to [`KeyExtractor::XForwardedFor`].
    pub fn key(mut self, k: KeyExtractor) -> Self {
        self.key = k;
        self
    }

    fn extract_key<B>(&self, req: &Request<B>) -> Option<String> {
        match &self.key {
            KeyExtractor::XForwardedFor => req
                .headers()
                .get("x-forwarded-for")
                .and_then(|v| v.to_str().ok())
                .map(|s| s.split(',').next().unwrap_or(s).trim().to_owned()),
            KeyExtractor::Header(name) => req
                .headers()
                .get(name)
                .and_then(|v| v.to_str().ok())
                .map(|s| s.to_owned()),
        }
    }
}

/// Tower [`Layer`] that applies fixed-window rate limiting per client.
///
/// All service clones produced by this layer share the same in-memory counter store via an
/// [`std::sync::Arc`], so the limit is enforced consistently across concurrent requests.
///
/// Apply with Axum's `.layer()` call. Use [`RateLimitLayer::default()`] for 100 req/60 s
/// per `X-Forwarded-For` IP, or [`RateLimitLayer::new()`] to supply a custom
/// [`RateLimitConfig`].
///
/// > [!NOTE]
/// > Requests that carry no extractable key (header absent) pass through uncounted.
/// > The counter store grows with the number of unique client keys and is never evicted;
/// > for long-running servers with many unique IPs consider restarting periodically or
/// > using a shared external store.
///
/// # Examples
///
/// ```rust,no_run
/// use axum::{routing::get, Router};
/// use rune_axum_ratelimit::RateLimitLayer;
///
/// let app: Router = Router::new()
///     .route("/api", get(|| async { "ok" }))
///     .layer(RateLimitLayer::default());
/// ```
#[derive(Clone, Debug)]
pub struct RateLimitLayer {
    config: RateLimitConfig,
    store: Arc<Mutex<HashMap<String, WindowState>>>,
}

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

impl RateLimitLayer {
    /// Creates a `RateLimitLayer` from a custom [`RateLimitConfig`].
    pub fn new(config: RateLimitConfig) -> Self {
        Self {
            config,
            store: Arc::new(Mutex::new(HashMap::new())),
        }
    }
}

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

    fn layer(&self, inner: S) -> Self::Service {
        RateLimitService {
            inner,
            config: self.config.clone(),
            store: Arc::clone(&self.store),
        }
    }
}

/// Tower [`Service`] produced by [`RateLimitLayer`].
#[derive(Clone, Debug)]
pub struct RateLimitService<S> {
    inner: S,
    config: RateLimitConfig,
    store: Arc<Mutex<HashMap<String, WindowState>>>,
}

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

    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 {
        // retry_after_secs: Some(n) means rate-limited with n-second retry hint
        let retry_after_secs: Option<u64> = self.config.extract_key(&req).and_then(|key| {
            let mut store = self.store.lock().expect("rate limit store lock poisoned");
            let now = Instant::now();
            let window = self.config.window;
            let limit = self.config.requests;

            let state = store.entry(key).or_insert_with(|| WindowState {
                count: 0,
                window_start: now,
            });

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

            if state.count >= limit {
                let remaining = window
                    .saturating_sub(now.duration_since(state.window_start))
                    .as_secs()
                    .max(1);
                return Some(remaining);
            }

            state.count += 1;
            None
        });

        if let Some(retry_secs) = retry_after_secs {
            let status = self.config.status;
            return Box::pin(async move {
                Ok(Response::builder()
                    .status(status)
                    .header("retry-after", retry_secs.to_string())
                    .body(ResBody::default())
                    .expect("error response is valid"))
            });
        }

        Box::pin(self.inner.call(req))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use axum::{body::Body, routing::get, Router};
    use http::StatusCode;
    use tower::ServiceExt;

    fn build_app(config: RateLimitConfig) -> Router {
        Router::new()
            .route("/", get(|| async { "ok" }))
            .layer(RateLimitLayer::new(config))
    }

    async fn send(app: Router, req: http::Request<Body>) -> http::Response<Body> {
        app.oneshot(req).await.unwrap()
    }

    fn req_with_ip(ip: &str) -> http::Request<Body> {
        http::Request::builder()
            .method("GET")
            .uri("/")
            .header("x-forwarded-for", ip)
            .body(Body::empty())
            .unwrap()
    }

    fn req_with_key(header: &str, value: &str) -> http::Request<Body> {
        http::Request::builder()
            .method("GET")
            .uri("/")
            .header(header, value)
            .body(Body::empty())
            .unwrap()
    }

    fn req_bare() -> http::Request<Body> {
        http::Request::builder()
            .method("GET")
            .uri("/")
            .body(Body::empty())
            .unwrap()
    }

    #[tokio::test]
    async fn passes_within_limit() {
        let app = build_app(RateLimitConfig::new().requests(3));
        for _ in 0..3 {
            let response = send(app.clone(), req_with_ip("1.2.3.4")).await;
            assert_eq!(response.status(), StatusCode::OK);
        }
    }

    #[tokio::test]
    async fn rejects_when_limit_exceeded() {
        let app = build_app(RateLimitConfig::new().requests(2));
        send(app.clone(), req_with_ip("1.2.3.4")).await;
        send(app.clone(), req_with_ip("1.2.3.4")).await;
        let response = send(app.clone(), req_with_ip("1.2.3.4")).await;
        assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS);
    }

    #[tokio::test]
    async fn retry_after_header_present_on_rejection() {
        let app = build_app(RateLimitConfig::new().requests(1));
        send(app.clone(), req_with_ip("1.2.3.4")).await;
        let response = send(app.clone(), req_with_ip("1.2.3.4")).await;
        assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS);
        let retry = response
            .headers()
            .get("retry-after")
            .and_then(|v| v.to_str().ok())
            .and_then(|s| s.parse::<u64>().ok())
            .unwrap();
        assert!(retry >= 1);
    }

    #[tokio::test]
    async fn different_ips_have_independent_counters() {
        let app = build_app(RateLimitConfig::new().requests(1));
        send(app.clone(), req_with_ip("1.1.1.1")).await; // exhausts 1.1.1.1
        let response = send(app.clone(), req_with_ip("2.2.2.2")).await; // 2.2.2.2 still fresh
        assert_eq!(response.status(), StatusCode::OK);
    }

    #[tokio::test]
    async fn zero_limit_rejects_immediately() {
        let app = build_app(RateLimitConfig::new().requests(0));
        let response = send(app, req_with_ip("1.2.3.4")).await;
        assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS);
    }

    #[tokio::test]
    async fn request_without_key_passes_uncounted() {
        let app = build_app(RateLimitConfig::new().requests(0)); // limit=0 would reject keyed reqs
        let response = send(app, req_bare()).await; // no X-Forwarded-For → no key → passes
        assert_eq!(response.status(), StatusCode::OK);
    }

    #[tokio::test]
    async fn header_key_extractor() {
        let config = RateLimitConfig::new()
            .requests(1)
            .key(KeyExtractor::Header(HeaderName::from_static("x-api-key")));
        let app = build_app(config);

        let response = send(app.clone(), req_with_key("x-api-key", "token-a")).await;
        assert_eq!(response.status(), StatusCode::OK);

        // Different key: independent counter
        let response = send(app.clone(), req_with_key("x-api-key", "token-b")).await;
        assert_eq!(response.status(), StatusCode::OK);

        // Same key again: now limited
        let response = send(app.clone(), req_with_key("x-api-key", "token-a")).await;
        assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS);
    }

    #[tokio::test]
    async fn custom_rejection_status() {
        let config = RateLimitConfig::new()
            .requests(0)
            .status(StatusCode::SERVICE_UNAVAILABLE);
        let response = send(build_app(config), req_with_ip("1.2.3.4")).await;
        assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
    }

    #[tokio::test]
    async fn window_resets_after_expiry() {
        let config = RateLimitConfig::new()
            .requests(1)
            .window(Duration::from_millis(50));
        let app = build_app(config);

        // First request: OK
        let r = send(app.clone(), req_with_ip("10.0.0.1")).await;
        assert_eq!(r.status(), StatusCode::OK);

        // Second request immediately: limited
        let r = send(app.clone(), req_with_ip("10.0.0.1")).await;
        assert_eq!(r.status(), StatusCode::TOO_MANY_REQUESTS);

        // Wait for window to expire
        tokio::time::sleep(Duration::from_millis(60)).await;

        // Now a new window: should pass again
        let r = send(app.clone(), req_with_ip("10.0.0.1")).await;
        assert_eq!(r.status(), StatusCode::OK);
    }

    #[tokio::test]
    async fn x_forwarded_for_uses_first_ip() {
        let app = build_app(RateLimitConfig::new().requests(1));
        // Multi-IP header: client IP is the leftmost
        send(
            app.clone(),
            http::Request::builder()
                .method("GET")
                .uri("/")
                .header("x-forwarded-for", "5.5.5.5, 10.0.0.1, 172.16.0.1")
                .body(Body::empty())
                .unwrap(),
        )
        .await;
        // Same leftmost IP should now be limited
        let response = send(
            app.clone(),
            http::Request::builder()
                .method("GET")
                .uri("/")
                .header("x-forwarded-for", "5.5.5.5, 192.168.1.1")
                .body(Body::empty())
                .unwrap(),
        )
        .await;
        assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS);
    }

    #[tokio::test]
    async fn default_layer_uses_100_req_per_60s() {
        let app = Router::new()
            .route("/", get(|| async { "ok" }))
            .layer(RateLimitLayer::default());

        // 100 requests should all pass
        for _ in 0..100 {
            let r = send(app.clone(), req_with_ip("9.9.9.9")).await;
            assert_eq!(r.status(), StatusCode::OK);
        }
        // The 101st should be rejected
        let r = send(app.clone(), req_with_ip("9.9.9.9")).await;
        assert_eq!(r.status(), StatusCode::TOO_MANY_REQUESTS);
    }
}