ullun/
rate_limit.rs

1//! In-memory rate limiting middleware
2//!
3//! Fixed window per key (the peer IP by default). State lives in this process,
4//! so with several instances each one enforces the limit on its own; put a
5//! shared limiter (reverse proxy, Redis) in front for a global limit.
6//!
7//! ```no_run
8//! use std::time::Duration;
9//! use ullun::prelude::*;
10//! use ullun::rate_limit::RateLimiter;
11//!
12//! # async fn run() {
13//! App::new()
14//!     .middleware(RateLimiter::new(100, Duration::from_secs(60)))
15//!     .get("/", |_req: Request| async { Ok(Response::text("ok")) })
16//!     .run("127.0.0.1:3000")
17//!     .await
18//!     .unwrap();
19//! # }
20//! ```
21
22use crate::error::{Error, Result};
23use crate::middleware::{Middleware, Next};
24use crate::request::Request;
25use crate::response::Response;
26use async_trait::async_trait;
27use std::collections::HashMap;
28use std::sync::{Arc, Mutex};
29use std::time::{Duration, Instant};
30
31type KeyFn = dyn Fn(&Request) -> Option<String> + Send + Sync;
32
33struct Window {
34    started: Instant,
35    count: u32,
36}
37
38/// Outcome of counting one request
39#[derive(Debug, Clone, Copy, PartialEq, Eq)]
40pub struct Decision {
41    pub allowed: bool,
42    pub limit: u32,
43    pub remaining: u32,
44    /// Seconds until the current window ends (at least 1)
45    pub reset_secs: u64,
46}
47
48/// Fixed-window rate limiter, usable as middleware
49pub struct RateLimiter {
50    limit: u32,
51    window: Duration,
52    max_keys: usize,
53    key_fn: Arc<KeyFn>,
54    state: Mutex<HashMap<String, Window>>,
55}
56
57impl RateLimiter {
58    /// Allow `limit` requests per `window` for each client IP
59    pub fn new(limit: u32, window: Duration) -> Self {
60        Self {
61            limit,
62            window,
63            max_keys: 100_000,
64            key_fn: Arc::new(|req| req.remote_addr.map(|a| a.ip().to_string())),
65            state: Mutex::new(HashMap::new()),
66        }
67    }
68
69    /// Choose the bucket per request, e.g. an API key header. Returning `None`
70    /// exempts the request from limiting.
71    pub fn key_by<F>(mut self, f: F) -> Self
72    where
73        F: Fn(&Request) -> Option<String> + Send + Sync + 'static,
74    {
75        self.key_fn = Arc::new(f);
76        self
77    }
78
79    /// Use the first `X-Forwarded-For` address, falling back to the peer IP.
80    ///
81    /// Only enable this behind a reverse proxy that sets the header: otherwise
82    /// clients choose their own bucket and can evade the limit.
83    pub fn trust_forwarded_for(self) -> Self {
84        self.key_by(|req| {
85            req.headers
86                .get("x-forwarded-for")
87                .and_then(|v| v.to_str().ok())
88                .and_then(|v| v.split(',').next())
89                .map(|ip| ip.trim().to_string())
90                .filter(|ip| !ip.is_empty())
91                .or_else(|| req.remote_addr.map(|a| a.ip().to_string()))
92        })
93    }
94
95    /// Cap on tracked keys (default 100 000). When full after dropping expired
96    /// windows, unseen keys are let through untracked rather than blocked.
97    pub fn max_keys(mut self, max_keys: usize) -> Self {
98        self.max_keys = max_keys;
99        self
100    }
101
102    /// Count one request for `key` at time `now`
103    pub fn check(&self, key: &str, now: Instant) -> Decision {
104        let mut map = self.state.lock().unwrap_or_else(|e| e.into_inner());
105
106        if map.len() >= self.max_keys && !map.contains_key(key) {
107            let window = self.window;
108            map.retain(|_, w| now.duration_since(w.started) < window);
109            if map.len() >= self.max_keys {
110                return Decision {
111                    allowed: true,
112                    limit: self.limit,
113                    remaining: self.limit,
114                    reset_secs: self.window.as_secs().max(1),
115                };
116            }
117        }
118
119        let w = map.entry(key.to_string()).or_insert(Window {
120            started: now,
121            count: 0,
122        });
123        if now.duration_since(w.started) >= self.window {
124            *w = Window {
125                started: now,
126                count: 0,
127            };
128        }
129        let elapsed = now.duration_since(w.started);
130        let reset_secs = self
131            .window
132            .saturating_sub(elapsed)
133            .as_secs_f64()
134            .ceil()
135            .max(1.0) as u64;
136
137        if w.count < self.limit {
138            w.count += 1;
139            Decision {
140                allowed: true,
141                limit: self.limit,
142                remaining: self.limit - w.count,
143                reset_secs,
144            }
145        } else {
146            Decision {
147                allowed: false,
148                limit: self.limit,
149                remaining: 0,
150                reset_secs,
151            }
152        }
153    }
154}
155
156#[async_trait]
157impl Middleware for RateLimiter {
158    async fn handle(&self, req: Request, next: Next) -> Result<Response> {
159        // CORS preflights carry no credentials and must not eat the budget
160        if req.method == http::Method::OPTIONS {
161            return next.run(req).await;
162        }
163        let Some(key) = (self.key_fn)(&req) else {
164            return next.run(req).await;
165        };
166
167        let decision = self.check(&key, Instant::now());
168        let res = if decision.allowed {
169            match next.run(req).await {
170                Ok(res) => res,
171                Err(err) => Response::from_error(err),
172            }
173        } else {
174            Response::from_error(Error::custom(429, "Too many requests"))
175                .with_header("retry-after", decision.reset_secs.to_string())
176        };
177
178        Ok(res
179            .with_header("x-ratelimit-limit", decision.limit.to_string())
180            .with_header("x-ratelimit-remaining", decision.remaining.to_string())
181            .with_header("x-ratelimit-reset", decision.reset_secs.to_string()))
182    }
183}
184
185#[cfg(test)]
186mod tests {
187    use super::*;
188
189    #[test]
190    fn counts_within_window_and_resets_after() {
191        let rl = RateLimiter::new(2, Duration::from_secs(10));
192        let t0 = Instant::now();
193        let d1 = rl.check("a", t0);
194        assert_eq!((d1.allowed, d1.remaining), (true, 1));
195        assert_eq!(rl.check("a", t0 + Duration::from_secs(1)).remaining, 0);
196
197        let blocked = rl.check("a", t0 + Duration::from_millis(2500));
198        assert!(!blocked.allowed);
199        assert_eq!((blocked.remaining, blocked.reset_secs), (0, 8));
200
201        let again = rl.check("a", t0 + Duration::from_secs(10));
202        assert!(again.allowed && again.remaining == 1);
203    }
204
205    #[test]
206    fn keys_are_independent() {
207        let rl = RateLimiter::new(1, Duration::from_secs(60));
208        let t = Instant::now();
209        assert!(rl.check("a", t).allowed);
210        assert!(!rl.check("a", t).allowed);
211        assert!(rl.check("b", t).allowed);
212    }
213
214    #[test]
215    fn full_table_drops_expired_then_fails_open() {
216        let rl = RateLimiter::new(1, Duration::from_secs(10)).max_keys(2);
217        let t = Instant::now();
218        rl.check("a", t);
219        rl.check("b", t);
220        // table full and nothing expired: untracked, allowed, not remembered
221        assert!(rl.check("c", t).allowed);
222        assert!(rl.check("c", t).allowed);
223        assert!(!rl.check("a", t).allowed);
224        // later the expired entries are swept and "c" gets tracked
225        let later = t + Duration::from_secs(11);
226        assert!(rl.check("c", later).allowed);
227        assert!(!rl.check("c", later).allowed);
228    }
229
230    #[test]
231    fn zero_limit_blocks_everything() {
232        let rl = RateLimiter::new(0, Duration::from_secs(1));
233        assert!(!rl.check("a", Instant::now()).allowed);
234    }
235}
ullun/rate_limit.rs

ullun/
rate_limit.rs
