use std::{
collections::HashMap,
net::IpAddr,
sync::Mutex,
time::{Duration, Instant},
};
pub struct RateLimiter {
clients: Mutex<HashMap<IpAddr, Window>>,
requests: u32,
window: Duration,
max_clients: usize,
}
struct Window {
started: Instant,
requests: u32,
}
impl RateLimiter {
pub fn new(requests: u32, window: Duration, max_clients: usize) -> Self {
Self {
clients: Mutex::new(HashMap::new()),
requests,
window,
max_clients,
}
}
pub fn check(&self, ip: IpAddr) -> Result<(), u64> {
self.check_at(ip, Instant::now())
}
pub(crate) fn check_at(&self, ip: IpAddr, now: Instant) -> Result<(), u64> {
let mut clients = self
.clients
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
if !clients.contains_key(&ip) && clients.len() >= self.max_clients {
clients.retain(|_, entry| now.duration_since(entry.started) < self.window);
if clients.len() >= self.max_clients {
return Err(self.window.as_secs().max(1));
}
}
let entry = clients.entry(ip).or_insert(Window {
started: now,
requests: 0,
});
if now.duration_since(entry.started) >= self.window {
*entry = Window {
started: now,
requests: 0,
};
}
if entry.requests >= self.requests {
return Err((self.window - now.duration_since(entry.started))
.as_secs()
.saturating_add(1));
}
entry.requests += 1;
Ok(())
}
}