1use std::collections::HashMap;
8use std::net::IpAddr;
9use std::sync::{Mutex, MutexGuard};
10use std::time::{Duration, Instant};
11
12pub trait KeyVerifier: Send + Sync {
14 fn extract(&self, x_api_key: Option<&str>, authorization: Option<&str>) -> Option<String>;
17
18 fn verify(&self, provided: Option<&str>, expected: &str) -> bool;
21}
22
23#[derive(Debug, Default, Clone, Copy)]
25pub struct SecurityKeyVerifier;
26
27impl KeyVerifier for SecurityKeyVerifier {
28 fn extract(&self, x_api_key: Option<&str>, authorization: Option<&str>) -> Option<String> {
29 talon_security::auth::extract_bearer(x_api_key, authorization)
30 }
31
32 fn verify(&self, provided: Option<&str>, expected: &str) -> bool {
33 talon_security::auth::verify_api_key(provided, expected)
34 }
35}
36
37const MAX_TRACKED_IPS: usize = 10_000;
39
40#[derive(Debug, Clone, Copy)]
41struct FailureWindow {
42 started: Instant,
43 failures: u32,
44}
45
46#[derive(Debug, Clone, Copy, PartialEq, Eq)]
48pub(crate) struct FailureRecord {
49 pub failures: u32,
51 pub now_throttled: bool,
53}
54
55pub(crate) struct AuthThrottle {
59 max_failures: u32,
60 window: Duration,
61 max_entries: usize,
62 inner: Mutex<HashMap<IpAddr, FailureWindow>>,
63}
64
65impl AuthThrottle {
66 pub(crate) fn new(max_failures: u32, window: Duration) -> Self {
67 Self::with_capacity(max_failures, window, MAX_TRACKED_IPS)
68 }
69
70 fn with_capacity(max_failures: u32, window: Duration, max_entries: usize) -> Self {
71 Self {
72 max_failures: max_failures.max(1),
73 window,
74 max_entries: max_entries.max(1),
75 inner: Mutex::new(HashMap::new()),
76 }
77 }
78
79 fn lock(&self) -> MutexGuard<'_, HashMap<IpAddr, FailureWindow>> {
80 self.inner.lock().unwrap_or_else(|e| e.into_inner())
81 }
82
83 pub(crate) fn blocked(&self, ip: IpAddr) -> Option<Duration> {
85 self.blocked_at(ip, Instant::now())
86 }
87
88 fn blocked_at(&self, ip: IpAddr, now: Instant) -> Option<Duration> {
89 let mut map = self.lock();
90 let entry = map.get(&ip).copied()?;
91 let elapsed = now.saturating_duration_since(entry.started);
92 if elapsed >= self.window {
93 map.remove(&ip);
94 return None;
95 }
96 (entry.failures >= self.max_failures).then(|| self.window - elapsed)
97 }
98
99 pub(crate) fn record_failure(&self, ip: IpAddr) -> FailureRecord {
101 self.record_failure_at(ip, Instant::now())
102 }
103
104 fn record_failure_at(&self, ip: IpAddr, now: Instant) -> FailureRecord {
105 let mut map = self.lock();
106 if !map.contains_key(&ip) && map.len() >= self.max_entries {
107 let window = self.window;
108 map.retain(|_, e| now.saturating_duration_since(e.started) < window);
109 if map.len() >= self.max_entries {
110 if let Some(oldest) = map.iter().min_by_key(|(_, e)| e.started).map(|(k, _)| *k) {
112 map.remove(&oldest);
113 }
114 }
115 }
116 let entry = map.entry(ip).or_insert(FailureWindow {
117 started: now,
118 failures: 0,
119 });
120 if now.saturating_duration_since(entry.started) >= self.window {
121 *entry = FailureWindow {
122 started: now,
123 failures: 0,
124 };
125 }
126 entry.failures = entry.failures.saturating_add(1);
127 FailureRecord {
128 failures: entry.failures,
129 now_throttled: entry.failures == self.max_failures,
130 }
131 }
132
133 pub(crate) fn record_success(&self, ip: IpAddr) {
135 let mut map = self.lock();
136 if !map.is_empty() {
137 map.remove(&ip);
138 }
139 }
140
141 #[cfg(test)]
142 fn tracked(&self) -> usize {
143 self.lock().len()
144 }
145}
146
147#[cfg(test)]
148pub(crate) mod test_support {
149 use super::KeyVerifier;
150
151 #[derive(Debug, Default, Clone, Copy)]
155 pub(crate) struct LocalVerifier;
156
157 impl KeyVerifier for LocalVerifier {
158 fn extract(&self, x_api_key: Option<&str>, authorization: Option<&str>) -> Option<String> {
159 if let Some(k) = x_api_key.map(str::trim).filter(|k| !k.is_empty()) {
160 return Some(k.to_string());
161 }
162 let auth = authorization?.trim();
163 let (scheme, rest) = auth.split_once(char::is_whitespace)?;
164 if !scheme.eq_ignore_ascii_case("bearer") {
165 return None;
166 }
167 let token = rest.trim();
168 (!token.is_empty()).then(|| token.to_string())
169 }
170
171 fn verify(&self, provided: Option<&str>, expected: &str) -> bool {
172 if expected.trim().is_empty() {
173 return false;
174 }
175 let Some(p) = provided.filter(|p| !p.is_empty()) else {
176 return false;
177 };
178 let (a, b) = (p.as_bytes(), expected.as_bytes());
179 if a.len() != b.len() {
180 return false;
181 }
182 a.iter().zip(b).fold(0u8, |acc, (x, y)| acc | (x ^ y)) == 0
183 }
184 }
185}
186
187#[cfg(test)]
188mod tests {
189 use super::test_support::LocalVerifier;
190 use super::*;
191 use std::net::Ipv4Addr;
192
193 fn ip(n: u8) -> IpAddr {
194 IpAddr::V4(Ipv4Addr::new(10, 0, 0, n))
195 }
196
197 #[test]
198 fn throttles_after_max_failures_and_recovers_after_the_window() {
199 let t = AuthThrottle::new(3, Duration::from_secs(60));
200 let start = Instant::now();
201 assert_eq!(t.blocked_at(ip(1), start), None);
202 assert!(!t.record_failure_at(ip(1), start).now_throttled);
203 assert!(!t.record_failure_at(ip(1), start).now_throttled);
204 let third = t.record_failure_at(ip(1), start);
205 assert!(third.now_throttled);
206 assert_eq!(third.failures, 3);
207 let retry = t.blocked_at(ip(1), start + Duration::from_secs(10));
208 assert_eq!(retry, Some(Duration::from_secs(50)));
209 assert_eq!(t.blocked_at(ip(2), start), None);
211 assert_eq!(t.blocked_at(ip(1), start + Duration::from_secs(61)), None);
213 assert_eq!(
214 t.record_failure_at(ip(1), start + Duration::from_secs(62))
215 .failures,
216 1
217 );
218 }
219
220 #[test]
221 fn success_resets_the_count() {
222 let t = AuthThrottle::new(2, Duration::from_secs(60));
223 t.record_failure(ip(1));
224 t.record_success(ip(1));
225 assert_eq!(t.record_failure(ip(1)).failures, 1);
226 assert_eq!(t.blocked(ip(1)), None);
227 }
228
229 #[test]
230 fn tracked_ip_set_is_bounded() {
231 let t = AuthThrottle::with_capacity(5, Duration::from_secs(60), 4);
232 for n in 0..50 {
233 t.record_failure(ip(n));
234 }
235 assert!(t.tracked() <= 4);
236 }
237
238 #[test]
239 fn local_verifier_contract() {
240 let v = LocalVerifier;
241 assert!(v.verify(Some("k3y"), "k3y"));
242 assert!(!v.verify(Some("k3y"), "k3z"));
243 assert!(!v.verify(Some(""), ""));
244 assert!(!v.verify(None, ""));
245 assert!(!v.verify(Some("x"), " "));
246 assert!(!v.verify(None, "k3y"));
247 assert_eq!(v.extract(Some("a"), Some("Bearer b")).as_deref(), Some("a"));
248 assert_eq!(v.extract(None, Some("bearer b ")).as_deref(), Some("b"));
249 assert_eq!(v.extract(None, Some("Basic b")), None);
250 assert_eq!(v.extract(None, None), None);
251 }
252
253 #[test]
255 fn security_crate_verifier_fails_closed() {
256 let v = SecurityKeyVerifier;
257 assert!(v.verify(Some("correct-horse-battery"), "correct-horse-battery"));
258 assert!(!v.verify(Some("wrong"), "correct-horse-battery"));
259 assert!(!v.verify(Some(""), ""));
260 assert!(!v.verify(None, ""));
261 assert_eq!(
262 v.extract(None, Some("Bearer correct-horse-battery"))
263 .as_deref(),
264 Some("correct-horse-battery")
265 );
266 assert_eq!(v.extract(Some("k"), None).as_deref(), Some("k"));
267 }
268}