use std::collections::HashMap;
use std::net::IpAddr;
use std::sync::{Mutex, MutexGuard};
use std::time::{Duration, Instant};
pub trait KeyVerifier: Send + Sync {
fn extract(&self, x_api_key: Option<&str>, authorization: Option<&str>) -> Option<String>;
fn verify(&self, provided: Option<&str>, expected: &str) -> bool;
}
#[derive(Debug, Default, Clone, Copy)]
pub struct SecurityKeyVerifier;
impl KeyVerifier for SecurityKeyVerifier {
fn extract(&self, x_api_key: Option<&str>, authorization: Option<&str>) -> Option<String> {
talon_security::auth::extract_bearer(x_api_key, authorization)
}
fn verify(&self, provided: Option<&str>, expected: &str) -> bool {
talon_security::auth::verify_api_key(provided, expected)
}
}
const MAX_TRACKED_IPS: usize = 10_000;
#[derive(Debug, Clone, Copy)]
struct FailureWindow {
started: Instant,
failures: u32,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct FailureRecord {
pub failures: u32,
pub now_throttled: bool,
}
pub(crate) struct AuthThrottle {
max_failures: u32,
window: Duration,
max_entries: usize,
inner: Mutex<HashMap<IpAddr, FailureWindow>>,
}
impl AuthThrottle {
pub(crate) fn new(max_failures: u32, window: Duration) -> Self {
Self::with_capacity(max_failures, window, MAX_TRACKED_IPS)
}
fn with_capacity(max_failures: u32, window: Duration, max_entries: usize) -> Self {
Self {
max_failures: max_failures.max(1),
window,
max_entries: max_entries.max(1),
inner: Mutex::new(HashMap::new()),
}
}
fn lock(&self) -> MutexGuard<'_, HashMap<IpAddr, FailureWindow>> {
self.inner.lock().unwrap_or_else(|e| e.into_inner())
}
pub(crate) fn blocked(&self, ip: IpAddr) -> Option<Duration> {
self.blocked_at(ip, Instant::now())
}
fn blocked_at(&self, ip: IpAddr, now: Instant) -> Option<Duration> {
let mut map = self.lock();
let entry = map.get(&ip).copied()?;
let elapsed = now.saturating_duration_since(entry.started);
if elapsed >= self.window {
map.remove(&ip);
return None;
}
(entry.failures >= self.max_failures).then(|| self.window - elapsed)
}
pub(crate) fn record_failure(&self, ip: IpAddr) -> FailureRecord {
self.record_failure_at(ip, Instant::now())
}
fn record_failure_at(&self, ip: IpAddr, now: Instant) -> FailureRecord {
let mut map = self.lock();
if !map.contains_key(&ip) && map.len() >= self.max_entries {
let window = self.window;
map.retain(|_, e| now.saturating_duration_since(e.started) < window);
if map.len() >= self.max_entries {
if let Some(oldest) = map.iter().min_by_key(|(_, e)| e.started).map(|(k, _)| *k) {
map.remove(&oldest);
}
}
}
let entry = map.entry(ip).or_insert(FailureWindow {
started: now,
failures: 0,
});
if now.saturating_duration_since(entry.started) >= self.window {
*entry = FailureWindow {
started: now,
failures: 0,
};
}
entry.failures = entry.failures.saturating_add(1);
FailureRecord {
failures: entry.failures,
now_throttled: entry.failures == self.max_failures,
}
}
pub(crate) fn record_success(&self, ip: IpAddr) {
let mut map = self.lock();
if !map.is_empty() {
map.remove(&ip);
}
}
#[cfg(test)]
fn tracked(&self) -> usize {
self.lock().len()
}
}
#[cfg(test)]
pub(crate) mod test_support {
use super::KeyVerifier;
#[derive(Debug, Default, Clone, Copy)]
pub(crate) struct LocalVerifier;
impl KeyVerifier for LocalVerifier {
fn extract(&self, x_api_key: Option<&str>, authorization: Option<&str>) -> Option<String> {
if let Some(k) = x_api_key.map(str::trim).filter(|k| !k.is_empty()) {
return Some(k.to_string());
}
let auth = authorization?.trim();
let (scheme, rest) = auth.split_once(char::is_whitespace)?;
if !scheme.eq_ignore_ascii_case("bearer") {
return None;
}
let token = rest.trim();
(!token.is_empty()).then(|| token.to_string())
}
fn verify(&self, provided: Option<&str>, expected: &str) -> bool {
if expected.trim().is_empty() {
return false;
}
let Some(p) = provided.filter(|p| !p.is_empty()) else {
return false;
};
let (a, b) = (p.as_bytes(), expected.as_bytes());
if a.len() != b.len() {
return false;
}
a.iter().zip(b).fold(0u8, |acc, (x, y)| acc | (x ^ y)) == 0
}
}
}
#[cfg(test)]
mod tests {
use super::test_support::LocalVerifier;
use super::*;
use std::net::Ipv4Addr;
fn ip(n: u8) -> IpAddr {
IpAddr::V4(Ipv4Addr::new(10, 0, 0, n))
}
#[test]
fn throttles_after_max_failures_and_recovers_after_the_window() {
let t = AuthThrottle::new(3, Duration::from_secs(60));
let start = Instant::now();
assert_eq!(t.blocked_at(ip(1), start), None);
assert!(!t.record_failure_at(ip(1), start).now_throttled);
assert!(!t.record_failure_at(ip(1), start).now_throttled);
let third = t.record_failure_at(ip(1), start);
assert!(third.now_throttled);
assert_eq!(third.failures, 3);
let retry = t.blocked_at(ip(1), start + Duration::from_secs(10));
assert_eq!(retry, Some(Duration::from_secs(50)));
assert_eq!(t.blocked_at(ip(2), start), None);
assert_eq!(t.blocked_at(ip(1), start + Duration::from_secs(61)), None);
assert_eq!(
t.record_failure_at(ip(1), start + Duration::from_secs(62))
.failures,
1
);
}
#[test]
fn success_resets_the_count() {
let t = AuthThrottle::new(2, Duration::from_secs(60));
t.record_failure(ip(1));
t.record_success(ip(1));
assert_eq!(t.record_failure(ip(1)).failures, 1);
assert_eq!(t.blocked(ip(1)), None);
}
#[test]
fn tracked_ip_set_is_bounded() {
let t = AuthThrottle::with_capacity(5, Duration::from_secs(60), 4);
for n in 0..50 {
t.record_failure(ip(n));
}
assert!(t.tracked() <= 4);
}
#[test]
fn local_verifier_contract() {
let v = LocalVerifier;
assert!(v.verify(Some("k3y"), "k3y"));
assert!(!v.verify(Some("k3y"), "k3z"));
assert!(!v.verify(Some(""), ""));
assert!(!v.verify(None, ""));
assert!(!v.verify(Some("x"), " "));
assert!(!v.verify(None, "k3y"));
assert_eq!(v.extract(Some("a"), Some("Bearer b")).as_deref(), Some("a"));
assert_eq!(v.extract(None, Some("bearer b ")).as_deref(), Some("b"));
assert_eq!(v.extract(None, Some("Basic b")), None);
assert_eq!(v.extract(None, None), None);
}
#[test]
fn security_crate_verifier_fails_closed() {
let v = SecurityKeyVerifier;
assert!(v.verify(Some("correct-horse-battery"), "correct-horse-battery"));
assert!(!v.verify(Some("wrong"), "correct-horse-battery"));
assert!(!v.verify(Some(""), ""));
assert!(!v.verify(None, ""));
assert_eq!(
v.extract(None, Some("Bearer correct-horse-battery"))
.as_deref(),
Some("correct-horse-battery")
);
assert_eq!(v.extract(Some("k"), None).as_deref(), Some("k"));
}
}