zzstoatzz.io
slack status without the slack
status.zzstoatzz.io
hatk
statusphere
1use actix_web::HttpRequest;
2use std::collections::HashMap;
3use std::sync::{Arc, Mutex};
4use std::time::{Duration, Instant};
5
6#[derive(Clone)]
7pub struct RateLimiter {
8 buckets: Arc<Mutex<HashMap<String, TokenBucket>>>,
9 max_tokens: u32,
10 refill_rate: Duration,
11}
12
13struct TokenBucket {
14 tokens: u32,
15 last_refill: Instant,
16}
17
18impl RateLimiter {
19 pub fn new(max_tokens: u32, refill_rate: Duration) -> Self {
20 Self {
21 buckets: Arc::new(Mutex::new(HashMap::new())),
22 max_tokens,
23 refill_rate,
24 }
25 }
26
27 pub fn check_rate_limit(&self, key: &str) -> bool {
28 let mut buckets = self.buckets.lock().unwrap();
29 let now = Instant::now();
30
31 let bucket = buckets.entry(key.to_string()).or_insert(TokenBucket {
32 tokens: self.max_tokens,
33 last_refill: now,
34 });
35
36 // Refill tokens based on elapsed time
37 let elapsed = now.duration_since(bucket.last_refill);
38 let tokens_to_add = (elapsed.as_secs_f64() / self.refill_rate.as_secs_f64()
39 * self.max_tokens as f64) as u32;
40
41 if tokens_to_add > 0 {
42 bucket.tokens = (bucket.tokens + tokens_to_add).min(self.max_tokens);
43 bucket.last_refill = now;
44 }
45
46 // Check if we have tokens available
47 if bucket.tokens > 0 {
48 bucket.tokens -= 1;
49 true
50 } else {
51 false
52 }
53 }
54
55 pub fn get_client_key(req: &HttpRequest) -> String {
56 // Use IP address as the key for rate limiting
57 req.connection_info()
58 .realip_remote_addr()
59 .unwrap_or("unknown")
60 .to_string()
61 }
62}
63
64#[cfg(test)]
65mod tests {
66 use super::*;
67 use std::thread;
68
69 #[test]
70 fn test_rate_limiter_basic() {
71 let limiter = RateLimiter::new(5, Duration::from_secs(1));
72
73 // Should allow first 5 requests
74 for _ in 0..5 {
75 assert!(limiter.check_rate_limit("test_client"));
76 }
77
78 // 6th request should be blocked
79 assert!(!limiter.check_rate_limit("test_client"));
80 }
81
82 #[test]
83 fn test_rate_limiter_refill() {
84 let limiter = RateLimiter::new(2, Duration::from_millis(100));
85
86 // Use up tokens
87 assert!(limiter.check_rate_limit("test_client"));
88 assert!(limiter.check_rate_limit("test_client"));
89 assert!(!limiter.check_rate_limit("test_client"));
90
91 // Wait for refill
92 thread::sleep(Duration::from_millis(150));
93
94 // Should have tokens again
95 assert!(limiter.check_rate_limit("test_client"));
96 }
97
98 #[test]
99 fn test_rate_limiter_different_clients() {
100 let limiter = RateLimiter::new(1, Duration::from_secs(1));
101
102 // Different clients should have separate buckets
103 assert!(limiter.check_rate_limit("client1"));
104 assert!(limiter.check_rate_limit("client2"));
105
106 // But same client should be limited
107 assert!(!limiter.check_rate_limit("client1"));
108 assert!(!limiter.check_rate_limit("client2"));
109 }
110}