velesdb_server/
rate_limit.rs

1//! Global per-IP rate limiting middleware backed by `tower-governor`.
2//!
3//! Provides a token-bucket rate limiter keyed by client IP address.
4//! When the bucket is exhausted the server replies with `429 Too Many Requests`
5//! and standard rate-limit headers (`x-ratelimit-limit`, `x-ratelimit-remaining`,
6//! `retry-after`).
7
8use std::sync::Arc;
9use std::time::Duration;
10use tower_governor::governor::{GovernorConfig, GovernorConfigBuilder};
11use tower_governor::key_extractor::SmartIpKeyExtractor;
12
13/// Re-export so callers can build a `GovernorLayer` from the config.
14pub use tower_governor::GovernorLayer;
15
16/// The middleware type used when `use_headers()` is enabled.
17type HeaderMiddleware = ::governor::middleware::StateInformationMiddleware;
18
19/// Concrete governor config type with per-IP keying and rate-limit headers.
20pub type RateLimitConfig = GovernorConfig<SmartIpKeyExtractor, HeaderMiddleware>;
21
22/// Build a [`GovernorConfig`] that enforces `burst` requests/second per IP.
23///
24/// Uses [`SmartIpKeyExtractor`] which inspects `x-forwarded-for`,
25/// `x-real-ip`, `forwarded` headers before falling back to the peer IP,
26/// making it safe behind reverse proxies.
27///
28/// A background thread periodically prunes stale entries from the
29/// governor limiter map (every 60 s).
30///
31/// # Errors
32///
33/// Returns an error if the governor configuration cannot be built.
34pub fn build_rate_limit_config(burst: u32) -> anyhow::Result<Arc<RateLimitConfig>> {
35    let mut builder = GovernorConfigBuilder::default();
36    builder.per_second(u64::from(burst));
37    builder.burst_size(burst);
38    let mut builder = builder.key_extractor(SmartIpKeyExtractor);
39    let mut builder = builder.use_headers();
40
41    let config = Arc::new(
42        builder
43            .finish()
44            .ok_or_else(|| anyhow::anyhow!("failed to build rate limiter configuration"))?,
45    );
46
47    spawn_limiter_cleanup(&config);
48
49    Ok(config)
50}
51
52/// Spawns a background thread that prunes stale rate-limiter entries every 60 s.
53fn spawn_limiter_cleanup(config: &Arc<RateLimitConfig>) {
54    let limiter = config.limiter().clone();
55    let interval = Duration::from_secs(60);
56    std::thread::spawn(move || loop {
57        std::thread::sleep(interval);
58        tracing::debug!("rate limiter cleanup: {} tracked IPs", limiter.len());
59        limiter.retain_recent();
60    });
61}
62
63#[cfg(test)]
64mod tests {
65    use super::*;
66
67    #[test]
68    fn test_build_rate_limit_config_succeeds() {
69        let config = build_rate_limit_config(100);
70        assert!(
71            config.is_ok(),
72            "governor config should build with burst=100"
73        );
74    }
75
76    #[test]
77    fn test_build_rate_limit_config_burst_one() {
78        let config = build_rate_limit_config(1);
79        assert!(config.is_ok(), "governor config should build with burst=1");
80    }
81
82    /// BUG-1 regression: replenishment rate must scale with `burst`.
83    ///
84    /// With `per_second(burst)`, the full bucket refills within 1 second.
85    /// Previously, `per_second(1)` hardcoded 1 token/sec regardless of burst.
86    #[test]
87    fn test_rate_limit_replenishment_scales_with_burst() {
88        // burst=100 → per_second(100) → 100 tokens/sec replenishment.
89        // The limiter should allow at least 2 requests in rapid succession
90        // after one token has been consumed and a brief wait.
91        let config = build_rate_limit_config(100).expect("config should build");
92        let limiter = config.limiter();
93
94        // Simulate a key (IP address represented as a simple value).
95        let key = std::net::IpAddr::V4(std::net::Ipv4Addr::LOCALHOST);
96
97        // First request should always succeed (bucket starts full).
98        assert!(
99            limiter.check_key(&key).is_ok(),
100            "first request should succeed"
101        );
102
103        // With burst=100 and per_second(100), we should have 99 remaining
104        // tokens immediately — second request must also succeed.
105        assert!(
106            limiter.check_key(&key).is_ok(),
107            "second request should succeed (burst=100 allows many concurrent)"
108        );
109    }
110}
velesdb_server/rate_limit.rs

velesdb_server/
rate_limit.rs
