arc_web/helpers/
rate_limit.rs

1use std::collections::HashMap;
2use std::env;
3use std::sync::{Arc, Mutex};
4use std::time::{Duration, Instant};
5use tracing::info;
6
7#[derive(Clone)]
8pub struct RateLimiter {
9    max_requests: usize,
10    period: Duration,
11    requests: Arc<Mutex<HashMap<String, Vec<Instant>>>>,
12}
13
14/// Wrapper type for login-specific rate limiter to distinguish from global rate limiter
15#[derive(Clone)]
16pub struct LoginRateLimiter(pub RateLimiter);
17
18impl std::ops::Deref for LoginRateLimiter {
19    type Target = RateLimiter;
20    fn deref(&self) -> &Self::Target {
21        &self.0
22    }
23}
24
25impl RateLimiter {
26    pub fn new(max_requests: usize, period: Duration) -> Self {
27        Self {
28            max_requests,
29            period,
30            requests: Arc::new(Mutex::new(HashMap::new())),
31        }
32    }
33
34    /// Check if a key is allowed to make a request.
35    /// Returns Ok(()) if allowed, Err(Duration) with retry_after time if rate limited.
36    pub fn check(&self, key: String) -> Result<(), Duration> {
37        let now = Instant::now();
38        let mut requests = self.requests.lock().unwrap();
39
40        // Get or create entry for this key
41        let key_requests = requests.entry(key).or_default();
42
43        // Remove old requests outside the time window
44        key_requests.retain(|&time| now.duration_since(time) < self.period);
45
46        // Check if we've exceeded the limit
47        if key_requests.len() >= self.max_requests {
48            // Calculate retry_after duration
49            if let Some(&oldest) = key_requests.first() {
50                let retry_after = self.period.saturating_sub(now.duration_since(oldest));
51                return Err(retry_after);
52            }
53            return Err(self.period);
54        }
55
56        // Add this request
57        key_requests.push(now);
58        Ok(())
59    }
60}
61
62/// Create a rate limiter for login endpoints from environment configuration.
63pub fn create_rate_limiter() -> LoginRateLimiter {
64    let max_requests = env::var("RATE_LIMIT_MAX_REQUESTS")
65        .ok()
66        .and_then(|s| s.parse::<usize>().ok())
67        .unwrap_or(5);
68
69    let period_secs = env::var("RATE_LIMIT_PERIOD_SECS")
70        .ok()
71        .and_then(|s| s.parse::<u64>().ok())
72        .unwrap_or(60);
73
74    info!(
75        max_requests = max_requests,
76        period_secs = period_secs,
77        "Configuring login rate limiter"
78    );
79
80    LoginRateLimiter(RateLimiter::new(
81        max_requests,
82        Duration::from_secs(period_secs),
83    ))
84}
85
86/// Create a global rate limiter for all endpoints.
87pub fn create_global_rate_limiter() -> RateLimiter {
88    let max_requests = env::var("GLOBAL_RATE_LIMIT_MAX_REQUESTS")
89        .ok()
90        .and_then(|s| s.parse::<usize>().ok())
91        .unwrap_or(100);
92
93    let period_secs = env::var("GLOBAL_RATE_LIMIT_PERIOD_SECS")
94        .ok()
95        .and_then(|s| s.parse::<u64>().ok())
96        .unwrap_or(60);
97
98    info!(
99        max_requests = max_requests,
100        period_secs = period_secs,
101        "Configuring global rate limiter"
102    );
103
104    RateLimiter::new(max_requests, Duration::from_secs(period_secs))
105}
106
107#[cfg(test)]
108mod tests {
109    use super::*;
110    use std::thread::sleep;
111
112    #[test]
113    fn test_rate_limiter_allows_within_limit() {
114        let limiter = RateLimiter::new(3, Duration::from_secs(1));
115
116        assert!(limiter.check("test_key".to_string()).is_ok());
117        assert!(limiter.check("test_key".to_string()).is_ok());
118        assert!(limiter.check("test_key".to_string()).is_ok());
119    }
120
121    #[test]
122    fn test_rate_limiter_blocks_over_limit() {
123        let limiter = RateLimiter::new(2, Duration::from_secs(1));
124
125        assert!(limiter.check("test_key".to_string()).is_ok());
126        assert!(limiter.check("test_key".to_string()).is_ok());
127        assert!(limiter.check("test_key".to_string()).is_err());
128    }
129
130    #[test]
131    fn test_rate_limiter_resets_after_period() {
132        let limiter = RateLimiter::new(2, Duration::from_millis(100));
133
134        assert!(limiter.check("test_key".to_string()).is_ok());
135        assert!(limiter.check("test_key".to_string()).is_ok());
136        assert!(limiter.check("test_key".to_string()).is_err());
137
138        sleep(Duration::from_millis(150));
139
140        // Should be allowed again after period expires
141        assert!(limiter.check("test_key".to_string()).is_ok());
142    }
143
144    #[test]
145    fn test_rate_limiter_separate_keys() {
146        let limiter = RateLimiter::new(1, Duration::from_secs(1));
147
148        assert!(limiter.check("key1".to_string()).is_ok());
149        assert!(limiter.check("key2".to_string()).is_ok());
150
151        // Each key should be rate limited independently
152        assert!(limiter.check("key1".to_string()).is_err());
153        assert!(limiter.check("key2".to_string()).is_err());
154    }
155}
arc_web/helpers/rate_limit.rs

arc_web/helpers/
rate_limit.rs
