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)]
134pub struct IpBanConfig {
135 pub enable_ip_banning: bool,
139 pub auto_ban_threshold: u32,
142 pub auto_ban_duration: u64,
144 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#[derive(Debug, Clone, PartialEq, Eq)]
163pub enum IpBanConfigError {
164 NonPositive {
166 field: &'static str,
168 },
169 UnknownCategory {
173 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#[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 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 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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
270pub enum BanError {
271 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#[derive(Debug, Clone, PartialEq)]
290pub struct BanRecord {
291 pub expiry: f64,
293 pub reason: String,
295}
296
297pub 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
309fn 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
317pub struct IpBanManager {
325 bans: Arc<Mutex<LruCache<IpAddr, BanRecord>>>,
326 trusted_proxies: Vec<IpNet>,
327 clock: Clock,
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 }
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 #[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 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 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 #[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 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 #[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 #[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 pub fn unban(&self, ip: IpAddr) {
505 self.bans.lock().expect("ban store").pop(&canonical(ip));
506 }
507
508 pub fn reset(&self) {
511 self.bans.lock().expect("ban store").clear();
512 }
513
514 #[must_use]
517 pub fn banned_count(&self) -> usize {
518 self.bans.lock().expect("ban store").len()
519 }
520
521 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
531pub struct ViolationCounters {
539 counts: Arc<Mutex<LruCache<IpAddr, HashMap<String, u64>>>>,
540}
541
542impl Clone for ViolationCounters {
543 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 #[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 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 for category in categories {
592 *ip_counts.entry((*category).to_owned()).or_insert(0) += 1;
593 }
594 drop(counts);
595 }
596
597 #[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 pub fn reset(&self) {
611 self.counts.lock().expect("violation counter store").clear();
612 }
613
614 #[must_use]
616 pub fn tracked_ips(&self) -> usize {
617 self.counts.lock().expect("violation counter store").len()
618 }
619}
620
621#[derive(Debug, Clone, PartialEq, Eq)]
624pub struct ResolvedBan {
625 pub duration: u64,
627 pub reason: String,
631 pub category: Option<String>,
633}
634
635#[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 #[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 #[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 fake.advance(59);
837 assert!(manager.is_banned(ip("192.0.2.9")));
838 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 assert!(
915 !manager
916 .ban_ip(ip("::ffff:127.0.0.1"), 60, "mistake")
917 .expect("refusal")
918 );
919 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 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 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 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 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 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 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 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}