forge
ExploreLog in

Tahinli/forge

public
0 watching
★ 0
7.5 KiB
1//! In-memory sliding-window failure counters for brute-force throttling.
2//!
3//! Each key keeps the timestamps of its recent failures. A key is limited while it has `limit`
4//! failures inside the window. Failures are only recorded for rejected credentials, so a locked
5//! key stops growing. Memory stays bounded: expired timestamps and empty keys are pruned
6//! periodically, and when the key cap is reached the key with the oldest failure is evicted.
7
8use std::collections::{HashMap, VecDeque};
9use std::net::IpAddr;
10use std::sync::{Mutex, MutexGuard};
11use std::time::{Duration, Instant};
12
13/// Failure window shared by login and git authentication.
14pub const WINDOW: Duration = Duration::from_secs(15 * 60);
15/// Failed logins allowed per folded username inside [`WINDOW`].
16pub const LOGIN_USER_LIMIT: usize = 10;
17/// Failed logins allowed per client IP inside [`WINDOW`].
18pub const LOGIN_IP_LIMIT: usize = 30;
19/// Failed git Basic/token authentications allowed per client IP inside [`WINDOW`].
20pub const GIT_IP_LIMIT: usize = 30;
21/// Failed second-factor checks (login and settings) allowed per user inside [`WINDOW`].
22pub const TWO_FACTOR_USER_LIMIT: usize = 10;
23/// Maximum number of distinct keys tracked at once.
24const MAX_KEYS: usize = 100_000;
25/// Run a full prune every this many recorded failures.
26const PRUNE_EVERY: u32 = 256;
27
28pub struct Throttle {
29 window: Duration,
30 max_keys: usize,
31 state: Mutex<State>,
32}
33
34#[derive(Default)]
35struct State {
36 failures: HashMap<String, VecDeque<Instant>>,
37 recorded: u32,
38}
39
40impl Default for Throttle {
41 fn default() -> Self {
42 Self::with_limits(WINDOW, MAX_KEYS)
43 }
44}
45
46impl Throttle {
47 pub fn with_limits(window: Duration, max_keys: usize) -> Self {
48 Self {
49 window,
50 max_keys: max_keys.max(1),
51 state: Mutex::new(State::default()),
52 }
53 }
54
55 /// Returns how long until `key` drops below `limit` failures, or `None` when it is not limited.
56 pub fn check(&self, key: &str, limit: usize, now: Instant) -> Option<Duration> {
57 let mut state = self.lock();
58 let timestamps = state.failures.get_mut(key)?;
59 drop_expired(timestamps, self.window, now);
60 if timestamps.len() < limit {
61 return None;
62 }
63 // The entry at this index expires first and brings the count below `limit`.
64 let index = timestamps.len().checked_sub(limit)?;
65 let oldest = *timestamps.get(index)?;
66 Some((oldest + self.window).saturating_duration_since(now))
67 }
68
69 /// Records one failure for `key` at `now`.
70 pub fn record_failure(&self, key: &str, now: Instant) {
71 let mut state = self.lock();
72 state.recorded = state.recorded.wrapping_add(1);
73 if state.recorded.is_multiple_of(PRUNE_EVERY) {
74 self.prune_locked(&mut state, now);
75 }
76 if !state.failures.contains_key(key) && state.failures.len() >= self.max_keys {
77 self.evict_oldest_locked(&mut state);
78 }
79 let timestamps = state.failures.entry(key.to_owned()).or_default();
80 drop_expired(timestamps, self.window, now);
81 timestamps.push_back(now);
82 }
83
84 /// Forgets all failures for `key`, for example after a successful login.
85 pub fn clear(&self, key: &str) {
86 self.lock().failures.remove(key);
87 }
88
89 /// Drops expired failures and empty keys.
90 #[cfg(test)]
91 pub fn prune(&self, now: Instant) {
92 let mut state = self.lock();
93 self.prune_locked(&mut state, now);
94 }
95
96 /// Number of keys currently tracked.
97 #[cfg(test)]
98 pub fn tracked_keys(&self) -> usize {
99 self.lock().failures.len()
100 }
101
102 fn prune_locked(&self, state: &mut State, now: Instant) {
103 let window = self.window;
104 state.failures.retain(|_, timestamps| {
105 drop_expired(timestamps, window, now);
106 !timestamps.is_empty()
107 });
108 }
109
110 /// Removes the key whose most recent failure is the oldest. Runs only when the cap is hit.
111 fn evict_oldest_locked(&self, state: &mut State) {
112 let victim = state
113 .failures
114 .iter()
115 .filter_map(|(key, timestamps)| timestamps.back().map(|last| (key, *last)))
116 .min_by_key(|(_, last)| *last)
117 .map(|(key, _)| key.clone());
118 if let Some(key) = victim {
119 state.failures.remove(&key);
120 }
121 }
122
123 fn lock(&self) -> MutexGuard<'_, State> {
124 // A poisoned lock still holds consistent counters, so keep using it.
125 self.state.lock().unwrap_or_else(|e| e.into_inner())
126 }
127}
128
129fn drop_expired(timestamps: &mut VecDeque<Instant>, window: Duration, now: Instant) {
130 while let Some(front) = timestamps.front() {
131 if now.saturating_duration_since(*front) < window {
132 break;
133 }
134 timestamps.pop_front();
135 }
136}
137
138/// Throttle key for a folded username.
139pub fn user_key(folded: &str) -> String {
140 format!("user:{folded}")
141}
142
143/// Throttle key for second-factor checks of one user, by user id.
144pub fn two_factor_key(user_id: &str) -> String {
145 format!("2fa:{user_id}")
146}
147
148/// Throttle key for an IP under a named scope, such as `login-ip` or `git-ip`.
149pub fn ip_key(scope: &str, ip: IpAddr) -> String {
150 format!("{scope}:{ip}")
151}
152
153#[cfg(test)]
154mod tests {
155 use super::*;
156
157 fn base() -> Instant {
158 Instant::now()
159 }
160
161 #[test]
162 fn limits_after_threshold_and_reports_retry_after() {
163 let throttle = Throttle::with_limits(Duration::from_secs(60), 16);
164 let t0 = base();
165 for i in 0..3 {
166 assert!(
167 throttle
168 .check("k", 3, t0 + Duration::from_secs(i))
169 .is_none()
170 );
171 throttle.record_failure("k", t0 + Duration::from_secs(i));
172 }
173 let retry = throttle
174 .check("k", 3, t0 + Duration::from_secs(10))
175 .expect("limited");
176 assert_eq!(retry, Duration::from_secs(50));
177 // Once the first failure leaves the window, the key is free again.
178 assert!(
179 throttle
180 .check("k", 3, t0 + Duration::from_secs(60))
181 .is_none()
182 );
183 }
184
185 #[test]
186 fn clear_resets_a_key() {
187 let throttle = Throttle::with_limits(Duration::from_secs(60), 16);
188 let t0 = base();
189 throttle.record_failure("k", t0);
190 throttle.clear("k");
191 assert!(throttle.check("k", 1, t0).is_none());
192 assert_eq!(throttle.tracked_keys(), 0);
193 }
194
195 #[test]
196 fn pruning_keeps_the_map_bounded() {
197 let throttle = Throttle::with_limits(Duration::from_secs(10), 8);
198 let t0 = base();
199 for i in 0..1000u64 {
200 let now = t0 + Duration::from_millis(i);
201 throttle.record_failure(&format!("key-{i}"), now);
202 assert!(throttle.tracked_keys() <= 8);
203 }
204 assert_eq!(throttle.tracked_keys(), 8);
205 }
206
207 #[test]
208 fn prune_drops_expired_keys() {
209 let throttle = Throttle::with_limits(Duration::from_secs(10), 8);
210 let t0 = base();
211 throttle.record_failure("a", t0);
212 throttle.record_failure("b", t0 + Duration::from_secs(5));
213 throttle.prune(t0 + Duration::from_secs(12));
214 assert_eq!(throttle.tracked_keys(), 1);
215 throttle.prune(t0 + Duration::from_secs(20));
216 assert_eq!(throttle.tracked_keys(), 0);
217 }
218
219 #[test]
220 fn keys_are_scoped_by_ip() {
221 let ip: IpAddr = "192.0.2.1".parse().unwrap();
222 assert_eq!(ip_key("git-ip", ip), "git-ip:192.0.2.1");
223 assert_eq!(user_key("alice"), "user:alice");
224 }
225}