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)]
64#[path = "rate_limit_tests.rs"]
65mod tests;
velesdb_server/rate_limit.rs

velesdb_server/
rate_limit.rs
