use std::collections::VecDeque;
use std::net::IpAddr;
use std::num::NonZeroUsize;
use std::sync::{Arc, Mutex};
use lru::LruCache;
use crate::ip_gate::canonical;
pub const DEFAULT_RATE_LIMIT: u32 = 10;
pub const DEFAULT_RATE_LIMIT_WINDOW: u64 = 60;
pub const MAX_TRACKED_RATE_LIMIT_KEYS: usize = 10_000;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RateLimitConfig {
pub enable_rate_limiting: bool,
pub rate_limit: u32,
pub rate_limit_window: u64,
pub enable_rate_limit_auto_ban: bool,
}
impl Default for RateLimitConfig {
fn default() -> Self {
Self {
enable_rate_limiting: false,
rate_limit: DEFAULT_RATE_LIMIT,
rate_limit_window: DEFAULT_RATE_LIMIT_WINDOW,
enable_rate_limit_auto_ban: false,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RateLimitConfigError {
pub field: &'static str,
pub reason: &'static str,
}
impl core::fmt::Display for RateLimitConfigError {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
write!(f, "invalid {}: {}", self.field, self.reason)
}
}
impl std::error::Error for RateLimitConfigError {}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct RateLimitDecision {
pub allowed: bool,
pub count: u64,
pub window: u64,
}
impl RateLimitDecision {
#[must_use]
pub const fn retry_after(self) -> u64 {
self.window
}
}
pub type Clock = Arc<dyn Fn() -> f64 + Send + Sync>;
#[derive(Debug, PartialEq, Eq, Hash)]
struct WindowKey {
ip: IpAddr,
endpoint: Option<String>,
}
pub struct RateLimiter {
config: RateLimitConfig,
windows: Arc<Mutex<LruCache<WindowKey, VecDeque<f64>>>>,
clock: Clock,
}
impl Clone for RateLimiter {
fn clone(&self) -> Self {
Self {
config: self.config.clone(),
windows: Arc::clone(&self.windows),
clock: Arc::clone(&self.clock),
}
}
}
impl core::fmt::Debug for RateLimiter {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
f.debug_struct("RateLimiter")
.field("config", &self.config)
.finish_non_exhaustive()
}
}
#[must_use]
pub fn system_clock() -> f64 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map_or(0.0, |d| d.as_secs_f64())
}
fn new_window_store() -> Arc<Mutex<LruCache<WindowKey, VecDeque<f64>>>> {
Arc::new(Mutex::new(LruCache::new(
NonZeroUsize::new(MAX_TRACKED_RATE_LIMIT_KEYS).expect("constant above zero"),
)))
}
impl RateLimiter {
pub fn new(config: RateLimitConfig) -> Result<Self, RateLimitConfigError> {
Self::with_config_and_clock(config, Arc::new(system_clock))
}
pub fn with_config_and_clock(
config: RateLimitConfig,
clock: Clock,
) -> Result<Self, RateLimitConfigError> {
if config.rate_limit == 0 {
return Err(RateLimitConfigError {
field: "rate_limit",
reason: "must be at least 1 request per window",
});
}
if config.rate_limit_window == 0 {
return Err(RateLimitConfigError {
field: "rate_limit_window",
reason: "must be at least 1 second",
});
}
Ok(Self {
config,
windows: new_window_store(),
clock,
})
}
#[must_use]
pub fn with_clock(clock: Clock) -> Self {
Self {
config: RateLimitConfig::default(),
windows: new_window_store(),
clock,
}
}
#[must_use]
pub const fn config(&self) -> &RateLimitConfig {
&self.config
}
#[allow(clippy::significant_drop_tightening)]
#[must_use]
pub fn check(&self, ip: IpAddr, endpoint: Option<&str>) -> RateLimitDecision {
if !self.config.enable_rate_limiting {
return RateLimitDecision {
allowed: true,
count: 0,
window: self.config.rate_limit_window,
};
}
#[allow(clippy::cast_precision_loss)]
let window = self.config.rate_limit_window as f64;
let now = (self.clock)();
let window_start = now - window;
let key = WindowKey {
ip: canonical(ip),
endpoint: endpoint.map(str::to_owned),
};
let inside = {
let mut windows = self.windows.lock().expect("rate window store");
let timestamps = windows
.try_get_or_insert_mut(key, || Ok::<_, core::convert::Infallible>(VecDeque::new()))
.expect("key capacity just reserved");
while timestamps
.front()
.is_some_and(|&recorded| recorded <= window_start)
{
timestamps.pop_front();
}
let inside = timestamps.len();
timestamps.push_back(now);
inside
};
RateLimitDecision {
allowed: inside < usize::try_from(self.config.rate_limit).unwrap_or(usize::MAX),
count: u64::try_from(inside + 1).unwrap_or(u64::MAX),
window: self.config.rate_limit_window,
}
}
pub fn reset(&self) {
self.windows.lock().expect("rate window store").clear();
}
#[must_use]
pub fn tracked_windows(&self) -> usize {
self.windows.lock().expect("rate window store").len()
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::str::FromStr;
use std::sync::atomic::{AtomicU64, Ordering};
#[derive(Clone, Default)]
struct FakeClock(Arc<AtomicU64>);
impl FakeClock {
fn advance(&self, seconds: u64) {
self.0.fetch_add(seconds, Ordering::Relaxed);
}
fn clock(&self) -> Clock {
let state = self.0.clone();
#[allow(clippy::cast_precision_loss)]
Arc::new(move || state.load(Ordering::Relaxed) as f64)
}
}
fn ip(text: &str) -> IpAddr {
IpAddr::from_str(text).expect("test address")
}
fn enabled_config(rate_limit: u32) -> RateLimitConfig {
RateLimitConfig {
enable_rate_limiting: true,
rate_limit,
..RateLimitConfig::default()
}
}
#[test]
fn default_config_is_off_with_reference_thresholds() {
let config = RateLimitConfig::default();
assert!(!config.enable_rate_limiting);
assert_eq!(config.rate_limit, 10);
assert_eq!(config.rate_limit_window, 60);
assert!(!config.enable_rate_limit_auto_ban);
}
#[test]
fn new_fails_closed_on_zero_limit_or_window() {
let error = RateLimiter::new(RateLimitConfig {
rate_limit: 0,
..RateLimitConfig::default()
})
.unwrap_err();
assert_eq!(error.field, "rate_limit");
assert_eq!(
error.to_string(),
"invalid rate_limit: must be at least 1 request per window"
);
let error = RateLimiter::new(RateLimitConfig {
enable_rate_limiting: true,
rate_limit: 5,
rate_limit_window: 0,
enable_rate_limit_auto_ban: false,
})
.unwrap_err();
assert_eq!(error.field, "rate_limit_window");
}
#[test]
fn disabled_limiter_allows_and_records_nothing() {
let limiter = RateLimiter::new(RateLimitConfig::default()).expect("default config");
for _ in 0..50 {
assert!(limiter.check(ip("192.0.2.1"), None).allowed);
}
assert_eq!(limiter.tracked_windows(), 0, "no window was recorded");
}
#[test]
fn blocks_at_the_crossing_and_reports_the_crossing_count() {
let limiter = RateLimiter::new(enabled_config(3)).expect("valid config");
for expected in 1..=3 {
let decision = limiter.check(ip("192.0.2.1"), None);
assert!(decision.allowed, "request {expected} must pass");
assert_eq!(decision.count, expected);
}
let decision = limiter.check(ip("192.0.2.1"), None);
assert!(
!decision.allowed,
"the 4th request inside the window blocks"
);
assert_eq!(decision.count, 4);
assert_eq!(decision.retry_after(), 60, "Retry-After is the window");
}
#[test]
fn window_slide_restores_the_budget() {
let fake = FakeClock::default();
let limiter = RateLimiter {
config: enabled_config(2),
windows: Arc::new(Mutex::new(LruCache::new(
NonZeroUsize::new(16).expect("above zero"),
))),
clock: fake.clock(),
};
assert!(limiter.check(ip("192.0.2.1"), None).allowed);
assert!(limiter.check(ip("192.0.2.1"), None).allowed);
assert!(!limiter.check(ip("192.0.2.1"), None).allowed);
fake.advance(30);
assert!(!limiter.check(ip("192.0.2.1"), None).allowed);
fake.advance(31);
assert!(limiter.check(ip("192.0.2.1"), None).allowed);
}
#[test]
fn windows_are_per_ip() {
let limiter = RateLimiter::new(enabled_config(1)).expect("valid config");
assert!(limiter.check(ip("192.0.2.1"), None).allowed);
assert!(!limiter.check(ip("192.0.2.1"), None).allowed);
assert!(
limiter.check(ip("192.0.2.2"), None).allowed,
"another IP has its own budget"
);
}
#[test]
fn endpoint_scope_is_isolated_from_the_global_scope() {
let limiter = RateLimiter::new(enabled_config(1)).expect("valid config");
assert!(limiter.check(ip("192.0.2.1"), None).allowed);
assert!(
limiter.check(ip("192.0.2.1"), Some("/login")).allowed,
"the endpoint window is a separate budget"
);
assert!(!limiter.check(ip("192.0.2.1"), Some("/login")).allowed);
assert!(
limiter.check(ip("192.0.2.1"), Some("/signup")).allowed,
"another endpoint has its own budget"
);
assert_eq!(limiter.tracked_windows(), 3);
}
#[test]
fn ipv4_mapped_requests_share_the_ipv4_budget() {
let limiter = RateLimiter::new(enabled_config(1)).expect("valid config");
assert!(limiter.check(ip("::ffff:192.0.2.1"), None).allowed);
assert!(
!limiter.check(ip("192.0.2.1"), None).allowed,
"the mapped form must count toward the same window"
);
}
#[test]
fn reset_drops_every_window() {
let limiter = RateLimiter::new(enabled_config(1)).expect("valid config");
assert!(limiter.check(ip("192.0.2.1"), None).allowed);
assert!(!limiter.check(ip("192.0.2.1"), None).allowed);
limiter.reset();
assert_eq!(limiter.tracked_windows(), 0);
assert!(limiter.check(ip("192.0.2.1"), None).allowed);
}
#[test]
fn large_windows_keep_their_exact_length() {
let limiter = RateLimiter::new(RateLimitConfig {
enable_rate_limiting: true,
rate_limit: 1,
rate_limit_window: u64::from(u32::MAX) + 1,
enable_rate_limit_auto_ban: false,
})
.expect("valid config");
assert!(limiter.check(ip("192.0.2.1"), None).allowed);
let decision = limiter.check(ip("192.0.2.1"), None);
assert!(!decision.allowed);
assert_eq!(decision.retry_after(), u64::from(u32::MAX) + 1);
}
}