ratelimit/
lib.rs

1//! Bounded sliding-window throttles for authentication attempts.
2//!
3//! Pure state plus an injected clock, so tests drive the time axis instead of
4//! sleeping. Deliberately in-process: a shared backend belongs behind
5//! [`Bucket`]'s interface, not woven through its callers, and until one exists
6//! the effective limit is `replicas × limit` — which is why the surfaces using
7//! this run single-replica.
8
9use std::collections::HashMap;
10use std::hash::Hash;
11use std::time::{Duration, Instant};
12
13#[cfg(test)]
14mod tests;
15
16/// "What time is it now", injected so tests need not sleep.
17pub trait Clock {
18    fn now(&self) -> Instant;
19}
20
21/// Delegates to [`Instant::now`].
22#[derive(Debug, Default, Clone, Copy)]
23pub struct SystemClock;
24
25impl Clock for SystemClock {
26    fn now(&self) -> Instant {
27        Instant::now()
28    }
29}
30
31#[derive(Debug, Clone, Copy, PartialEq, Eq)]
32pub enum Decision {
33    Allow,
34    Deny,
35}
36
37impl Decision {
38    #[must_use]
39    pub const fn is_denied(self) -> bool {
40        matches!(self, Self::Deny)
41    }
42}
43
44/// Throttle shape: how many attempts, over what window, across how many keys.
45#[derive(Debug, Clone, Copy)]
46pub struct Limits {
47    /// Attempts one key may make inside `window` before being refused.
48    pub max_attempts: u32,
49    /// Width of the sliding window, measured from each attempt.
50    pub window: Duration,
51    /// Distinct keys held at once. Bounding this is not an optimization: the
52    /// key is derived from the peer, so an unbounded map is a memory DoS
53    /// reachable by anyone who can open connections.
54    pub capacity: usize,
55}
56
57struct Entry {
58    attempts: Vec<Instant>,
59    last_seen: Instant,
60}
61
62/// A sliding-window counter per key, holding at most [`Limits::capacity`] keys.
63pub struct Bucket<K, C = SystemClock> {
64    entries: HashMap<K, Entry>,
65    limits: Limits,
66    clock: C,
67}
68
69impl<K: Eq + Hash + Clone> Bucket<K, SystemClock> {
70    #[must_use]
71    pub fn new(limits: Limits) -> Self {
72        Self::with_clock(limits, SystemClock)
73    }
74}
75
76impl<K: Eq + Hash + Clone, C: Clock> Bucket<K, C> {
77    pub fn with_clock(limits: Limits, clock: C) -> Self {
78        Self {
79            entries: HashMap::new(),
80            limits,
81            clock,
82        }
83    }
84
85    /// Records an attempt for `key` and says whether it may proceed.
86    pub fn check(&mut self, key: &K) -> Decision {
87        let now = self.clock.now();
88        let window = self.limits.window;
89
90        if !self.entries.contains_key(key) {
91            self.make_room(now);
92        }
93
94        let entry = self.entries.entry(key.clone()).or_insert_with(|| Entry {
95            attempts: Vec::new(),
96            last_seen: now,
97        });
98        entry.last_seen = now;
99        entry.attempts.retain(|t| now.duration_since(*t) < window);
100
101        let seen = u32::try_from(entry.attempts.len()).unwrap_or(u32::MAX);
102        if seen >= self.limits.max_attempts {
103            return Decision::Deny;
104        }
105
106        entry.attempts.push(now);
107        Decision::Allow
108    }
109
110    /// Frees a slot for a newcomer: drops keys whose window has fully elapsed,
111    /// then, if that was not enough, the least-recently-seen key.
112    ///
113    /// Evicting rather than refusing is deliberate. Refusing when full would
114    /// let anyone who can fill the map deny service to every other client — a
115    /// far cheaper attack than the credential stuffing this exists to stop.
116    fn make_room(&mut self, now: Instant) {
117        if self.entries.len() < self.limits.capacity {
118            return;
119        }
120
121        let window = self.limits.window;
122        self.entries
123            .retain(|_, e| now.duration_since(e.last_seen) < window);
124
125        if self.entries.len() < self.limits.capacity {
126            return;
127        }
128
129        if let Some(stale) = self
130            .entries
131            .iter()
132            .min_by_key(|(_, e)| e.last_seen)
133            .map(|(k, _)| k.clone())
134        {
135            self.entries.remove(&stale);
136        }
137    }
138
139    /// Keys currently tracked. Exposed for metrics and tests.
140    #[must_use]
141    pub fn tracked(&self) -> usize {
142        self.entries.len()
143    }
144}
ratelimit/lib.rs

ratelimit/
lib.rs
