doido_controller/
rate_limit.rs

1//! Fixed-window rate limiting (Rails 8 `rate_limit to:, within:`), backed by any
2//! [`doido_cache::CacheStore`].
3//!
4//! Each key counts requests within a window; the stored entry carries the window
5//! start, so the window resets correctly regardless of the cache's TTL support.
6
7use doido_cache::CacheStore;
8use std::sync::Arc;
9use std::time::{SystemTime, UNIX_EPOCH};
10
11/// Allow at most `limit` requests per `window_secs` per key.
12pub struct RateLimiter {
13    store: Arc<dyn CacheStore>,
14    limit: i64,
15    window_secs: u64,
16}
17
18impl RateLimiter {
19    pub fn new(store: Arc<dyn CacheStore>, limit: i64, window_secs: u64) -> Self {
20        Self {
21            store,
22            limit,
23            window_secs,
24        }
25    }
26
27    fn now() -> u64 {
28        SystemTime::now()
29            .duration_since(UNIX_EPOCH)
30            .map(|d| d.as_secs())
31            .unwrap_or(0)
32    }
33
34    /// Record a hit for `key`, returning `true` if it is within the limit.
35    pub async fn check(&self, key: &str) -> bool {
36        let full_key = format!("ratelimit:{key}");
37        let now = Self::now();
38
39        let entry = self
40            .store
41            .get(&full_key)
42            .await
43            .ok()
44            .flatten()
45            .and_then(|v| serde_json::from_value::<(i64, u64)>(v).ok());
46
47        // Reset the window if it has elapsed (or there is no entry yet).
48        let (count, start) = match entry {
49            Some((count, start)) if now.saturating_sub(start) < self.window_secs => (count, start),
50            _ => (0, now),
51        };
52
53        if count >= self.limit {
54            return false;
55        }
56
57        let value = serde_json::to_value((count + 1, start)).unwrap_or(serde_json::Value::Null);
58        let _ = self
59            .store
60            .set(&full_key, value, Some(self.window_secs))
61            .await;
62        true
63    }
64}
doido_controller/rate_limit.rs

doido_controller/
rate_limit.rs
