oris_evolution_network/
rate_limiter.rs

1use std::collections::{HashMap, VecDeque};
2use std::sync::Mutex;
3
4use chrono::Utc;
5use serde::{Deserialize, Serialize};
6
7#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
8pub struct PeerRateLimitConfig {
9    pub max_capsules_per_hour: usize,
10    pub window_secs: i64,
11}
12
13impl Default for PeerRateLimitConfig {
14    fn default() -> Self {
15        Self {
16            max_capsules_per_hour: 100,
17            window_secs: 3600,
18        }
19    }
20}
21
22pub struct PeerRateLimiter {
23    config: PeerRateLimitConfig,
24    windows: Mutex<HashMap<String, VecDeque<i64>>>,
25}
26
27impl PeerRateLimiter {
28    pub fn new(config: PeerRateLimitConfig) -> Self {
29        Self {
30            config,
31            windows: Mutex::new(HashMap::new()),
32        }
33    }
34
35    pub fn config(&self) -> &PeerRateLimitConfig {
36        &self.config
37    }
38
39    pub fn check(&self, peer_id: &str) -> bool {
40        self.check_at(peer_id, Utc::now().timestamp())
41    }
42
43    pub fn check_at(&self, peer_id: &str, timestamp: i64) -> bool {
44        let mut windows = self.windows.lock().unwrap();
45        let entry = windows.entry(peer_id.to_string()).or_default();
46        while let Some(front) = entry.front().copied() {
47            if timestamp - front < self.config.window_secs {
48                break;
49            }
50            entry.pop_front();
51        }
52
53        if entry.len() >= self.config.max_capsules_per_hour {
54            return false;
55        }
56
57        entry.push_back(timestamp);
58        true
59    }
60}
61
62#[cfg(test)]
63mod tests {
64    use super::*;
65
66    #[test]
67    fn rate_limiter_blocks_after_capacity() {
68        let limiter = PeerRateLimiter::new(PeerRateLimitConfig {
69            max_capsules_per_hour: 2,
70            window_secs: 3600,
71        });
72
73        assert!(limiter.check_at("peer-a", 10));
74        assert!(limiter.check_at("peer-a", 11));
75        assert!(!limiter.check_at("peer-a", 12));
76    }
77
78    #[test]
79    fn rate_limiter_allows_after_window_expires() {
80        let limiter = PeerRateLimiter::new(PeerRateLimitConfig {
81            max_capsules_per_hour: 1,
82            window_secs: 5,
83        });
84
85        assert!(limiter.check_at("peer-a", 10));
86        assert!(!limiter.check_at("peer-a", 11));
87        assert!(limiter.check_at("peer-a", 16));
88    }
89}
oris_evolution_network/rate_limiter.rs

oris_evolution_network/
rate_limiter.rs
