use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
#[derive(Clone)]
pub struct LoginGuard {
inner: Arc<Mutex<HashMap<String, Entry>>>,
}
struct Entry {
failures: Vec<Instant>,
last_seen: Instant,
}
impl Entry {
fn new(now: Instant) -> Self {
Self { failures: Vec::new(), last_seen: now }
}
fn prune(&mut self, now: Instant) {
self.failures.retain(|t| now.duration_since(*t) < WINDOW);
self.last_seen = now;
}
fn failure_count(&self) -> usize {
self.failures.len()
}
}
const MAX_FAILURES: usize = 5;
const WINDOW: Duration = Duration::from_secs(300);
pub const MAX_TRACKED_IPS: usize = 4096;
impl LoginGuard {
pub fn new() -> Self {
Self { inner: Arc::new(Mutex::new(HashMap::new())) }
}
pub fn is_blocked(&self, ip: &str) -> bool {
let now = Instant::now();
let mut map = self.inner.lock().unwrap();
let Some(entry) = map.get_mut(ip) else { return false };
entry.prune(now);
entry.failure_count() >= MAX_FAILURES
}
pub fn record_failure(&self, ip: &str) {
let now = Instant::now();
let mut map = self.inner.lock().unwrap();
let entry = map.entry(ip.to_string()).or_insert_with(|| Entry::new(now));
entry.prune(now);
entry.failures.push(now);
cap_tracked_ips(&mut map);
}
pub fn record_success(&self, ip: &str) {
self.inner.lock().unwrap().remove(ip);
}
#[cfg(test)]
pub fn tracked_len(&self) -> usize {
self.inner.lock().unwrap().len()
}
}
fn cap_tracked_ips(map: &mut HashMap<String, Entry>) {
if map.len() <= MAX_TRACKED_IPS {
return;
}
let Some(oldest) = map.iter().min_by_key(|(_, entry)| entry.last_seen).map(|(k, _)| k.clone())
else {
return;
};
map.remove(&oldest);
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn blocks_after_max_failures() {
let g = LoginGuard::new();
for _ in 0..MAX_FAILURES {
assert!(!g.is_blocked("1.2.3.4")); g.record_failure("1.2.3.4");
}
assert!(g.is_blocked("1.2.3.4"));
assert!(g.is_blocked("1.2.3.4"));
}
#[test]
fn success_resets_counter() {
let g = LoginGuard::new();
for _ in 0..MAX_FAILURES {
g.record_failure("1.2.3.4");
}
assert!(g.is_blocked("1.2.3.4"));
g.record_success("1.2.3.4");
assert!(!g.is_blocked("1.2.3.4"));
}
#[test]
fn window_expiry_resets_counter() {
let g = LoginGuard::new();
for _ in 0..MAX_FAILURES {
g.record_failure("1.2.3.4");
}
assert!(g.is_blocked("1.2.3.4"));
{
let mut map = g.inner.lock().unwrap();
let entry = map.get_mut("1.2.3.4").unwrap();
for t in entry.failures.iter_mut() {
*t = Instant::now() - WINDOW - Duration::from_secs(1);
}
}
assert!(!g.is_blocked("1.2.3.4"));
}
#[test]
fn ips_are_isolated() {
let g = LoginGuard::new();
for _ in 0..MAX_FAILURES {
g.record_failure("1.2.3.4");
}
assert!(g.is_blocked("1.2.3.4"));
assert!(!g.is_blocked("5.6.7.8"));
}
#[test]
fn is_blocked_does_not_insert_unknown_ip() {
let g = LoginGuard::new();
for i in 0..10 {
assert!(!g.is_blocked(&format!("10.0.0.{i}")));
}
assert_eq!(g.tracked_len(), 0);
}
#[test]
fn tracked_ip_count_never_exceeds_cap() {
let g = LoginGuard::new();
for i in 0..MAX_TRACKED_IPS + 200 {
g.record_failure(&format!("10.{}.{}.{}", i / 65536, (i / 256) % 256, i % 256));
}
assert!(
g.tracked_len() <= MAX_TRACKED_IPS,
"tracked {} IPs, cap is {}",
g.tracked_len(),
MAX_TRACKED_IPS
);
assert!(g.inner.lock().unwrap().values().all(|e| e.failure_count() == 1));
}
#[test]
fn one_insert_over_cap_evicts_exactly_one_entry() {
let g = LoginGuard::new();
for i in 0..MAX_TRACKED_IPS - 1 {
g.record_failure(&format!("10.0.0.{:04}", i));
}
assert_eq!(g.tracked_len(), MAX_TRACKED_IPS - 1);
g.record_failure("10.0.0.9999"); assert_eq!(g.tracked_len(), MAX_TRACKED_IPS);
}
#[test]
fn cap_evicts_least_recently_seen_first() {
let base = Instant::now() - Duration::from_secs(10);
let g = LoginGuard::new();
for i in 0..MAX_TRACKED_IPS {
g.record_failure(&format!("10.0.0.{:04}", i));
}
fn restamp(g: &LoginGuard, base: Instant) {
let mut map = g.inner.lock().unwrap();
for (k, entry) in map.iter_mut() {
let n: u64 =
k.rsplit('.').next().and_then(|t| t.parse().ok()).expect("fill key parses");
entry.last_seen = match k.as_str() {
"10.0.0.0000" | "10.0.0.9999" => base + Duration::from_secs(3600),
_ => base + Duration::from_millis(n),
};
}
}
restamp(&g, base);
let floor = {
let map = g.inner.lock().unwrap();
map.values().map(|e| e.last_seen).min().unwrap()
};
let keys_before: std::collections::HashSet<String> = {
let map = g.inner.lock().unwrap();
map.keys().cloned().collect()
};
g.record_failure("10.0.0.9999"); restamp(&g, base);
let map = g.inner.lock().unwrap();
assert_eq!(map.len(), MAX_TRACKED_IPS, "cap must be re-established exactly");
assert!(map.contains_key("10.0.0.0000"), "protected IP was evicted");
assert!(map.contains_key("10.0.0.9999"), "the newly inserted IP must survive");
let gone: Vec<&String> =
keys_before.iter().filter(|k| !map.contains_key(k.as_str())).collect();
assert_eq!(gone.len(), 1, "exactly one entry may be evicted; gone={gone:?}");
let min_surviving = map.values().map(|e| e.last_seen).min().unwrap();
assert!(
min_surviving > floor,
"the evicted entry must have been the least recently seen one: \
surviving minimum {min_surviving:?} should be newer than floor {floor:?}"
);
}
#[test]
fn a_read_refreshes_recency_of_an_existing_entry() {
let base = Instant::now();
let g = LoginGuard::new();
g.record_failure("10.0.0.1");
{
let mut map = g.inner.lock().unwrap();
map.get_mut("10.0.0.1").unwrap().last_seen = base;
}
assert!(!g.is_blocked("10.0.0.1"));
let after = g.inner.lock().unwrap();
assert!(
after.get("10.0.0.1").unwrap().last_seen > base,
"is_blocked did not refresh last_seen; read path is not part of LRU"
);
}
#[test]
fn eviction_keeps_limiter_functional() {
let g = LoginGuard::new();
for i in 0..MAX_TRACKED_IPS + 50 {
g.record_failure(&format!("10.1.{}.{}", i / 256, i % 256));
}
for _ in 0..MAX_FAILURES {
g.record_failure("10.2.0.1");
}
assert!(g.is_blocked("10.2.0.1"));
}
}