postrust_proxy/ratelimit/
limiter.rs

1//! Rate limiter with per-key tracking.
2
3use crate::config::RateLimitDefaults;
4use crate::ratelimit::TokenBucket;
5use dashmap::DashMap;
6use std::net::IpAddr;
7use std::sync::Arc;
8use std::time::{Duration, Instant};
9use tokio_util::sync::CancellationToken;
10use tracing::{debug, info};
11
12/// Key for rate limiting.
13#[derive(Clone, Debug, Hash, PartialEq, Eq)]
14pub enum RateLimitKey {
15    /// Rate limit by IP address
16    Ip(IpAddr),
17    /// Rate limit by custom header value
18    Header(String),
19    /// Rate limit by route ID
20    Route(uuid::Uuid),
21    /// Global rate limit
22    Global,
23}
24
25/// Entry in the rate limiter cache.
26struct RateLimitEntry {
27    bucket: TokenBucket,
28    last_access: Instant,
29}
30
31/// Rate limiter with per-key token buckets.
32pub struct RateLimiter {
33    /// Per-key token buckets
34    buckets: DashMap<RateLimitKey, RateLimitEntry>,
35    /// Default configuration
36    defaults: RateLimitDefaults,
37    /// Entry TTL for cleanup
38    entry_ttl: Duration,
39}
40
41impl RateLimiter {
42    /// Create a new rate limiter.
43    pub fn new(defaults: RateLimitDefaults) -> Self {
44        Self {
45            buckets: DashMap::new(),
46            defaults,
47            entry_ttl: Duration::from_secs(3600), // 1 hour TTL
48        }
49    }
50
51    /// Check if a request should be allowed.
52    ///
53    /// Returns `true` if the request is allowed, `false` if rate limited.
54    pub fn check(&self, key: RateLimitKey) -> bool {
55        // Convert requests per window to requests per second
56        let rps = self
57            .defaults
58            .requests
59            .checked_div(self.defaults.window_secs)
60            .unwrap_or(self.defaults.requests);
61        self.check_with_config(key, rps, self.defaults.burst)
62    }
63
64    /// Check with custom rate limit configuration.
65    pub fn check_with_config(&self, key: RateLimitKey, rps: u32, burst: u32) -> bool {
66        let mut entry = self.buckets.entry(key).or_insert_with(|| RateLimitEntry {
67            bucket: TokenBucket::new(burst as u64, rps as f64),
68            last_access: Instant::now(),
69        });
70
71        entry.last_access = Instant::now();
72        entry.bucket.try_acquire()
73    }
74
75    /// Get remaining tokens for a key (approximate).
76    pub fn remaining(&self, key: &RateLimitKey) -> Option<u64> {
77        self.buckets.get(key).map(|entry| entry.bucket.available())
78    }
79
80    /// Start background cleanup task.
81    pub async fn start_cleanup(self: Arc<Self>, cancel_token: CancellationToken) {
82        let cleanup_interval = Duration::from_secs(300); // 5 minutes
83
84        info!("Rate limiter cleanup task started");
85
86        loop {
87            tokio::select! {
88                _ = cancel_token.cancelled() => {
89                    info!("Rate limiter cleanup task stopped");
90                    break;
91                }
92                _ = tokio::time::sleep(cleanup_interval) => {
93                    self.cleanup_expired();
94                }
95            }
96        }
97    }
98
99    /// Remove expired entries.
100    fn cleanup_expired(&self) {
101        let now = Instant::now();
102        let before = self.buckets.len();
103
104        self.buckets
105            .retain(|_, entry| now.duration_since(entry.last_access) < self.entry_ttl);
106
107        let removed = before - self.buckets.len();
108        if removed > 0 {
109            debug!("Rate limiter cleanup: removed {} expired entries", removed);
110        }
111    }
112}
113
114impl Default for RateLimiter {
115    fn default() -> Self {
116        Self::new(RateLimitDefaults::default())
117    }
118}
119
120#[cfg(test)]
121mod tests {
122    use super::*;
123    use std::net::Ipv4Addr;
124
125    #[test]
126    fn test_rate_limiter_by_ip() {
127        let limiter = RateLimiter::new(RateLimitDefaults {
128            requests: 600, // 10 per second
129            window_secs: 60,
130            burst: 5,
131        });
132
133        let ip1 = RateLimitKey::Ip(IpAddr::V4(Ipv4Addr::new(192, 168, 1, 1)));
134        let ip2 = RateLimitKey::Ip(IpAddr::V4(Ipv4Addr::new(192, 168, 1, 2)));
135
136        // Each IP gets its own bucket
137        for _ in 0..5 {
138            assert!(limiter.check(ip1.clone()));
139            assert!(limiter.check(ip2.clone()));
140        }
141
142        // Both exhausted
143        assert!(!limiter.check(ip1.clone()));
144        assert!(!limiter.check(ip2.clone()));
145    }
146
147    #[test]
148    fn test_rate_limiter_remaining() {
149        let limiter = RateLimiter::new(RateLimitDefaults {
150            requests: 600, // 10 per second
151            window_secs: 60,
152            burst: 10,
153        });
154
155        let key = RateLimitKey::Global;
156        assert!(limiter.check(key.clone()));
157
158        let remaining = limiter.remaining(&key).unwrap();
159        assert!(remaining < 10);
160    }
161}
postrust_proxy/ratelimit/limiter.rs

postrust_proxy/ratelimit/
limiter.rs
