use std::net::IpAddr;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct StoreError(pub String);
impl core::fmt::Display for StoreError {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
write!(f, "distributed store error: {}", self.0)
}
}
impl std::error::Error for StoreError {}
pub trait SlidingWindowStore: Send + Sync {
fn record_hit(&self, key: &str, now: f64, window: u64) -> Result<u64, StoreError>;
}
pub trait BanStore: Send + Sync {
fn set_ban(&self, key: &str, expiry: f64, ttl_seconds: u64) -> Result<(), StoreError>;
fn get_ban(&self, key: &str) -> Result<Option<f64>, StoreError>;
fn delete_ban(&self, key: &str) -> Result<(), StoreError>;
}
#[must_use]
pub fn rate_window_key(prefix: &str, ip: IpAddr) -> String {
format!("{prefix}rate_limit:rate:{ip}")
}
#[must_use]
pub fn rate_window_key_endpoint(prefix: &str, ip: IpAddr, endpoint_path: &str) -> String {
format!(
"{prefix}rate_limit:rate:{ip}:{}",
identity_hash(endpoint_path)
)
}
#[must_use]
pub fn ban_key(prefix: &str, ip: IpAddr) -> String {
format!("{prefix}banned_ips:{ip}")
}
#[must_use]
pub fn identity_hash(value: &str) -> String {
use sha2::{Digest, Sha256};
let mut hasher = Sha256::new();
hasher.update(value.as_bytes());
let digest = hasher.finalize();
digest
.iter()
.fold(String::with_capacity(64), |mut out, byte| {
use core::fmt::Write as _;
let _ = write!(out, "{byte:02x}");
out
})
}
#[cfg(test)]
mod tests {
use super::*;
use std::str::FromStr;
#[test]
fn key_layouts_match_the_reference() {
let ip = IpAddr::from_str("192.0.2.1").unwrap();
assert_eq!(
rate_window_key("guard_core:", ip),
"guard_core:rate_limit:rate:192.0.2.1"
);
assert_eq!(
ban_key("guard_core:", ip),
"guard_core:banned_ips:192.0.2.1"
);
}
#[test]
fn endpoint_keys_hash_the_path_segment() {
let ip = IpAddr::from_str("192.0.2.1").unwrap();
let key = rate_window_key_endpoint("guard_core:", ip, "/login");
assert!(key.starts_with("guard_core:rate_limit:rate:192.0.2.1:"));
let digest = &key["guard_core:rate_limit:rate:192.0.2.1:".len()..];
assert_eq!(digest.len(), 64, "sha256 hex digest");
}
#[test]
fn identity_hash_matches_the_reference_vectors() {
assert_eq!(
identity_hash(""),
"e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
);
assert_eq!(
identity_hash("/login"),
"7e93fba0bc7adda858a2500090e639e4f563a3965c7cf4c719fb56eb6b12c666"
);
}
}
#[derive(Default)]
pub struct MemoryStore {
pub windows: std::sync::Mutex<std::collections::HashMap<String, Vec<(f64, f64)>>>,
pub bans: std::sync::Mutex<std::collections::HashMap<String, f64>>,
pub fail: std::sync::atomic::AtomicBool,
}
impl SlidingWindowStore for MemoryStore {
fn record_hit(&self, key: &str, now: f64, window: u64) -> Result<u64, StoreError> {
if self.fail.load(std::sync::atomic::Ordering::Relaxed) {
return Err(StoreError(String::from("backend down")));
}
let count;
{
let mut windows = self.windows.lock().expect("windows");
let entries = windows.entry(key.to_owned()).or_default();
entries.push((now, now));
let window_start = now - f64::from(u32::try_from(window).unwrap_or(u32::MAX));
entries.retain(|&(score, _)| score > window_start);
count = entries.len();
drop(windows);
}
#[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
Ok(count as u64)
}
}
impl BanStore for MemoryStore {
fn set_ban(&self, key: &str, expiry: f64, _ttl_seconds: u64) -> Result<(), StoreError> {
if self.fail.load(std::sync::atomic::Ordering::Relaxed) {
return Err(StoreError(String::from("backend down")));
}
self.bans
.lock()
.expect("bans")
.insert(key.to_owned(), expiry);
Ok(())
}
fn get_ban(&self, key: &str) -> Result<Option<f64>, StoreError> {
if self.fail.load(std::sync::atomic::Ordering::Relaxed) {
return Err(StoreError(String::from("backend down")));
}
Ok(self.bans.lock().expect("bans").get(key).copied())
}
fn delete_ban(&self, key: &str) -> Result<(), StoreError> {
if self.fail.load(std::sync::atomic::Ordering::Relaxed) {
return Err(StoreError(String::from("backend down")));
}
self.bans.lock().expect("bans").remove(key);
Ok(())
}
}
#[cfg(test)]
mod distributed_tests {
use super::*;
use std::str::FromStr;
use std::sync::Arc;
use crate::ip_ban::IpBanManager;
use crate::rate_limit::{RateLimitConfig, RateLimiter};
fn ip(text: &str) -> IpAddr {
IpAddr::from_str(text).expect("test address")
}
fn store_config(limit: u32) -> RateLimitConfig {
RateLimitConfig {
enable_rate_limiting: true,
rate_limit: limit,
..RateLimitConfig::default()
}
}
#[test]
fn distributed_windows_share_one_budget_across_limiter_clones() {
let store = Arc::new(MemoryStore::default());
let limiter = RateLimiter::new(store_config(3))
.expect("config")
.with_distributed_store(
Arc::clone(&store) as Arc<dyn SlidingWindowStore>,
"guard_core:",
false,
);
let visitor = ip("192.0.2.81");
assert!(
limiter
.check_distributed(visitor, None)
.expect("ok")
.allowed
);
assert!(
limiter
.check_distributed(visitor, None)
.expect("ok")
.allowed
);
assert!(
limiter
.check_distributed(visitor, None)
.expect("ok")
.allowed
);
let blocked = limiter.check_distributed(visitor, None).expect("ok");
assert!(!blocked.allowed);
assert_eq!(blocked.count, 4);
assert!(
store
.windows
.lock()
.expect("windows")
.contains_key("guard_core:rate_limit:rate:192.0.2.81")
);
}
#[test]
fn fail_closed_surfaces_the_backend_error() {
let store = Arc::new(MemoryStore::default());
store.fail.store(true, std::sync::atomic::Ordering::Relaxed);
let limiter = RateLimiter::new(store_config(3))
.expect("config")
.with_distributed_store(
Arc::clone(&store) as Arc<dyn SlidingWindowStore>,
"guard_core:",
false,
);
let error = limiter
.check_distributed(ip("192.0.2.82"), None)
.expect_err("fail closed");
assert!(error.to_string().contains("backend down"));
}
#[test]
fn fail_open_degrades_to_the_in_memory_window() {
let store = Arc::new(MemoryStore::default());
store.fail.store(true, std::sync::atomic::Ordering::Relaxed);
let limiter = RateLimiter::new(store_config(1))
.expect("config")
.with_distributed_store(
Arc::clone(&store) as Arc<dyn SlidingWindowStore>,
"guard_core:",
true,
);
let visitor = ip("192.0.2.83");
assert!(
limiter
.check_distributed(visitor, None)
.expect("ok")
.allowed
);
let blocked = limiter.check_distributed(visitor, None).expect("ok");
assert!(!blocked.allowed);
assert!(store.windows.lock().expect("windows").is_empty());
}
#[test]
fn distributed_bans_read_through_and_expire() {
let store = Arc::new(MemoryStore::default());
let bans = IpBanManager::with_clock(Arc::new(|| 1_000.0))
.with_distributed_store(Arc::clone(&store) as Arc<dyn BanStore>, "guard_core:");
let attacker = ip("192.0.2.84");
assert!(
bans.ban_ip(attacker, 60, "penetration_attempt")
.expect("ban")
);
assert_eq!(
store
.bans
.lock()
.expect("bans")
.get("guard_core:banned_ips:192.0.2.84"),
Some(&1_060.0)
);
let other_worker = IpBanManager::with_clock(Arc::new(|| 1_010.0))
.with_distributed_store(Arc::clone(&store) as Arc<dyn BanStore>, "guard_core:");
assert!(other_worker.is_banned(attacker));
let expired_worker = IpBanManager::with_clock(Arc::new(|| 2_000.0))
.with_distributed_store(Arc::clone(&store) as Arc<dyn BanStore>, "guard_core:");
assert!(!expired_worker.is_banned(attacker));
assert!(
!store
.bans
.lock()
.expect("bans")
.contains_key("guard_core:banned_ips:192.0.2.84")
);
}
}
#[cfg(test)]
mod coverage_tests {
use super::*;
#[test]
fn a_failing_backend_surfaces_the_store_error() {
let store = MemoryStore::default();
store.fail.store(true, std::sync::atomic::Ordering::Relaxed);
assert!(store.record_hit("k", 1_000.0, 60).is_err());
}
}
#[cfg(test)]
mod gap_tests {
use super::*;
#[test]
fn a_failing_backend_surfaces_the_store_error_for_ban_operations() {
let store = MemoryStore::default();
store.fail.store(true, std::sync::atomic::Ordering::Relaxed);
assert!(store.set_ban("k", 1_000.0, 60).is_err());
assert!(store.get_ban("k").is_err());
assert!(store.delete_ban("k").is_err());
}
}