Skip to main content

max / makenotwork

2.8 KB · 103 lines History Blame Raw
1 //! Per-IP authentication rate limiting.
2 //!
3 //! russh's `auth_rejection_time` only delays within a single connection.
4 //! Parallel connections bypass it. This module tracks failed auth attempts
5 //! per IP and rejects early when a threshold is exceeded.
6
7 use std::collections::HashMap;
8 use std::net::IpAddr;
9 use std::sync::Mutex;
10 use std::time::Instant;
11
12 const MAX_FAILURES: usize = 10;
13 const WINDOW_SECS: u64 = 60;
14 const PRUNE_THRESHOLD: usize = 1000;
15
16 pub(crate) struct AuthRateLimiter {
17 failures: Mutex<HashMap<IpAddr, Vec<Instant>>>,
18 }
19
20 impl AuthRateLimiter {
21 pub(crate) fn new() -> Self {
22 Self {
23 failures: Mutex::new(HashMap::new()),
24 }
25 }
26
27 /// Returns `true` if the IP is allowed to attempt auth.
28 /// Returns `false` if the IP has exceeded the failure threshold.
29 pub(crate) fn check(&self, ip: IpAddr) -> bool {
30 let mut map = self.failures.lock().unwrap();
31 let cutoff = Instant::now()
32 .checked_sub(std::time::Duration::from_secs(WINDOW_SECS))
33 .unwrap();
34
35 if let Some(times) = map.get_mut(&ip) {
36 times.retain(|t| *t > cutoff);
37 times.len() < MAX_FAILURES
38 } else {
39 true
40 }
41 }
42
43 /// Record a failed auth attempt for the given IP.
44 pub(crate) fn record_failure(&self, ip: IpAddr) {
45 let mut map = self.failures.lock().unwrap();
46
47 // Prune stale entries when map grows large
48 if map.len() > PRUNE_THRESHOLD {
49 let cutoff = Instant::now()
50 .checked_sub(std::time::Duration::from_secs(WINDOW_SECS))
51 .unwrap();
52 map.retain(|_, times| {
53 times.retain(|t| *t > cutoff);
54 !times.is_empty()
55 });
56 }
57
58 map.entry(ip).or_default().push(Instant::now());
59 }
60 }
61
62 #[cfg(test)]
63 mod tests {
64 use super::*;
65 use std::net::Ipv4Addr;
66
67 #[test]
68 fn allows_under_threshold() {
69 let limiter = AuthRateLimiter::new();
70 let ip = IpAddr::V4(Ipv4Addr::new(192, 168, 1, 1));
71
72 for _ in 0..MAX_FAILURES - 1 {
73 assert!(limiter.check(ip));
74 limiter.record_failure(ip);
75 }
76 assert!(limiter.check(ip));
77 }
78
79 #[test]
80 fn blocks_at_threshold() {
81 let limiter = AuthRateLimiter::new();
82 let ip = IpAddr::V4(Ipv4Addr::new(10, 0, 0, 1));
83
84 for _ in 0..MAX_FAILURES {
85 limiter.record_failure(ip);
86 }
87 assert!(!limiter.check(ip));
88 }
89
90 #[test]
91 fn independent_ips() {
92 let limiter = AuthRateLimiter::new();
93 let ip_a = IpAddr::V4(Ipv4Addr::new(10, 0, 0, 1));
94 let ip_b = IpAddr::V4(Ipv4Addr::new(10, 0, 0, 2));
95
96 for _ in 0..MAX_FAILURES {
97 limiter.record_failure(ip_a);
98 }
99 assert!(!limiter.check(ip_a));
100 assert!(limiter.check(ip_b));
101 }
102 }
103