use std::{
collections::HashMap,
hash::{BuildHasher, RandomState},
net::IpAddr,
sync::Mutex,
time::{Duration, Instant},
};
const MAX_TRACKED: usize = 10_000;
pub struct LoginThrottle {
limit: Option<Limit>,
hasher: RandomState,
counts: Mutex<HashMap<u64, Count>>,
}
#[derive(Debug, Clone, Copy)]
struct Limit {
attempts: u32,
window: Duration,
}
#[derive(Debug, Clone, Copy)]
struct Count {
attempts: u32,
since: Instant,
}
impl LoginThrottle {
#[must_use]
pub fn new(attempts: u32, window: Duration) -> Self {
assert!(attempts > 0, "a login throttle allows at least one attempt");
assert!(!window.is_zero(), "a login throttle's window is not zero");
Self::with(Some(Limit { attempts, window }))
}
#[must_use]
pub fn off() -> Self {
Self::with(None)
}
fn with(limit: Option<Limit>) -> Self {
Self {
limit,
hasher: RandomState::new(),
counts: Mutex::new(HashMap::new()),
}
}
pub(crate) fn attempt(&self, login: &str, client: Option<IpAddr>) -> bool {
self.attempt_at(login, client, Instant::now())
}
fn attempt_at(&self, login: &str, client: Option<IpAddr>, now: Instant) -> bool {
let Some(limit) = self.limit else {
return true;
};
let key = self.key(login, client);
let mut counts = self
.counts
.lock()
.unwrap_or_else(|poison| poison.into_inner());
if !counts.contains_key(&key) && counts.len() >= MAX_TRACKED {
counts.retain(|_, count| now.duration_since(count.since) < limit.window);
if counts.len() >= MAX_TRACKED
&& let Some(oldest) = counts
.iter()
.min_by_key(|(_, count)| count.since)
.map(|(key, _)| *key)
{
counts.remove(&oldest);
}
}
let count = counts.entry(key).or_insert(Count {
attempts: 0,
since: now,
});
if now.duration_since(count.since) >= limit.window {
*count = Count {
attempts: 0,
since: now,
};
}
if count.attempts >= limit.attempts {
return false;
}
count.attempts += 1;
true
}
pub(crate) fn clear(&self, login: &str, client: Option<IpAddr>) {
if self.limit.is_some() {
let key = self.key(login, client);
self.counts
.lock()
.unwrap_or_else(|poison| poison.into_inner())
.remove(&key);
}
}
fn key(&self, login: &str, client: Option<IpAddr>) -> u64 {
self.hasher.hash_one((login.trim().to_lowercase(), client))
}
}
impl Default for LoginThrottle {
fn default() -> Self {
Self::new(5, Duration::from_secs(60))
}
}
impl std::fmt::Debug for LoginThrottle {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self.limit {
Some(Limit { attempts, window }) => f
.debug_struct("LoginThrottle")
.field("attempts", &attempts)
.field("window", &window)
.finish(),
None => f.write_str("LoginThrottle::off"),
}
}
}
#[cfg(test)]
mod tests;