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    trust_proxy: bool,
58}
59
60impl RateLimiter {
61    pub fn new(requests_per_second: f64, burst_size: u32, max_tracked_ips: usize) -> Self {
62        Self {
63            buckets: Mutex::new(HashMap::new()),
64            rate: requests_per_second,
65            burst: burst_size as f64,
66            max_ips: max_tracked_ips.max(1),
67            trust_proxy: false,
68        }
69    }
70
71    pub fn with_trust_proxy(mut self, trust: bool) -> Self {
72        self.trust_proxy = trust;
73        self
74    }
75
76    async fn allow(&self, ip: IpAddr) -> bool {
77        let mut buckets = self.buckets.lock().await;
78        if !buckets.contains_key(&ip) && buckets.len() >= self.max_ips {
79            evict_oldest(&mut buckets);
80        }
81        let bucket = buckets
82            .entry(ip)
83            .or_insert_with(|| TokenBucket::new(self.burst));
84        bucket.try_consume(self.rate, self.burst)
85    }
86
87    /// Current number of tracked client IPs (for tests).
88    pub async fn tracked_ip_count(&self) -> usize {
89        self.buckets.lock().await.len()
90    }
91}
92
93fn evict_oldest(buckets: &mut HashMap<IpAddr, TokenBucket>) {
94    if buckets.is_empty() {
95        return;
96    }
97    let oldest = buckets
98        .iter()
99        .min_by_key(|(_, b)| b.last_refill)
100        .map(|(ip, _)| *ip);
101    if let Some(ip) = oldest {
102        buckets.remove(&ip);
103    }
104}
105
106pub struct RateLimitMiddleware {
107    limiter: RateLimiter,
108    trust_proxy: bool,
109}
110
111impl RateLimitMiddleware {
112    pub fn new(requests_per_second: f64, burst_size: u32) -> Self {
113        Self::new_with_max_ips(requests_per_second, burst_size, 10_000)
114    }
115
116    pub fn new_with_max_ips(requests_per_second: f64, burst_size: u32, max_tracked_ips: usize) -> Self {
117        Self {
118            limiter: RateLimiter::new(requests_per_second, burst_size, max_tracked_ips),
119            trust_proxy: false,
120        }
121    }
122
123    pub fn from_config(cfg: &RateLimitSection) -> Self {
124        Self {
125            limiter: RateLimiter::new(
126                cfg.requests_per_second,
127                cfg.burst_size,
128                cfg.max_tracked_ips,
129            ),
130            trust_proxy: cfg.trust_proxy,
131        }
132    }
133}
134
135#[async_trait::async_trait]
136impl IMiddleware for RateLimitMiddleware {
137    async fn invoke(&self, ctx: &mut dyn IHttpContext) -> Result<ControlFlow<()>> {
138        let ip = extract_client_ip(ctx, self.trust_proxy);
139
140        if !self.limiter.allow(ip).await {
141            write_problem(ctx, 429, "Too Many Requests").await;
142            return Ok(ControlFlow::Break(()));
143        }
144
145        Ok(ControlFlow::Continue(()))
146    }
147}
148
149fn extract_client_ip(ctx: &dyn IHttpContext, trust_proxy: bool) -> IpAddr {
150    if trust_proxy {
151        if let Some(fwd) = ctx.request().header("x-forwarded-for") {
152            let first = fwd.split(',').next().unwrap_or("").trim();
153            if let Ok(ip) = IpAddr::from_str(first) {
154                return ip;
155            }
156        }
157
158        if let Some(real) = ctx.request().header("x-real-ip") {
159            if let Ok(ip) = IpAddr::from_str(real.trim()) {
160                return ip;
161            }
162        }
163    }
164
165    IpAddr::from_str("127.0.0.1").unwrap()
166}
167
168#[cfg(test)]
169mod tests {
170    use super::*;
171    use std::net::Ipv4Addr;
172
173    #[tokio::test]
174    async fn rate_limiter_evicts_oldest_when_max_ips_reached() {
175        let limiter = RateLimiter::new(100.0, 10, 2);
176        assert!(limiter.allow(IpAddr::V4(Ipv4Addr::new(1, 0, 0, 1))).await);
177        assert!(limiter.allow(IpAddr::V4(Ipv4Addr::new(2, 0, 0, 1))).await);
178        assert_eq!(limiter.tracked_ip_count().await, 2);
179
180        assert!(limiter.allow(IpAddr::V4(Ipv4Addr::new(3, 0, 0, 1))).await);
181        assert_eq!(limiter.tracked_ip_count().await, 2);
182    }
183}
rust_webx_host/rate_limit.rs

rust_webx_host/
rate_limit.rs
