shared_framework/middleware/

rate.rs

//! Fixed-window rate limiting keyed by caller and route.
//!
//! The [`RateLimiter`] trait answers whether a `key` is still within `limit` requests
//! for the current minute. [`InMemoryRateLimiter`] keeps per-minute buckets in process
//! memory, [`RedisBackedRateLimiter`] shares buckets in Redis via `INCR` plus a 60s
//! expiry (falling back to memory on Redis errors), and [`create_from_env`] picks the
//! Redis variant when `REDIS_URL` or `REDIS_URI` is set, else the in-memory one.
//! Routers build keys as `<user id or x-forwarded-for>:<route path>`.
//! ```ignore
//! if !limiter.is_allowed("user-1:/v1/users/list", 100).await { return Err(too_many()); }
//! ```

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

use crate::data::cache::RedisStorage;

/// Per-minute request gate: returns `true` when `key` is still within `limit`
/// requests for the current minute, `false` when the caller must be rejected.
#[async_trait::async_trait]
pub trait RateLimiter: Send + Sync {
    /// Returns whether `key` may proceed under `limit` requests for the current minute.
    async fn is_allowed(&self, key: &str, limit: u32) -> bool;
}

// ── In-memory ────────────────────────────────────────────────────────────────

/// In-process per-minute limiter: buckets are keyed by `"{key}:{minute}"` with a
/// 60-second window and expired buckets are pruned by periodic cleanup.
pub struct InMemoryRateLimiter {
    inner: Mutex<HashMap<String, (u32, Instant)>>,
    window: Duration,
}

impl InMemoryRateLimiter {
    /// Creates a limiter with a 60-second window and no background cleanup task.
    pub fn new() -> Self {
        Self {
            inner: Mutex::new(HashMap::new()),
            window: Duration::from_secs(60),
        }
    }

    /// Creates a limiter wrapped in `Arc` with a background task pruning expired buckets every 60s.
    /// Hold the returned `Arc`; the task exits when the last clone is dropped.
    pub fn new_with_cleanup() -> Arc<Self> {
        let arc = Arc::new(Self {
            inner: Mutex::new(HashMap::new()),
            window: Duration::from_secs(60),
        });
        let weak = Arc::downgrade(&arc);
        tokio::spawn(async move {
            let mut interval = tokio::time::interval(Duration::from_secs(60));
            loop {
                interval.tick().await;
                if let Some(limiter) = weak.upgrade() {
                    limiter.cleanup();
                } else {
                    break;
                }
            }
        });
        arc
    }

    /// Synchronous check-and-increment for `key` against `limit` in the current minute bucket.
    pub fn is_allowed_sync(&self, key: &str, limit: u32) -> bool {
        let minute = chrono::Utc::now().timestamp() / 60;
        let bucket_key = format!("{}:{}", key, minute);
        let mut map = self.inner.lock().unwrap();
        let now = Instant::now();
        let entry = map.entry(bucket_key).or_insert((0, now + self.window));
        if now > entry.1 {
            *entry = (1, now + self.window);
            return true;
        }
        if entry.0 < limit {
            entry.0 += 1;
            true
        } else {
            false
        }
    }

    fn cleanup(&self) {
        let now = Instant::now();
        let mut map = self.inner.lock().unwrap();
        map.retain(|_, (_, expiry)| *expiry > now);
    }
}

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

#[async_trait::async_trait]
impl RateLimiter for InMemoryRateLimiter {
    async fn is_allowed(&self, key: &str, limit: u32) -> bool {
        self.is_allowed_sync(key, limit)
    }
}

// ── Redis-backed ─────────────────────────────────────────────────────────────

/// Redis-shared per-minute limiter: increments `rate_limit:{key}:{minute}` with a 60s
/// expiry on first use. Uses the in-memory limiter whenever Redis is unconfigured or fails.
pub struct RedisBackedRateLimiter {
    storage: Option<RedisStorage>,
    fallback: Arc<InMemoryRateLimiter>,
}

impl RedisBackedRateLimiter {
    /// Builds a limiter using Redis configuration from the environment when available,
    /// otherwise operating on the in-memory fallback only.
    pub fn new() -> Self {
        let storage = RedisStorage::from_env().ok();
        Self {
            storage,
            fallback: InMemoryRateLimiter::new_with_cleanup(),
        }
    }

    /// Builds a limiter backed by the given Redis storage plus an in-memory fallback.
    pub fn with_storage(storage: RedisStorage) -> Self {
        Self {
            storage: Some(storage),
            fallback: InMemoryRateLimiter::new_with_cleanup(),
        }
    }

    /// Builds a limiter connecting to `url`, falling back to memory when unreachable or on errors.
    pub fn from_url(url: &str) -> Self {
        let storage = RedisStorage::new(url).ok();
        Self {
            storage,
            fallback: InMemoryRateLimiter::new_with_cleanup(),
        }
    }
}

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

#[async_trait::async_trait]
impl RateLimiter for RedisBackedRateLimiter {
    async fn is_allowed(&self, key: &str, limit: u32) -> bool {
        let Some(storage) = &self.storage else {
            return self.fallback.is_allowed(key, limit).await;
        };

        let minute = chrono::Utc::now().timestamp() / 60;
        let redis_key = format!("rate_limit:{}:{}", key, minute);

        // INCR
        match storage.increment_value(&redis_key).await {
            Ok(count) => {
                if count == 1 {
                    // First request in this minute — set 60s TTL.
                    let _ = storage.set_expiration(&redis_key, 60).await;
                    return true;
                }
                count <= limit as i64
            }
            Err(e) => {
                tracing::warn!(backend = "redis", fallback_backend = "in_memory", error = %e, "Rate limiter failed; using fallback");
                self.fallback.is_allowed(key, limit).await
            }
        }
    }
}

// ── Factory ──────────────────────────────────────────────────────────────────

/// Selects which limiter backend [`create`] builds.
pub enum LimiterKind {
    /// In-process memory buckets.
    InMemory,
    /// Redis-shared buckets with an in-memory fallback.
    Redis,
}

/// Builds the backend selected by `kind` as a shared trait object.
pub fn create(kind: LimiterKind) -> Arc<dyn RateLimiter> {
    match kind {
        LimiterKind::InMemory => InMemoryRateLimiter::new_with_cleanup(),
        LimiterKind::Redis => Arc::new(RedisBackedRateLimiter::new()),
    }
}

/// Builds a shared limiter from the environment: Redis-backed when `REDIS_URL` or
/// `REDIS_URI` is set, otherwise in-memory.
pub fn create_from_env() -> Arc<dyn RateLimiter> {
    if std::env::var("REDIS_URL").is_ok() || std::env::var("REDIS_URI").is_ok() {
        tracing::info!(component = "rate_limiter", backend = "redis", "Configured rate limiter");
        Arc::new(RedisBackedRateLimiter::new())
    } else {
        tracing::info!(component = "rate_limiter", backend = "in_memory", "Configured rate limiter");
        InMemoryRateLimiter::new_with_cleanup()
    }
}