aa_auth/
rate_limit.rs

1//! In-memory per-key rate limiter using token bucket algorithm.
2
3use std::time::Instant;
4
5use dashmap::DashMap;
6
7/// A token bucket that refills at a fixed rate.
8#[derive(Debug)]
9struct TokenBucket {
10    /// Current number of available tokens.
11    tokens: f64,
12    /// Maximum number of tokens (bucket capacity).
13    capacity: f64,
14    /// Tokens added per second.
15    refill_rate: f64,
16    /// Last time the bucket was refilled.
17    last_refill: Instant,
18}
19
20impl TokenBucket {
21    /// Create a new full bucket with the given capacity (requests per minute).
22    #[cfg(test)]
23    fn new(rpm: u32) -> Self {
24        Self::new_with_window(rpm, 60)
25    }
26
27    /// Create a new full bucket with an explicit refill window (seconds per full cycle).
28    fn new_with_window(rpm: u32, window_secs: u64) -> Self {
29        let capacity = f64::from(rpm);
30        let window = window_secs.max(1) as f64;
31        Self {
32            tokens: capacity,
33            capacity,
34            refill_rate: capacity / window,
35            last_refill: Instant::now(),
36        }
37    }
38
39    /// Try to consume one token. Returns `Ok(())` if allowed, or
40    /// `Err(seconds)` with the number of seconds until the next token
41    /// is available.
42    fn try_consume(&mut self) -> Result<(), u64> {
43        self.refill();
44        if self.tokens >= 1.0 {
45            self.tokens -= 1.0;
46            Ok(())
47        } else {
48            // Calculate seconds until one token is available.
49            let deficit = 1.0 - self.tokens;
50            let wait_secs = (deficit / self.refill_rate).ceil() as u64;
51            Err(wait_secs.max(1))
52        }
53    }
54
55    /// Refill tokens based on elapsed time since last refill.
56    fn refill(&mut self) {
57        let now = Instant::now();
58        let elapsed = now.duration_since(self.last_refill).as_secs_f64();
59        self.tokens = (self.tokens + elapsed * self.refill_rate).min(self.capacity);
60        self.last_refill = now;
61    }
62}
63
64/// Concurrent per-key rate limiter.
65///
66/// Each API key ID gets its own [`TokenBucket`]. Buckets are created
67/// on first access and cleaned up when stale.
68pub struct RateLimiter {
69    buckets: DashMap<String, TokenBucket>,
70    rpm: u32,
71    /// Refill window in seconds.  Production default: 60 (1 rpm = 1 token per
72    /// 60 s).  Override via `AA_RATE_LIMIT_WINDOW_SECS` env-var or by
73    /// constructing with [`RateLimiter::new_with_window`].
74    window_secs: u64,
75}
76
77impl RateLimiter {
78    /// Create a new rate limiter with the given per-key requests-per-minute limit.
79    ///
80    /// The refill window defaults to 60 seconds unless `AA_RATE_LIMIT_WINDOW_SECS`
81    /// is set in the environment.
82    pub fn new(rpm: u32) -> Self {
83        let window_secs = std::env::var("AA_RATE_LIMIT_WINDOW_SECS")
84            .ok()
85            .and_then(|v| v.parse::<u64>().ok())
86            .unwrap_or(60);
87        Self::new_with_window(rpm, window_secs)
88    }
89
90    /// Create a new rate limiter with an explicit refill window.
91    ///
92    /// Intended for tests that need a short window without setting an env-var.
93    pub fn new_with_window(rpm: u32, window_secs: u64) -> Self {
94        Self {
95            buckets: DashMap::new(),
96            rpm,
97            window_secs: window_secs.max(1),
98        }
99    }
100
101    /// Check whether a request from the given key ID is allowed.
102    ///
103    /// Returns `Ok(())` if the request is within the rate limit, or
104    /// `Err(retry_after_secs)` if the rate limit is exceeded.
105    pub fn check(&self, key_id: &str) -> Result<(), u64> {
106        let mut bucket = self
107            .buckets
108            .entry(key_id.to_string())
109            .or_insert_with(|| TokenBucket::new_with_window(self.rpm, self.window_secs));
110        bucket.try_consume()
111    }
112}
113
114#[cfg(test)]
115mod tests {
116    use super::*;
117
118    #[test]
119    fn test_token_bucket_allows_within_limit() {
120        let mut bucket = TokenBucket::new(100);
121        for _ in 0..100 {
122            assert!(bucket.try_consume().is_ok());
123        }
124    }
125
126    #[test]
127    fn test_token_bucket_rejects_over_limit() {
128        let mut bucket = TokenBucket::new(10);
129        for _ in 0..10 {
130            bucket.try_consume().unwrap();
131        }
132        let result = bucket.try_consume();
133        assert!(result.is_err(), "should reject when tokens exhausted");
134        let retry_after = result.unwrap_err();
135        assert!(retry_after >= 1, "retry_after should be at least 1 second");
136    }
137
138    #[test]
139    fn test_token_bucket_refills_over_time() {
140        let mut bucket = TokenBucket::new(60); // 1 token per second
141                                               // Exhaust all tokens.
142        for _ in 0..60 {
143            bucket.try_consume().unwrap();
144        }
145        assert!(bucket.try_consume().is_err());
146
147        // Simulate time passing by adjusting last_refill.
148        bucket.last_refill = Instant::now() - std::time::Duration::from_secs(2);
149        assert!(bucket.try_consume().is_ok(), "should have tokens after refill");
150    }
151
152    #[test]
153    fn test_rate_limiter_per_key_isolation() {
154        let limiter = RateLimiter::new(5);
155
156        // Exhaust key-a.
157        for _ in 0..5 {
158            limiter.check("key-a").unwrap();
159        }
160        assert!(limiter.check("key-a").is_err(), "key-a should be exhausted");
161
162        // key-b should still have tokens.
163        assert!(limiter.check("key-b").is_ok(), "key-b should be independent");
164    }
165}
aa_auth/rate_limit.rs

aa_auth/
rate_limit.rs
