1use std::collections::HashMap;
8use std::net::IpAddr;
9use std::time::Instant;
10
11const MAX_ATTEMPTS: usize = 5;
13
14const WINDOW_SECS: u64 = 60;
16
17const BAN_DURATION_SECS: u64 = 300;
19
20#[derive(Debug, Clone)]
22struct RateLimitEntry {
23 attempts: usize,
25 window_start: Instant,
27 banned_until: Option<Instant>,
29}
30
31#[derive(Debug, Default)]
33pub struct RateLimiter {
34 entries: HashMap<IpAddr, RateLimitEntry>,
35}
36
37impl RateLimiter {
38 pub fn new() -> Self {
40 Self {
41 entries: HashMap::new(),
42 }
43 }
44
45 pub fn check_attempt(&mut self, ip: IpAddr) -> bool {
50 let now = Instant::now();
51 self.cleanup_expired(now);
52
53 let entry = self.entries.entry(ip).or_insert(RateLimitEntry {
54 attempts: 0,
55 window_start: now,
56 banned_until: None,
57 });
58
59 if let Some(banned_until) = entry.banned_until {
61 if now < banned_until {
62 log::warn!("Rejected auth attempt from banned IP: {}", ip);
63 return false;
64 }
65 entry.banned_until = None;
67 entry.attempts = 0;
68 entry.window_start = now;
69 }
70
71 if now.duration_since(entry.window_start).as_secs() > WINDOW_SECS {
73 entry.attempts = 0;
74 entry.window_start = now;
75 }
76
77 entry.attempts += 1;
78
79 if entry.attempts > MAX_ATTEMPTS {
81 entry.banned_until = Some(now + std::time::Duration::from_secs(BAN_DURATION_SECS));
82 log::warn!(
83 "IP {} banned for {} seconds after {} failed attempts",
84 ip,
85 BAN_DURATION_SECS,
86 entry.attempts
87 );
88 return false;
89 }
90
91 true
93 }
94
95 pub fn record_success(&mut self, ip: IpAddr) {
97 if let Some(entry) = self.entries.get_mut(&ip) {
98 entry.attempts = 0;
99 entry.banned_until = None;
100 }
101 }
102
103 fn cleanup_expired(&mut self, now: Instant) {
105 self.entries.retain(|_ip, entry| {
106 let expired = now.duration_since(entry.window_start).as_secs() > WINDOW_SECS + BAN_DURATION_SECS;
107 if let Some(banned_until) = entry.banned_until {
108 !expired || now < banned_until
109 } else {
110 !expired
111 }
112 });
113 }
114
115 pub fn tracked_count(&self) -> usize {
117 self.entries.len()
118 }
119}
120
121#[cfg(test)]
122mod tests {
123 use super::*;
124
125 fn test_ip() -> IpAddr {
126 "192.168.1.100".parse().unwrap()
127 }
128
129 #[test]
130 fn test_allows_first_five_attempts() {
131 let mut limiter = RateLimiter::new();
132 let ip = test_ip();
133
134 for _ in 0..5 {
135 assert!(limiter.check_attempt(ip));
136 }
137 }
138
139 #[test]
140 fn test_blocks_sixth_attempt() {
141 let mut limiter = RateLimiter::new();
142 let ip = test_ip();
143
144 for _ in 0..5 {
145 assert!(limiter.check_attempt(ip));
146 }
147 assert!(!limiter.check_attempt(ip));
149 }
150
151 #[test]
152 fn test_success_resets_counter() {
153 let mut limiter = RateLimiter::new();
154 let ip = test_ip();
155
156 for _ in 0..3 {
158 limiter.check_attempt(ip);
159 }
160 limiter.record_success(ip);
162
163 for _ in 0..5 {
165 assert!(limiter.check_attempt(ip));
166 }
167 }
168
169 #[test]
170 fn test_different_ips_independent() {
171 let mut limiter = RateLimiter::new();
172 let ip1: IpAddr = "10.0.0.1".parse().unwrap();
173 let ip2: IpAddr = "10.0.0.2".parse().unwrap();
174
175 for _ in 0..5 {
177 limiter.check_attempt(ip1);
178 }
179 assert!(!limiter.check_attempt(ip1));
180
181 assert!(limiter.check_attempt(ip2));
183 }
184}