use std::collections::HashMap;
use std::net::IpAddr;
use std::sync::Mutex;
use std::time::{Duration, Instant};
const ADDRESS_FACTOR: u32 = 4;
const MAX_TRACKED: usize = 100_000;
#[derive(PartialEq, Eq, Hash)]
enum Key {
Account(String),
Address(IpAddr),
}
struct Window {
attempts: u32,
started: Instant,
}
pub(crate) struct LoginLimiter {
max_attempts: u32,
window: Duration,
windows: Mutex<HashMap<Key, Window>>,
}
impl LoginLimiter {
pub(crate) fn new(max_attempts: u32, window_secs: u64) -> Self {
Self {
max_attempts,
window: Duration::from_secs(window_secs),
windows: Mutex::new(HashMap::new()),
}
}
pub(crate) fn attempt(&self, username: &str, address: Option<IpAddr>) -> Result<(), u64> {
self.attempt_at(username, address, Instant::now())
}
fn attempt_at(&self, username: &str, address: Option<IpAddr>, now: Instant) -> Result<(), u64> {
if self.max_attempts == 0 {
return Ok(());
}
let mut windows = self.lock();
if windows.len() >= MAX_TRACKED {
windows.retain(|_, w| now.duration_since(w.started) < self.window);
if windows.len() >= MAX_TRACKED {
windows.clear();
}
}
let keys = [
Some((Key::Account(username.to_owned()), self.max_attempts)),
address.map(|a| {
(
Key::Address(a),
self.max_attempts.saturating_mul(ADDRESS_FACTOR),
)
}),
];
for (key, limit) in keys.iter().flatten() {
if let Some(window) = windows.get(key) {
let age = now.duration_since(window.started);
if age < self.window && window.attempts >= *limit {
return Err((self.window - age).as_secs().max(1));
}
}
}
for (key, _) in keys.into_iter().flatten() {
let window = windows.entry(key).or_insert(Window {
attempts: 0,
started: now,
});
if now.duration_since(window.started) >= self.window {
*window = Window {
attempts: 0,
started: now,
};
}
window.attempts += 1;
}
Ok(())
}
pub(crate) fn succeeded(&self, username: &str, address: Option<IpAddr>) {
let mut windows = self.lock();
windows.remove(&Key::Account(username.to_owned()));
if let Some(window) = address.and_then(|a| windows.get_mut(&Key::Address(a))) {
window.attempts = window.attempts.saturating_sub(1);
}
}
fn lock(&self) -> std::sync::MutexGuard<'_, HashMap<Key, Window>> {
self.windows
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
}
}
#[cfg(test)]
mod tests {
use super::*;
fn address(last: u8) -> Option<IpAddr> {
Some(IpAddr::from([10, 0, 0, last]))
}
#[test]
fn an_account_is_locked_after_too_many_attempts_and_released_when_the_window_ends() {
let limiter = LoginLimiter::new(3, 60);
let start = Instant::now();
for _ in 0..3 {
assert_eq!(limiter.attempt_at("ann", None, start), Ok(()));
}
assert_eq!(limiter.attempt_at("ann", None, start), Err(60));
let later = start + Duration::from_secs(45);
assert_eq!(limiter.attempt_at("ann", None, later), Err(15));
assert_eq!(limiter.attempt_at("bob", None, later), Ok(()));
let after = start + Duration::from_secs(60);
assert_eq!(limiter.attempt_at("ann", None, after), Ok(()));
}
#[test]
fn logging_in_clears_the_count() {
let limiter = LoginLimiter::new(3, 60);
let now = Instant::now();
for _ in 0..2 {
limiter.attempt_at("ann", address(1), now).unwrap();
}
limiter.attempt_at("ann", address(1), now).unwrap();
limiter.succeeded("ann", address(1));
for _ in 0..3 {
assert_eq!(limiter.attempt_at("ann", address(1), now), Ok(()));
}
assert!(limiter.attempt_at("ann", address(1), now).is_err());
}
#[test]
fn an_address_is_limited_across_accounts() {
let limiter = LoginLimiter::new(2, 60);
let now = Instant::now();
for i in 0..8 {
assert_eq!(
limiter.attempt_at(&format!("user{i}"), address(1), now),
Ok(())
);
}
assert!(limiter.attempt_at("fresh", address(1), now).is_err());
assert_eq!(limiter.attempt_at("fresh", address(2), now), Ok(()));
}
#[test]
fn zero_attempts_turns_it_off() {
let limiter = LoginLimiter::new(0, 60);
for _ in 0..1000 {
assert_eq!(limiter.attempt("ann", address(1)), Ok(()));
}
}
}