Skip to main content

guard_core_engine/
ip_ban.rs

1//! The dynamic IP ban store and the auto-ban engine: `ban_ip` with a
2//! duration, expiry-honoring `is_banned`, and the per-IP violation counters
3//! with the threshold resolution both are fed by.
4//!
5//! This is the Rust family's in-memory port of the reference engine's
6//! dynamic ban layer (`guard_core/handlers/ipban_handler.py` and the
7//! threshold helper `core/checks/helpers.py:_resolve_and_apply_threshold_ban`,
8//! via the Go port's `ipban.go`). Redis is out of scope: the store is
9//! process-local, exactly what both references fall back to when Redis is
10//! off, and the follow-up that adds Redis adopts the full distributed
11//! semantics (shared bans, expiry refresh, legacy key migration).
12//!
13//! ## Bans
14//!
15//! [`IpBanManager::ban_ip`] records an expiry (`now + duration`) and
16//! [`IpBanManager::is_banned`] honors it: a lookup past the expiry returns
17//! `false` and drops the entry. Mirroring the references' local store, an
18//! entry lives at most
19//! [`LOCAL_CACHE_TTL_CAP_SECONDS`] (3600) seconds - the reference clamps
20//! longer durations to the local cap when Redis is not configured
21//! (`clampToLocalCap`, cause `not configured`) and the Python `TTLCache`
22//! caps every local entry at the same TTL - so `ban_ip` with a longer
23//! duration records the clamped expiry. The default `auto_ban_duration`
24//! (3600) sits exactly at the cap, so default autobans are unaffected.
25//!
26//! The store is an LRU capped at [`LOCAL_CACHE_MAX_SIZE`] (10 000) entries
27//! (`localCacheMaxSize` / the reference `maxsize=10000`): inserting into a
28//! full store silently evicts the least recently banned entry, the same
29//! silent-overflow behavior the references log a warning for.
30//!
31//! ## Self-ban refusal
32//!
33//! `ban_ip` refuses (`Ok(false)`, no entry recorded) when the target
34//! overlaps the loopback space (`127.0.0.0/8`, `::1/128`) or a configured
35//! trusted proxy - the references' self-DoS guard, so a deployment can never
36//! ban its own ingress and go dark on itself.
37//!
38//! ## The auto-ban engine
39//!
40//! Detection violations are counted per IP per category in
41//! [`ViolationCounters`] (the reference's `suspicious_request_counts`). One
42//! call, [`IpBanManager::register_violations`], records the categories and
43//! applies the reference threshold resolution
44//! (`_resolve_and_apply_threshold_ban`, same shape as the TypeScript port's
45//! `resolveThresholdBan`):
46//!
47//! ```text
48//! banning disabled:                          no ban
49//! for each category in the recorded order:
50//!   threat_ban_config[category] exists and
51//!   count[category] >= its threshold:        ban (entry duration, reason:category)
52//! total of all counted categories >=
53//! auto_ban_threshold:                        ban (auto_ban_duration, reason)
54//! otherwise:                                 no ban
55//! ```
56//!
57//! The per-category entries win first (in the order the caller lists the
58//! categories), the flat threshold is the fallback measured against the
59//! total. The rate-limit stage feeds the same engine with the pseudo-category
60//! `rate_limit` when `enable_rate_limit_auto_ban` is on, so
61//! `threat_ban_config["rate_limit"]` overrides and the flat threshold backs
62//! it up, exactly as the reference pipeline does.
63//!
64//! # Example
65//!
66//! ```
67//! use std::net::IpAddr;
68//! use std::str::FromStr;
69//!
70//! use guard_core_engine::ip_ban::{IpBanConfig, IpBanManager};
71//!
72//! let manager = IpBanManager::new();
73//! let attacker = IpAddr::from_str("192.0.2.9").unwrap();
74//!
75//! assert!(manager.ban_ip(attacker, 60, "threshold_exceeded").unwrap());
76//! assert!(manager.is_banned(attacker));
77//!
78//! // Loopback is refused: banning it would self-DoS the deployment.
79//! let loopback = IpAddr::from_str("127.0.0.1").unwrap();
80//! assert!(!manager.ban_ip(loopback, 60, "mistake").unwrap());
81//! assert!(!manager.is_banned(loopback));
82//!
83//! manager.unban(attacker);
84//! assert!(!manager.is_banned(attacker));
85//! ```
86
87use std::collections::{HashMap, HashSet};
88use std::net::IpAddr;
89use std::num::NonZeroUsize;
90use std::sync::{Arc, Mutex};
91
92use lru::LruCache;
93
94use crate::ip_gate::{IpGateError, IpNet, canonical, parse_network_entry};
95
96/// The flat auto-ban fallback threshold (`auto_ban_threshold`, reference
97/// default).
98pub const DEFAULT_AUTO_BAN_THRESHOLD: u32 = 10;
99/// The flat auto-ban fallback duration in seconds (`auto_ban_duration`,
100/// reference default: one hour).
101pub const DEFAULT_AUTO_BAN_DURATION: u64 = 3600;
102/// The local store's TTL cap in seconds (`LOCAL_CACHE_TTL_CAP_SECONDS` in
103/// both references): no in-memory ban outlives it.
104pub const LOCAL_CACHE_TTL_CAP_SECONDS: u64 = 3600;
105/// The local ban store's entry cap (`localCacheMaxSize` /
106/// the reference `maxsize=10000`).
107pub const LOCAL_CACHE_MAX_SIZE: usize = 10_000;
108/// The violation counter store's IP cap (`_MAX_TRACKED_SUSPICIOUS_IPS`).
109pub const MAX_TRACKED_SUSPICIOUS_IPS: usize = 10_000;
110
111/// The pseudo-category the rate-limit stage feeds the auto-ban engine under
112/// (`enable_rate_limit_auto_ban`).
113pub const RATE_LIMIT_CATEGORY: &str = "rate_limit";
114
115/// One `threat_ban_config` entry: the violation count per category that
116/// triggers a ban, and how long the ban lasts.
117#[derive(Debug, Clone, Copy, PartialEq, Eq)]
118pub struct ThreatBanEntry {
119    /// Violations of this category after which the IP is banned.
120    pub threshold: u32,
121    /// Ban length in seconds.
122    pub duration: u64,
123}
124
125/// The IP-banning knobs (the reference `enable_ip_banning` /
126/// `auto_ban_threshold` / `auto_ban_duration` / `threat_ban_config` group).
127///
128/// Default is the opt-in pair: banning off (`false`, so a stage upgrade
129/// never starts banning uninvited - note the Python `SecurityConfig`
130/// defaults `enable_ip_banning` to `true`; the Rust family pins the
131/// conservative `false` like the Go port does) with the reference
132/// thresholds for when it is turned on.
133#[derive(Debug, Clone, PartialEq, Eq)]
134pub struct IpBanConfig {
135    /// `enable_ip_banning`. `false` makes every ban resolution a no-op;
136    /// violations still count (the reference counts regardless so enabling
137    /// banning later starts from observed history).
138    pub enable_ip_banning: bool,
139    /// `auto_ban_threshold`: the flat fallback, measured against the total
140    /// of all counted categories.
141    pub auto_ban_threshold: u32,
142    /// `auto_ban_duration`: the flat fallback ban length (seconds).
143    pub auto_ban_duration: u64,
144    /// `threat_ban_config`: per-category threshold/duration overrides,
145    /// consulted before the flat pair.
146    pub threat_ban_config: HashMap<String, ThreatBanEntry>,
147}
148
149impl Default for IpBanConfig {
150    fn default() -> Self {
151        Self {
152            enable_ip_banning: false,
153            auto_ban_threshold: DEFAULT_AUTO_BAN_THRESHOLD,
154            auto_ban_duration: DEFAULT_AUTO_BAN_DURATION,
155            threat_ban_config: HashMap::new(),
156        }
157    }
158}
159
160/// An invalid [`IpBanConfig`]: the config error [`IpBanConfig::new`] fails
161/// closed with.
162#[derive(Debug, Clone, PartialEq, Eq)]
163pub enum IpBanConfigError {
164    /// A threshold or duration that must be at least 1 is not.
165    NonPositive {
166        /// The rejected field name.
167        field: &'static str,
168    },
169    /// A `threat_ban_config` key that is neither a detection category nor
170    /// the `rate_limit` pseudo-category (the reference and the TypeScript
171    /// port both reject unknown categories).
172    UnknownCategory {
173        /// The rejected key.
174        category: String,
175    },
176}
177
178impl core::fmt::Display for IpBanConfigError {
179    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
180        match self {
181            Self::NonPositive { field } => write!(f, "invalid {field}: must be at least 1"),
182            Self::UnknownCategory { category } => write!(
183                f,
184                "unknown threat category '{category}': expected a detection category or 'rate_limit'"
185            ),
186        }
187    }
188}
189
190impl std::error::Error for IpBanConfigError {}
191
192/// The detection categories a `threat_ban_config` key may name: every
193/// pattern-table category plus the `rate_limit` pseudo-category (the
194/// reference's `ALL_DETECTION_CATEGORIES | {'rate_limit'}`).
195#[must_use]
196pub fn valid_threat_categories() -> HashSet<&'static str> {
197    let mut categories: HashSet<&'static str> = crate::patterns::table::PATTERN_DEFINITIONS
198        .iter()
199        .map(|entry| entry.category)
200        .collect();
201    categories.insert(RATE_LIMIT_CATEGORY);
202    categories
203}
204
205impl IpBanConfig {
206    /// Build and validate a config, failing closed on a non-positive
207    /// threshold/duration or an unknown `threat_ban_config` category.
208    ///
209    /// # Errors
210    ///
211    /// [`IpBanConfigError::NonPositive`] for a zero `auto_ban_threshold` or
212    /// `auto_ban_duration` (the reference rejects both with `ge=1`),
213    /// [`IpBanConfigError::UnknownCategory`] for a key outside
214    /// [`valid_threat_categories`].
215    pub fn new<I, K, V>(
216        enable_ip_banning: bool,
217        auto_ban_threshold: u32,
218        auto_ban_duration: u64,
219        threat_ban_config: I,
220    ) -> Result<Self, IpBanConfigError>
221    where
222        I: IntoIterator<Item = (K, V)>,
223        K: AsRef<str>,
224        V: Into<ThreatBanEntry>,
225    {
226        let config = Self {
227            enable_ip_banning,
228            auto_ban_threshold,
229            auto_ban_duration,
230            threat_ban_config: threat_ban_config
231                .into_iter()
232                .map(|(key, entry)| (key.as_ref().to_owned(), entry.into()))
233                .collect(),
234        };
235        config.validate()?;
236        Ok(config)
237    }
238
239    /// Validate this config. Struct-literal builders can run the same
240    /// checks the constructor applies.
241    ///
242    /// # Errors
243    ///
244    /// See [`IpBanConfig::new`].
245    pub fn validate(&self) -> Result<(), IpBanConfigError> {
246        if self.auto_ban_threshold == 0 {
247            return Err(IpBanConfigError::NonPositive {
248                field: "auto_ban_threshold",
249            });
250        }
251        if self.auto_ban_duration == 0 {
252            return Err(IpBanConfigError::NonPositive {
253                field: "auto_ban_duration",
254            });
255        }
256        let valid = valid_threat_categories();
257        for category in self.threat_ban_config.keys() {
258            if !valid.contains(category.as_str()) {
259                return Err(IpBanConfigError::UnknownCategory {
260                    category: category.to_owned(),
261                });
262            }
263        }
264        Ok(())
265    }
266}
267
268/// Why [`IpBanManager::ban_ip`] refused a ban.
269#[derive(Debug, Clone, Copy, PartialEq, Eq)]
270pub enum BanError {
271    /// A ban duration must be at least one second (the reference raises;
272    /// the Go port answers with a 400-shaped error).
273    NonPositiveDuration,
274}
275
276impl core::fmt::Display for BanError {
277    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
278        match self {
279            Self::NonPositiveDuration => write!(f, "ban duration must be positive"),
280        }
281    }
282}
283
284impl std::error::Error for BanError {}
285
286/// One recorded ban: when it expires and why it was issued. The reason is
287/// metadata for observability (the references persist it on the Redis
288/// side); the in-memory store only acts on the expiry.
289#[derive(Debug, Clone, PartialEq)]
290pub struct BanRecord {
291    /// Unix seconds after which the ban no longer applies.
292    pub expiry: f64,
293    /// The reason `ban_ip` was called with.
294    pub reason: String,
295}
296
297/// The monotonic-ish wall clock, in seconds since the Unix epoch.
298/// Injectable so expiry behavior is testable without sleeping.
299pub type Clock = Arc<dyn Fn() -> f64 + Send + Sync>;
300
301fn system_clock() -> f64 {
302    std::time::SystemTime::now()
303        .duration_since(std::time::UNIX_EPOCH)
304        .map_or(0.0, |d| d.as_secs_f64())
305}
306
307const LOOPBACK_NETWORKS: [&str; 2] = ["127.0.0.0/8", "::1/128"];
308
309/// A fresh shared ban store (the store lives behind an `Arc` so clones of
310/// the manager share it, the references' singleton semantics).
311fn new_ban_store() -> Arc<Mutex<LruCache<IpAddr, BanRecord>>> {
312    Arc::new(Mutex::new(LruCache::new(
313        NonZeroUsize::new(LOCAL_CACHE_MAX_SIZE).expect("constant above zero"),
314    )))
315}
316
317/// The dynamic IP ban store over one shared in-memory map.
318///
319/// Build it once at startup and share the handle across requests; the
320/// internal store is a mutex-guarded LRU, safe for concurrent services.
321/// Trusted proxies are configured at construction
322/// ([`IpBanManager::with_trusted_proxies`]) because they shape the
323/// self-ban refusal, not per-request behavior.
324pub struct IpBanManager {
325    bans: Arc<Mutex<LruCache<IpAddr, BanRecord>>>,
326    trusted_proxies: Vec<IpNet>,
327    clock: Clock,
328}
329
330impl Clone for IpBanManager {
331    /// A clone shares the ban store and the clock, exactly the references'
332    /// singleton semantics: bans recorded through any handle are visible to
333    /// every handle. Trusted proxies are copied (they shape the self-DoS
334    /// refusal only).
335    fn clone(&self) -> Self {
336        Self {
337            bans: Arc::clone(&self.bans),
338            trusted_proxies: self.trusted_proxies.clone(),
339            clock: Arc::clone(&self.clock),
340        }
341    }
342}
343
344impl core::fmt::Debug for IpBanManager {
345    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
346        f.debug_struct("IpBanManager")
347            .field("trusted_proxies", &self.trusted_proxies.len())
348            .finish_non_exhaustive()
349    }
350}
351
352impl Default for IpBanManager {
353    fn default() -> Self {
354        Self::new()
355    }
356}
357
358impl IpBanManager {
359    /// Build a ban store with no trusted proxies (loopback refusal only).
360    #[must_use]
361    pub fn new() -> Self {
362        Self {
363            bans: new_ban_store(),
364            trusted_proxies: Vec::new(),
365            clock: Arc::new(system_clock),
366        }
367    }
368
369    /// Build a ban store whose self-ban refusal also covers the given
370    /// trusted proxy networks (a bare IP or a CIDR range per entry, same
371    /// matching semantics as the IP gate's lists).
372    ///
373    /// # Errors
374    ///
375    /// [`IpGateError`] naming the offending entry, exactly like
376    /// `IpGateConfig::new` fails closed.
377    pub fn with_trusted_proxies<I>(entries: I) -> Result<Self, IpGateError>
378    where
379        I: IntoIterator,
380        I::Item: AsRef<str>,
381    {
382        Self::with_trusted_proxies_and_clock(entries, Arc::new(system_clock))
383    }
384
385    /// Build a ban store whose self-ban refusal also covers the given
386    /// trusted-proxy networks and whose wall clock is injected: the combined
387    /// seam [`with_trusted_proxies`](Self::with_trusted_proxies) and
388    /// [`with_clock`](Self::with_clock) cover separately. Adapters that run
389    /// deterministic expiry tests under a production-shaped proxy list use
390    /// this one.
391    ///
392    /// # Errors
393    ///
394    /// [`IpGateError`] naming the offending entry, exactly like
395    /// `IpBanManager::with_trusted_proxies` fails closed.
396    pub fn with_trusted_proxies_and_clock<I>(entries: I, clock: Clock) -> Result<Self, IpGateError>
397    where
398        I: IntoIterator,
399        I::Item: AsRef<str>,
400    {
401        let mut trusted_proxies = Vec::new();
402        for entry in entries {
403            let entry = entry.as_ref();
404            match parse_network_entry(entry) {
405                Some(network) => trusted_proxies.push(network),
406                None => {
407                    return Err(IpGateError {
408                        list: "trusted_proxies",
409                        entry: entry.to_owned(),
410                    });
411                }
412            }
413        }
414        Ok(Self {
415            bans: new_ban_store(),
416            trusted_proxies,
417            clock,
418        })
419    }
420
421    /// Swap the wall clock. Test seam: production builds use the system
422    /// clock; deterministic expiry coverage injects a fake.
423    #[must_use]
424    pub fn with_clock(clock: Clock) -> Self {
425        Self {
426            bans: new_ban_store(),
427            trusted_proxies: Vec::new(),
428            clock,
429        }
430    }
431
432    /// Ban `ip` for `duration` seconds, recording `reason`.
433    ///
434    /// Returns `Ok(true)` when the ban was recorded, `Ok(false)` when it
435    /// was refused by the self-DoS guard (loopback or a trusted proxy:
436    /// banning those would take the deployment's own view of client IPs
437    /// offline), and [`BanError::NonPositiveDuration`] for a zero duration.
438    /// Re-banning a live IP overwrites the record: the new expiry and
439    /// reason replace the old ones.
440    ///
441    /// Durations beyond [`LOCAL_CACHE_TTL_CAP_SECONDS`] are clamped to it:
442    /// this is the in-memory store, and the references cap their local
443    /// store the same way (only the Redis backend honors longer bans).
444    ///
445    /// # Errors
446    ///
447    /// [`BanError::NonPositiveDuration`] when `duration` is zero.
448    pub fn ban_ip(&self, ip: IpAddr, duration: u64, reason: &str) -> Result<bool, BanError> {
449        if duration == 0 {
450            return Err(BanError::NonPositiveDuration);
451        }
452        let ip = canonical(ip);
453        if self.self_ban_refusal(&ip) {
454            return Ok(false);
455        }
456        let clamped = duration.min(LOCAL_CACHE_TTL_CAP_SECONDS);
457        #[allow(clippy::cast_precision_loss)]
458        let expiry = (self.clock)() + clamped as f64;
459        self.bans.lock().expect("ban store").put(
460            ip,
461            BanRecord {
462                expiry,
463                reason: reason.to_owned(),
464            },
465        );
466        Ok(true)
467    }
468
469    /// Whether `ip` currently carries a live ban. An entry past its expiry
470    /// reads as unbanned and is dropped from the store.
471    #[must_use]
472    pub fn is_banned(&self, ip: IpAddr) -> bool {
473        let ip = canonical(ip);
474        let now = (self.clock)();
475        let mut bans = self.bans.lock().expect("ban store");
476        match bans.get(&ip) {
477            Some(record) if now <= record.expiry => true,
478            Some(_) => {
479                bans.pop(&ip);
480                false
481            }
482            None => false,
483        }
484    }
485
486    /// The recorded ban for `ip`, if one is still live (test and
487    /// observability seam: the reason a ban was issued).
488    #[must_use]
489    pub fn ban_record(&self, ip: IpAddr) -> Option<BanRecord> {
490        let ip = canonical(ip);
491        let now = (self.clock)();
492        let mut bans = self.bans.lock().expect("ban store");
493        match bans.get(&ip) {
494            Some(record) if now <= record.expiry => Some(record.clone()),
495            Some(_) => {
496                bans.pop(&ip);
497                None
498            }
499            None => None,
500        }
501    }
502
503    /// Lift `ip`'s ban, live or expired (`unban` in the references).
504    pub fn unban(&self, ip: IpAddr) {
505        self.bans.lock().expect("ban store").pop(&canonical(ip));
506    }
507
508    /// Drop every ban and every violation count (`reset` in the
509    /// references; test harnesses use it to isolate cases).
510    pub fn reset(&self) {
511        self.bans.lock().expect("ban store").clear();
512    }
513
514    /// How many bans the store currently holds, live or expired (test and
515    /// observability seam for the LRU cap).
516    #[must_use]
517    pub fn banned_count(&self) -> usize {
518        self.bans.lock().expect("ban store").len()
519    }
520
521    /// The self-DoS guard: loopback and trusted-proxy targets are refused.
522    fn self_ban_refusal(&self, ip: &IpAddr) -> bool {
523        LOOPBACK_NETWORKS
524            .iter()
525            .filter_map(|entry| parse_network_entry(entry))
526            .chain(self.trusted_proxies.iter().copied())
527            .any(|network| network.contains(*ip))
528    }
529}
530
531/// Per-IP per-category violation counts: the reference's
532/// `suspicious_request_counts`, the store the auto-ban engine resolves
533/// thresholds against.
534///
535/// The store is an LRU capped at [`MAX_TRACKED_SUSPICIOUS_IPS`] IPs (the
536/// reference evicts the oldest IP at the same cap), so attacker IP
537/// cardinality cannot grow it without bound.
538pub struct ViolationCounters {
539    counts: Arc<Mutex<LruCache<IpAddr, HashMap<String, u64>>>>,
540}
541
542impl Clone for ViolationCounters {
543    /// A clone shares the count store: violations recorded through any
544    /// handle are visible to every handle.
545    fn clone(&self) -> Self {
546        Self {
547            counts: Arc::clone(&self.counts),
548        }
549    }
550}
551
552impl core::fmt::Debug for ViolationCounters {
553    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
554        f.debug_struct("ViolationCounters").finish_non_exhaustive()
555    }
556}
557
558impl Default for ViolationCounters {
559    fn default() -> Self {
560        Self::new()
561    }
562}
563
564impl ViolationCounters {
565    /// Build an empty counter store.
566    #[must_use]
567    pub fn new() -> Self {
568        Self {
569            counts: Arc::new(Mutex::new(LruCache::new(
570                NonZeroUsize::new(MAX_TRACKED_SUSPICIOUS_IPS).expect("constant above zero"),
571            ))),
572        }
573    }
574
575    /// Record one violation per category for `ip` (each listed category
576    /// gains `+1`; the reference maps an empty category list to the
577    /// `uncategorized` pseudo-category, mirrored here).
578    pub fn record(&self, ip: IpAddr, categories: &[&str]) {
579        let categories: &[&str] = if categories.is_empty() {
580            &["uncategorized"]
581        } else {
582            categories
583        };
584        let mut counts = self.counts.lock().expect("violation counter store");
585        let ip_counts = counts
586            .try_get_or_insert_mut(canonical(ip), || {
587                Ok::<_, core::convert::Infallible>(HashMap::new())
588            })
589            .expect("ip capacity just reserved");
590        // The write loop is the guard's only lifetime reason: it ends here.
591        for category in categories {
592            *ip_counts.entry((*category).to_owned()).or_insert(0) += 1;
593        }
594        drop(counts);
595    }
596
597    /// A copy of `ip`'s counts (the resolution reads a snapshot while the
598    /// store stays locked for the next writer).
599    #[must_use]
600    pub fn snapshot(&self, ip: IpAddr) -> HashMap<String, u64> {
601        self.counts
602            .lock()
603            .expect("violation counter store")
604            .peek(&canonical(ip))
605            .cloned()
606            .unwrap_or_default()
607    }
608
609    /// Drop every count.
610    pub fn reset(&self) {
611        self.counts.lock().expect("violation counter store").clear();
612    }
613
614    /// How many IPs the store currently tracks (test/observability seam).
615    #[must_use]
616    pub fn tracked_ips(&self) -> usize {
617        self.counts.lock().expect("violation counter store").len()
618    }
619}
620
621/// The ban a threshold resolution decided on: how long, under what reason,
622/// and which category crossed (none for the flat fallback).
623#[derive(Debug, Clone, PartialEq, Eq)]
624pub struct ResolvedBan {
625    /// Ban length in seconds (the crossed entry's or the flat default's).
626    pub duration: u64,
627    /// The reason to record (`"reason"` for the flat fallback,
628    /// `"reason:category"` for a crossed entry, exactly as the references
629    /// build it).
630    pub reason: String,
631    /// The category whose entry crossed, `None` for the flat fallback.
632    pub category: Option<String>,
633}
634
635/// The reference threshold resolution against one IP's current counts.
636///
637/// This is `_resolve_and_apply_threshold_ban` (the TypeScript port's
638/// `resolveThresholdBan`): per-category entries first, in the order the
639/// categories are listed, then the flat threshold against the total.
640///
641/// Banning disabled always resolves to no ban.
642#[must_use]
643pub fn resolve_threshold_ban<S: core::hash::BuildHasher>(
644    counts: &HashMap<String, u64, S>,
645    config: &IpBanConfig,
646    threat_categories: &[&str],
647    reason: &str,
648) -> Option<ResolvedBan> {
649    if !config.enable_ip_banning {
650        return None;
651    }
652    for category in threat_categories {
653        let Some(entry) = config.threat_ban_config.get(*category) else {
654            continue;
655        };
656        let count = counts.get(*category).copied().unwrap_or(0);
657        if count < u64::from(entry.threshold) {
658            continue;
659        }
660        return Some(ResolvedBan {
661            duration: entry.duration,
662            reason: format!("{reason}:{category}"),
663            category: Some((*category).to_owned()),
664        });
665    }
666    let total: u64 = counts.values().sum();
667    if total < u64::from(config.auto_ban_threshold) {
668        return None;
669    }
670    Some(ResolvedBan {
671        duration: config.auto_ban_duration,
672        reason: reason.to_owned(),
673        category: None,
674    })
675}
676
677impl IpBanManager {
678    /// The auto-ban engine's one-call shape: record the categories for
679    /// `ip`, resolve the thresholds, and apply the ban when one crossed.
680    ///
681    /// `Some(..)` means a ban now stands (the stage answers with the
682    /// family's banned shape), `None` means nothing fired - banning
683    /// disabled, no threshold crossed, or the self-DoS guard refusing the
684    /// target. Violations are counted regardless of the banning switch,
685    /// mirroring the reference: enabling banning later starts from observed
686    /// history.
687    #[must_use]
688    pub fn register_violations(
689        &self,
690        counters: &ViolationCounters,
691        ip: IpAddr,
692        categories: &[&str],
693        config: &IpBanConfig,
694        reason: &str,
695    ) -> Option<ResolvedBan> {
696        counters.record(ip, categories);
697        let counts = counters.snapshot(ip);
698        let resolved = resolve_threshold_ban(&counts, config, categories, reason)?;
699        match self.ban_ip(ip, resolved.duration, &resolved.reason) {
700            Ok(true) => Some(resolved),
701            _ => None,
702        }
703    }
704}
705
706#[cfg(test)]
707mod tests {
708    use super::*;
709    use std::str::FromStr;
710    use std::sync::atomic::{AtomicU64, Ordering};
711
712    /// A fake clock: f64 unix seconds starting at `1_000.0`, advanced by
713    /// `advance`.
714    #[derive(Clone, Default)]
715    struct FakeClock(Arc<AtomicU64>);
716
717    impl FakeClock {
718        fn advance(&self, seconds: u64) {
719            self.0.fetch_add(seconds, Ordering::Relaxed);
720        }
721
722        fn clock(&self) -> Clock {
723            let state = self.0.clone();
724            #[allow(clippy::cast_precision_loss)]
725            Arc::new(move || state.load(Ordering::Relaxed) as f64)
726        }
727    }
728
729    fn ip(text: &str) -> IpAddr {
730        IpAddr::from_str(text).expect("test address")
731    }
732
733    #[test]
734    fn default_config_is_off_with_reference_thresholds() {
735        let config = IpBanConfig::default();
736        assert!(!config.enable_ip_banning);
737        assert_eq!(config.auto_ban_threshold, 10);
738        assert_eq!(config.auto_ban_duration, 3600);
739        assert!(config.threat_ban_config.is_empty());
740    }
741
742    #[test]
743    fn config_validation_fails_closed() {
744        let error =
745            IpBanConfig::new(true, 0, 3600, Vec::<(String, ThreatBanEntry)>::new()).unwrap_err();
746        assert_eq!(
747            error,
748            IpBanConfigError::NonPositive {
749                field: "auto_ban_threshold"
750            }
751        );
752        assert_eq!(
753            error.to_string(),
754            "invalid auto_ban_threshold: must be at least 1"
755        );
756
757        let error =
758            IpBanConfig::new(true, 10, 0, Vec::<(String, ThreatBanEntry)>::new()).unwrap_err();
759        assert_eq!(
760            error,
761            IpBanConfigError::NonPositive {
762                field: "auto_ban_duration"
763            }
764        );
765
766        let error = IpBanConfig::new(
767            true,
768            10,
769            3600,
770            [(
771                "not_a_category",
772                ThreatBanEntry {
773                    threshold: 1,
774                    duration: 1,
775                },
776            )],
777        )
778        .unwrap_err();
779        assert_eq!(
780            error,
781            IpBanConfigError::UnknownCategory {
782                category: "not_a_category".to_owned()
783            }
784        );
785        assert_eq!(
786            error.to_string(),
787            "unknown threat category 'not_a_category': expected a detection category or 'rate_limit'"
788        );
789    }
790
791    #[test]
792    fn config_accepts_table_categories_and_rate_limit() {
793        IpBanConfig::new(
794            true,
795            10,
796            3600,
797            [
798                (
799                    "sqli",
800                    ThreatBanEntry {
801                        threshold: 3,
802                        duration: 60,
803                    },
804                ),
805                (
806                    RATE_LIMIT_CATEGORY,
807                    ThreatBanEntry {
808                        threshold: 5,
809                        duration: 30,
810                    },
811                ),
812            ],
813        )
814        .expect("both categories are valid");
815        assert!(valid_threat_categories().contains("sqli"));
816        assert!(valid_threat_categories().contains(RATE_LIMIT_CATEGORY));
817    }
818
819    #[test]
820    fn ban_then_lookup_then_expiry() {
821        let fake = FakeClock::default();
822        let manager = IpBanManager::with_clock(fake.clock());
823        manager
824            .ban_ip(ip("192.0.2.9"), 60, "threshold_exceeded")
825            .expect("ban");
826        assert!(manager.is_banned(ip("192.0.2.9")));
827        assert_eq!(
828            manager
829                .ban_record(ip("192.0.2.9"))
830                .expect("live record")
831                .reason,
832            "threshold_exceeded"
833        );
834
835        // Just before expiry the ban still stands.
836        fake.advance(59);
837        assert!(manager.is_banned(ip("192.0.2.9")));
838        // At the expiry second (`now <= expiry`) it still stands; strictly
839        // past it the reference boundary (`now > expiry` evicts) releases.
840        fake.advance(1);
841        assert!(manager.is_banned(ip("192.0.2.9")));
842        fake.advance(1);
843        assert!(!manager.is_banned(ip("192.0.2.9")));
844        assert_eq!(manager.banned_count(), 0, "the expired entry was dropped");
845    }
846
847    #[test]
848    fn rebanning_overwrites_expiry_and_reason() {
849        let manager = IpBanManager::new();
850        manager.ban_ip(ip("192.0.2.9"), 60, "first").expect("ban");
851        manager.ban_ip(ip("192.0.2.9"), 120, "second").expect("ban");
852        let record = manager.ban_record(ip("192.0.2.9")).expect("live record");
853        assert_eq!(record.reason, "second");
854    }
855
856    #[test]
857    fn zero_duration_is_an_error() {
858        let manager = IpBanManager::new();
859        assert_eq!(
860            manager.ban_ip(ip("192.0.2.9"), 0, "nope").unwrap_err(),
861            BanError::NonPositiveDuration
862        );
863        assert!(!manager.is_banned(ip("192.0.2.9")));
864        assert_eq!(
865            BanError::NonPositiveDuration.to_string(),
866            "ban duration must be positive"
867        );
868    }
869
870    #[test]
871    fn loopback_is_refused_for_both_families() {
872        let manager = IpBanManager::new();
873        for loopback in ["127.0.0.1", "127.8.8.8", "::1"] {
874            assert!(
875                !manager
876                    .ban_ip(ip(loopback), 60, "mistake")
877                    .expect("refusal is not an error"),
878                "{loopback} must be refused"
879            );
880            assert!(!manager.is_banned(ip(loopback)));
881        }
882    }
883
884    #[test]
885    fn trusted_proxies_are_refused_and_parsed_fail_closed() {
886        let manager = IpBanManager::with_trusted_proxies(["10.0.0.0/8", "2001:db8:ffff::/48"])
887            .expect("valid");
888        assert!(
889            !manager
890                .ban_ip(ip("10.1.2.3"), 60, "proxy")
891                .expect("refusal")
892        );
893        assert!(
894            !manager
895                .ban_ip(ip("2001:db8:ffff::1"), 60, "proxy")
896                .expect("refusal")
897        );
898        assert!(
899            manager
900                .ban_ip(ip("11.1.2.3"), 60, "outside the proxy range")
901                .expect("ban"),
902            "an address outside the proxy range is bannable"
903        );
904
905        let error = IpBanManager::with_trusted_proxies(["not-an-ip"]).unwrap_err();
906        assert_eq!(error.list, "trusted_proxies");
907        assert_eq!(error.entry, "not-an-ip");
908    }
909
910    #[test]
911    fn ipv4_mapped_targets_canonicalize() {
912        let manager = IpBanManager::new();
913        // The mapped loopback form is the loopback: refused.
914        assert!(
915            !manager
916                .ban_ip(ip("::ffff:127.0.0.1"), 60, "mistake")
917                .expect("refusal")
918        );
919        // A mapped ban and its IPv4 form are one entry.
920        manager
921            .ban_ip(ip("::ffff:192.0.2.9"), 60, "x")
922            .expect("ban");
923        assert!(manager.is_banned(ip("192.0.2.9")));
924        assert_eq!(manager.banned_count(), 1);
925    }
926
927    #[test]
928    fn unban_and_reset() {
929        let manager = IpBanManager::new();
930        manager.ban_ip(ip("192.0.2.9"), 60, "x").expect("ban");
931        manager.ban_ip(ip("192.0.2.8"), 60, "x").expect("ban");
932        manager.unban(ip("192.0.2.9"));
933        assert!(!manager.is_banned(ip("192.0.2.9")));
934        assert!(manager.is_banned(ip("192.0.2.8")));
935        manager.reset();
936        assert!(!manager.is_banned(ip("192.0.2.8")));
937        assert_eq!(manager.banned_count(), 0);
938    }
939
940    #[test]
941    fn durations_beyond_the_local_cap_are_clamped() {
942        let fake = FakeClock::default();
943        let manager = IpBanManager::with_clock(fake.clock());
944        manager.ban_ip(ip("192.0.2.9"), 7200, "long").expect("ban");
945        // The expiry is the local cap, not the requested duration.
946        fake.advance(3600);
947        assert!(manager.is_banned(ip("192.0.2.9")));
948        fake.advance(1);
949        assert!(!manager.is_banned(ip("192.0.2.9")));
950    }
951
952    #[test]
953    fn counters_accumulate_per_ip_per_category() {
954        let counters = ViolationCounters::new();
955        let attacker = ip("192.0.2.9");
956        counters.record(attacker, &["sqli"]);
957        counters.record(attacker, &["sqli", "xss"]);
958        counters.record(ip("192.0.2.8"), &["xss"]);
959
960        let counts = counters.snapshot(attacker);
961        assert_eq!(counts.get("sqli"), Some(&2));
962        assert_eq!(counts.get("xss"), Some(&1));
963        assert_eq!(counters.snapshot(ip("192.0.2.8")).get("xss"), Some(&1));
964        assert_eq!(counters.tracked_ips(), 2);
965
966        // The empty category list records the uncategorized pseudo-category,
967        // exactly like the reference's `_increment_suspicious_counts`.
968        counters.record(attacker, &[]);
969        assert_eq!(counters.snapshot(attacker).get("uncategorized"), Some(&1));
970
971        counters.reset();
972        assert_eq!(counters.tracked_ips(), 0);
973    }
974
975    #[test]
976    fn ipv4_mapped_violations_share_the_ipv4_counts() {
977        let counters = ViolationCounters::new();
978        counters.record(ip("::ffff:192.0.2.9"), &["sqli"]);
979        assert_eq!(counters.snapshot(ip("192.0.2.9")).get("sqli"), Some(&1));
980    }
981
982    #[test]
983    fn resolution_disabled_never_bans() {
984        let counts = HashMap::from([("sqli".to_owned(), 99)]);
985        let config = IpBanConfig {
986            enable_ip_banning: false,
987            ..IpBanConfig::default()
988        };
989        assert!(
990            resolve_threshold_ban(&counts, &config, &["sqli"], "penetration_attempt").is_none()
991        );
992    }
993
994    #[test]
995    fn resolution_category_entry_wins_in_order() {
996        let counts = HashMap::from([("sqli".to_owned(), 3), ("xss".to_owned(), 10)]);
997        let config = IpBanConfig::new(
998            true,
999            10,
1000            3600,
1001            [
1002                (
1003                    "sqli",
1004                    ThreatBanEntry {
1005                        threshold: 3,
1006                        duration: 60,
1007                    },
1008                ),
1009                (
1010                    "xss",
1011                    ThreatBanEntry {
1012                        threshold: 5,
1013                        duration: 30,
1014                    },
1015                ),
1016            ],
1017        )
1018        .expect("valid config");
1019
1020        let resolved =
1021            resolve_threshold_ban(&counts, &config, &["sqli", "xss"], "penetration_attempt")
1022                .expect("sqli crosses first in list order");
1023        assert_eq!(resolved.category.as_deref(), Some("sqli"));
1024        assert_eq!(resolved.duration, 60);
1025        assert_eq!(resolved.reason, "penetration_attempt:sqli");
1026
1027        // Listed the other way round, the xss entry wins.
1028        let resolved =
1029            resolve_threshold_ban(&counts, &config, &["xss", "sqli"], "penetration_attempt")
1030                .expect("xss crosses first in list order");
1031        assert_eq!(resolved.category.as_deref(), Some("xss"));
1032        assert_eq!(resolved.duration, 30);
1033    }
1034
1035    #[test]
1036    fn resolution_below_every_entry_falls_back_to_the_flat_total() {
1037        let config = IpBanConfig::new(
1038            true,
1039            10,
1040            3600,
1041            [(
1042                "sqli",
1043                ThreatBanEntry {
1044                    threshold: 5,
1045                    duration: 60,
1046                },
1047            )],
1048        )
1049        .expect("valid config");
1050
1051        // No category at its entry threshold, total 8 < 10: nothing.
1052        let counts = HashMap::from([("sqli".to_owned(), 4), ("xss".to_owned(), 4)]);
1053        assert!(
1054            resolve_threshold_ban(&counts, &config, &["sqli"], "penetration_attempt").is_none()
1055        );
1056
1057        // Still no entry crossing, but the total reaches the flat threshold.
1058        let counts = HashMap::from([("sqli".to_owned(), 4), ("xss".to_owned(), 6)]);
1059        let resolved = resolve_threshold_ban(&counts, &config, &["sqli"], "penetration_attempt")
1060            .expect("the flat threshold fires on the total");
1061        assert_eq!(resolved.category, None);
1062        assert_eq!(resolved.duration, 3600);
1063        assert_eq!(resolved.reason, "penetration_attempt");
1064    }
1065
1066    #[test]
1067    fn resolution_ignores_uncounted_entries() {
1068        let config = IpBanConfig::new(
1069            true,
1070            10,
1071            3600,
1072            [(
1073                "xss",
1074                ThreatBanEntry {
1075                    threshold: 1,
1076                    duration: 60,
1077                },
1078            )],
1079        )
1080        .expect("valid config");
1081        // sqli has no entry and no count: it neither bans nor blocks the
1082        // flat fallback.
1083        let counts = HashMap::from([("sqli".to_owned(), 4)]);
1084        assert!(
1085            resolve_threshold_ban(&counts, &config, &["sqli"], "penetration_attempt").is_none()
1086        );
1087    }
1088
1089    #[test]
1090    fn register_violations_counts_then_bans_at_the_threshold() {
1091        let fake = FakeClock::default();
1092        let manager = IpBanManager::with_clock(fake.clock());
1093        let counters = ViolationCounters::new();
1094        let attacker = ip("192.0.2.9");
1095        let config = IpBanConfig::new(
1096            true,
1097            100,
1098            3600,
1099            [(
1100                "sqli",
1101                ThreatBanEntry {
1102                    threshold: 3,
1103                    duration: 120,
1104                },
1105            )],
1106        )
1107        .expect("valid config");
1108
1109        assert!(
1110            manager
1111                .register_violations(
1112                    &counters,
1113                    attacker,
1114                    &["sqli"],
1115                    &config,
1116                    "penetration_attempt"
1117                )
1118                .is_none()
1119        );
1120        assert!(
1121            manager
1122                .register_violations(
1123                    &counters,
1124                    attacker,
1125                    &["sqli"],
1126                    &config,
1127                    "penetration_attempt"
1128                )
1129                .is_none()
1130        );
1131        assert!(
1132            !manager.is_banned(attacker),
1133            "below the threshold nobody is banned"
1134        );
1135
1136        let resolved = manager
1137            .register_violations(
1138                &counters,
1139                attacker,
1140                &["sqli"],
1141                &config,
1142                "penetration_attempt",
1143            )
1144            .expect("the third violation crosses the entry");
1145        assert_eq!(resolved.category.as_deref(), Some("sqli"));
1146        assert_eq!(resolved.reason, "penetration_attempt:sqli");
1147        assert!(manager.is_banned(attacker));
1148        assert_eq!(
1149            manager.ban_record(attacker).expect("record").reason,
1150            "penetration_attempt:sqli"
1151        );
1152    }
1153
1154    #[test]
1155    fn register_violations_without_banning_still_counts() {
1156        let manager = IpBanManager::new();
1157        let counters = ViolationCounters::new();
1158        let attacker = ip("192.0.2.9");
1159        let config = IpBanConfig::default();
1160
1161        for _ in 0..(config.auto_ban_threshold + 5) {
1162            assert!(
1163                manager
1164                    .register_violations(
1165                        &counters,
1166                        attacker,
1167                        &["xss"],
1168                        &config,
1169                        "penetration_attempt"
1170                    )
1171                    .is_none()
1172            );
1173        }
1174        assert!(!manager.is_banned(attacker), "banning is off");
1175        assert!(
1176            counters.snapshot(attacker).values().sum::<u64>()
1177                >= u64::from(config.auto_ban_threshold),
1178            "the violations were counted anyway"
1179        );
1180    }
1181
1182    #[test]
1183    fn register_violations_refusal_counts_but_does_not_ban() {
1184        let manager = IpBanManager::new();
1185        let counters = ViolationCounters::new();
1186        let config = IpBanConfig::new(true, 1, 3600, Vec::<(String, ThreatBanEntry)>::new())
1187            .expect("valid config");
1188        // The flat threshold is 1, but the target is loopback: refused.
1189        assert!(
1190            manager
1191                .register_violations(
1192                    &counters,
1193                    ip("127.0.0.1"),
1194                    &["xss"],
1195                    &config,
1196                    "penetration_attempt"
1197                )
1198                .is_none()
1199        );
1200        assert!(!manager.is_banned(ip("127.0.0.1")));
1201        assert_eq!(counters.snapshot(ip("127.0.0.1")).get("xss"), Some(&1));
1202    }
1203}