use std::collections::{HashMap, 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 DEFAULT_ROUTE_RATE_LIMIT_WINDOW: u64 = 60;
pub const MAX_TRACKED_RATE_LIMIT_KEYS: usize = 10_000;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct RateLimitEntry {
pub requests: u32,
pub window: u64,
}
impl RateLimitEntry {
pub fn new(requests: u32, window: u64) -> Result<Self, RateLimitConfigError> {
if requests == 0 {
return Err(RateLimitConfigError {
field: "requests".into(),
reason: "must be at least 1 request per window",
});
}
if window == 0 {
return Err(RateLimitConfigError {
field: "window".into(),
reason: "must be at least 1 second",
});
}
Ok(Self { requests, window })
}
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct RouteRateLimits {
rate_limit: Option<u32>,
rate_limit_window: Option<u64>,
geo_rate_limits: Option<HashMap<String, RateLimitEntry>>,
}
impl RouteRateLimits {
pub fn new(
rate_limit: Option<u32>,
rate_limit_window: Option<u64>,
geo_rate_limits: Option<HashMap<String, RateLimitEntry>>,
) -> Result<Self, RateLimitConfigError> {
if rate_limit.is_some_and(|limit| limit == 0) {
return Err(RateLimitConfigError {
field: "rate_limit".into(),
reason: "must be at least 1 request per window",
});
}
if rate_limit.is_some() && rate_limit_window.is_some_and(|window| window == 0) {
return Err(RateLimitConfigError {
field: "rate_limit_window".into(),
reason: "must be at least 1 second",
});
}
if let Some(geo) = &geo_rate_limits {
for (country, entry) in geo {
if entry.requests == 0 || entry.window == 0 {
return Err(RateLimitConfigError {
field: format!("geo_rate_limits[{country}]").into(),
reason: "every entry needs at least 1 request per at least 1 second",
});
}
}
}
Ok(Self {
rate_limit,
rate_limit_window,
geo_rate_limits,
})
}
#[must_use]
pub const fn rate_limit(&self) -> Option<u32> {
self.rate_limit
}
#[must_use]
pub const fn rate_limit_window(&self) -> Option<u64> {
self.rate_limit_window
}
#[must_use]
pub const fn geo_rate_limits(&self) -> Option<&HashMap<String, RateLimitEntry>> {
self.geo_rate_limits.as_ref()
}
}
#[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,
pub endpoint_rate_limits: HashMap<String, RateLimitEntry>,
}
impl Default for RateLimitConfig {
fn default() -> Self {
Self {
enable_rate_limiting: true,
rate_limit: DEFAULT_RATE_LIMIT,
rate_limit_window: DEFAULT_RATE_LIMIT_WINDOW,
enable_rate_limit_auto_ban: false,
endpoint_rate_limits: HashMap::new(),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RateLimitConfigError {
pub field: std::borrow::Cow<'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
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RateLimitTier {
Endpoint,
Route,
Geo,
Global,
}
impl RateLimitTier {
#[must_use]
pub const fn name(self) -> &'static str {
match self {
Self::Endpoint => "endpoint",
Self::Route => "route",
Self::Geo => "geo",
Self::Global => "global",
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct TierDecision {
allowed: bool,
count: u64,
tier: RateLimitTier,
window: u64,
}
impl TierDecision {
#[must_use]
pub const fn allowed(&self) -> bool {
self.allowed
}
#[must_use]
pub const fn count(&self) -> u64 {
self.count
}
#[must_use]
pub const fn tier(&self) -> RateLimitTier {
self.tier
}
#[must_use]
pub const fn tier_name(&self) -> &'static str {
self.tier.name()
}
#[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,
store: Option<Arc<dyn crate::distributed::SlidingWindowStore>>,
redis_prefix: String,
redis_fail_open: bool,
}
impl Clone for RateLimiter {
fn clone(&self) -> Self {
Self {
config: self.config.clone(),
windows: Arc::clone(&self.windows),
clock: Arc::clone(&self.clock),
store: self.store.clone(),
redis_prefix: self.redis_prefix.clone(),
redis_fail_open: self.redis_fail_open,
}
}
}
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".into(),
reason: "must be at least 1 request per window",
});
}
if config.rate_limit_window == 0 {
return Err(RateLimitConfigError {
field: "rate_limit_window".into(),
reason: "must be at least 1 second",
});
}
for (path, entry) in &config.endpoint_rate_limits {
if entry.requests == 0 || entry.window == 0 {
return Err(RateLimitConfigError {
field: format!("endpoint_rate_limits[{path}]").into(),
reason: "every entry needs at least 1 request per at least 1 second",
});
}
}
Ok(Self {
config,
windows: new_window_store(),
clock,
store: None,
redis_prefix: String::from("guard_core:"),
redis_fail_open: false,
})
}
#[must_use]
pub fn with_distributed_store(
mut self,
store: Arc<dyn crate::distributed::SlidingWindowStore>,
redis_prefix: &str,
redis_fail_open: bool,
) -> Self {
self.store = Some(store);
redis_prefix.clone_into(&mut self.redis_prefix);
self.redis_fail_open = redis_fail_open;
self
}
pub fn check_distributed(
&self,
ip: IpAddr,
endpoint: Option<&str>,
) -> Result<RateLimitDecision, crate::distributed::StoreError> {
if !self.config.enable_rate_limiting {
return Ok(RateLimitDecision {
allowed: true,
count: 0,
window: self.config.rate_limit_window,
});
}
let (allowed, count, window) = self.record_and_decide_distributed(
ip,
endpoint,
self.config.rate_limit,
self.config.rate_limit_window,
)?;
Ok(RateLimitDecision {
allowed,
count,
window,
})
}
pub fn check_tiers_distributed(
&self,
ip: IpAddr,
url_path: Option<&str>,
route: Option<&RouteRateLimits>,
country_of_ip: Option<&dyn Fn(IpAddr) -> Option<String>>,
) -> Result<TierDecision, crate::distributed::StoreError> {
let inert = TierDecision {
allowed: true,
count: 0,
tier: RateLimitTier::Global,
window: self.config.rate_limit_window,
};
if !self.config.enable_rate_limiting {
return Ok(inert);
}
if let Some(path) = url_path {
if let Some(entry) = self.config.endpoint_rate_limits.get(path) {
let (allowed, count, window) = self.record_and_decide_distributed(
ip,
Some(path),
entry.requests,
entry.window,
)?;
if !allowed {
return Ok(TierDecision {
allowed,
count,
tier: RateLimitTier::Endpoint,
window,
});
}
}
if let Some(route) = route {
if let Some(limit) = route.rate_limit() {
let window = route
.rate_limit_window()
.unwrap_or(DEFAULT_ROUTE_RATE_LIMIT_WINDOW);
let (allowed, count, window) =
self.record_and_decide_distributed(ip, Some(path), limit, window)?;
if !allowed {
return Ok(TierDecision {
allowed,
count,
tier: RateLimitTier::Route,
window,
});
}
}
if let (Some(limits), Some(country_of_ip)) =
(route.geo_rate_limits(), country_of_ip)
{
let country = country_of_ip(ip);
let entry = country
.as_ref()
.and_then(|code| limits.get(code))
.or_else(|| limits.get("*"));
if let Some(entry) = entry {
let (allowed, count, window) = self.record_and_decide_distributed(
ip,
Some(path),
entry.requests,
entry.window,
)?;
if !allowed {
return Ok(TierDecision {
allowed,
count,
tier: RateLimitTier::Geo,
window,
});
}
}
}
}
}
let (allowed, count, window) = self.record_and_decide_distributed(
ip,
None,
self.config.rate_limit,
self.config.rate_limit_window,
)?;
Ok(TierDecision {
allowed,
count,
tier: RateLimitTier::Global,
window,
})
}
pub fn with_clock(clock: Clock) -> Self {
Self {
config: RateLimitConfig::default(),
windows: new_window_store(),
store: None,
redis_prefix: String::from("guard_core:"),
redis_fail_open: false,
clock,
}
}
#[must_use]
pub const fn config(&self) -> &RateLimitConfig {
&self.config
}
#[allow(clippy::significant_drop_tightening)]
fn record_and_decide_distributed(
&self,
ip: IpAddr,
endpoint: Option<&str>,
limit: u32,
window_seconds: u64,
) -> Result<(bool, u64, u64), crate::distributed::StoreError> {
let Some(store) = &self.store else {
return Ok(self.record_and_decide(ip, endpoint, limit, window_seconds));
};
let now = (self.clock)();
let key = endpoint.map_or_else(
|| crate::distributed::rate_window_key(&self.redis_prefix, ip),
|path| crate::distributed::rate_window_key_endpoint(&self.redis_prefix, ip, path),
);
match store.record_hit(&key, now, window_seconds) {
Ok(count) => {
let allowed = count <= u64::from(limit);
Ok((allowed, count, window_seconds))
}
Err(error) if self.redis_fail_open => {
let _ = error;
Ok(self.record_and_decide(ip, endpoint, limit, window_seconds))
}
Err(error) => Err(error),
}
}
fn record_and_decide(
&self,
ip: IpAddr,
endpoint: Option<&str>,
limit: u32,
window_seconds: u64,
) -> (bool, u64, u64) {
#[allow(clippy::cast_precision_loss)]
let window = window_seconds 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);
drop(windows);
inside
};
let allowed = inside < usize::try_from(limit).unwrap_or(usize::MAX);
let count = u64::try_from(inside + 1).unwrap_or(u64::MAX);
(allowed, count, window_seconds)
}
#[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,
};
}
let (allowed, count, window) = self.record_and_decide(
ip,
endpoint,
self.config.rate_limit,
self.config.rate_limit_window,
);
RateLimitDecision {
allowed,
count,
window,
}
}
#[must_use]
pub fn check_tiers(
&self,
ip: IpAddr,
url_path: Option<&str>,
route: Option<&RouteRateLimits>,
country_of_ip: Option<&dyn Fn(IpAddr) -> Option<String>>,
) -> TierDecision {
let inert = TierDecision {
allowed: true,
count: 0,
tier: RateLimitTier::Global,
window: self.config.rate_limit_window,
};
if !self.config.enable_rate_limiting {
return inert;
}
if let Some(path) = url_path {
if let Some(entry) = self.config.endpoint_rate_limits.get(path) {
let (allowed, count, window) =
self.record_and_decide(ip, Some(path), entry.requests, entry.window);
if !allowed {
return TierDecision {
allowed,
count,
tier: RateLimitTier::Endpoint,
window,
};
}
}
if let Some(route) = route {
if let Some(limit) = route.rate_limit() {
let window = route
.rate_limit_window()
.unwrap_or(DEFAULT_ROUTE_RATE_LIMIT_WINDOW);
let (allowed, count, window) =
self.record_and_decide(ip, Some(path), limit, window);
if !allowed {
return TierDecision {
allowed,
count,
tier: RateLimitTier::Route,
window,
};
}
}
if let (Some(limits), Some(country_of_ip)) =
(route.geo_rate_limits(), country_of_ip)
{
let country = country_of_ip(ip);
let entry = country
.as_ref()
.and_then(|code| limits.get(code))
.or_else(|| limits.get("*"));
if let Some(entry) = entry {
let (allowed, count, window) =
self.record_and_decide(ip, Some(path), entry.requests, entry.window);
if !allowed {
return TierDecision {
allowed,
count,
tier: RateLimitTier::Geo,
window,
};
}
}
}
}
}
let (allowed, count, window) = self.record_and_decide(
ip,
None,
self.config.rate_limit,
self.config.rate_limit_window,
);
TierDecision {
allowed,
count,
tier: RateLimitTier::Global,
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::collections::HashMap;
use std::str::FromStr;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use crate::distributed::MemoryStore;
#[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_matches_the_reference() {
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.field, std::borrow::Cow::Borrowed("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,
endpoint_rate_limits: std::collections::HashMap::new(),
})
.unwrap_err();
assert_eq!(error.field, "rate_limit_window");
}
#[test]
fn new_fails_closed_on_invalid_endpoint_entries() {
let error = RateLimiter::new(RateLimitConfig {
endpoint_rate_limits: std::iter::once((
"/login".to_owned(),
RateLimitEntry {
requests: 0,
window: 60,
},
))
.collect(),
..RateLimitConfig::default()
})
.unwrap_err();
assert_eq!(error.field, "endpoint_rate_limits[/login]");
assert_eq!(
error.to_string(),
"invalid endpoint_rate_limits[/login]: every entry needs at least 1 request \
per at least 1 second"
);
let error = RateLimitEntry::new(1, 0).unwrap_err();
assert_eq!(error.field, "window");
let error = RouteRateLimits::new(Some(0), Some(60), None).unwrap_err();
assert_eq!(error.field, "rate_limit");
let error = RouteRateLimits::new(Some(10), Some(0), None).unwrap_err();
assert_eq!(error.field, "rate_limit_window");
let error = RouteRateLimits::new(
None,
None,
Some(
std::iter::once((
"RU".to_owned(),
RateLimitEntry {
requests: 1,
window: 0,
},
))
.collect(),
),
)
.unwrap_err();
assert_eq!(error.field, "geo_rate_limits[RU]");
assert!(RouteRateLimits::new(None, None, None).is_ok());
assert!(RouteRateLimits::new(None, Some(30), None).is_ok());
}
#[test]
fn disabled_limiter_allows_and_records_nothing() {
let limiter = RateLimiter::new(RateLimitConfig {
enable_rate_limiting: false,
..RateLimitConfig::default()
})
.expect("valid 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(),
store: None,
redis_prefix: String::from("guard_core:"),
redis_fail_open: false,
};
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,
..RateLimitConfig::default()
})
.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);
}
fn country_of(code: &'static str) -> impl Fn(IpAddr) -> Option<String> {
move |_ip| Some(code.to_owned())
}
#[test]
fn tierless_check_tiers_collapses_to_the_global_tier() {
let limiter = RateLimiter::new(enabled_config(2)).expect("valid config");
let visitor = ip("192.0.2.40");
assert!(
limiter
.check_tiers(visitor, Some("/x"), None, None)
.allowed()
);
assert!(
limiter
.check_tiers(visitor, Some("/x"), None, None)
.allowed()
);
let blocked = limiter.check_tiers(visitor, Some("/x"), None, None);
assert!(!blocked.allowed());
assert_eq!(blocked.tier(), RateLimitTier::Global);
assert_eq!(blocked.tier_name(), "global");
assert_eq!(blocked.retry_after(), 60);
assert_eq!(limiter.tracked_windows(), 1, "one global window");
}
#[test]
fn endpoint_tier_matches_the_exact_path_and_keys_its_own_window() {
let limiter = RateLimiter::new(RateLimitConfig {
endpoint_rate_limits: std::iter::once((
"/login".to_owned(),
RateLimitEntry::new(1, 60).expect("valid entry"),
))
.collect(),
..RateLimitConfig::default()
})
.expect("valid config");
let visitor = ip("192.0.2.41");
assert!(
limiter
.check_tiers(visitor, Some("/login"), None, None)
.allowed()
);
let blocked = limiter.check_tiers(visitor, Some("/login"), None, None);
assert!(!blocked.allowed());
assert_eq!(blocked.tier(), RateLimitTier::Endpoint);
assert_eq!(blocked.retry_after(), 60);
assert!(
limiter
.check_tiers(visitor, Some("/other"), None, None)
.allowed()
);
assert_eq!(
limiter.tracked_windows(),
2,
"the login endpoint window + the shared global window"
);
}
#[test]
fn route_tier_uses_the_decorator_window_default_of_60() {
let fake = FakeClock::default();
let limiter = RateLimiter {
config: enabled_config(100),
windows: Arc::new(Mutex::new(LruCache::new(
NonZeroUsize::new(16).expect("above zero"),
))),
clock: fake.clock(),
store: None,
redis_prefix: String::from("guard_core:"),
redis_fail_open: false,
};
let route = RouteRateLimits::new(Some(1), None, None).expect("valid route");
let visitor = ip("192.0.2.42");
assert!(
limiter
.check_tiers(visitor, Some("/x"), Some(&route), None)
.allowed()
);
let blocked = limiter.check_tiers(visitor, Some("/x"), Some(&route), None);
assert!(!blocked.allowed());
assert_eq!(blocked.tier(), RateLimitTier::Route);
assert_eq!(
blocked.retry_after(),
60,
"the reference default window (rate_limit_window or 60)"
);
let route = RouteRateLimits::new(Some(1), Some(30), None).expect("valid route");
let blocked = limiter.check_tiers(visitor, Some("/x"), Some(&route), None);
assert!(!blocked.allowed());
assert_eq!(blocked.retry_after(), 30);
}
#[test]
fn geo_tier_resolves_the_country_then_the_star_fallback() {
let limiter = RateLimiter::new(RateLimitConfig {
rate_limit: 100,
..RateLimitConfig::default()
})
.expect("valid config");
let route = RouteRateLimits::new(
None,
None,
Some(
[
("RU".to_owned(), RateLimitEntry::new(1, 60).expect("valid")),
("*".to_owned(), RateLimitEntry::new(2, 45).expect("valid")),
]
.into_iter()
.collect(),
),
)
.expect("valid route");
let russian = ip("192.0.2.43");
assert!(
limiter
.check_tiers(russian, Some("/x"), Some(&route), Some(&country_of("RU")))
.allowed()
);
let blocked =
limiter.check_tiers(russian, Some("/x"), Some(&route), Some(&country_of("RU")));
assert!(!blocked.allowed());
assert_eq!(blocked.tier(), RateLimitTier::Geo);
assert_eq!(blocked.retry_after(), 60);
let other = ip("192.0.2.44");
let resolver = country_of("DE");
assert!(
limiter
.check_tiers(other, Some("/x"), Some(&route), Some(&resolver))
.allowed()
);
assert!(
limiter
.check_tiers(other, Some("/x"), Some(&route), Some(&resolver))
.allowed()
);
let blocked = limiter.check_tiers(other, Some("/x"), Some(&route), Some(&resolver));
assert!(!blocked.allowed());
assert_eq!(blocked.tier(), RateLimitTier::Geo);
assert_eq!(blocked.retry_after(), 45);
let unknown: fn(IpAddr) -> Option<String> = |_ip| None;
let third = ip("192.0.2.45");
assert!(
limiter
.check_tiers(third, Some("/x"), Some(&route), Some(&unknown))
.allowed()
);
assert!(
limiter
.check_tiers(third, Some("/x"), Some(&route), Some(&unknown))
.allowed()
);
let blocked = limiter.check_tiers(third, Some("/x"), Some(&route), Some(&unknown));
assert!(!blocked.allowed());
assert_eq!(blocked.tier(), RateLimitTier::Geo);
}
#[test]
fn geo_tier_needs_a_configured_map_and_country_lookup() {
let limiter = RateLimiter::new(RateLimitConfig {
rate_limit: 1,
..RateLimitConfig::default()
})
.expect("valid config");
let route = RouteRateLimits::new(None, None, Some(std::collections::HashMap::new()))
.expect("valid route");
let visitor = ip("192.0.2.46");
assert!(
limiter
.check_tiers(visitor, Some("/x"), Some(&route), Some(&country_of("RU")))
.allowed()
);
let blocked =
limiter.check_tiers(visitor, Some("/x"), Some(&route), Some(&country_of("RU")));
assert!(!blocked.allowed());
assert_eq!(blocked.tier(), RateLimitTier::Global);
let route = RouteRateLimits::new(
None,
None,
Some(
std::iter::once(("*".to_owned(), RateLimitEntry::new(1, 60).expect("valid")))
.collect(),
),
)
.expect("valid route");
let fresh = ip("192.0.2.51");
assert!(
limiter
.check_tiers(fresh, Some("/x"), Some(&route), None)
.allowed()
);
let blocked = limiter.check_tiers(fresh, Some("/x"), Some(&route), None);
assert!(!blocked.allowed());
assert_eq!(
blocked.tier(),
RateLimitTier::Global,
"no handler, no geo tier"
);
}
#[test]
fn every_configured_tier_records_so_the_budgets_stay_independent() {
let limiter = RateLimiter::new(RateLimitConfig {
rate_limit: 2,
..RateLimitConfig::default()
})
.expect("valid config");
let route = RouteRateLimits::new(Some(1_000), Some(60), None).expect("valid route");
let visitor = ip("192.0.2.47");
assert!(
limiter
.check_tiers(visitor, Some("/x"), Some(&route), None)
.allowed()
);
assert!(
limiter
.check_tiers(visitor, Some("/x"), Some(&route), None)
.allowed()
);
let blocked = limiter.check_tiers(visitor, Some("/x"), Some(&route), None);
assert!(!blocked.allowed());
assert_eq!(
blocked.tier(),
RateLimitTier::Global,
"the global tier crossed first"
);
assert_eq!(blocked.count(), 3);
assert_eq!(limiter.tracked_windows(), 2, "route + global windows");
}
#[test]
fn first_blocked_tier_wins_the_response_shape() {
let limiter = RateLimiter::new(RateLimitConfig {
rate_limit: 1,
rate_limit_window: 120,
endpoint_rate_limits: std::iter::once((
"/login".to_owned(),
RateLimitEntry::new(1, 30).expect("valid"),
))
.collect(),
..RateLimitConfig::default()
})
.expect("valid config");
let visitor = ip("192.0.2.48");
assert!(
limiter
.check_tiers(visitor, Some("/login"), None, None)
.allowed()
);
let blocked = limiter.check_tiers(visitor, Some("/login"), None, None);
assert!(!blocked.allowed());
assert_eq!(blocked.tier(), RateLimitTier::Endpoint);
assert_eq!(blocked.retry_after(), 30, "the endpoint tier's window");
let blocked = limiter.check_tiers(visitor, Some("/other"), None, None);
assert!(!blocked.allowed());
assert_eq!(blocked.tier(), RateLimitTier::Global);
assert_eq!(blocked.retry_after(), 120);
}
#[test]
fn tiers_without_a_request_path_skip_the_path_keyed_tiers() {
let limiter = RateLimiter::new(RateLimitConfig {
rate_limit: 1,
endpoint_rate_limits: std::iter::once((
"/login".to_owned(),
RateLimitEntry::new(1, 60).expect("valid"),
))
.collect(),
..RateLimitConfig::default()
})
.expect("valid config");
let route = RouteRateLimits::new(Some(1), Some(60), None).expect("valid route");
let visitor = ip("192.0.2.49");
assert!(
limiter
.check_tiers(visitor, None, Some(&route), Some(&country_of("RU")))
.allowed()
);
let blocked = limiter.check_tiers(visitor, None, Some(&route), Some(&country_of("RU")));
assert!(!blocked.allowed());
assert_eq!(blocked.tier(), RateLimitTier::Global);
}
#[test]
fn disabled_limiter_skips_the_tiers_without_recording() {
let limiter = RateLimiter::new(RateLimitConfig {
enable_rate_limiting: false,
endpoint_rate_limits: std::iter::once((
"/login".to_owned(),
RateLimitEntry::new(1, 60).expect("valid"),
))
.collect(),
..RateLimitConfig::default()
})
.expect("valid config");
let route = RouteRateLimits::new(Some(1), Some(60), None).expect("valid route");
let visitor = ip("192.0.2.50");
for _ in 0..10 {
let decision = limiter.check_tiers(visitor, Some("/login"), Some(&route), None);
assert!(decision.allowed());
}
assert_eq!(limiter.tracked_windows(), 0, "nothing was recorded");
}
#[test]
fn zero_request_entries_are_rejected() {
assert_eq!(RateLimitEntry::new(0, 60).err().unwrap().field, "requests");
}
#[test]
fn tier_names_read_like_the_reference_reasons() {
assert_eq!(RateLimitTier::Endpoint.name(), "endpoint");
assert_eq!(RateLimitTier::Route.name(), "route");
assert_eq!(RateLimitTier::Geo.name(), "geo");
assert_eq!(RateLimitTier::Global.name(), "global");
}
#[test]
fn limiter_debug_carries_the_config() {
let limiter = RateLimiter::with_clock(Arc::new(system_clock));
assert!(format!("{limiter:?}").contains("RateLimiter"));
assert_eq!(limiter.config().rate_limit_window, 60);
}
#[test]
fn distributed_checks_short_circuit_when_disabled() {
let limiter = RateLimiter::with_config_and_clock(
RateLimitConfig {
enable_rate_limiting: false,
..RateLimitConfig::default()
},
Arc::new(system_clock),
)
.expect("config");
let decision = limiter
.check_distributed(ip("192.0.2.51"), None)
.expect("inert");
assert!(decision.allowed);
let tiers = limiter
.check_tiers_distributed(ip("192.0.2.51"), Some("/x"), None, None)
.expect("inert");
assert!(tiers.allowed());
}
#[test]
fn distributed_tiers_cover_endpoint_route_and_geo_windows() {
let fake = FakeClock::default();
let store = Arc::new(MemoryStore::default());
let geo_limits = HashMap::from([
("RU".to_owned(), RateLimitEntry::new(2, 60).unwrap()),
("*".to_owned(), RateLimitEntry::new(5, 60).unwrap()),
]);
let route = RouteRateLimits::new(Some(3), Some(30), Some(geo_limits)).expect("valid route");
let limiter = RateLimiter::with_config_and_clock(
RateLimitConfig {
enable_rate_limiting: true,
rate_limit: 100,
endpoint_rate_limits: std::iter::once((
"/api".to_owned(),
RateLimitEntry::new(2, 60).unwrap(),
))
.collect(),
..RateLimitConfig::default()
},
fake.clock(),
)
.expect("config")
.with_distributed_store(
Arc::clone(&store) as Arc<dyn crate::distributed::SlidingWindowStore>,
"gc:",
false,
);
let country_ru = |probe: IpAddr| -> Option<String> {
if probe == ip("192.0.2.52") {
Some("RU".to_owned())
} else {
None
}
};
let first = limiter
.check_tiers_distributed(
ip("192.0.2.52"),
Some("/api"),
Some(&route),
Some(&country_ru),
)
.expect("store");
assert!(!first.allowed());
assert_eq!(first.tier(), RateLimitTier::Geo);
assert_eq!(first.count(), 3);
let second = limiter
.check_tiers_distributed(
ip("192.0.2.53"),
Some("/other"),
Some(&route),
Some(&country_ru),
)
.expect("store");
assert!(second.allowed());
let third = limiter
.check_tiers_distributed(
ip("192.0.2.52"),
Some("/other"),
Some(&route),
Some(&country_ru),
)
.expect("store");
assert!(third.allowed());
}
#[test]
fn a_store_error_propagates_from_the_endpoint_tier() {
let store = Arc::new(MemoryStore::default());
let limiter = RateLimiter::with_config_and_clock(
RateLimitConfig {
enable_rate_limiting: true,
rate_limit: 100,
endpoint_rate_limits: std::iter::once((
"/api".to_owned(),
RateLimitEntry::new(2, 60).unwrap(),
))
.collect(),
..RateLimitConfig::default()
},
FakeClock::default().clock(),
)
.expect("config")
.with_distributed_store(
Arc::clone(&store) as Arc<dyn crate::distributed::SlidingWindowStore>,
"gc:",
false,
);
store.fail.store(true, std::sync::atomic::Ordering::Relaxed);
assert!(
limiter
.check_tiers_distributed(ip("192.0.2.60"), Some("/api"), None, None)
.is_err()
);
}
#[test]
fn a_store_error_propagates_from_the_route_tier() {
let store = Arc::new(MemoryStore::default());
let route = RouteRateLimits::new(Some(3), Some(60), None).expect("valid route");
let limiter = RateLimiter::with_config_and_clock(
RateLimitConfig {
enable_rate_limiting: true,
..RateLimitConfig::default()
},
FakeClock::default().clock(),
)
.expect("config")
.with_distributed_store(
Arc::clone(&store) as Arc<dyn crate::distributed::SlidingWindowStore>,
"gc:",
false,
);
store.fail.store(true, std::sync::atomic::Ordering::Relaxed);
assert!(
limiter
.check_tiers_distributed(ip("192.0.2.59"), Some("/r"), Some(&route), None)
.is_err()
);
}
#[test]
fn a_store_error_propagates_from_the_geo_tier() {
let store = Arc::new(MemoryStore::default());
let geo_limits = HashMap::from([("RU".to_owned(), RateLimitEntry::new(2, 60).unwrap())]);
let route = RouteRateLimits::new(None, None, Some(geo_limits)).expect("valid route");
let limiter = RateLimiter::with_config_and_clock(
RateLimitConfig {
enable_rate_limiting: true,
..RateLimitConfig::default()
},
FakeClock::default().clock(),
)
.expect("config")
.with_distributed_store(
Arc::clone(&store) as Arc<dyn crate::distributed::SlidingWindowStore>,
"gc:",
false,
);
store.fail.store(true, std::sync::atomic::Ordering::Relaxed);
assert!(
limiter
.check_tiers_distributed(
ip("192.0.2.61"),
Some("/other"),
Some(&route),
Some(&|_| Some("RU".to_owned())),
)
.is_err()
);
}
#[test]
fn a_geo_limit_miss_falls_through_to_the_global_tier() {
let store = Arc::new(MemoryStore::default());
let geo_limits = HashMap::from([("CN".to_owned(), RateLimitEntry::new(2, 60).unwrap())]);
let route = RouteRateLimits::new(None, None, Some(geo_limits)).expect("valid route");
let limiter = RateLimiter::with_config_and_clock(
RateLimitConfig {
enable_rate_limiting: true,
..RateLimitConfig::default()
},
FakeClock::default().clock(),
)
.expect("config")
.with_distributed_store(
Arc::clone(&store) as Arc<dyn crate::distributed::SlidingWindowStore>,
"gc:",
false,
);
let decision = limiter
.check_tiers_distributed(
ip("192.0.2.62"),
Some("/other"),
Some(&route),
Some(&|_| Some("US".to_owned())),
)
.expect("store");
assert!(decision.allowed());
assert_eq!(decision.tier(), RateLimitTier::Global);
}
}
#[cfg(test)]
mod unit_twins {
use super::*;
use crate::distributed::MemoryStore;
use std::str::FromStr;
fn ip(text: &str) -> IpAddr {
IpAddr::from_str(text).expect("test address")
}
fn failing_store_limiter(endpoint: Option<(&str, RateLimitEntry)>) -> RateLimiter {
let store = MemoryStore::default();
store.fail.store(true, std::sync::atomic::Ordering::Relaxed);
let mut endpoint_rate_limits = HashMap::new();
if let Some((path, entry)) = endpoint {
endpoint_rate_limits.insert(path.to_owned(), entry);
}
RateLimiter::with_config_and_clock(
RateLimitConfig {
enable_rate_limiting: true,
rate_limit: 100,
endpoint_rate_limits,
..RateLimitConfig::default()
},
Arc::new(system_clock),
)
.expect("config")
.with_distributed_store(
Arc::new(store) as Arc<dyn crate::distributed::SlidingWindowStore>,
"gc:",
false,
)
}
#[test]
fn the_endpoint_tier_surfaces_a_failing_backend() {
let limiter = failing_store_limiter(Some((
"/api",
RateLimitEntry::new(2, 60).expect("valid entry"),
)));
assert!(
limiter
.check_tiers_distributed(ip("192.0.2.60"), Some("/api"), None, None)
.is_err()
);
}
#[test]
fn the_geo_tier_surfaces_a_failing_backend() {
let geo_limits = HashMap::from([(
"RU".to_owned(),
RateLimitEntry::new(2, 60).expect("valid entry"),
)]);
let route = RouteRateLimits::new(None, None, Some(geo_limits)).expect("valid route");
let limiter = failing_store_limiter(None);
let country = |probe: IpAddr| -> Option<String> {
if probe == ip("192.0.2.61") {
Some("RU".to_owned())
} else {
None
}
};
assert!(
limiter
.check_tiers_distributed(
ip("192.0.2.61"),
Some("/other"),
Some(&route),
Some(&country),
)
.is_err()
);
assert!(
limiter
.check_tiers_distributed(
ip("192.0.2.62"),
Some("/other"),
Some(&route),
Some(&country),
)
.is_err()
);
}
#[test]
fn exhausted_tiers_answer_their_denials() {
let store = Arc::new(MemoryStore::default());
let geo_limits = HashMap::from([(
"RU".to_owned(),
RateLimitEntry::new(1, 60).expect("valid entry"),
)]);
let limiter = RateLimiter::with_config_and_clock(
RateLimitConfig {
enable_rate_limiting: true,
rate_limit: 1,
endpoint_rate_limits: std::iter::once((
"/api".to_owned(),
RateLimitEntry::new(1, 60).expect("valid entry"),
))
.collect(),
..RateLimitConfig::default()
},
Arc::new(system_clock),
)
.expect("config")
.with_distributed_store(
Arc::clone(&store) as Arc<dyn crate::distributed::SlidingWindowStore>,
"gc:",
false,
);
let country_ru = |probe: IpAddr| -> Option<String> {
if probe == ip("192.0.2.63") || probe == ip("192.0.2.67") {
Some("RU".to_owned())
} else {
None
}
};
let route = RouteRateLimits::new(Some(1), Some(30), None).expect("valid route");
let geo_route = RouteRateLimits::new(None, None, Some(geo_limits)).expect("valid route");
let endpoint = limiter
.check_tiers_distributed(ip("192.0.2.63"), Some("/api"), None, None)
.expect("store");
assert!(endpoint.allowed());
let endpoint = limiter
.check_tiers_distributed(ip("192.0.2.63"), Some("/api"), None, None)
.expect("store");
assert!(!endpoint.allowed());
let route_hit = limiter
.check_tiers_distributed(ip("192.0.2.64"), Some("/route"), Some(&route), None)
.expect("store");
assert!(route_hit.allowed());
let route_hit = limiter
.check_tiers_distributed(ip("192.0.2.64"), Some("/route"), Some(&route), None)
.expect("store");
assert!(!route_hit.allowed());
let geo = limiter
.check_tiers_distributed(
ip("192.0.2.67"),
Some("/geo"),
Some(&geo_route),
Some(&country_ru),
)
.expect("store");
assert!(geo.allowed(), "first geo: {geo:?}");
let skipped = limiter
.check_tiers_distributed(
ip("192.0.2.68"),
Some("/geo"),
Some(&geo_route),
Some(&country_ru),
)
.expect("store");
assert!(skipped.allowed(), "skipped geo: {skipped:?}");
let geo = limiter
.check_tiers_distributed(
ip("192.0.2.67"),
Some("/geo"),
Some(&geo_route),
Some(&country_ru),
)
.expect("store");
assert!(!geo.allowed(), "second geo: {geo:?}");
assert_eq!(geo.tier(), RateLimitTier::Geo);
let global = limiter
.check_tiers_distributed(ip("192.0.2.65"), Some("/other"), None, None)
.expect("store");
assert!(global.allowed());
let global = limiter
.check_tiers_distributed(ip("192.0.2.65"), Some("/other"), None, None)
.expect("store");
assert!(!global.allowed());
assert_eq!(global.tier(), RateLimitTier::Global);
}
#[test]
fn the_global_tier_surfaces_a_failing_backend() {
let limiter = failing_store_limiter(None);
assert!(
limiter
.check_tiers_distributed(ip("192.0.2.66"), Some("/other"), None, None)
.is_err()
);
assert!(
limiter
.check_tiers_distributed(ip("192.0.2.66"), None, None, None)
.is_err()
);
}
#[test]
fn a_geo_tier_without_a_matching_country_falls_through_to_the_global_tier() {
let store = Arc::new(MemoryStore::default());
let geo_limits = HashMap::from([(
"RU".to_owned(),
RateLimitEntry::new(2, 60).expect("valid entry"),
)]);
let route = RouteRateLimits::new(Some(3), Some(30), Some(geo_limits)).expect("valid route");
let limiter = RateLimiter::with_config_and_clock(
RateLimitConfig {
enable_rate_limiting: true,
rate_limit: 100,
..RateLimitConfig::default()
},
Arc::new(system_clock),
)
.expect("config")
.with_distributed_store(
Arc::clone(&store) as Arc<dyn crate::distributed::SlidingWindowStore>,
"gc:",
false,
);
let no_country = |_: IpAddr| -> Option<String> { None };
let decision = limiter
.check_tiers_distributed(
ip("192.0.2.62"),
Some("/other"),
Some(&route),
Some(&no_country),
)
.expect("store");
assert!(decision.allowed());
assert_eq!(decision.tier(), RateLimitTier::Global);
}
}