1use 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
96pub const DEFAULT_AUTO_BAN_THRESHOLD: u32 = 10;
99pub const DEFAULT_AUTO_BAN_DURATION: u64 = 3600;
102pub const LOCAL_CACHE_TTL_CAP_SECONDS: u64 = 3600;
105pub const LOCAL_CACHE_MAX_SIZE: usize = 10_000;
108pub const MAX_TRACKED_SUSPICIOUS_IPS: usize = 10_000;
110
111pub const RATE_LIMIT_CATEGORY: &str = "rate_limit";
114
115#[derive(Debug, Clone, Copy, PartialEq, Eq)]
118pub struct ThreatBanEntry {
119 pub threshold: u32,
121 pub duration: u64,
123}
124
125#[derive(Debug, Clone, PartialEq, Eq)]
132pub struct IpBanConfig {
133 pub enable_ip_banning: bool,
137 pub auto_ban_threshold: u32,
140 pub auto_ban_duration: u64,
142 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#[derive(Debug, Clone, PartialEq, Eq)]
161pub enum IpBanConfigError {
162 NonPositive {
164 field: &'static str,
166 },
167 UnknownCategory {
171 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#[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 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 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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
268pub enum BanError {
269 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#[derive(Debug, Clone, PartialEq)]
288pub struct BanRecord {
289 pub expiry: f64,
291 pub reason: String,
293}
294
295pub 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
307fn 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
315pub 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 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 #[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 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 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 #[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 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 #[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 #[must_use]
517 pub fn is_banned(&self, ip: IpAddr) -> bool {
518 self.try_is_banned(ip).unwrap_or(false)
519 }
520
521 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 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 #[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 pub fn unban(&self, ip: IpAddr) {
599 self.bans.lock().expect("ban store").pop(&canonical(ip));
600 }
601
602 pub fn reset(&self) {
605 self.bans.lock().expect("ban store").clear();
606 }
607
608 #[must_use]
611 pub fn banned_count(&self) -> usize {
612 self.bans.lock().expect("ban store").len()
613 }
614
615 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
625pub struct ViolationCounters {
633 counts: Arc<Mutex<LruCache<IpAddr, HashMap<String, u64>>>>,
634}
635
636impl Clone for ViolationCounters {
637 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 #[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 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 for category in categories {
686 *ip_counts.entry((*category).to_owned()).or_insert(0) += 1;
687 }
688 drop(counts);
689 }
690
691 #[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 pub fn reset(&self) {
705 self.counts.lock().expect("violation counter store").clear();
706 }
707
708 #[must_use]
710 pub fn tracked_ips(&self) -> usize {
711 self.counts.lock().expect("violation counter store").len()
712 }
713}
714
715#[derive(Debug, Clone, PartialEq, Eq)]
718pub struct ResolvedBan {
719 pub duration: u64,
721 pub reason: String,
725 pub category: Option<String>,
727}
728
729#[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 #[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 #[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 fake.advance(59);
931 assert!(manager.is_banned(ip("192.0.2.9")));
932 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 assert!(
1009 !manager
1010 .ban_ip(ip("::ffff:127.0.0.1"), 60, "mistake")
1011 .expect("refusal")
1012 );
1013 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 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 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 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 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 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 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 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 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 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 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 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 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 live.ban_ip(ip("192.0.2.72"), 60, "penetration_attempt")
1407 .expect("ban");
1408 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 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 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 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 let error = manager.try_is_banned(stranger).expect_err("fail secure");
1460 assert!(error.to_string().contains("backend down"));
1461 assert!(!manager.is_banned(stranger));
1464 }
1465
1466 #[test]
1467 fn live_local_ban_answers_ok_even_when_the_store_is_down() {
1468 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 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}