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/// The defaults are the reference `SecurityConfig` defaults: banning on
129/// (`enable_ip_banning = true`, `guard_core/_security_config_fields.py`)
130/// with the reference threshold/duration pair.
131#[derive(Debug, Clone, PartialEq, Eq)]
132pub struct IpBanConfig {
133    /// `enable_ip_banning`. `false` makes every ban resolution a no-op;
134    /// violations still count (the reference counts regardless so enabling
135    /// banning later starts from observed history).
136    pub enable_ip_banning: bool,
137    /// `auto_ban_threshold`: the flat fallback, measured against the total
138    /// of all counted categories.
139    pub auto_ban_threshold: u32,
140    /// `auto_ban_duration`: the flat fallback ban length (seconds).
141    pub auto_ban_duration: u64,
142    /// `threat_ban_config`: per-category threshold/duration overrides,
143    /// consulted before the flat pair.
144    pub threat_ban_config: HashMap<String, ThreatBanEntry>,
145}
146
147impl Default for IpBanConfig {
148    fn default() -> Self {
149        Self {
150            enable_ip_banning: true,
151            auto_ban_threshold: DEFAULT_AUTO_BAN_THRESHOLD,
152            auto_ban_duration: DEFAULT_AUTO_BAN_DURATION,
153            threat_ban_config: HashMap::new(),
154        }
155    }
156}
157
158/// An invalid [`IpBanConfig`]: the config error [`IpBanConfig::new`] fails
159/// closed with.
160#[derive(Debug, Clone, PartialEq, Eq)]
161pub enum IpBanConfigError {
162    /// A threshold or duration that must be at least 1 is not.
163    NonPositive {
164        /// The rejected field name.
165        field: &'static str,
166    },
167    /// A `threat_ban_config` key that is neither a detection category nor
168    /// the `rate_limit` pseudo-category (the reference and the TypeScript
169    /// port both reject unknown categories).
170    UnknownCategory {
171        /// The rejected key.
172        category: String,
173    },
174}
175
176impl core::fmt::Display for IpBanConfigError {
177    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
178        match self {
179            Self::NonPositive { field } => write!(f, "invalid {field}: must be at least 1"),
180            Self::UnknownCategory { category } => write!(
181                f,
182                "unknown threat category '{category}': expected a detection category or 'rate_limit'"
183            ),
184        }
185    }
186}
187
188impl std::error::Error for IpBanConfigError {}
189
190/// The detection categories a `threat_ban_config` key may name: every
191/// pattern-table category plus the `rate_limit` pseudo-category (the
192/// reference's `ALL_DETECTION_CATEGORIES | {'rate_limit'}`).
193#[must_use]
194pub fn valid_threat_categories() -> HashSet<&'static str> {
195    let mut categories: HashSet<&'static str> = crate::patterns::table::PATTERN_DEFINITIONS
196        .iter()
197        .map(|entry| entry.category)
198        .collect();
199    categories.insert(RATE_LIMIT_CATEGORY);
200    categories
201}
202
203impl IpBanConfig {
204    /// Build and validate a config, failing closed on a non-positive
205    /// threshold/duration or an unknown `threat_ban_config` category.
206    ///
207    /// # Errors
208    ///
209    /// [`IpBanConfigError::NonPositive`] for a zero `auto_ban_threshold` or
210    /// `auto_ban_duration` (the reference rejects both with `ge=1`),
211    /// [`IpBanConfigError::UnknownCategory`] for a key outside
212    /// [`valid_threat_categories`].
213    pub fn new<I, K, V>(
214        enable_ip_banning: bool,
215        auto_ban_threshold: u32,
216        auto_ban_duration: u64,
217        threat_ban_config: I,
218    ) -> Result<Self, IpBanConfigError>
219    where
220        I: IntoIterator<Item = (K, V)>,
221        K: AsRef<str>,
222        V: Into<ThreatBanEntry>,
223    {
224        let config = Self {
225            enable_ip_banning,
226            auto_ban_threshold,
227            auto_ban_duration,
228            threat_ban_config: threat_ban_config
229                .into_iter()
230                .map(|(key, entry)| (key.as_ref().to_owned(), entry.into()))
231                .collect(),
232        };
233        config.validate()?;
234        Ok(config)
235    }
236
237    /// Validate this config. Struct-literal builders can run the same
238    /// checks the constructor applies.
239    ///
240    /// # Errors
241    ///
242    /// See [`IpBanConfig::new`].
243    pub fn validate(&self) -> Result<(), IpBanConfigError> {
244        if self.auto_ban_threshold == 0 {
245            return Err(IpBanConfigError::NonPositive {
246                field: "auto_ban_threshold",
247            });
248        }
249        if self.auto_ban_duration == 0 {
250            return Err(IpBanConfigError::NonPositive {
251                field: "auto_ban_duration",
252            });
253        }
254        let valid = valid_threat_categories();
255        for category in self.threat_ban_config.keys() {
256            if !valid.contains(category.as_str()) {
257                return Err(IpBanConfigError::UnknownCategory {
258                    category: category.to_owned(),
259                });
260            }
261        }
262        Ok(())
263    }
264}
265
266/// Why [`IpBanManager::ban_ip`] refused a ban.
267#[derive(Debug, Clone, Copy, PartialEq, Eq)]
268pub enum BanError {
269    /// A ban duration must be at least one second (the reference raises;
270    /// the Go port answers with a 400-shaped error).
271    NonPositiveDuration,
272}
273
274impl core::fmt::Display for BanError {
275    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
276        match self {
277            Self::NonPositiveDuration => write!(f, "ban duration must be positive"),
278        }
279    }
280}
281
282impl std::error::Error for BanError {}
283
284/// One recorded ban: when it expires and why it was issued. The reason is
285/// metadata for observability (the references persist it on the Redis
286/// side); the in-memory store only acts on the expiry.
287#[derive(Debug, Clone, PartialEq)]
288pub struct BanRecord {
289    /// Unix seconds after which the ban no longer applies.
290    pub expiry: f64,
291    /// The reason `ban_ip` was called with.
292    pub reason: String,
293}
294
295/// The monotonic-ish wall clock, in seconds since the Unix epoch.
296/// Injectable so expiry behavior is testable without sleeping.
297pub type Clock = Arc<dyn Fn() -> f64 + Send + Sync>;
298
299fn system_clock() -> f64 {
300    std::time::SystemTime::now()
301        .duration_since(std::time::UNIX_EPOCH)
302        .map_or(0.0, |d| d.as_secs_f64())
303}
304
305const LOOPBACK_NETWORKS: [&str; 2] = ["127.0.0.0/8", "::1/128"];
306
307/// A fresh shared ban store (the store lives behind an `Arc` so clones of
308/// the manager share it, the references' singleton semantics).
309fn new_ban_store() -> Arc<Mutex<LruCache<IpAddr, BanRecord>>> {
310    Arc::new(Mutex::new(LruCache::new(
311        NonZeroUsize::new(LOCAL_CACHE_MAX_SIZE).expect("constant above zero"),
312    )))
313}
314
315/// The dynamic IP ban store over one shared in-memory map.
316///
317/// Build it once at startup and share the handle across requests; the
318/// internal store is a mutex-guarded LRU, safe for concurrent services.
319/// Trusted proxies are configured at construction
320/// ([`IpBanManager::with_trusted_proxies`]) because they shape the
321/// self-ban refusal, not per-request behavior.
322pub struct IpBanManager {
323    bans: Arc<Mutex<LruCache<IpAddr, BanRecord>>>,
324    trusted_proxies: Vec<IpNet>,
325    clock: Clock,
326    ban_store: Option<Arc<dyn crate::distributed::BanStore>>,
327    redis_prefix: String,
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            ban_store: self.ban_store.clone(),
341            redis_prefix: self.redis_prefix.clone(),
342        }
343    }
344}
345
346impl core::fmt::Debug for IpBanManager {
347    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
348        f.debug_struct("IpBanManager")
349            .field("trusted_proxies", &self.trusted_proxies.len())
350            .finish_non_exhaustive()
351    }
352}
353
354impl Default for IpBanManager {
355    fn default() -> Self {
356        Self::new()
357    }
358}
359
360impl IpBanManager {
361    /// Build a ban store with no trusted proxies (loopback refusal only).
362    #[must_use]
363    pub fn new() -> Self {
364        Self {
365            bans: new_ban_store(),
366            ban_store: None,
367            redis_prefix: String::from("guard_core:"),
368            trusted_proxies: Vec::new(),
369            clock: Arc::new(system_clock),
370        }
371    }
372
373    /// Build a ban store whose self-ban refusal also covers the given
374    /// trusted proxy networks (a bare IP or a CIDR range per entry, same
375    /// matching semantics as the IP gate's lists).
376    ///
377    /// # Errors
378    ///
379    /// [`IpGateError`] naming the offending entry, exactly like
380    /// `IpGateConfig::new` fails closed.
381    pub fn with_trusted_proxies<I>(entries: I) -> Result<Self, IpGateError>
382    where
383        I: IntoIterator,
384        I::Item: AsRef<str>,
385    {
386        Self::with_trusted_proxies_and_clock(entries, Arc::new(system_clock))
387    }
388
389    /// Build a ban store whose self-ban refusal also covers the given
390    /// trusted-proxy networks and whose wall clock is injected: the combined
391    /// seam [`with_trusted_proxies`](Self::with_trusted_proxies) and
392    /// [`with_clock`](Self::with_clock) cover separately. Adapters that run
393    /// deterministic expiry tests under a production-shaped proxy list use
394    /// this one.
395    ///
396    /// # Errors
397    ///
398    /// [`IpGateError`] naming the offending entry, exactly like
399    /// `IpBanManager::with_trusted_proxies` fails closed.
400    pub fn with_trusted_proxies_and_clock<I>(entries: I, clock: Clock) -> Result<Self, IpGateError>
401    where
402        I: IntoIterator,
403        I::Item: AsRef<str>,
404    {
405        let mut trusted_proxies = Vec::new();
406        for entry in entries {
407            let entry = entry.as_ref();
408            match parse_network_entry(entry) {
409                Some(network) => trusted_proxies.push(network),
410                None => {
411                    return Err(IpGateError {
412                        list: "trusted_proxies",
413                        entry: entry.to_owned(),
414                    });
415                }
416            }
417        }
418        Ok(Self {
419            bans: new_ban_store(),
420            ban_store: None,
421            redis_prefix: String::from("guard_core:"),
422            trusted_proxies,
423            clock,
424        })
425    }
426
427    /// Swap the wall clock. Test seam: production builds use the system
428    /// clock; deterministic expiry coverage injects a fake.
429    /// Run the ban engine over a distributed store (the reference
430    /// `redis_handler` seam): ban writes go to
431    /// `{prefix}banned_ips:{ip}` with the ban TTL (a failed write
432    /// degrades to the local store, the reference `except` branch), and
433    /// lookups fall through to the store on a local miss, caching the
434    /// answer (the reference `_check_redis_exact`). A failing lookup
435    /// surfaces the backend error through [`IpBanManager::try_is_banned`]
436    /// so a caller can fail secure the way the reference's
437    /// `GuardRedisError(503)` propagates into the pipeline; the lossy
438    /// [`IpBanManager::is_banned`] keeps reading an error as unbanned.
439    #[must_use]
440    pub fn with_distributed_store(
441        mut self,
442        store: Arc<dyn crate::distributed::BanStore>,
443        redis_prefix: &str,
444    ) -> Self {
445        self.ban_store = Some(store);
446        redis_prefix.clone_into(&mut self.redis_prefix);
447        self
448    }
449
450    pub fn with_clock(clock: Clock) -> Self {
451        Self {
452            bans: new_ban_store(),
453            ban_store: None,
454            redis_prefix: String::from("guard_core:"),
455            trusted_proxies: Vec::new(),
456            clock,
457        }
458    }
459
460    /// Ban `ip` for `duration` seconds, recording `reason`.
461    ///
462    /// Returns `Ok(true)` when the ban was recorded, `Ok(false)` when it
463    /// was refused by the self-DoS guard (loopback or a trusted proxy:
464    /// banning those would take the deployment's own view of client IPs
465    /// offline), and [`BanError::NonPositiveDuration`] for a zero duration.
466    /// Re-banning a live IP overwrites the record: the new expiry and
467    /// reason replace the old ones.
468    ///
469    /// Durations beyond [`LOCAL_CACHE_TTL_CAP_SECONDS`] are clamped to it:
470    /// this is the in-memory store, and the references cap their local
471    /// store the same way (only the Redis backend honors longer bans).
472    ///
473    /// # Errors
474    ///
475    /// [`BanError::NonPositiveDuration`] when `duration` is zero.
476    pub fn ban_ip(&self, ip: IpAddr, duration: u64, reason: &str) -> Result<bool, BanError> {
477        if duration == 0 {
478            return Err(BanError::NonPositiveDuration);
479        }
480        let ip = canonical(ip);
481        if self.self_ban_refusal(&ip) {
482            return Ok(false);
483        }
484        let clamped = duration.min(LOCAL_CACHE_TTL_CAP_SECONDS);
485        #[allow(clippy::cast_precision_loss)]
486        let expiry = (self.clock)() + clamped as f64;
487        self.bans.lock().expect("ban store").put(
488            ip,
489            BanRecord {
490                expiry,
491                reason: reason.to_owned(),
492            },
493        );
494        if let Some(store) = &self.ban_store {
495            // The distributed store honors the unclamped TTL (only the
496            // local cache is capped, exactly the reference). A failed
497            // write degrades to the local store above.
498            #[allow(clippy::cast_precision_loss)]
499            let distributed_expiry = (self.clock)() + duration as f64;
500            let _ = store.set_ban(
501                &crate::distributed::ban_key(&self.redis_prefix, ip),
502                distributed_expiry,
503                duration,
504            );
505        }
506        Ok(true)
507    }
508
509    /// Whether `ip` currently carries a live ban. An entry past its expiry
510    /// reads as unbanned and is dropped from the store.
511    ///
512    /// This is the lossy form: a failing distributed lookup reads as
513    /// unbanned. Fail-secure callers use [`IpBanManager::try_is_banned`],
514    /// which surfaces the backend error the way the reference's
515    /// `GuardRedisError(503)` propagates into its pipeline.
516    #[must_use]
517    pub fn is_banned(&self, ip: IpAddr) -> bool {
518        self.try_is_banned(ip).unwrap_or(false)
519    }
520
521    /// [`IpBanManager::is_banned`] with the backend error surfaced: a
522    /// failing distributed ban lookup returns [`Err`] instead of reading
523    /// as unbanned, so a caller can fail secure (the reference
524    /// `redis_handler.safe_operation` raising `GuardRedisError(503, "Redis
525    /// operation failed")` out of `_check_redis_exact`'s `get_key`).
526    ///
527    /// # Errors
528    ///
529    /// [`crate::distributed::StoreError`] when the distributed lookup
530    /// fails and the local cache holds no live answer.
531    pub fn try_is_banned(&self, ip: IpAddr) -> Result<bool, crate::distributed::StoreError> {
532        let ip = canonical(ip);
533        let now = (self.clock)();
534        {
535            let mut bans = self.bans.lock().expect("ban store");
536            match bans.get(&ip) {
537                Some(record) if now <= record.expiry => return Ok(true),
538                Some(_) => {
539                    bans.pop(&ip);
540                }
541                None => {}
542            }
543        }
544        self.check_distributed_ban(ip, now)
545    }
546
547    /// The distributed lookup on a local miss (the reference
548    /// `_check_redis_exact`): a stored expiry at or after `now` counts as
549    /// banned and is cached locally, a stale one is deleted, a backend
550    /// error is returned to the caller (the fail-secure seam; the lossy
551    /// [`IpBanManager::is_banned`] flattens it to unbanned).
552    fn check_distributed_ban(
553        &self,
554        ip: IpAddr,
555        now: f64,
556    ) -> Result<bool, crate::distributed::StoreError> {
557        let Some(store) = &self.ban_store else {
558            return Ok(false);
559        };
560        match store.get_ban(&crate::distributed::ban_key(&self.redis_prefix, ip)) {
561            Ok(Some(expiry)) if now <= expiry => {
562                self.bans.lock().expect("ban store").put(
563                    ip,
564                    BanRecord {
565                        expiry,
566                        reason: String::from("distributed"),
567                    },
568                );
569                Ok(true)
570            }
571            Ok(Some(_)) => {
572                let _ = store.delete_ban(&crate::distributed::ban_key(&self.redis_prefix, ip));
573                Ok(false)
574            }
575            Ok(None) => Ok(false),
576            Err(error) => Err(error),
577        }
578    }
579
580    /// The recorded ban for `ip`, if one is still live (test and
581    /// observability seam: the reason a ban was issued).
582    #[must_use]
583    pub fn ban_record(&self, ip: IpAddr) -> Option<BanRecord> {
584        let ip = canonical(ip);
585        let now = (self.clock)();
586        let mut bans = self.bans.lock().expect("ban store");
587        match bans.get(&ip) {
588            Some(record) if now <= record.expiry => Some(record.clone()),
589            Some(_) => {
590                bans.pop(&ip);
591                None
592            }
593            None => None,
594        }
595    }
596
597    /// Lift `ip`'s ban, live or expired (`unban` in the references).
598    pub fn unban(&self, ip: IpAddr) {
599        self.bans.lock().expect("ban store").pop(&canonical(ip));
600    }
601
602    /// Drop every ban and every violation count (`reset` in the
603    /// references; test harnesses use it to isolate cases).
604    pub fn reset(&self) {
605        self.bans.lock().expect("ban store").clear();
606    }
607
608    /// How many bans the store currently holds, live or expired (test and
609    /// observability seam for the LRU cap).
610    #[must_use]
611    pub fn banned_count(&self) -> usize {
612        self.bans.lock().expect("ban store").len()
613    }
614
615    /// The self-DoS guard: loopback and trusted-proxy targets are refused.
616    fn self_ban_refusal(&self, ip: &IpAddr) -> bool {
617        LOOPBACK_NETWORKS
618            .iter()
619            .filter_map(|entry| parse_network_entry(entry))
620            .chain(self.trusted_proxies.iter().copied())
621            .any(|network| network.contains(*ip))
622    }
623}
624
625/// Per-IP per-category violation counts: the reference's
626/// `suspicious_request_counts`, the store the auto-ban engine resolves
627/// thresholds against.
628///
629/// The store is an LRU capped at [`MAX_TRACKED_SUSPICIOUS_IPS`] IPs (the
630/// reference evicts the oldest IP at the same cap), so attacker IP
631/// cardinality cannot grow it without bound.
632pub struct ViolationCounters {
633    counts: Arc<Mutex<LruCache<IpAddr, HashMap<String, u64>>>>,
634}
635
636impl Clone for ViolationCounters {
637    /// A clone shares the count store: violations recorded through any
638    /// handle are visible to every handle.
639    fn clone(&self) -> Self {
640        Self {
641            counts: Arc::clone(&self.counts),
642        }
643    }
644}
645
646impl core::fmt::Debug for ViolationCounters {
647    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
648        f.debug_struct("ViolationCounters").finish_non_exhaustive()
649    }
650}
651
652impl Default for ViolationCounters {
653    fn default() -> Self {
654        Self::new()
655    }
656}
657
658impl ViolationCounters {
659    /// Build an empty counter store.
660    #[must_use]
661    pub fn new() -> Self {
662        Self {
663            counts: Arc::new(Mutex::new(LruCache::new(
664                NonZeroUsize::new(MAX_TRACKED_SUSPICIOUS_IPS).expect("constant above zero"),
665            ))),
666        }
667    }
668
669    /// Record one violation per category for `ip` (each listed category
670    /// gains `+1`; the reference maps an empty category list to the
671    /// `uncategorized` pseudo-category, mirrored here).
672    pub fn record(&self, ip: IpAddr, categories: &[&str]) {
673        let categories: &[&str] = if categories.is_empty() {
674            &["uncategorized"]
675        } else {
676            categories
677        };
678        let mut counts = self.counts.lock().expect("violation counter store");
679        let ip_counts = counts
680            .try_get_or_insert_mut(canonical(ip), || {
681                Ok::<_, core::convert::Infallible>(HashMap::new())
682            })
683            .expect("ip capacity just reserved");
684        // The write loop is the guard's only lifetime reason: it ends here.
685        for category in categories {
686            *ip_counts.entry((*category).to_owned()).or_insert(0) += 1;
687        }
688        drop(counts);
689    }
690
691    /// A copy of `ip`'s counts (the resolution reads a snapshot while the
692    /// store stays locked for the next writer).
693    #[must_use]
694    pub fn snapshot(&self, ip: IpAddr) -> HashMap<String, u64> {
695        self.counts
696            .lock()
697            .expect("violation counter store")
698            .peek(&canonical(ip))
699            .cloned()
700            .unwrap_or_default()
701    }
702
703    /// Drop every count.
704    pub fn reset(&self) {
705        self.counts.lock().expect("violation counter store").clear();
706    }
707
708    /// How many IPs the store currently tracks (test/observability seam).
709    #[must_use]
710    pub fn tracked_ips(&self) -> usize {
711        self.counts.lock().expect("violation counter store").len()
712    }
713}
714
715/// The ban a threshold resolution decided on: how long, under what reason,
716/// and which category crossed (none for the flat fallback).
717#[derive(Debug, Clone, PartialEq, Eq)]
718pub struct ResolvedBan {
719    /// Ban length in seconds (the crossed entry's or the flat default's).
720    pub duration: u64,
721    /// The reason to record (`"reason"` for the flat fallback,
722    /// `"reason:category"` for a crossed entry, exactly as the references
723    /// build it).
724    pub reason: String,
725    /// The category whose entry crossed, `None` for the flat fallback.
726    pub category: Option<String>,
727}
728
729/// The reference threshold resolution against one IP's current counts.
730///
731/// This is `_resolve_and_apply_threshold_ban` (the TypeScript port's
732/// `resolveThresholdBan`): per-category entries first, in the order the
733/// categories are listed, then the flat threshold against the total.
734///
735/// Banning disabled always resolves to no ban.
736#[must_use]
737pub fn resolve_threshold_ban<S: core::hash::BuildHasher>(
738    counts: &HashMap<String, u64, S>,
739    config: &IpBanConfig,
740    threat_categories: &[&str],
741    reason: &str,
742) -> Option<ResolvedBan> {
743    if !config.enable_ip_banning {
744        return None;
745    }
746    for category in threat_categories {
747        let Some(entry) = config.threat_ban_config.get(*category) else {
748            continue;
749        };
750        let count = counts.get(*category).copied().unwrap_or(0);
751        if count < u64::from(entry.threshold) {
752            continue;
753        }
754        return Some(ResolvedBan {
755            duration: entry.duration,
756            reason: format!("{reason}:{category}"),
757            category: Some((*category).to_owned()),
758        });
759    }
760    let total: u64 = counts.values().sum();
761    if total < u64::from(config.auto_ban_threshold) {
762        return None;
763    }
764    Some(ResolvedBan {
765        duration: config.auto_ban_duration,
766        reason: reason.to_owned(),
767        category: None,
768    })
769}
770
771impl IpBanManager {
772    /// The auto-ban engine's one-call shape: record the categories for
773    /// `ip`, resolve the thresholds, and apply the ban when one crossed.
774    ///
775    /// `Some(..)` means a ban now stands (the stage answers with the
776    /// family's banned shape), `None` means nothing fired - banning
777    /// disabled, no threshold crossed, or the self-DoS guard refusing the
778    /// target. Violations are counted regardless of the banning switch,
779    /// mirroring the reference: enabling banning later starts from observed
780    /// history.
781    #[must_use]
782    pub fn register_violations(
783        &self,
784        counters: &ViolationCounters,
785        ip: IpAddr,
786        categories: &[&str],
787        config: &IpBanConfig,
788        reason: &str,
789    ) -> Option<ResolvedBan> {
790        counters.record(ip, categories);
791        let counts = counters.snapshot(ip);
792        let resolved = resolve_threshold_ban(&counts, config, categories, reason)?;
793        match self.ban_ip(ip, resolved.duration, &resolved.reason) {
794            Ok(true) => Some(resolved),
795            _ => None,
796        }
797    }
798}
799
800#[cfg(test)]
801mod tests {
802    use super::*;
803    use std::str::FromStr;
804    use std::sync::atomic::{AtomicU64, Ordering};
805
806    /// A fake clock: f64 unix seconds starting at `1_000.0`, advanced by
807    /// `advance`.
808    #[derive(Clone, Default)]
809    struct FakeClock(Arc<AtomicU64>);
810
811    impl FakeClock {
812        fn advance(&self, seconds: u64) {
813            self.0.fetch_add(seconds, Ordering::Relaxed);
814        }
815
816        fn clock(&self) -> Clock {
817            let state = self.0.clone();
818            #[allow(clippy::cast_precision_loss)]
819            Arc::new(move || state.load(Ordering::Relaxed) as f64)
820        }
821    }
822
823    fn ip(text: &str) -> IpAddr {
824        IpAddr::from_str(text).expect("test address")
825    }
826
827    #[test]
828    fn default_config_matches_the_reference() {
829        let config = IpBanConfig::default();
830        assert!(config.enable_ip_banning);
831        assert_eq!(config.auto_ban_threshold, 10);
832        assert_eq!(config.auto_ban_duration, 3600);
833        assert!(config.threat_ban_config.is_empty());
834    }
835
836    #[test]
837    fn config_validation_fails_closed() {
838        let error =
839            IpBanConfig::new(true, 0, 3600, Vec::<(String, ThreatBanEntry)>::new()).unwrap_err();
840        assert_eq!(
841            error,
842            IpBanConfigError::NonPositive {
843                field: "auto_ban_threshold"
844            }
845        );
846        assert_eq!(
847            error.to_string(),
848            "invalid auto_ban_threshold: must be at least 1"
849        );
850
851        let error =
852            IpBanConfig::new(true, 10, 0, Vec::<(String, ThreatBanEntry)>::new()).unwrap_err();
853        assert_eq!(
854            error,
855            IpBanConfigError::NonPositive {
856                field: "auto_ban_duration"
857            }
858        );
859
860        let error = IpBanConfig::new(
861            true,
862            10,
863            3600,
864            [(
865                "not_a_category",
866                ThreatBanEntry {
867                    threshold: 1,
868                    duration: 1,
869                },
870            )],
871        )
872        .unwrap_err();
873        assert_eq!(
874            error,
875            IpBanConfigError::UnknownCategory {
876                category: "not_a_category".to_owned()
877            }
878        );
879        assert_eq!(
880            error.to_string(),
881            "unknown threat category 'not_a_category': expected a detection category or 'rate_limit'"
882        );
883    }
884
885    #[test]
886    fn config_accepts_table_categories_and_rate_limit() {
887        IpBanConfig::new(
888            true,
889            10,
890            3600,
891            [
892                (
893                    "sqli",
894                    ThreatBanEntry {
895                        threshold: 3,
896                        duration: 60,
897                    },
898                ),
899                (
900                    RATE_LIMIT_CATEGORY,
901                    ThreatBanEntry {
902                        threshold: 5,
903                        duration: 30,
904                    },
905                ),
906            ],
907        )
908        .expect("both categories are valid");
909        assert!(valid_threat_categories().contains("sqli"));
910        assert!(valid_threat_categories().contains(RATE_LIMIT_CATEGORY));
911    }
912
913    #[test]
914    fn ban_then_lookup_then_expiry() {
915        let fake = FakeClock::default();
916        let manager = IpBanManager::with_clock(fake.clock());
917        manager
918            .ban_ip(ip("192.0.2.9"), 60, "threshold_exceeded")
919            .expect("ban");
920        assert!(manager.is_banned(ip("192.0.2.9")));
921        assert_eq!(
922            manager
923                .ban_record(ip("192.0.2.9"))
924                .expect("live record")
925                .reason,
926            "threshold_exceeded"
927        );
928
929        // Just before expiry the ban still stands.
930        fake.advance(59);
931        assert!(manager.is_banned(ip("192.0.2.9")));
932        // At the expiry second (`now <= expiry`) it still stands; strictly
933        // past it the reference boundary (`now > expiry` evicts) releases.
934        fake.advance(1);
935        assert!(manager.is_banned(ip("192.0.2.9")));
936        fake.advance(1);
937        assert!(!manager.is_banned(ip("192.0.2.9")));
938        assert_eq!(manager.banned_count(), 0, "the expired entry was dropped");
939    }
940
941    #[test]
942    fn rebanning_overwrites_expiry_and_reason() {
943        let manager = IpBanManager::new();
944        manager.ban_ip(ip("192.0.2.9"), 60, "first").expect("ban");
945        manager.ban_ip(ip("192.0.2.9"), 120, "second").expect("ban");
946        let record = manager.ban_record(ip("192.0.2.9")).expect("live record");
947        assert_eq!(record.reason, "second");
948    }
949
950    #[test]
951    fn zero_duration_is_an_error() {
952        let manager = IpBanManager::new();
953        assert_eq!(
954            manager.ban_ip(ip("192.0.2.9"), 0, "nope").unwrap_err(),
955            BanError::NonPositiveDuration
956        );
957        assert!(!manager.is_banned(ip("192.0.2.9")));
958        assert_eq!(
959            BanError::NonPositiveDuration.to_string(),
960            "ban duration must be positive"
961        );
962    }
963
964    #[test]
965    fn loopback_is_refused_for_both_families() {
966        let manager = IpBanManager::new();
967        for loopback in ["127.0.0.1", "127.8.8.8", "::1"] {
968            assert!(
969                !manager
970                    .ban_ip(ip(loopback), 60, "mistake")
971                    .expect("refusal is not an error"),
972                "{loopback} must be refused"
973            );
974            assert!(!manager.is_banned(ip(loopback)));
975        }
976    }
977
978    #[test]
979    fn trusted_proxies_are_refused_and_parsed_fail_closed() {
980        let manager = IpBanManager::with_trusted_proxies(["10.0.0.0/8", "2001:db8:ffff::/48"])
981            .expect("valid");
982        assert!(
983            !manager
984                .ban_ip(ip("10.1.2.3"), 60, "proxy")
985                .expect("refusal")
986        );
987        assert!(
988            !manager
989                .ban_ip(ip("2001:db8:ffff::1"), 60, "proxy")
990                .expect("refusal")
991        );
992        assert!(
993            manager
994                .ban_ip(ip("11.1.2.3"), 60, "outside the proxy range")
995                .expect("ban"),
996            "an address outside the proxy range is bannable"
997        );
998
999        let error = IpBanManager::with_trusted_proxies(["not-an-ip"]).unwrap_err();
1000        assert_eq!(error.list, "trusted_proxies");
1001        assert_eq!(error.entry, "not-an-ip");
1002    }
1003
1004    #[test]
1005    fn ipv4_mapped_targets_canonicalize() {
1006        let manager = IpBanManager::new();
1007        // The mapped loopback form is the loopback: refused.
1008        assert!(
1009            !manager
1010                .ban_ip(ip("::ffff:127.0.0.1"), 60, "mistake")
1011                .expect("refusal")
1012        );
1013        // A mapped ban and its IPv4 form are one entry.
1014        manager
1015            .ban_ip(ip("::ffff:192.0.2.9"), 60, "x")
1016            .expect("ban");
1017        assert!(manager.is_banned(ip("192.0.2.9")));
1018        assert_eq!(manager.banned_count(), 1);
1019    }
1020
1021    #[test]
1022    fn unban_and_reset() {
1023        let manager = IpBanManager::new();
1024        manager.ban_ip(ip("192.0.2.9"), 60, "x").expect("ban");
1025        manager.ban_ip(ip("192.0.2.8"), 60, "x").expect("ban");
1026        manager.unban(ip("192.0.2.9"));
1027        assert!(!manager.is_banned(ip("192.0.2.9")));
1028        assert!(manager.is_banned(ip("192.0.2.8")));
1029        manager.reset();
1030        assert!(!manager.is_banned(ip("192.0.2.8")));
1031        assert_eq!(manager.banned_count(), 0);
1032    }
1033
1034    #[test]
1035    fn durations_beyond_the_local_cap_are_clamped() {
1036        let fake = FakeClock::default();
1037        let manager = IpBanManager::with_clock(fake.clock());
1038        manager.ban_ip(ip("192.0.2.9"), 7200, "long").expect("ban");
1039        // The expiry is the local cap, not the requested duration.
1040        fake.advance(3600);
1041        assert!(manager.is_banned(ip("192.0.2.9")));
1042        fake.advance(1);
1043        assert!(!manager.is_banned(ip("192.0.2.9")));
1044    }
1045
1046    #[test]
1047    fn counters_accumulate_per_ip_per_category() {
1048        let counters = ViolationCounters::new();
1049        let attacker = ip("192.0.2.9");
1050        counters.record(attacker, &["sqli"]);
1051        counters.record(attacker, &["sqli", "xss"]);
1052        counters.record(ip("192.0.2.8"), &["xss"]);
1053
1054        let counts = counters.snapshot(attacker);
1055        assert_eq!(counts.get("sqli"), Some(&2));
1056        assert_eq!(counts.get("xss"), Some(&1));
1057        assert_eq!(counters.snapshot(ip("192.0.2.8")).get("xss"), Some(&1));
1058        assert_eq!(counters.tracked_ips(), 2);
1059
1060        // The empty category list records the uncategorized pseudo-category,
1061        // exactly like the reference's `_increment_suspicious_counts`.
1062        counters.record(attacker, &[]);
1063        assert_eq!(counters.snapshot(attacker).get("uncategorized"), Some(&1));
1064
1065        counters.reset();
1066        assert_eq!(counters.tracked_ips(), 0);
1067    }
1068
1069    #[test]
1070    fn ipv4_mapped_violations_share_the_ipv4_counts() {
1071        let counters = ViolationCounters::new();
1072        counters.record(ip("::ffff:192.0.2.9"), &["sqli"]);
1073        assert_eq!(counters.snapshot(ip("192.0.2.9")).get("sqli"), Some(&1));
1074    }
1075
1076    #[test]
1077    fn resolution_disabled_never_bans() {
1078        let counts = HashMap::from([("sqli".to_owned(), 99)]);
1079        let config = IpBanConfig {
1080            enable_ip_banning: false,
1081            ..IpBanConfig::default()
1082        };
1083        assert!(
1084            resolve_threshold_ban(&counts, &config, &["sqli"], "penetration_attempt").is_none()
1085        );
1086    }
1087
1088    #[test]
1089    fn resolution_category_entry_wins_in_order() {
1090        let counts = HashMap::from([("sqli".to_owned(), 3), ("xss".to_owned(), 10)]);
1091        let config = IpBanConfig::new(
1092            true,
1093            10,
1094            3600,
1095            [
1096                (
1097                    "sqli",
1098                    ThreatBanEntry {
1099                        threshold: 3,
1100                        duration: 60,
1101                    },
1102                ),
1103                (
1104                    "xss",
1105                    ThreatBanEntry {
1106                        threshold: 5,
1107                        duration: 30,
1108                    },
1109                ),
1110            ],
1111        )
1112        .expect("valid config");
1113
1114        let resolved =
1115            resolve_threshold_ban(&counts, &config, &["sqli", "xss"], "penetration_attempt")
1116                .expect("sqli crosses first in list order");
1117        assert_eq!(resolved.category.as_deref(), Some("sqli"));
1118        assert_eq!(resolved.duration, 60);
1119        assert_eq!(resolved.reason, "penetration_attempt:sqli");
1120
1121        // Listed the other way round, the xss entry wins.
1122        let resolved =
1123            resolve_threshold_ban(&counts, &config, &["xss", "sqli"], "penetration_attempt")
1124                .expect("xss crosses first in list order");
1125        assert_eq!(resolved.category.as_deref(), Some("xss"));
1126        assert_eq!(resolved.duration, 30);
1127    }
1128
1129    #[test]
1130    fn resolution_below_every_entry_falls_back_to_the_flat_total() {
1131        let config = IpBanConfig::new(
1132            true,
1133            10,
1134            3600,
1135            [(
1136                "sqli",
1137                ThreatBanEntry {
1138                    threshold: 5,
1139                    duration: 60,
1140                },
1141            )],
1142        )
1143        .expect("valid config");
1144
1145        // No category at its entry threshold, total 8 < 10: nothing.
1146        let counts = HashMap::from([("sqli".to_owned(), 4), ("xss".to_owned(), 4)]);
1147        assert!(
1148            resolve_threshold_ban(&counts, &config, &["sqli"], "penetration_attempt").is_none()
1149        );
1150
1151        // Still no entry crossing, but the total reaches the flat threshold.
1152        let counts = HashMap::from([("sqli".to_owned(), 4), ("xss".to_owned(), 6)]);
1153        let resolved = resolve_threshold_ban(&counts, &config, &["sqli"], "penetration_attempt")
1154            .expect("the flat threshold fires on the total");
1155        assert_eq!(resolved.category, None);
1156        assert_eq!(resolved.duration, 3600);
1157        assert_eq!(resolved.reason, "penetration_attempt");
1158    }
1159
1160    #[test]
1161    fn resolution_ignores_uncounted_entries() {
1162        let config = IpBanConfig::new(
1163            true,
1164            10,
1165            3600,
1166            [(
1167                "xss",
1168                ThreatBanEntry {
1169                    threshold: 1,
1170                    duration: 60,
1171                },
1172            )],
1173        )
1174        .expect("valid config");
1175        // sqli has no entry and no count: it neither bans nor blocks the
1176        // flat fallback.
1177        let counts = HashMap::from([("sqli".to_owned(), 4)]);
1178        assert!(
1179            resolve_threshold_ban(&counts, &config, &["sqli"], "penetration_attempt").is_none()
1180        );
1181    }
1182
1183    #[test]
1184    fn register_violations_counts_then_bans_at_the_threshold() {
1185        let fake = FakeClock::default();
1186        let manager = IpBanManager::with_clock(fake.clock());
1187        let counters = ViolationCounters::new();
1188        let attacker = ip("192.0.2.9");
1189        let config = IpBanConfig::new(
1190            true,
1191            100,
1192            3600,
1193            [(
1194                "sqli",
1195                ThreatBanEntry {
1196                    threshold: 3,
1197                    duration: 120,
1198                },
1199            )],
1200        )
1201        .expect("valid config");
1202
1203        assert!(
1204            manager
1205                .register_violations(
1206                    &counters,
1207                    attacker,
1208                    &["sqli"],
1209                    &config,
1210                    "penetration_attempt"
1211                )
1212                .is_none()
1213        );
1214        assert!(
1215            manager
1216                .register_violations(
1217                    &counters,
1218                    attacker,
1219                    &["sqli"],
1220                    &config,
1221                    "penetration_attempt"
1222                )
1223                .is_none()
1224        );
1225        assert!(
1226            !manager.is_banned(attacker),
1227            "below the threshold nobody is banned"
1228        );
1229
1230        let resolved = manager
1231            .register_violations(
1232                &counters,
1233                attacker,
1234                &["sqli"],
1235                &config,
1236                "penetration_attempt",
1237            )
1238            .expect("the third violation crosses the entry");
1239        assert_eq!(resolved.category.as_deref(), Some("sqli"));
1240        assert_eq!(resolved.reason, "penetration_attempt:sqli");
1241        assert!(manager.is_banned(attacker));
1242        assert_eq!(
1243            manager.ban_record(attacker).expect("record").reason,
1244            "penetration_attempt:sqli"
1245        );
1246    }
1247
1248    #[test]
1249    fn register_violations_without_banning_still_counts() {
1250        let manager = IpBanManager::new();
1251        let counters = ViolationCounters::new();
1252        let attacker = ip("192.0.2.9");
1253        let config = IpBanConfig {
1254            enable_ip_banning: false,
1255            ..IpBanConfig::default()
1256        };
1257
1258        for _ in 0..(config.auto_ban_threshold + 5) {
1259            assert!(
1260                manager
1261                    .register_violations(
1262                        &counters,
1263                        attacker,
1264                        &["xss"],
1265                        &config,
1266                        "penetration_attempt"
1267                    )
1268                    .is_none()
1269            );
1270        }
1271        assert!(!manager.is_banned(attacker), "banning is off");
1272        assert!(
1273            counters.snapshot(attacker).values().sum::<u64>()
1274                >= u64::from(config.auto_ban_threshold),
1275            "the violations were counted anyway"
1276        );
1277    }
1278
1279    #[test]
1280    fn register_violations_refusal_counts_but_does_not_ban() {
1281        let manager = IpBanManager::new();
1282        let counters = ViolationCounters::new();
1283        let config = IpBanConfig::new(true, 1, 3600, Vec::<(String, ThreatBanEntry)>::new())
1284            .expect("valid config");
1285        // The flat threshold is 1, but the target is loopback: refused.
1286        assert!(
1287            manager
1288                .register_violations(
1289                    &counters,
1290                    ip("127.0.0.1"),
1291                    &["xss"],
1292                    &config,
1293                    "penetration_attempt"
1294                )
1295                .is_none()
1296        );
1297        assert!(!manager.is_banned(ip("127.0.0.1")));
1298        assert_eq!(counters.snapshot(ip("127.0.0.1")).get("xss"), Some(&1));
1299    }
1300
1301    #[test]
1302    fn manager_debug_and_default_render_and_construct() {
1303        let manager = IpBanManager::default();
1304        assert!(format!("{manager:?}").contains("IpBanManager"));
1305        assert!(format!("{:?}", ViolationCounters::default()).contains("ViolationCounters"));
1306        assert!(manager.ban_record(ip("192.0.2.8")).is_none());
1307    }
1308
1309    #[test]
1310    fn trusted_proxies_parse_cidr_entries() {
1311        let manager = IpBanManager::with_trusted_proxies_and_clock(
1312            ["10.0.0.0/8", "192.168.1.1"],
1313            Arc::new(system_clock),
1314        )
1315        .expect("valid proxies");
1316        assert!(format!("{manager:?}").contains("IpBanManager"));
1317        // a trusted-proxy IP is not treated as the client
1318        assert!(!manager.is_banned(ip("10.1.2.3")));
1319    }
1320
1321    #[test]
1322    fn expired_bans_are_evicted_on_read() {
1323        let fake = FakeClock::default();
1324        let manager = IpBanManager::with_clock(fake.clock());
1325        manager
1326            .ban_ip(ip("192.0.2.9"), 60, "rate_limit")
1327            .expect("ban recorded");
1328        assert!(manager.ban_record(ip("192.0.2.9")).is_some());
1329        fake.advance(120);
1330        // the expired ban is popped on read
1331        assert!(manager.ban_record(ip("192.0.2.9")).is_none());
1332        assert!(!manager.is_banned(ip("192.0.2.9")));
1333    }
1334}
1335
1336#[cfg(test)]
1337mod unit_twins {
1338    use super::*;
1339    use std::str::FromStr;
1340
1341    fn ip(text: &str) -> IpAddr {
1342        IpAddr::from_str(text).expect("test address")
1343    }
1344
1345    /// A ban store the test scripts per call: one expiry or a backend
1346    /// failure, so every distributed-lookup arm runs.
1347    struct ScriptedBanStore {
1348        answer: Result<Option<f64>, crate::distributed::StoreError>,
1349    }
1350
1351    impl crate::distributed::BanStore for ScriptedBanStore {
1352        fn set_ban(
1353            &self,
1354            _key: &str,
1355            _expiry: f64,
1356            _ttl_seconds: u64,
1357        ) -> Result<(), crate::distributed::StoreError> {
1358            Ok(())
1359        }
1360
1361        fn get_ban(&self, _key: &str) -> Result<Option<f64>, crate::distributed::StoreError> {
1362            self.answer.clone()
1363        }
1364
1365        fn delete_ban(&self, _key: &str) -> Result<(), crate::distributed::StoreError> {
1366            Ok(())
1367        }
1368    }
1369
1370    #[test]
1371    fn trusted_proxy_lists_accept_and_reject_entries_fail_closed() {
1372        // one call site drives both outcomes: valid CIDR and bare-IP entries
1373        // land in the trusted list, an unparsable entry fails closed exactly
1374        // like the gate
1375        for (entry, valid) in [
1376            ("10.0.0.0/8", true),
1377            ("192.168.1.1", true),
1378            ("not-a-network", false),
1379        ] {
1380            let built =
1381                IpBanManager::with_trusted_proxies_and_clock([entry], Arc::new(system_clock));
1382            assert_eq!(built.is_ok(), valid, "entry {entry}");
1383            if let Ok(manager) = built {
1384                assert!(!manager.is_banned(ip("10.1.2.3")));
1385            }
1386        }
1387    }
1388
1389    #[test]
1390    fn distributed_lookups_answer_live_stale_missing_and_failed_backends() {
1391        let now = 1_000.0_f64;
1392        // a live distributed ban reads as banned and caches locally
1393        let live = IpBanManager::with_trusted_proxies_and_clock(
1394            Vec::<String>::new(),
1395            Arc::new(move || now),
1396        )
1397        .expect("manager")
1398        .with_distributed_store(
1399            Arc::new(ScriptedBanStore {
1400                answer: Ok(Some(2_000.0)),
1401            }),
1402            "gc:",
1403        );
1404        assert!(live.is_banned(ip("192.0.2.71")));
1405        // banning through the manager writes the expiry into the store
1406        live.ban_ip(ip("192.0.2.72"), 60, "penetration_attempt")
1407            .expect("ban");
1408        // a stale expiry deletes the key and reads unbanned
1409        let stale = IpBanManager::with_trusted_proxies_and_clock(
1410            Vec::<String>::new(),
1411            Arc::new(move || now),
1412        )
1413        .expect("manager")
1414        .with_distributed_store(
1415            Arc::new(ScriptedBanStore {
1416                answer: Ok(Some(500.0)),
1417            }),
1418            "gc:",
1419        );
1420        assert!(!stale.is_banned(ip("192.0.2.71")));
1421        // a missing key reads unbanned
1422        let missing = IpBanManager::with_trusted_proxies_and_clock(
1423            Vec::<String>::new(),
1424            Arc::new(move || now),
1425        )
1426        .expect("manager")
1427        .with_distributed_store(Arc::new(ScriptedBanStore { answer: Ok(None) }), "gc:");
1428        assert!(!missing.is_banned(ip("192.0.2.71")));
1429        // a backend error reads unbanned
1430        let failed = IpBanManager::with_trusted_proxies_and_clock(
1431            Vec::<String>::new(),
1432            Arc::new(move || now),
1433        )
1434        .expect("manager")
1435        .with_distributed_store(
1436            Arc::new(ScriptedBanStore {
1437                answer: Err(crate::distributed::StoreError(String::from("down"))),
1438            }),
1439            "gc:",
1440        );
1441        assert!(!failed.is_banned(ip("192.0.2.71")));
1442    }
1443
1444    #[test]
1445    fn failing_distributed_ban_lookup_fails_secure() {
1446        // The reference `_check_redis_exact` -> `redis_handler.get_key` ->
1447        // `safe_operation` raises `GuardRedisError(503, "Redis operation
1448        // failed")` on a backend failure: the pipeline never reads a
1449        // failing lookup as unbanned.
1450        let store = Arc::new(crate::distributed::MemoryStore::default());
1451        store.fail.store(true, std::sync::atomic::Ordering::Relaxed);
1452        let manager = IpBanManager::with_clock(Arc::new(|| 1_000.0)).with_distributed_store(
1453            Arc::clone(&store) as Arc<dyn crate::distributed::BanStore>,
1454            "guard_core:",
1455        );
1456        let stranger = ip("192.0.2.77");
1457
1458        // The fail-secure seam surfaces the backend error.
1459        let error = manager.try_is_banned(stranger).expect_err("fail secure");
1460        assert!(error.to_string().contains("backend down"));
1461        // The lossy form keeps its documented shape (reads unbanned) so
1462        // existing callers stay total.
1463        assert!(!manager.is_banned(stranger));
1464    }
1465
1466    #[test]
1467    fn live_local_ban_answers_ok_even_when_the_store_is_down() {
1468        // The reference checks the local TTLCache before Redis, so a live
1469        // local ban short-circuits and a backend failure never surfaces.
1470        let store = Arc::new(crate::distributed::MemoryStore::default());
1471        let manager = IpBanManager::with_clock(Arc::new(|| 1_000.0)).with_distributed_store(
1472            Arc::clone(&store) as Arc<dyn crate::distributed::BanStore>,
1473            "guard_core:",
1474        );
1475        let attacker = ip("192.0.2.78");
1476        manager
1477            .ban_ip(attacker, 60, "threshold_exceeded")
1478            .expect("ban");
1479
1480        store.fail.store(true, std::sync::atomic::Ordering::Relaxed);
1481        assert!(manager.try_is_banned(attacker).expect("local hit"));
1482    }
1483
1484    #[test]
1485    fn distributed_ban_lookup_error_names_the_backend() {
1486        // Without a store at all the fail-secure seam still answers: no
1487        // backend, no error, just the local store's verdict.
1488        let manager = IpBanManager::with_clock(Arc::new(|| 1_000.0));
1489        assert!(!manager.try_is_banned(ip("192.0.2.79")).expect("no store"));
1490    }
1491}