recursive/http/
rate_limit.rs

1//! Token-bucket rate limiting and metrics middleware.
2
3use axum::http::StatusCode;
4use std::collections::HashMap;
5use std::sync::atomic::Ordering;
6use std::sync::Arc;
7use std::time::Instant;
8use tokio::sync::Mutex;
9
10use super::Metrics;
11
12/// Token-bucket rate limiter keyed by client identifier (API key or remote IP).
13#[derive(Clone)]
14pub struct RateLimiter {
15    /// Tokens remaining per client key.
16    buckets: Arc<Mutex<HashMap<String, TokenBucket>>>,
17    /// Max tokens per bucket.
18    capacity: u32,
19    /// Tokens refilled per second.
20    refill_rate: f64,
21}
22
23/// A single token bucket for one client.
24struct TokenBucket {
25    tokens: f64,
26    last_refill: Instant,
27}
28
29impl RateLimiter {
30    /// Create a new rate limiter with the given capacity and refill rate.
31    ///
32    /// - `capacity`: maximum number of tokens (burst size).
33    /// - `refill_rate`: tokens added per second.
34    ///
35    /// External callers (tests, custom embedders) can construct a
36    /// `RateLimiter` directly and inject it via
37    /// [`build_router_with_auth_and_rate_limit`] when env-driven
38    /// configuration is undesirable.
39    pub fn new(capacity: u32, refill_rate: f64) -> Self {
40        Self {
41            buckets: Arc::new(Mutex::new(HashMap::new())),
42            capacity,
43            refill_rate,
44        }
45    }
46
47    /// Check if a request from `key` is allowed.
48    ///
49    /// Returns `true` if the request is within the rate limit, `false` if it
50    /// should be rejected (429).
51    async fn check(&self, key: &str) -> bool {
52        let mut buckets = self.buckets.lock().await;
53        let now = Instant::now();
54
55        let bucket = buckets.entry(key.to_string()).or_insert_with(|| {
56            // New client gets a full bucket
57            let tokens = self.capacity as f64;
58            TokenBucket {
59                tokens,
60                last_refill: now,
61            }
62        });
63
64        // Refill tokens based on elapsed time
65        let elapsed = now.duration_since(bucket.last_refill).as_secs_f64();
66        let refill = elapsed * self.refill_rate;
67        bucket.tokens = (bucket.tokens + refill).min(self.capacity as f64);
68        bucket.last_refill = now;
69
70        if bucket.tokens >= 1.0 {
71            bucket.tokens -= 1.0;
72            true
73        } else {
74            false
75        }
76    }
77}
78
79/// Build a `RateLimiter` from environment variables.
80///
81/// - `RECURSIVE_RATE_LIMIT_RPM`: requests per minute (default: 60)
82/// - `RECURSIVE_RATE_LIMIT_BURST`: burst capacity (default: 10)
83pub(super) fn rate_limiter_from_env() -> RateLimiter {
84    let rpm = std::env::var("RECURSIVE_RATE_LIMIT_RPM")
85        .ok()
86        .and_then(|v| v.parse::<f64>().ok())
87        .unwrap_or(60.0);
88    let burst = std::env::var("RECURSIVE_RATE_LIMIT_BURST")
89        .ok()
90        .and_then(|v| v.parse::<u32>().ok())
91        .unwrap_or(10);
92    // Convert RPM to per-second refill rate
93    let refill_rate = rpm / 60.0;
94    RateLimiter::new(burst, refill_rate)
95}
96
97/// Extract a client key from the request for rate limiting.
98///
99/// Uses the `X-API-Key` header if present, otherwise falls back to the
100/// remote IP address.
101pub(super) fn extract_client_key(req: &axum::extract::Request) -> String {
102    if let Some(api_key) = req.headers().get("x-api-key") {
103        if let Ok(key) = api_key.to_str() {
104            return format!("apikey:{}", key);
105        }
106    }
107    // Fall back to remote IP
108    req.extensions()
109        .get::<axum::extract::ConnectInfo<std::net::SocketAddr>>()
110        .map(|info| format!("ip:{}", info.ip()))
111        .unwrap_or_else(|| "ip:unknown".to_string())
112}
113
114/// Middleware that increments request counters.
115pub(super) async fn metrics_middleware(
116    axum::extract::State(metrics): axum::extract::State<Arc<Metrics>>,
117    req: axum::extract::Request,
118    next: axum::middleware::Next,
119) -> axum::response::Response {
120    metrics.requests_total.fetch_add(1, Ordering::Relaxed);
121    metrics.requests_active.fetch_add(1, Ordering::Relaxed);
122    let response = next.run(req).await;
123    metrics.requests_active.fetch_sub(1, Ordering::Relaxed);
124    response
125}
126
127/// Middleware that enforces rate limits on all API requests.
128pub(super) async fn rate_limit_middleware(
129    axum::extract::State(limiter): axum::extract::State<RateLimiter>,
130    req: axum::extract::Request,
131    next: axum::middleware::Next,
132) -> axum::response::Response {
133    let key = extract_client_key(&req);
134    if !limiter.check(&key).await {
135        let mut resp = axum::response::Response::new(axum::body::Body::from("rate limit exceeded"));
136        *resp.status_mut() = StatusCode::TOO_MANY_REQUESTS;
137        return resp;
138    }
139    next.run(req).await
140}
141
142#[cfg(test)]
143mod tests {
144    use super::*;
145    use std::time::Duration;
146
147    /// Helper: create a rate limiter with very small capacity for testing.
148    fn test_limiter(capacity: u32, rpm: f64) -> RateLimiter {
149        RateLimiter::new(capacity, rpm / 60.0)
150    }
151
152    #[tokio::test]
153    async fn test_requests_within_limit_succeed() {
154        let limiter = test_limiter(5, 60.0); // 5 burst, 60 RPM
155        for _ in 0..5 {
156            assert!(limiter.check("client-a").await, "request should be allowed");
157        }
158    }
159
160    #[tokio::test]
161    async fn test_requests_exceeding_limit_get_429() {
162        let limiter = test_limiter(3, 60.0); // 3 burst, 60 RPM
163        for _ in 0..3 {
164            assert!(limiter.check("client-b").await, "request should be allowed");
165        }
166        // Fourth request should be denied
167        assert!(
168            !limiter.check("client-b").await,
169            "request should be rate limited"
170        );
171    }
172
173    #[tokio::test]
174    async fn test_tokens_refill_after_waiting() {
175        let limiter = test_limiter(2, 60.0); // 2 burst, 60 RPM = 1 per second
176                                             // Exhaust the bucket
177        assert!(limiter.check("client-c").await);
178        assert!(limiter.check("client-c").await);
179        assert!(!limiter.check("client-c").await, "should be denied");
180
181        // Wait for refill (1 token per second, wait 1.1s to be safe)
182        tokio::time::sleep(Duration::from_millis(1100)).await;
183
184        // Should have at least 1 token now
185        assert!(
186            limiter.check("client-c").await,
187            "should be allowed after refill"
188        );
189    }
190
191    #[tokio::test]
192    async fn test_different_clients_have_independent_buckets() {
193        let limiter = test_limiter(2, 60.0); // 2 burst
194
195        // Exhaust client-d
196        assert!(limiter.check("client-d").await);
197        assert!(limiter.check("client-d").await);
198        assert!(
199            !limiter.check("client-d").await,
200            "client-d should be denied"
201        );
202
203        // client-e should still have a full bucket
204        assert!(
205            limiter.check("client-e").await,
206            "client-e should be allowed"
207        );
208        assert!(
209            limiter.check("client-e").await,
210            "client-e should be allowed"
211        );
212    }
213
214    #[tokio::test]
215    async fn test_extract_client_key_with_api_key() {
216        let req = axum::http::Request::builder()
217            .header("x-api-key", "test-key-123")
218            .body(axum::body::Body::empty())
219            .unwrap();
220        let key = extract_client_key(&req);
221        assert_eq!(key, "apikey:test-key-123");
222    }
223
224    #[tokio::test]
225    async fn test_extract_client_key_without_api_key() {
226        let req = axum::http::Request::builder()
227            .body(axum::body::Body::empty())
228            .unwrap();
229        let key = extract_client_key(&req);
230        // No ConnectInfo extension, so falls back to "ip:unknown"
231        assert_eq!(key, "ip:unknown");
232    }
233
234    #[tokio::test]
235    async fn test_rate_limiter_from_env_defaults() {
236        // Unset the env vars to test defaults
237        std::env::remove_var("RECURSIVE_RATE_LIMIT_RPM");
238        std::env::remove_var("RECURSIVE_RATE_LIMIT_BURST");
239        let limiter = rate_limiter_from_env();
240        // Default: 60 RPM, 10 burst
241        for _ in 0..10 {
242            assert!(limiter.check("default-client").await);
243        }
244        assert!(!limiter.check("default-client").await, "burst exceeded");
245    }
246}
recursive/http/rate_limit.rs

recursive/http/
rate_limit.rs
