1use std::collections::HashMap;
45use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
46use std::sync::Mutex;
47use std::time::{Duration, Instant};
48
49#[derive(Debug, Clone, Copy, PartialEq, Eq)]
50pub struct Policy {
51 pub free_failures: u32,
54 pub window: Duration,
57 pub first_lockout: Duration,
59 pub max_lockout: Duration,
60}
61
62impl Policy {
63 pub const ENROLL: Policy = Policy {
65 free_failures: 5,
66 window: Duration::from_secs(10 * 60),
67 first_lockout: Duration::from_secs(60),
68 max_lockout: Duration::from_secs(60 * 60),
69 };
70
71 pub const LOGIN: Policy = Policy {
74 free_failures: 5,
75 window: Duration::from_secs(60),
76 first_lockout: Duration::from_secs(60),
77 max_lockout: Duration::from_secs(60 * 60),
78 };
79
80 pub const PASSKEY_LOGIN: Policy = Policy {
86 free_failures: 30,
87 window: Duration::from_secs(60),
88 first_lockout: Duration::from_secs(60),
89 max_lockout: Duration::from_secs(60 * 60),
90 };
91}
92
93pub const DEFAULT_CAPACITY: usize = 4096;
95
96#[derive(Debug, Clone, Copy)]
97struct Entry {
98 failures: u32,
99 last_failure: Instant,
100 locked_until: Option<Instant>,
101}
102
103impl Entry {
104 fn locked(&self, now: Instant) -> bool {
105 self.locked_until.is_some_and(|until| now < until)
106 }
107
108 fn forgotten(&self, policy: &Policy, now: Instant) -> bool {
111 let since = self
112 .locked_until
113 .map_or(self.last_failure, |u| u.max(self.last_failure));
114 !self.locked(now) && now.saturating_duration_since(since) >= policy.window
115 }
116}
117
118pub fn key(addr: IpAddr) -> IpAddr {
122 match addr.to_canonical() {
123 IpAddr::V4(v4) if v4.is_loopback() => IpAddr::V4(Ipv4Addr::LOCALHOST),
124 IpAddr::V4(v4) => IpAddr::V4(v4),
125 IpAddr::V6(v6) if v6.is_loopback() => IpAddr::V6(v6),
127 IpAddr::V6(v6) => {
128 let s = v6.segments();
129 IpAddr::V6(Ipv6Addr::new(s[0], s[1], s[2], s[3], 0, 0, 0, 0))
130 }
131 }
132}
133
134struct State {
137 entries: HashMap<IpAddr, Entry>,
138 overflow: Option<Entry>,
139}
140
141pub struct Limiter {
142 policy: Policy,
143 capacity: usize,
144 state: Mutex<State>,
145}
146
147impl Limiter {
148 pub fn new(policy: Policy) -> Self {
149 Self::with_capacity(policy, DEFAULT_CAPACITY)
150 }
151
152 pub fn with_capacity(policy: Policy, capacity: usize) -> Self {
153 Self {
154 policy,
155 capacity: capacity.max(1),
156 state: Mutex::new(State {
157 entries: HashMap::new(),
158 overflow: None,
159 }),
160 }
161 }
162
163 pub fn attempt(&self, addr: IpAddr, now: Instant) -> Result<(), Duration> {
167 let policy = self.policy;
168 let addr = key(addr);
169 let mut state = self.state.lock().expect("limiter lock");
170 if let Some(entry) = state.entries.get_mut(&addr) {
171 return record_attempt(entry, &policy, now);
172 }
173 state.entries.retain(|_, e| !e.forgotten(&policy, now));
177 if state.entries.len() < self.capacity || addr.is_loopback() {
178 let entry = state.entries.entry(addr).or_insert(Entry {
179 failures: 0,
180 last_failure: now,
181 locked_until: None,
182 });
183 return record_attempt(entry, &policy, now);
184 }
185 let overflow = state.overflow.get_or_insert(Entry {
188 failures: 0,
189 last_failure: now,
190 locked_until: None,
191 });
192 record_attempt(overflow, &policy, now)
193 }
194
195 pub fn succeeded(&self, addr: IpAddr) {
200 self.state.lock().expect("limiter lock").entries.remove(&key(addr));
201 }
202
203 pub fn clear(&self) {
205 let mut state = self.state.lock().expect("limiter lock");
206 state.entries.clear();
207 state.overflow = None;
208 }
209
210 pub fn tracked(&self) -> usize {
213 self.state.lock().expect("limiter lock").entries.len()
214 }
215}
216
217fn record_attempt(entry: &mut Entry, policy: &Policy, now: Instant) -> Result<(), Duration> {
223 if let Some(until) = entry.locked_until
224 && now < until
225 {
226 return Err(until - now);
227 }
228 if entry.forgotten(policy, now) {
229 entry.failures = 0;
230 entry.locked_until = None;
231 }
232 entry.failures += 1;
233 entry.last_failure = now;
234 if entry.failures >= policy.free_failures {
235 let doublings = (entry.failures - policy.free_failures).min(20);
236 let lockout = policy
237 .first_lockout
238 .saturating_mul(1 << doublings)
239 .min(policy.max_lockout);
240 entry.locked_until = Some(now + lockout);
241 }
242 Ok(())
243}
244
245#[cfg(test)]
246mod tests {
247 use super::*;
248 use std::sync::Arc;
249
250 const A: IpAddr = IpAddr::V4(std::net::Ipv4Addr::new(192, 0, 2, 1));
251 const B: IpAddr = IpAddr::V4(std::net::Ipv4Addr::new(192, 0, 2, 2));
252
253 fn secs(s: u64) -> Duration {
254 Duration::from_secs(s)
255 }
256
257 #[test]
258 fn the_fifth_attempt_locks_out_and_each_further_one_doubles_the_lockout() {
259 let limiter = Limiter::new(Policy::ENROLL);
260 let t0 = Instant::now();
261 for _ in 0..5 {
262 assert_eq!(limiter.attempt(A, t0), Ok(()));
263 }
264 assert_eq!(limiter.attempt(A, t0), Err(secs(60)));
265 assert_eq!(limiter.attempt(A, t0 + secs(59)), Err(secs(1)));
266 assert_eq!(limiter.attempt(A, t0 + secs(60)), Ok(()));
267 assert_eq!(limiter.attempt(A, t0 + secs(60)), Err(secs(120)));
268 assert_eq!(limiter.attempt(A, t0 + secs(180)), Ok(()));
269 assert_eq!(limiter.attempt(A, t0 + secs(180)), Err(secs(240)));
270 assert_eq!(limiter.attempt(B, t0), Ok(()));
272 }
273
274 #[test]
275 fn a_parallel_burst_gets_no_more_than_the_free_attempts() {
276 let limiter = Arc::new(Limiter::new(Policy::ENROLL));
277 let now = Instant::now();
278 let barrier = Arc::new(std::sync::Barrier::new(32));
279 let threads: Vec<_> = (0..32)
280 .map(|_| {
281 let limiter = limiter.clone();
282 let barrier = barrier.clone();
283 std::thread::spawn(move || {
284 barrier.wait();
285 limiter.attempt(A, now).is_ok()
286 })
287 })
288 .collect();
289 let allowed = threads.into_iter().map(|t| t.join().unwrap()).filter(|ok| *ok).count();
290 assert_eq!(allowed, 5);
291 }
292
293 #[test]
294 fn failures_older_than_the_window_are_forgotten() {
295 let limiter = Limiter::new(Policy::ENROLL);
296 let t0 = Instant::now();
297 for _ in 0..4 {
298 limiter.attempt(A, t0).unwrap();
299 }
300 let later = t0 + secs(10 * 60);
301 for _ in 0..4 {
302 assert_eq!(limiter.attempt(A, later), Ok(()));
303 }
304 assert_eq!(limiter.attempt(A, later), Ok(()));
305 assert!(limiter.attempt(A, later).is_err());
306 }
307
308 #[test]
309 fn a_success_clears_the_count_and_the_lockout_is_capped() {
310 let limiter = Limiter::new(Policy::ENROLL);
311 let t0 = Instant::now();
312 for _ in 0..4 {
313 limiter.attempt(A, t0).unwrap();
314 }
315 limiter.succeeded(A);
316 for _ in 0..4 {
317 assert_eq!(limiter.attempt(A, t0), Ok(()));
318 }
319 let mut t = t0;
322 let mut longest = Duration::ZERO;
323 for _ in 0..40 {
324 if let Err(wait) = limiter.attempt(A, t) {
325 longest = longest.max(wait);
326 t += wait;
327 }
328 }
329 assert_eq!(longest, secs(60 * 60));
330 }
331
332 #[test]
333 fn an_ipv6_64_is_one_address_and_a_mapped_ipv4_is_its_ipv4() {
334 let limiter = Limiter::new(Policy::ENROLL);
335 let t0 = Instant::now();
336 let v6 = |last: u16| IpAddr::V6(Ipv6Addr::new(0x2001, 0xdb8, 1, 2, 0, 0, 0, last));
337 for last in 1..=5 {
338 limiter.attempt(v6(last), t0).unwrap();
339 }
340 assert!(limiter.attempt(v6(99), t0).is_err(), "same /64");
341 let other_64 = IpAddr::V6(Ipv6Addr::new(0x2001, 0xdb8, 1, 3, 0, 0, 0, 1));
342 assert_eq!(limiter.attempt(other_64, t0), Ok(()));
343
344 for _ in 0..5 {
345 limiter.attempt(A, t0).unwrap();
346 }
347 let mapped = IpAddr::V6(Ipv6Addr::new(0, 0, 0, 0, 0, 0xffff, 0xc000, 0x0201));
348 assert_eq!(key(mapped), A);
349 assert!(limiter.attempt(mapped, t0).is_err());
350 }
351
352 #[test]
353 fn past_its_capacity_forgotten_entries_are_pruned_before_anything_else() {
354 let limiter = Limiter::with_capacity(Policy::ENROLL, 3);
355 let t0 = Instant::now();
356 let addr = |n: u8| IpAddr::V4(std::net::Ipv4Addr::new(198, 51, 100, n));
357 for n in 1..=3 {
358 limiter.attempt(addr(n), t0).unwrap();
359 }
360 assert_eq!(limiter.tracked(), 3);
361 let later = t0 + secs(10 * 60);
364 assert_eq!(limiter.attempt(addr(4), later), Ok(()));
365 assert_eq!(limiter.tracked(), 1);
366 }
367
368 #[test]
374 fn once_the_table_is_full_of_live_entries_new_addresses_share_one_overflow_budget() {
375 let limiter = Limiter::with_capacity(Policy::ENROLL, 3);
376 let t0 = Instant::now();
377 let live = |n: u8| IpAddr::V4(std::net::Ipv4Addr::new(203, 0, 113, n));
378 for n in 1..=3 {
379 assert_eq!(limiter.attempt(live(n), t0), Ok(()));
382 }
383 assert_eq!(limiter.tracked(), 3);
384
385 let flood = |n: u8| IpAddr::V4(std::net::Ipv4Addr::new(198, 51, 100, n));
388 for n in 1..=5 {
389 assert_eq!(limiter.attempt(flood(n), t0), Ok(()), "overflow attempt {n}");
390 }
391 assert_eq!(limiter.tracked(), 3);
394 assert_eq!(limiter.attempt(flood(6), t0), Err(secs(60)));
398
399 for n in 1..=3 {
405 for _ in 0..4 {
406 assert_eq!(
407 limiter.attempt(live(n), t0),
408 Ok(()),
409 "address {n} was reset by the flood"
410 );
411 }
412 assert!(
413 limiter.attempt(live(n), t0).is_err(),
414 "address {n} was reset by the flood (should be locked after 5 total failures)"
415 );
416 }
417 assert_eq!(limiter.tracked(), 3);
418 }
419
420 #[test]
425 fn loopback_keeps_one_entry_of_its_own_past_capacity() {
426 let limiter = Limiter::with_capacity(Policy::ENROLL, 1);
427 let t0 = Instant::now();
428 limiter.attempt(A, t0).unwrap();
429 for _ in 0..5 {
431 limiter.attempt(B, t0).unwrap();
432 }
433 assert!(limiter.attempt(B, t0).is_err(), "the overflow budget is used up");
434 let lo = |d: u8| IpAddr::V4(std::net::Ipv4Addr::new(127, d, 0, 1));
435 assert_eq!(key(lo(9)), IpAddr::V4(std::net::Ipv4Addr::LOCALHOST));
436 for d in 0..5 {
437 assert_eq!(
438 limiter.attempt(lo(d), t0),
439 Ok(()),
440 "loopback shared the overflow budget"
441 );
442 }
443 assert!(limiter.attempt(lo(200), t0).is_err(), "127/8 is one address");
444 assert_eq!(limiter.tracked(), 2);
445 assert_eq!(limiter.attempt(IpAddr::V6(Ipv6Addr::LOCALHOST), t0), Ok(()));
447 }
448
449 #[test]
454 fn locked_entries_are_never_evicted_even_by_a_large_flood() {
455 let limiter = Limiter::with_capacity(Policy::ENROLL, 2);
456 let t0 = Instant::now();
457 let locked = |n: u8| IpAddr::V4(std::net::Ipv4Addr::new(203, 0, 113, n));
458 for n in 1..=2 {
459 for _ in 0..5 {
460 limiter.attempt(locked(n), t0).unwrap();
461 }
462 assert!(limiter.attempt(locked(n), t0).is_err());
463 }
464 assert_eq!(limiter.tracked(), 2);
465
466 let flood = |n: u16| IpAddr::V4(std::net::Ipv4Addr::new(198, 51, 100, (n % 256) as u8));
467 for n in 1..=20 {
468 let _ = limiter.attempt(flood(n), t0);
469 assert_eq!(limiter.tracked(), 2, "a locked entry was evicted at flood attempt {n}");
470 }
471 for n in 1..=2 {
473 assert!(limiter.attempt(locked(n), t0).is_err(), "address {n} was reset");
474 }
475 }
476}