rust_webx_host/
rate_limit.rs

1//! Per-IP token-bucket rate-limiting middleware.
2//!
3//! Configure via `appsettings.json` `RateLimit` section or manually:
4//!
5//! ```ignore
6//! .use_middleware_with(|| {
7//!     Arc::new(RateLimitMiddleware::from_config(&options.rate_limit))
8//! })
9//! ```
10
11use crate::problem_response::write_problem;
12use rust_webx_core::config::RateLimitSection;
13use rust_webx_core::error::Result;
14use rust_webx_core::http::IHttpContext;
15use rust_webx_core::middleware::IMiddleware;
16use std::collections::HashMap;
17use std::net::IpAddr;
18use std::ops::ControlFlow;
19use std::str::FromStr;
20use std::time::Instant;
21use tokio::sync::Mutex;
22
23struct TokenBucket {
24    tokens: f64,
25    last_refill: Instant,
26}
27
28impl TokenBucket {
29    fn new(burst_size: f64) -> Self {
30        Self {
31            tokens: burst_size,
32            last_refill: Instant::now(),
33        }
34    }
35
36    fn try_consume(&mut self, rate: f64, burst: f64) -> bool {
37        let now = Instant::now();
38        let elapsed = now.duration_since(self.last_refill).as_secs_f64();
39        self.tokens = (self.tokens + elapsed * rate).min(burst);
40        self.last_refill = now;
41
42        if self.tokens >= 1.0 {
43            self.tokens -= 1.0;
44            true
45        } else {
46            false
47        }
48    }
49}
50
51/// Inner rate-limiting state shared behind a [`Mutex`].
52pub struct RateLimiter {
53    buckets: Mutex<HashMap<IpAddr, TokenBucket>>,
54    rate: f64,
55    burst: f64,
56    max_ips: usize,
57}
58
59impl RateLimiter {
60    pub fn new(requests_per_second: f64, burst_size: u32, max_tracked_ips: usize) -> Self {
61        Self {
62            buckets: Mutex::new(HashMap::new()),
63            rate: requests_per_second,
64            burst: burst_size as f64,
65            max_ips: max_tracked_ips.max(1),
66        }
67    }
68
69    async fn allow(&self, ip: IpAddr) -> bool {
70        let mut buckets = self.buckets.lock().await;
71        if !buckets.contains_key(&ip) && buckets.len() >= self.max_ips {
72            evict_oldest(&mut buckets);
73        }
74        let bucket = buckets
75            .entry(ip)
76            .or_insert_with(|| TokenBucket::new(self.burst));
77        bucket.try_consume(self.rate, self.burst)
78    }
79
80    /// Current number of tracked client IPs (for tests).
81    pub async fn tracked_ip_count(&self) -> usize {
82        self.buckets.lock().await.len()
83    }
84}
85
86fn evict_oldest(buckets: &mut HashMap<IpAddr, TokenBucket>) {
87    if buckets.is_empty() {
88        return;
89    }
90    let oldest = buckets
91        .iter()
92        .min_by_key(|(_, b)| b.last_refill)
93        .map(|(ip, _)| *ip);
94    if let Some(ip) = oldest {
95        buckets.remove(&ip);
96    }
97}
98
99pub struct RateLimitMiddleware {
100    limiter: RateLimiter,
101}
102
103impl RateLimitMiddleware {
104    pub fn new(requests_per_second: f64, burst_size: u32) -> Self {
105        Self::new_with_max_ips(requests_per_second, burst_size, 10_000)
106    }
107
108    pub fn new_with_max_ips(requests_per_second: f64, burst_size: u32, max_tracked_ips: usize) -> Self {
109        Self {
110            limiter: RateLimiter::new(requests_per_second, burst_size, max_tracked_ips),
111        }
112    }
113
114    pub fn from_config(cfg: &RateLimitSection) -> Self {
115        Self::new_with_max_ips(
116            cfg.requests_per_second,
117            cfg.burst_size,
118            cfg.max_tracked_ips,
119        )
120    }
121}
122
123#[async_trait::async_trait]
124impl IMiddleware for RateLimitMiddleware {
125    async fn invoke(&self, ctx: &mut dyn IHttpContext) -> Result<ControlFlow<()>> {
126        let ip = extract_client_ip(ctx);
127
128        if !self.limiter.allow(ip).await {
129            write_problem(ctx, 429, "Too Many Requests").await;
130            return Ok(ControlFlow::Break(()));
131        }
132
133        Ok(ControlFlow::Continue(()))
134    }
135}
136
137fn extract_client_ip(ctx: &dyn IHttpContext) -> IpAddr {
138    if let Some(fwd) = ctx.request().header("x-forwarded-for") {
139        let first = fwd.split(',').next().unwrap_or("").trim();
140        if let Ok(ip) = IpAddr::from_str(first) {
141            return ip;
142        }
143    }
144
145    if let Some(real) = ctx.request().header("x-real-ip") {
146        if let Ok(ip) = IpAddr::from_str(real.trim()) {
147            return ip;
148        }
149    }
150
151    IpAddr::from_str("127.0.0.1").unwrap()
152}
153
154#[cfg(test)]
155mod tests {
156    use super::*;
157    use std::net::Ipv4Addr;
158
159    #[tokio::test]
160    async fn rate_limiter_evicts_oldest_when_max_ips_reached() {
161        let limiter = RateLimiter::new(100.0, 10, 2);
162        assert!(limiter.allow(IpAddr::V4(Ipv4Addr::new(1, 0, 0, 1))).await);
163        assert!(limiter.allow(IpAddr::V4(Ipv4Addr::new(2, 0, 0, 1))).await);
164        assert_eq!(limiter.tracked_ip_count().await, 2);
165
166        assert!(limiter.allow(IpAddr::V4(Ipv4Addr::new(3, 0, 0, 1))).await);
167        assert_eq!(limiter.tracked_ip_count().await, 2);
168    }
169}
rust_webx_host/rate_limit.rs

rust_webx_host/
rate_limit.rs
