use std::collections::{HashMap, HashSet};
use std::net::IpAddr;
use std::num::NonZeroUsize;
use std::sync::{Arc, Mutex};
use lru::LruCache;
use crate::ip_gate::{IpGateError, IpNet, canonical, parse_network_entry};
pub const DEFAULT_AUTO_BAN_THRESHOLD: u32 = 10;
pub const DEFAULT_AUTO_BAN_DURATION: u64 = 3600;
pub const LOCAL_CACHE_TTL_CAP_SECONDS: u64 = 3600;
pub const LOCAL_CACHE_MAX_SIZE: usize = 10_000;
pub const MAX_TRACKED_SUSPICIOUS_IPS: usize = 10_000;
pub const RATE_LIMIT_CATEGORY: &str = "rate_limit";
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ThreatBanEntry {
pub threshold: u32,
pub duration: u64,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct IpBanConfig {
pub enable_ip_banning: bool,
pub auto_ban_threshold: u32,
pub auto_ban_duration: u64,
pub threat_ban_config: HashMap<String, ThreatBanEntry>,
}
impl Default for IpBanConfig {
fn default() -> Self {
Self {
enable_ip_banning: true,
auto_ban_threshold: DEFAULT_AUTO_BAN_THRESHOLD,
auto_ban_duration: DEFAULT_AUTO_BAN_DURATION,
threat_ban_config: HashMap::new(),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum IpBanConfigError {
NonPositive {
field: &'static str,
},
UnknownCategory {
category: String,
},
}
impl core::fmt::Display for IpBanConfigError {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
match self {
Self::NonPositive { field } => write!(f, "invalid {field}: must be at least 1"),
Self::UnknownCategory { category } => write!(
f,
"unknown threat category '{category}': expected a detection category or 'rate_limit'"
),
}
}
}
impl std::error::Error for IpBanConfigError {}
#[must_use]
pub fn valid_threat_categories() -> HashSet<&'static str> {
let mut categories: HashSet<&'static str> = crate::patterns::table::PATTERN_DEFINITIONS
.iter()
.map(|entry| entry.category)
.collect();
categories.insert(RATE_LIMIT_CATEGORY);
categories
}
impl IpBanConfig {
pub fn new<I, K, V>(
enable_ip_banning: bool,
auto_ban_threshold: u32,
auto_ban_duration: u64,
threat_ban_config: I,
) -> Result<Self, IpBanConfigError>
where
I: IntoIterator<Item = (K, V)>,
K: AsRef<str>,
V: Into<ThreatBanEntry>,
{
let config = Self {
enable_ip_banning,
auto_ban_threshold,
auto_ban_duration,
threat_ban_config: threat_ban_config
.into_iter()
.map(|(key, entry)| (key.as_ref().to_owned(), entry.into()))
.collect(),
};
config.validate()?;
Ok(config)
}
pub fn validate(&self) -> Result<(), IpBanConfigError> {
if self.auto_ban_threshold == 0 {
return Err(IpBanConfigError::NonPositive {
field: "auto_ban_threshold",
});
}
if self.auto_ban_duration == 0 {
return Err(IpBanConfigError::NonPositive {
field: "auto_ban_duration",
});
}
let valid = valid_threat_categories();
for category in self.threat_ban_config.keys() {
if !valid.contains(category.as_str()) {
return Err(IpBanConfigError::UnknownCategory {
category: category.to_owned(),
});
}
}
Ok(())
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BanError {
NonPositiveDuration,
}
impl core::fmt::Display for BanError {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
match self {
Self::NonPositiveDuration => write!(f, "ban duration must be positive"),
}
}
}
impl std::error::Error for BanError {}
#[derive(Debug, Clone, PartialEq)]
pub struct BanRecord {
pub expiry: f64,
pub reason: String,
}
pub type Clock = Arc<dyn Fn() -> f64 + Send + Sync>;
fn system_clock() -> f64 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map_or(0.0, |d| d.as_secs_f64())
}
const LOOPBACK_NETWORKS: [&str; 2] = ["127.0.0.0/8", "::1/128"];
fn new_ban_store() -> Arc<Mutex<LruCache<IpAddr, BanRecord>>> {
Arc::new(Mutex::new(LruCache::new(
NonZeroUsize::new(LOCAL_CACHE_MAX_SIZE).expect("constant above zero"),
)))
}
pub struct IpBanManager {
bans: Arc<Mutex<LruCache<IpAddr, BanRecord>>>,
trusted_proxies: Vec<IpNet>,
clock: Clock,
ban_store: Option<Arc<dyn crate::distributed::BanStore>>,
redis_prefix: String,
}
impl Clone for IpBanManager {
fn clone(&self) -> Self {
Self {
bans: Arc::clone(&self.bans),
trusted_proxies: self.trusted_proxies.clone(),
clock: Arc::clone(&self.clock),
ban_store: self.ban_store.clone(),
redis_prefix: self.redis_prefix.clone(),
}
}
}
impl core::fmt::Debug for IpBanManager {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
f.debug_struct("IpBanManager")
.field("trusted_proxies", &self.trusted_proxies.len())
.finish_non_exhaustive()
}
}
impl Default for IpBanManager {
fn default() -> Self {
Self::new()
}
}
impl IpBanManager {
#[must_use]
pub fn new() -> Self {
Self {
bans: new_ban_store(),
ban_store: None,
redis_prefix: String::from("guard_core:"),
trusted_proxies: Vec::new(),
clock: Arc::new(system_clock),
}
}
pub fn with_trusted_proxies<I>(entries: I) -> Result<Self, IpGateError>
where
I: IntoIterator,
I::Item: AsRef<str>,
{
Self::with_trusted_proxies_and_clock(entries, Arc::new(system_clock))
}
pub fn with_trusted_proxies_and_clock<I>(entries: I, clock: Clock) -> Result<Self, IpGateError>
where
I: IntoIterator,
I::Item: AsRef<str>,
{
let mut trusted_proxies = Vec::new();
for entry in entries {
let entry = entry.as_ref();
match parse_network_entry(entry) {
Some(network) => trusted_proxies.push(network),
None => {
return Err(IpGateError {
list: "trusted_proxies",
entry: entry.to_owned(),
});
}
}
}
Ok(Self {
bans: new_ban_store(),
ban_store: None,
redis_prefix: String::from("guard_core:"),
trusted_proxies,
clock,
})
}
#[must_use]
pub fn with_distributed_store(
mut self,
store: Arc<dyn crate::distributed::BanStore>,
redis_prefix: &str,
) -> Self {
self.ban_store = Some(store);
redis_prefix.clone_into(&mut self.redis_prefix);
self
}
pub fn with_clock(clock: Clock) -> Self {
Self {
bans: new_ban_store(),
ban_store: None,
redis_prefix: String::from("guard_core:"),
trusted_proxies: Vec::new(),
clock,
}
}
pub fn ban_ip(&self, ip: IpAddr, duration: u64, reason: &str) -> Result<bool, BanError> {
if duration == 0 {
return Err(BanError::NonPositiveDuration);
}
let ip = canonical(ip);
if self.self_ban_refusal(&ip) {
return Ok(false);
}
let clamped = duration.min(LOCAL_CACHE_TTL_CAP_SECONDS);
#[allow(clippy::cast_precision_loss)]
let expiry = (self.clock)() + clamped as f64;
self.bans.lock().expect("ban store").put(
ip,
BanRecord {
expiry,
reason: reason.to_owned(),
},
);
if let Some(store) = &self.ban_store {
#[allow(clippy::cast_precision_loss)]
let distributed_expiry = (self.clock)() + duration as f64;
let _ = store.set_ban(
&crate::distributed::ban_key(&self.redis_prefix, ip),
distributed_expiry,
duration,
);
}
Ok(true)
}
#[must_use]
pub fn is_banned(&self, ip: IpAddr) -> bool {
self.try_is_banned(ip).unwrap_or(false)
}
pub fn try_is_banned(&self, ip: IpAddr) -> Result<bool, crate::distributed::StoreError> {
let ip = canonical(ip);
let now = (self.clock)();
{
let mut bans = self.bans.lock().expect("ban store");
match bans.get(&ip) {
Some(record) if now <= record.expiry => return Ok(true),
Some(_) => {
bans.pop(&ip);
}
None => {}
}
}
self.check_distributed_ban(ip, now)
}
fn check_distributed_ban(
&self,
ip: IpAddr,
now: f64,
) -> Result<bool, crate::distributed::StoreError> {
let Some(store) = &self.ban_store else {
return Ok(false);
};
match store.get_ban(&crate::distributed::ban_key(&self.redis_prefix, ip)) {
Ok(Some(expiry)) if now <= expiry => {
self.bans.lock().expect("ban store").put(
ip,
BanRecord {
expiry,
reason: String::from("distributed"),
},
);
Ok(true)
}
Ok(Some(_)) => {
let _ = store.delete_ban(&crate::distributed::ban_key(&self.redis_prefix, ip));
Ok(false)
}
Ok(None) => Ok(false),
Err(error) => Err(error),
}
}
#[must_use]
pub fn ban_record(&self, ip: IpAddr) -> Option<BanRecord> {
let ip = canonical(ip);
let now = (self.clock)();
let mut bans = self.bans.lock().expect("ban store");
match bans.get(&ip) {
Some(record) if now <= record.expiry => Some(record.clone()),
Some(_) => {
bans.pop(&ip);
None
}
None => None,
}
}
pub fn unban(&self, ip: IpAddr) {
self.bans.lock().expect("ban store").pop(&canonical(ip));
}
pub fn reset(&self) {
self.bans.lock().expect("ban store").clear();
}
#[must_use]
pub fn banned_count(&self) -> usize {
self.bans.lock().expect("ban store").len()
}
fn self_ban_refusal(&self, ip: &IpAddr) -> bool {
LOOPBACK_NETWORKS
.iter()
.filter_map(|entry| parse_network_entry(entry))
.chain(self.trusted_proxies.iter().copied())
.any(|network| network.contains(*ip))
}
}
pub struct ViolationCounters {
counts: Arc<Mutex<LruCache<IpAddr, HashMap<String, u64>>>>,
}
impl Clone for ViolationCounters {
fn clone(&self) -> Self {
Self {
counts: Arc::clone(&self.counts),
}
}
}
impl core::fmt::Debug for ViolationCounters {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
f.debug_struct("ViolationCounters").finish_non_exhaustive()
}
}
impl Default for ViolationCounters {
fn default() -> Self {
Self::new()
}
}
impl ViolationCounters {
#[must_use]
pub fn new() -> Self {
Self {
counts: Arc::new(Mutex::new(LruCache::new(
NonZeroUsize::new(MAX_TRACKED_SUSPICIOUS_IPS).expect("constant above zero"),
))),
}
}
pub fn record(&self, ip: IpAddr, categories: &[&str]) {
let categories: &[&str] = if categories.is_empty() {
&["uncategorized"]
} else {
categories
};
let mut counts = self.counts.lock().expect("violation counter store");
let ip_counts = counts
.try_get_or_insert_mut(canonical(ip), || {
Ok::<_, core::convert::Infallible>(HashMap::new())
})
.expect("ip capacity just reserved");
for category in categories {
*ip_counts.entry((*category).to_owned()).or_insert(0) += 1;
}
drop(counts);
}
#[must_use]
pub fn snapshot(&self, ip: IpAddr) -> HashMap<String, u64> {
self.counts
.lock()
.expect("violation counter store")
.peek(&canonical(ip))
.cloned()
.unwrap_or_default()
}
pub fn reset(&self) {
self.counts.lock().expect("violation counter store").clear();
}
#[must_use]
pub fn tracked_ips(&self) -> usize {
self.counts.lock().expect("violation counter store").len()
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ResolvedBan {
pub duration: u64,
pub reason: String,
pub category: Option<String>,
}
#[must_use]
pub fn resolve_threshold_ban<S: core::hash::BuildHasher>(
counts: &HashMap<String, u64, S>,
config: &IpBanConfig,
threat_categories: &[&str],
reason: &str,
) -> Option<ResolvedBan> {
if !config.enable_ip_banning {
return None;
}
for category in threat_categories {
let Some(entry) = config.threat_ban_config.get(*category) else {
continue;
};
let count = counts.get(*category).copied().unwrap_or(0);
if count < u64::from(entry.threshold) {
continue;
}
return Some(ResolvedBan {
duration: entry.duration,
reason: format!("{reason}:{category}"),
category: Some((*category).to_owned()),
});
}
let total: u64 = counts.values().sum();
if total < u64::from(config.auto_ban_threshold) {
return None;
}
Some(ResolvedBan {
duration: config.auto_ban_duration,
reason: reason.to_owned(),
category: None,
})
}
impl IpBanManager {
#[must_use]
pub fn register_violations(
&self,
counters: &ViolationCounters,
ip: IpAddr,
categories: &[&str],
config: &IpBanConfig,
reason: &str,
) -> Option<ResolvedBan> {
counters.record(ip, categories);
let counts = counters.snapshot(ip);
let resolved = resolve_threshold_ban(&counts, config, categories, reason)?;
match self.ban_ip(ip, resolved.duration, &resolved.reason) {
Ok(true) => Some(resolved),
_ => None,
}
}
}
#[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")
}
#[test]
fn default_config_matches_the_reference() {
let config = IpBanConfig::default();
assert!(config.enable_ip_banning);
assert_eq!(config.auto_ban_threshold, 10);
assert_eq!(config.auto_ban_duration, 3600);
assert!(config.threat_ban_config.is_empty());
}
#[test]
fn config_validation_fails_closed() {
let error =
IpBanConfig::new(true, 0, 3600, Vec::<(String, ThreatBanEntry)>::new()).unwrap_err();
assert_eq!(
error,
IpBanConfigError::NonPositive {
field: "auto_ban_threshold"
}
);
assert_eq!(
error.to_string(),
"invalid auto_ban_threshold: must be at least 1"
);
let error =
IpBanConfig::new(true, 10, 0, Vec::<(String, ThreatBanEntry)>::new()).unwrap_err();
assert_eq!(
error,
IpBanConfigError::NonPositive {
field: "auto_ban_duration"
}
);
let error = IpBanConfig::new(
true,
10,
3600,
[(
"not_a_category",
ThreatBanEntry {
threshold: 1,
duration: 1,
},
)],
)
.unwrap_err();
assert_eq!(
error,
IpBanConfigError::UnknownCategory {
category: "not_a_category".to_owned()
}
);
assert_eq!(
error.to_string(),
"unknown threat category 'not_a_category': expected a detection category or 'rate_limit'"
);
}
#[test]
fn config_accepts_table_categories_and_rate_limit() {
IpBanConfig::new(
true,
10,
3600,
[
(
"sqli",
ThreatBanEntry {
threshold: 3,
duration: 60,
},
),
(
RATE_LIMIT_CATEGORY,
ThreatBanEntry {
threshold: 5,
duration: 30,
},
),
],
)
.expect("both categories are valid");
assert!(valid_threat_categories().contains("sqli"));
assert!(valid_threat_categories().contains(RATE_LIMIT_CATEGORY));
}
#[test]
fn ban_then_lookup_then_expiry() {
let fake = FakeClock::default();
let manager = IpBanManager::with_clock(fake.clock());
manager
.ban_ip(ip("192.0.2.9"), 60, "threshold_exceeded")
.expect("ban");
assert!(manager.is_banned(ip("192.0.2.9")));
assert_eq!(
manager
.ban_record(ip("192.0.2.9"))
.expect("live record")
.reason,
"threshold_exceeded"
);
fake.advance(59);
assert!(manager.is_banned(ip("192.0.2.9")));
fake.advance(1);
assert!(manager.is_banned(ip("192.0.2.9")));
fake.advance(1);
assert!(!manager.is_banned(ip("192.0.2.9")));
assert_eq!(manager.banned_count(), 0, "the expired entry was dropped");
}
#[test]
fn rebanning_overwrites_expiry_and_reason() {
let manager = IpBanManager::new();
manager.ban_ip(ip("192.0.2.9"), 60, "first").expect("ban");
manager.ban_ip(ip("192.0.2.9"), 120, "second").expect("ban");
let record = manager.ban_record(ip("192.0.2.9")).expect("live record");
assert_eq!(record.reason, "second");
}
#[test]
fn zero_duration_is_an_error() {
let manager = IpBanManager::new();
assert_eq!(
manager.ban_ip(ip("192.0.2.9"), 0, "nope").unwrap_err(),
BanError::NonPositiveDuration
);
assert!(!manager.is_banned(ip("192.0.2.9")));
assert_eq!(
BanError::NonPositiveDuration.to_string(),
"ban duration must be positive"
);
}
#[test]
fn loopback_is_refused_for_both_families() {
let manager = IpBanManager::new();
for loopback in ["127.0.0.1", "127.8.8.8", "::1"] {
assert!(
!manager
.ban_ip(ip(loopback), 60, "mistake")
.expect("refusal is not an error"),
"{loopback} must be refused"
);
assert!(!manager.is_banned(ip(loopback)));
}
}
#[test]
fn trusted_proxies_are_refused_and_parsed_fail_closed() {
let manager = IpBanManager::with_trusted_proxies(["10.0.0.0/8", "2001:db8:ffff::/48"])
.expect("valid");
assert!(
!manager
.ban_ip(ip("10.1.2.3"), 60, "proxy")
.expect("refusal")
);
assert!(
!manager
.ban_ip(ip("2001:db8:ffff::1"), 60, "proxy")
.expect("refusal")
);
assert!(
manager
.ban_ip(ip("11.1.2.3"), 60, "outside the proxy range")
.expect("ban"),
"an address outside the proxy range is bannable"
);
let error = IpBanManager::with_trusted_proxies(["not-an-ip"]).unwrap_err();
assert_eq!(error.list, "trusted_proxies");
assert_eq!(error.entry, "not-an-ip");
}
#[test]
fn ipv4_mapped_targets_canonicalize() {
let manager = IpBanManager::new();
assert!(
!manager
.ban_ip(ip("::ffff:127.0.0.1"), 60, "mistake")
.expect("refusal")
);
manager
.ban_ip(ip("::ffff:192.0.2.9"), 60, "x")
.expect("ban");
assert!(manager.is_banned(ip("192.0.2.9")));
assert_eq!(manager.banned_count(), 1);
}
#[test]
fn unban_and_reset() {
let manager = IpBanManager::new();
manager.ban_ip(ip("192.0.2.9"), 60, "x").expect("ban");
manager.ban_ip(ip("192.0.2.8"), 60, "x").expect("ban");
manager.unban(ip("192.0.2.9"));
assert!(!manager.is_banned(ip("192.0.2.9")));
assert!(manager.is_banned(ip("192.0.2.8")));
manager.reset();
assert!(!manager.is_banned(ip("192.0.2.8")));
assert_eq!(manager.banned_count(), 0);
}
#[test]
fn durations_beyond_the_local_cap_are_clamped() {
let fake = FakeClock::default();
let manager = IpBanManager::with_clock(fake.clock());
manager.ban_ip(ip("192.0.2.9"), 7200, "long").expect("ban");
fake.advance(3600);
assert!(manager.is_banned(ip("192.0.2.9")));
fake.advance(1);
assert!(!manager.is_banned(ip("192.0.2.9")));
}
#[test]
fn counters_accumulate_per_ip_per_category() {
let counters = ViolationCounters::new();
let attacker = ip("192.0.2.9");
counters.record(attacker, &["sqli"]);
counters.record(attacker, &["sqli", "xss"]);
counters.record(ip("192.0.2.8"), &["xss"]);
let counts = counters.snapshot(attacker);
assert_eq!(counts.get("sqli"), Some(&2));
assert_eq!(counts.get("xss"), Some(&1));
assert_eq!(counters.snapshot(ip("192.0.2.8")).get("xss"), Some(&1));
assert_eq!(counters.tracked_ips(), 2);
counters.record(attacker, &[]);
assert_eq!(counters.snapshot(attacker).get("uncategorized"), Some(&1));
counters.reset();
assert_eq!(counters.tracked_ips(), 0);
}
#[test]
fn ipv4_mapped_violations_share_the_ipv4_counts() {
let counters = ViolationCounters::new();
counters.record(ip("::ffff:192.0.2.9"), &["sqli"]);
assert_eq!(counters.snapshot(ip("192.0.2.9")).get("sqli"), Some(&1));
}
#[test]
fn resolution_disabled_never_bans() {
let counts = HashMap::from([("sqli".to_owned(), 99)]);
let config = IpBanConfig {
enable_ip_banning: false,
..IpBanConfig::default()
};
assert!(
resolve_threshold_ban(&counts, &config, &["sqli"], "penetration_attempt").is_none()
);
}
#[test]
fn resolution_category_entry_wins_in_order() {
let counts = HashMap::from([("sqli".to_owned(), 3), ("xss".to_owned(), 10)]);
let config = IpBanConfig::new(
true,
10,
3600,
[
(
"sqli",
ThreatBanEntry {
threshold: 3,
duration: 60,
},
),
(
"xss",
ThreatBanEntry {
threshold: 5,
duration: 30,
},
),
],
)
.expect("valid config");
let resolved =
resolve_threshold_ban(&counts, &config, &["sqli", "xss"], "penetration_attempt")
.expect("sqli crosses first in list order");
assert_eq!(resolved.category.as_deref(), Some("sqli"));
assert_eq!(resolved.duration, 60);
assert_eq!(resolved.reason, "penetration_attempt:sqli");
let resolved =
resolve_threshold_ban(&counts, &config, &["xss", "sqli"], "penetration_attempt")
.expect("xss crosses first in list order");
assert_eq!(resolved.category.as_deref(), Some("xss"));
assert_eq!(resolved.duration, 30);
}
#[test]
fn resolution_below_every_entry_falls_back_to_the_flat_total() {
let config = IpBanConfig::new(
true,
10,
3600,
[(
"sqli",
ThreatBanEntry {
threshold: 5,
duration: 60,
},
)],
)
.expect("valid config");
let counts = HashMap::from([("sqli".to_owned(), 4), ("xss".to_owned(), 4)]);
assert!(
resolve_threshold_ban(&counts, &config, &["sqli"], "penetration_attempt").is_none()
);
let counts = HashMap::from([("sqli".to_owned(), 4), ("xss".to_owned(), 6)]);
let resolved = resolve_threshold_ban(&counts, &config, &["sqli"], "penetration_attempt")
.expect("the flat threshold fires on the total");
assert_eq!(resolved.category, None);
assert_eq!(resolved.duration, 3600);
assert_eq!(resolved.reason, "penetration_attempt");
}
#[test]
fn resolution_ignores_uncounted_entries() {
let config = IpBanConfig::new(
true,
10,
3600,
[(
"xss",
ThreatBanEntry {
threshold: 1,
duration: 60,
},
)],
)
.expect("valid config");
let counts = HashMap::from([("sqli".to_owned(), 4)]);
assert!(
resolve_threshold_ban(&counts, &config, &["sqli"], "penetration_attempt").is_none()
);
}
#[test]
fn register_violations_counts_then_bans_at_the_threshold() {
let fake = FakeClock::default();
let manager = IpBanManager::with_clock(fake.clock());
let counters = ViolationCounters::new();
let attacker = ip("192.0.2.9");
let config = IpBanConfig::new(
true,
100,
3600,
[(
"sqli",
ThreatBanEntry {
threshold: 3,
duration: 120,
},
)],
)
.expect("valid config");
assert!(
manager
.register_violations(
&counters,
attacker,
&["sqli"],
&config,
"penetration_attempt"
)
.is_none()
);
assert!(
manager
.register_violations(
&counters,
attacker,
&["sqli"],
&config,
"penetration_attempt"
)
.is_none()
);
assert!(
!manager.is_banned(attacker),
"below the threshold nobody is banned"
);
let resolved = manager
.register_violations(
&counters,
attacker,
&["sqli"],
&config,
"penetration_attempt",
)
.expect("the third violation crosses the entry");
assert_eq!(resolved.category.as_deref(), Some("sqli"));
assert_eq!(resolved.reason, "penetration_attempt:sqli");
assert!(manager.is_banned(attacker));
assert_eq!(
manager.ban_record(attacker).expect("record").reason,
"penetration_attempt:sqli"
);
}
#[test]
fn register_violations_without_banning_still_counts() {
let manager = IpBanManager::new();
let counters = ViolationCounters::new();
let attacker = ip("192.0.2.9");
let config = IpBanConfig {
enable_ip_banning: false,
..IpBanConfig::default()
};
for _ in 0..(config.auto_ban_threshold + 5) {
assert!(
manager
.register_violations(
&counters,
attacker,
&["xss"],
&config,
"penetration_attempt"
)
.is_none()
);
}
assert!(!manager.is_banned(attacker), "banning is off");
assert!(
counters.snapshot(attacker).values().sum::<u64>()
>= u64::from(config.auto_ban_threshold),
"the violations were counted anyway"
);
}
#[test]
fn register_violations_refusal_counts_but_does_not_ban() {
let manager = IpBanManager::new();
let counters = ViolationCounters::new();
let config = IpBanConfig::new(true, 1, 3600, Vec::<(String, ThreatBanEntry)>::new())
.expect("valid config");
assert!(
manager
.register_violations(
&counters,
ip("127.0.0.1"),
&["xss"],
&config,
"penetration_attempt"
)
.is_none()
);
assert!(!manager.is_banned(ip("127.0.0.1")));
assert_eq!(counters.snapshot(ip("127.0.0.1")).get("xss"), Some(&1));
}
#[test]
fn manager_debug_and_default_render_and_construct() {
let manager = IpBanManager::default();
assert!(format!("{manager:?}").contains("IpBanManager"));
assert!(format!("{:?}", ViolationCounters::default()).contains("ViolationCounters"));
assert!(manager.ban_record(ip("192.0.2.8")).is_none());
}
#[test]
fn trusted_proxies_parse_cidr_entries() {
let manager = IpBanManager::with_trusted_proxies_and_clock(
["10.0.0.0/8", "192.168.1.1"],
Arc::new(system_clock),
)
.expect("valid proxies");
assert!(format!("{manager:?}").contains("IpBanManager"));
assert!(!manager.is_banned(ip("10.1.2.3")));
}
#[test]
fn expired_bans_are_evicted_on_read() {
let fake = FakeClock::default();
let manager = IpBanManager::with_clock(fake.clock());
manager
.ban_ip(ip("192.0.2.9"), 60, "rate_limit")
.expect("ban recorded");
assert!(manager.ban_record(ip("192.0.2.9")).is_some());
fake.advance(120);
assert!(manager.ban_record(ip("192.0.2.9")).is_none());
assert!(!manager.is_banned(ip("192.0.2.9")));
}
}
#[cfg(test)]
mod unit_twins {
use super::*;
use std::str::FromStr;
fn ip(text: &str) -> IpAddr {
IpAddr::from_str(text).expect("test address")
}
struct ScriptedBanStore {
answer: Result<Option<f64>, crate::distributed::StoreError>,
}
impl crate::distributed::BanStore for ScriptedBanStore {
fn set_ban(
&self,
_key: &str,
_expiry: f64,
_ttl_seconds: u64,
) -> Result<(), crate::distributed::StoreError> {
Ok(())
}
fn get_ban(&self, _key: &str) -> Result<Option<f64>, crate::distributed::StoreError> {
self.answer.clone()
}
fn delete_ban(&self, _key: &str) -> Result<(), crate::distributed::StoreError> {
Ok(())
}
}
#[test]
fn trusted_proxy_lists_accept_and_reject_entries_fail_closed() {
for (entry, valid) in [
("10.0.0.0/8", true),
("192.168.1.1", true),
("not-a-network", false),
] {
let built =
IpBanManager::with_trusted_proxies_and_clock([entry], Arc::new(system_clock));
assert_eq!(built.is_ok(), valid, "entry {entry}");
if let Ok(manager) = built {
assert!(!manager.is_banned(ip("10.1.2.3")));
}
}
}
#[test]
fn distributed_lookups_answer_live_stale_missing_and_failed_backends() {
let now = 1_000.0_f64;
let live = IpBanManager::with_trusted_proxies_and_clock(
Vec::<String>::new(),
Arc::new(move || now),
)
.expect("manager")
.with_distributed_store(
Arc::new(ScriptedBanStore {
answer: Ok(Some(2_000.0)),
}),
"gc:",
);
assert!(live.is_banned(ip("192.0.2.71")));
live.ban_ip(ip("192.0.2.72"), 60, "penetration_attempt")
.expect("ban");
let stale = IpBanManager::with_trusted_proxies_and_clock(
Vec::<String>::new(),
Arc::new(move || now),
)
.expect("manager")
.with_distributed_store(
Arc::new(ScriptedBanStore {
answer: Ok(Some(500.0)),
}),
"gc:",
);
assert!(!stale.is_banned(ip("192.0.2.71")));
let missing = IpBanManager::with_trusted_proxies_and_clock(
Vec::<String>::new(),
Arc::new(move || now),
)
.expect("manager")
.with_distributed_store(Arc::new(ScriptedBanStore { answer: Ok(None) }), "gc:");
assert!(!missing.is_banned(ip("192.0.2.71")));
let failed = IpBanManager::with_trusted_proxies_and_clock(
Vec::<String>::new(),
Arc::new(move || now),
)
.expect("manager")
.with_distributed_store(
Arc::new(ScriptedBanStore {
answer: Err(crate::distributed::StoreError(String::from("down"))),
}),
"gc:",
);
assert!(!failed.is_banned(ip("192.0.2.71")));
}
#[test]
fn failing_distributed_ban_lookup_fails_secure() {
let store = Arc::new(crate::distributed::MemoryStore::default());
store.fail.store(true, std::sync::atomic::Ordering::Relaxed);
let manager = IpBanManager::with_clock(Arc::new(|| 1_000.0)).with_distributed_store(
Arc::clone(&store) as Arc<dyn crate::distributed::BanStore>,
"guard_core:",
);
let stranger = ip("192.0.2.77");
let error = manager.try_is_banned(stranger).expect_err("fail secure");
assert!(error.to_string().contains("backend down"));
assert!(!manager.is_banned(stranger));
}
#[test]
fn live_local_ban_answers_ok_even_when_the_store_is_down() {
let store = Arc::new(crate::distributed::MemoryStore::default());
let manager = IpBanManager::with_clock(Arc::new(|| 1_000.0)).with_distributed_store(
Arc::clone(&store) as Arc<dyn crate::distributed::BanStore>,
"guard_core:",
);
let attacker = ip("192.0.2.78");
manager
.ban_ip(attacker, 60, "threshold_exceeded")
.expect("ban");
store.fail.store(true, std::sync::atomic::Ordering::Relaxed);
assert!(manager.try_is_banned(attacker).expect("local hit"));
}
#[test]
fn distributed_ban_lookup_error_names_the_backend() {
let manager = IpBanManager::with_clock(Arc::new(|| 1_000.0));
assert!(!manager.try_is_banned(ip("192.0.2.79")).expect("no store"));
}
}