1use std::collections::{HashMap, VecDeque};
130use std::net::IpAddr;
131use std::num::NonZeroUsize;
132use std::sync::{Arc, Mutex};
133
134use lru::LruCache;
135
136use crate::ip_gate::canonical;
137
138pub const DEFAULT_RATE_LIMIT: u32 = 10;
140pub const DEFAULT_RATE_LIMIT_WINDOW: u64 = 60;
142pub const DEFAULT_ROUTE_RATE_LIMIT_WINDOW: u64 = 60;
145pub const MAX_TRACKED_RATE_LIMIT_KEYS: usize = 10_000;
148
149#[derive(Debug, Clone, Copy, PartialEq, Eq)]
157pub struct RateLimitEntry {
158 pub requests: u32,
160 pub window: u64,
162}
163
164impl RateLimitEntry {
165 pub fn new(requests: u32, window: u64) -> Result<Self, RateLimitConfigError> {
172 if requests == 0 {
173 return Err(RateLimitConfigError {
174 field: "requests".into(),
175 reason: "must be at least 1 request per window",
176 });
177 }
178 if window == 0 {
179 return Err(RateLimitConfigError {
180 field: "window".into(),
181 reason: "must be at least 1 second",
182 });
183 }
184 Ok(Self { requests, window })
185 }
186}
187
188#[derive(Debug, Clone, Default, PartialEq, Eq)]
197pub struct RouteRateLimits {
198 rate_limit: Option<u32>,
200 rate_limit_window: Option<u64>,
203 geo_rate_limits: Option<HashMap<String, RateLimitEntry>>,
207}
208
209impl RouteRateLimits {
210 pub fn new(
220 rate_limit: Option<u32>,
221 rate_limit_window: Option<u64>,
222 geo_rate_limits: Option<HashMap<String, RateLimitEntry>>,
223 ) -> Result<Self, RateLimitConfigError> {
224 if rate_limit.is_some_and(|limit| limit == 0) {
225 return Err(RateLimitConfigError {
226 field: "rate_limit".into(),
227 reason: "must be at least 1 request per window",
228 });
229 }
230 if rate_limit.is_some() && rate_limit_window.is_some_and(|window| window == 0) {
231 return Err(RateLimitConfigError {
232 field: "rate_limit_window".into(),
233 reason: "must be at least 1 second",
234 });
235 }
236 if let Some(geo) = &geo_rate_limits {
237 for (country, entry) in geo {
238 if entry.requests == 0 || entry.window == 0 {
239 return Err(RateLimitConfigError {
240 field: format!("geo_rate_limits[{country}]").into(),
241 reason: "every entry needs at least 1 request per at least 1 second",
242 });
243 }
244 }
245 }
246 Ok(Self {
247 rate_limit,
248 rate_limit_window,
249 geo_rate_limits,
250 })
251 }
252
253 #[must_use]
255 pub const fn rate_limit(&self) -> Option<u32> {
256 self.rate_limit
257 }
258
259 #[must_use]
262 pub const fn rate_limit_window(&self) -> Option<u64> {
263 self.rate_limit_window
264 }
265
266 #[must_use]
268 pub const fn geo_rate_limits(&self) -> Option<&HashMap<String, RateLimitEntry>> {
269 self.geo_rate_limits.as_ref()
270 }
271}
272
273#[derive(Debug, Clone, PartialEq, Eq)]
280pub struct RateLimitConfig {
281 pub enable_rate_limiting: bool,
284 pub rate_limit: u32,
287 pub rate_limit_window: u64,
289 pub enable_rate_limit_auto_ban: bool,
293 pub endpoint_rate_limits: HashMap<String, RateLimitEntry>,
298}
299
300impl Default for RateLimitConfig {
301 fn default() -> Self {
302 Self {
303 enable_rate_limiting: true,
304 rate_limit: DEFAULT_RATE_LIMIT,
305 rate_limit_window: DEFAULT_RATE_LIMIT_WINDOW,
306 enable_rate_limit_auto_ban: false,
307 endpoint_rate_limits: HashMap::new(),
308 }
309 }
310}
311
312#[derive(Debug, Clone, PartialEq, Eq)]
315pub struct RateLimitConfigError {
316 pub field: std::borrow::Cow<'static, str>,
319 pub reason: &'static str,
321}
322
323impl core::fmt::Display for RateLimitConfigError {
324 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
325 write!(f, "invalid {}: {}", self.field, self.reason)
326 }
327}
328
329impl std::error::Error for RateLimitConfigError {}
330
331#[derive(Debug, Clone, Copy, PartialEq, Eq)]
333pub struct RateLimitDecision {
334 pub allowed: bool,
338 pub count: u64,
341 pub window: u64,
344}
345
346impl RateLimitDecision {
347 #[must_use]
350 pub const fn retry_after(self) -> u64 {
351 self.window
352 }
353}
354
355#[derive(Debug, Clone, Copy, PartialEq, Eq)]
358pub enum RateLimitTier {
359 Endpoint,
361 Route,
363 Geo,
365 Global,
367}
368
369impl RateLimitTier {
370 #[must_use]
375 pub const fn name(self) -> &'static str {
376 match self {
377 Self::Endpoint => "endpoint",
378 Self::Route => "route",
379 Self::Geo => "geo",
380 Self::Global => "global",
381 }
382 }
383}
384
385#[derive(Debug, Clone, Copy, PartialEq, Eq)]
388pub struct TierDecision {
389 allowed: bool,
393 count: u64,
397 tier: RateLimitTier,
399 window: u64,
401}
402
403impl TierDecision {
404 #[must_use]
406 pub const fn allowed(&self) -> bool {
407 self.allowed
408 }
409
410 #[must_use]
412 pub const fn count(&self) -> u64 {
413 self.count
414 }
415
416 #[must_use]
418 pub const fn tier(&self) -> RateLimitTier {
419 self.tier
420 }
421
422 #[must_use]
424 pub const fn tier_name(&self) -> &'static str {
425 self.tier.name()
426 }
427
428 #[must_use]
431 pub const fn retry_after(&self) -> u64 {
432 self.window
433 }
434}
435
436pub type Clock = Arc<dyn Fn() -> f64 + Send + Sync>;
439
440#[derive(Debug, PartialEq, Eq, Hash)]
441struct WindowKey {
442 ip: IpAddr,
443 endpoint: Option<String>,
444}
445
446pub struct RateLimiter {
452 config: RateLimitConfig,
453 windows: Arc<Mutex<LruCache<WindowKey, VecDeque<f64>>>>,
454 clock: Clock,
455 store: Option<Arc<dyn crate::distributed::SlidingWindowStore>>,
456 redis_prefix: String,
457 redis_fail_open: bool,
458}
459
460impl Clone for RateLimiter {
461 fn clone(&self) -> Self {
465 Self {
466 config: self.config.clone(),
467 windows: Arc::clone(&self.windows),
468 clock: Arc::clone(&self.clock),
469 store: self.store.clone(),
470 redis_prefix: self.redis_prefix.clone(),
471 redis_fail_open: self.redis_fail_open,
472 }
473 }
474}
475
476impl core::fmt::Debug for RateLimiter {
477 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
478 f.debug_struct("RateLimiter")
479 .field("config", &self.config)
480 .finish_non_exhaustive()
481 }
482}
483
484#[must_use]
488pub fn system_clock() -> f64 {
489 std::time::SystemTime::now()
490 .duration_since(std::time::UNIX_EPOCH)
491 .map_or(0.0, |d| d.as_secs_f64())
492}
493
494fn new_window_store() -> Arc<Mutex<LruCache<WindowKey, VecDeque<f64>>>> {
497 Arc::new(Mutex::new(LruCache::new(
498 NonZeroUsize::new(MAX_TRACKED_RATE_LIMIT_KEYS).expect("constant above zero"),
499 )))
500}
501
502impl RateLimiter {
503 pub fn new(config: RateLimitConfig) -> Result<Self, RateLimitConfigError> {
512 Self::with_config_and_clock(config, Arc::new(system_clock))
513 }
514
515 pub fn with_config_and_clock(
526 config: RateLimitConfig,
527 clock: Clock,
528 ) -> Result<Self, RateLimitConfigError> {
529 if config.rate_limit == 0 {
530 return Err(RateLimitConfigError {
531 field: "rate_limit".into(),
532 reason: "must be at least 1 request per window",
533 });
534 }
535 if config.rate_limit_window == 0 {
536 return Err(RateLimitConfigError {
537 field: "rate_limit_window".into(),
538 reason: "must be at least 1 second",
539 });
540 }
541 for (path, entry) in &config.endpoint_rate_limits {
545 if entry.requests == 0 || entry.window == 0 {
546 return Err(RateLimitConfigError {
547 field: format!("endpoint_rate_limits[{path}]").into(),
548 reason: "every entry needs at least 1 request per at least 1 second",
549 });
550 }
551 }
552 Ok(Self {
553 config,
554 windows: new_window_store(),
555 clock,
556 store: None,
557 redis_prefix: String::from("guard_core:"),
558 redis_fail_open: false,
559 })
560 }
561
562 #[must_use]
578 pub fn with_distributed_store(
579 mut self,
580 store: Arc<dyn crate::distributed::SlidingWindowStore>,
581 redis_prefix: &str,
582 redis_fail_open: bool,
583 ) -> Self {
584 self.store = Some(store);
585 redis_prefix.clone_into(&mut self.redis_prefix);
586 self.redis_fail_open = redis_fail_open;
587 self
588 }
589
590 pub fn check_distributed(
598 &self,
599 ip: IpAddr,
600 endpoint: Option<&str>,
601 ) -> Result<RateLimitDecision, crate::distributed::StoreError> {
602 if !self.config.enable_rate_limiting {
603 return Ok(RateLimitDecision {
604 allowed: true,
605 count: 0,
606 window: self.config.rate_limit_window,
607 });
608 }
609 let (allowed, count, window) = self.record_and_decide_distributed(
610 ip,
611 endpoint,
612 self.config.rate_limit,
613 self.config.rate_limit_window,
614 )?;
615 Ok(RateLimitDecision {
616 allowed,
617 count,
618 window,
619 })
620 }
621
622 pub fn check_tiers_distributed(
631 &self,
632 ip: IpAddr,
633 url_path: Option<&str>,
634 route: Option<&RouteRateLimits>,
635 country_of_ip: Option<&dyn Fn(IpAddr) -> Option<String>>,
636 ) -> Result<TierDecision, crate::distributed::StoreError> {
637 let inert = TierDecision {
638 allowed: true,
639 count: 0,
640 tier: RateLimitTier::Global,
641 window: self.config.rate_limit_window,
642 };
643 if !self.config.enable_rate_limiting {
644 return Ok(inert);
645 }
646 if let Some(path) = url_path {
647 if let Some(entry) = self.config.endpoint_rate_limits.get(path) {
648 let (allowed, count, window) = self.record_and_decide_distributed(
649 ip,
650 Some(path),
651 entry.requests,
652 entry.window,
653 )?;
654 if !allowed {
655 return Ok(TierDecision {
656 allowed,
657 count,
658 tier: RateLimitTier::Endpoint,
659 window,
660 });
661 }
662 }
663 if let Some(route) = route {
664 if let Some(limit) = route.rate_limit() {
665 let window = route
666 .rate_limit_window()
667 .unwrap_or(DEFAULT_ROUTE_RATE_LIMIT_WINDOW);
668 let (allowed, count, window) =
669 self.record_and_decide_distributed(ip, Some(path), limit, window)?;
670 if !allowed {
671 return Ok(TierDecision {
672 allowed,
673 count,
674 tier: RateLimitTier::Route,
675 window,
676 });
677 }
678 }
679 if let (Some(limits), Some(country_of_ip)) =
680 (route.geo_rate_limits(), country_of_ip)
681 {
682 let country = country_of_ip(ip);
683 let entry = country
684 .as_ref()
685 .and_then(|code| limits.get(code))
686 .or_else(|| limits.get("*"));
687 if let Some(entry) = entry {
688 let (allowed, count, window) = self.record_and_decide_distributed(
689 ip,
690 Some(path),
691 entry.requests,
692 entry.window,
693 )?;
694 if !allowed {
695 return Ok(TierDecision {
696 allowed,
697 count,
698 tier: RateLimitTier::Geo,
699 window,
700 });
701 }
702 }
703 }
704 }
705 }
706 let (allowed, count, window) = self.record_and_decide_distributed(
707 ip,
708 None,
709 self.config.rate_limit,
710 self.config.rate_limit_window,
711 )?;
712 Ok(TierDecision {
713 allowed,
714 count,
715 tier: RateLimitTier::Global,
716 window,
717 })
718 }
719
720 pub fn with_clock(clock: Clock) -> Self {
721 Self {
722 config: RateLimitConfig::default(),
723 windows: new_window_store(),
724 store: None,
725 redis_prefix: String::from("guard_core:"),
726 redis_fail_open: false,
727 clock,
728 }
729 }
730
731 #[must_use]
733 pub const fn config(&self) -> &RateLimitConfig {
734 &self.config
735 }
736
737 #[allow(clippy::significant_drop_tightening)]
745 fn record_and_decide_distributed(
750 &self,
751 ip: IpAddr,
752 endpoint: Option<&str>,
753 limit: u32,
754 window_seconds: u64,
755 ) -> Result<(bool, u64, u64), crate::distributed::StoreError> {
756 let Some(store) = &self.store else {
757 return Ok(self.record_and_decide(ip, endpoint, limit, window_seconds));
758 };
759 let now = (self.clock)();
760 let key = endpoint.map_or_else(
761 || crate::distributed::rate_window_key(&self.redis_prefix, ip),
762 |path| crate::distributed::rate_window_key_endpoint(&self.redis_prefix, ip, path),
763 );
764 match store.record_hit(&key, now, window_seconds) {
765 Ok(count) => {
766 let allowed = count <= u64::from(limit);
767 Ok((allowed, count, window_seconds))
768 }
769 Err(error) if self.redis_fail_open => {
770 let _ = error;
775 Ok(self.record_and_decide(ip, endpoint, limit, window_seconds))
776 }
777 Err(error) => Err(error),
778 }
779 }
780
781 fn record_and_decide(
782 &self,
783 ip: IpAddr,
784 endpoint: Option<&str>,
785 limit: u32,
786 window_seconds: u64,
787 ) -> (bool, u64, u64) {
788 #[allow(clippy::cast_precision_loss)]
791 let window = window_seconds as f64;
792 let now = (self.clock)();
793 let window_start = now - window;
794 let key = WindowKey {
795 ip: canonical(ip),
796 endpoint: endpoint.map(str::to_owned),
797 };
798 let inside = {
800 let mut windows = self.windows.lock().expect("rate window store");
801 let timestamps = windows
802 .try_get_or_insert_mut(key, || Ok::<_, core::convert::Infallible>(VecDeque::new()))
803 .expect("key capacity just reserved");
804 while timestamps
805 .front()
806 .is_some_and(|&recorded| recorded <= window_start)
807 {
808 timestamps.pop_front();
809 }
810 let inside = timestamps.len();
811 timestamps.push_back(now);
812 drop(windows);
813 inside
814 };
815 let allowed = inside < usize::try_from(limit).unwrap_or(usize::MAX);
816 let count = u64::try_from(inside + 1).unwrap_or(u64::MAX);
817 (allowed, count, window_seconds)
818 }
819
820 #[must_use]
828 pub fn check(&self, ip: IpAddr, endpoint: Option<&str>) -> RateLimitDecision {
829 if !self.config.enable_rate_limiting {
830 return RateLimitDecision {
831 allowed: true,
832 count: 0,
833 window: self.config.rate_limit_window,
834 };
835 }
836 let (allowed, count, window) = self.record_and_decide(
837 ip,
838 endpoint,
839 self.config.rate_limit,
840 self.config.rate_limit_window,
841 );
842 RateLimitDecision {
843 allowed,
844 count,
845 window,
846 }
847 }
848
849 #[must_use]
867 pub fn check_tiers(
868 &self,
869 ip: IpAddr,
870 url_path: Option<&str>,
871 route: Option<&RouteRateLimits>,
872 country_of_ip: Option<&dyn Fn(IpAddr) -> Option<String>>,
873 ) -> TierDecision {
874 let inert = TierDecision {
875 allowed: true,
876 count: 0,
877 tier: RateLimitTier::Global,
878 window: self.config.rate_limit_window,
879 };
880 if !self.config.enable_rate_limiting {
881 return inert;
882 }
883
884 if let Some(path) = url_path {
887 if let Some(entry) = self.config.endpoint_rate_limits.get(path) {
888 let (allowed, count, window) =
889 self.record_and_decide(ip, Some(path), entry.requests, entry.window);
890 if !allowed {
891 return TierDecision {
892 allowed,
893 count,
894 tier: RateLimitTier::Endpoint,
895 window,
896 };
897 }
898 }
899
900 if let Some(route) = route {
903 if let Some(limit) = route.rate_limit() {
904 let window = route
905 .rate_limit_window()
906 .unwrap_or(DEFAULT_ROUTE_RATE_LIMIT_WINDOW);
907 let (allowed, count, window) =
908 self.record_and_decide(ip, Some(path), limit, window);
909 if !allowed {
910 return TierDecision {
911 allowed,
912 count,
913 tier: RateLimitTier::Route,
914 window,
915 };
916 }
917 }
918
919 if let (Some(limits), Some(country_of_ip)) =
923 (route.geo_rate_limits(), country_of_ip)
924 {
925 let country = country_of_ip(ip);
926 let entry = country
927 .as_ref()
928 .and_then(|code| limits.get(code))
929 .or_else(|| limits.get("*"));
930 if let Some(entry) = entry {
931 let (allowed, count, window) =
932 self.record_and_decide(ip, Some(path), entry.requests, entry.window);
933 if !allowed {
934 return TierDecision {
935 allowed,
936 count,
937 tier: RateLimitTier::Geo,
938 window,
939 };
940 }
941 }
942 }
943 }
944 }
945
946 let (allowed, count, window) = self.record_and_decide(
949 ip,
950 None,
951 self.config.rate_limit,
952 self.config.rate_limit_window,
953 );
954 TierDecision {
955 allowed,
956 count,
957 tier: RateLimitTier::Global,
958 window,
959 }
960 }
961
962 pub fn reset(&self) {
965 self.windows.lock().expect("rate window store").clear();
966 }
967
968 #[must_use]
971 pub fn tracked_windows(&self) -> usize {
972 self.windows.lock().expect("rate window store").len()
973 }
974}
975
976#[cfg(test)]
977mod tests {
978 use super::*;
979 use std::str::FromStr;
980 use std::sync::atomic::{AtomicU64, Ordering};
981
982 #[derive(Clone, Default)]
985 struct FakeClock(Arc<AtomicU64>);
986
987 impl FakeClock {
988 fn advance(&self, seconds: u64) {
989 self.0.fetch_add(seconds, Ordering::Relaxed);
990 }
991
992 fn clock(&self) -> Clock {
993 let state = self.0.clone();
994 #[allow(clippy::cast_precision_loss)]
995 Arc::new(move || state.load(Ordering::Relaxed) as f64)
996 }
997 }
998
999 fn ip(text: &str) -> IpAddr {
1000 IpAddr::from_str(text).expect("test address")
1001 }
1002
1003 fn enabled_config(rate_limit: u32) -> RateLimitConfig {
1004 RateLimitConfig {
1005 enable_rate_limiting: true,
1006 rate_limit,
1007 ..RateLimitConfig::default()
1008 }
1009 }
1010
1011 #[test]
1012 fn default_config_matches_the_reference() {
1013 let config = RateLimitConfig::default();
1014 assert!(config.enable_rate_limiting);
1015 assert_eq!(config.rate_limit, 10);
1016 assert_eq!(config.rate_limit_window, 60);
1017 assert!(!config.enable_rate_limit_auto_ban);
1018 }
1019
1020 #[test]
1021 fn new_fails_closed_on_zero_limit_or_window() {
1022 let error = RateLimiter::new(RateLimitConfig {
1023 rate_limit: 0,
1024 ..RateLimitConfig::default()
1025 })
1026 .unwrap_err();
1027 assert_eq!(error.field, "rate_limit");
1028 assert_eq!(error.field, std::borrow::Cow::Borrowed("rate_limit"));
1030 assert_eq!(
1031 error.to_string(),
1032 "invalid rate_limit: must be at least 1 request per window"
1033 );
1034
1035 let error = RateLimiter::new(RateLimitConfig {
1036 enable_rate_limiting: true,
1037 rate_limit: 5,
1038 rate_limit_window: 0,
1039 enable_rate_limit_auto_ban: false,
1040 endpoint_rate_limits: std::collections::HashMap::new(),
1041 })
1042 .unwrap_err();
1043 assert_eq!(error.field, "rate_limit_window");
1044 }
1045
1046 #[test]
1047 fn new_fails_closed_on_invalid_endpoint_entries() {
1048 let error = RateLimiter::new(RateLimitConfig {
1049 endpoint_rate_limits: std::iter::once((
1050 "/login".to_owned(),
1051 RateLimitEntry {
1052 requests: 0,
1053 window: 60,
1054 },
1055 ))
1056 .collect(),
1057 ..RateLimitConfig::default()
1058 })
1059 .unwrap_err();
1060 assert_eq!(error.field, "endpoint_rate_limits[/login]");
1061 assert_eq!(
1062 error.to_string(),
1063 "invalid endpoint_rate_limits[/login]: every entry needs at least 1 request \
1064 per at least 1 second"
1065 );
1066
1067 let error = RateLimitEntry::new(1, 0).unwrap_err();
1068 assert_eq!(error.field, "window");
1069
1070 let error = RouteRateLimits::new(Some(0), Some(60), None).unwrap_err();
1071 assert_eq!(error.field, "rate_limit");
1072
1073 let error = RouteRateLimits::new(Some(10), Some(0), None).unwrap_err();
1074 assert_eq!(error.field, "rate_limit_window");
1075
1076 let error = RouteRateLimits::new(
1077 None,
1078 None,
1079 Some(
1080 std::iter::once((
1081 "RU".to_owned(),
1082 RateLimitEntry {
1083 requests: 1,
1084 window: 0,
1085 },
1086 ))
1087 .collect(),
1088 ),
1089 )
1090 .unwrap_err();
1091 assert_eq!(error.field, "geo_rate_limits[RU]");
1092
1093 assert!(RouteRateLimits::new(None, None, None).is_ok());
1096 assert!(RouteRateLimits::new(None, Some(30), None).is_ok());
1097 }
1098
1099 #[test]
1100 fn disabled_limiter_allows_and_records_nothing() {
1101 let limiter = RateLimiter::new(RateLimitConfig {
1102 enable_rate_limiting: false,
1103 ..RateLimitConfig::default()
1104 })
1105 .expect("valid config");
1106 for _ in 0..50 {
1107 assert!(limiter.check(ip("192.0.2.1"), None).allowed);
1108 }
1109 assert_eq!(limiter.tracked_windows(), 0, "no window was recorded");
1110 }
1111
1112 #[test]
1113 fn blocks_at_the_crossing_and_reports_the_crossing_count() {
1114 let limiter = RateLimiter::new(enabled_config(3)).expect("valid config");
1115 for expected in 1..=3 {
1116 let decision = limiter.check(ip("192.0.2.1"), None);
1117 assert!(decision.allowed, "request {expected} must pass");
1118 assert_eq!(decision.count, expected);
1119 }
1120 let decision = limiter.check(ip("192.0.2.1"), None);
1121 assert!(
1122 !decision.allowed,
1123 "the 4th request inside the window blocks"
1124 );
1125 assert_eq!(decision.count, 4);
1126 assert_eq!(decision.retry_after(), 60, "Retry-After is the window");
1127 }
1128
1129 #[test]
1130 fn window_slide_restores_the_budget() {
1131 let fake = FakeClock::default();
1132 let limiter = RateLimiter {
1133 config: enabled_config(2),
1134 windows: Arc::new(Mutex::new(LruCache::new(
1135 NonZeroUsize::new(16).expect("above zero"),
1136 ))),
1137 clock: fake.clock(),
1138 store: None,
1139 redis_prefix: String::from("guard_core:"),
1140 redis_fail_open: false,
1141 };
1142 assert!(limiter.check(ip("192.0.2.1"), None).allowed);
1143 assert!(limiter.check(ip("192.0.2.1"), None).allowed);
1144 assert!(!limiter.check(ip("192.0.2.1"), None).allowed);
1145
1146 fake.advance(30);
1148 assert!(!limiter.check(ip("192.0.2.1"), None).allowed);
1149 fake.advance(31);
1152 assert!(limiter.check(ip("192.0.2.1"), None).allowed);
1153 }
1154
1155 #[test]
1156 fn windows_are_per_ip() {
1157 let limiter = RateLimiter::new(enabled_config(1)).expect("valid config");
1158 assert!(limiter.check(ip("192.0.2.1"), None).allowed);
1159 assert!(!limiter.check(ip("192.0.2.1"), None).allowed);
1160 assert!(
1161 limiter.check(ip("192.0.2.2"), None).allowed,
1162 "another IP has its own budget"
1163 );
1164 }
1165
1166 #[test]
1167 fn endpoint_scope_is_isolated_from_the_global_scope() {
1168 let limiter = RateLimiter::new(enabled_config(1)).expect("valid config");
1169 assert!(limiter.check(ip("192.0.2.1"), None).allowed);
1170 assert!(
1171 limiter.check(ip("192.0.2.1"), Some("/login")).allowed,
1172 "the endpoint window is a separate budget"
1173 );
1174 assert!(!limiter.check(ip("192.0.2.1"), Some("/login")).allowed);
1175 assert!(
1176 limiter.check(ip("192.0.2.1"), Some("/signup")).allowed,
1177 "another endpoint has its own budget"
1178 );
1179 assert_eq!(limiter.tracked_windows(), 3);
1180 }
1181
1182 #[test]
1183 fn ipv4_mapped_requests_share_the_ipv4_budget() {
1184 let limiter = RateLimiter::new(enabled_config(1)).expect("valid config");
1185 assert!(limiter.check(ip("::ffff:192.0.2.1"), None).allowed);
1186 assert!(
1187 !limiter.check(ip("192.0.2.1"), None).allowed,
1188 "the mapped form must count toward the same window"
1189 );
1190 }
1191
1192 #[test]
1193 fn reset_drops_every_window() {
1194 let limiter = RateLimiter::new(enabled_config(1)).expect("valid config");
1195 assert!(limiter.check(ip("192.0.2.1"), None).allowed);
1196 assert!(!limiter.check(ip("192.0.2.1"), None).allowed);
1197 limiter.reset();
1198 assert_eq!(limiter.tracked_windows(), 0);
1199 assert!(limiter.check(ip("192.0.2.1"), None).allowed);
1200 }
1201
1202 #[test]
1203 fn large_windows_keep_their_exact_length() {
1204 let limiter = RateLimiter::new(RateLimitConfig {
1207 enable_rate_limiting: true,
1208 rate_limit: 1,
1209 rate_limit_window: u64::from(u32::MAX) + 1,
1210 ..RateLimitConfig::default()
1211 })
1212 .expect("valid config");
1213 assert!(limiter.check(ip("192.0.2.1"), None).allowed);
1214 let decision = limiter.check(ip("192.0.2.1"), None);
1215 assert!(!decision.allowed);
1216 assert_eq!(decision.retry_after(), u64::from(u32::MAX) + 1);
1217 }
1218
1219 fn country_of(code: &'static str) -> impl Fn(IpAddr) -> Option<String> {
1221 move |_ip| Some(code.to_owned())
1222 }
1223
1224 #[test]
1225 fn tierless_check_tiers_collapses_to_the_global_tier() {
1226 let limiter = RateLimiter::new(enabled_config(2)).expect("valid config");
1230 let visitor = ip("192.0.2.40");
1231 assert!(
1232 limiter
1233 .check_tiers(visitor, Some("/x"), None, None)
1234 .allowed()
1235 );
1236 assert!(
1237 limiter
1238 .check_tiers(visitor, Some("/x"), None, None)
1239 .allowed()
1240 );
1241 let blocked = limiter.check_tiers(visitor, Some("/x"), None, None);
1242 assert!(!blocked.allowed());
1243 assert_eq!(blocked.tier(), RateLimitTier::Global);
1244 assert_eq!(blocked.tier_name(), "global");
1245 assert_eq!(blocked.retry_after(), 60);
1246 assert_eq!(limiter.tracked_windows(), 1, "one global window");
1247 }
1248
1249 #[test]
1250 fn endpoint_tier_matches_the_exact_path_and_keys_its_own_window() {
1251 let limiter = RateLimiter::new(RateLimitConfig {
1252 endpoint_rate_limits: std::iter::once((
1253 "/login".to_owned(),
1254 RateLimitEntry::new(1, 60).expect("valid entry"),
1255 ))
1256 .collect(),
1257 ..RateLimitConfig::default()
1258 })
1259 .expect("valid config");
1260 let visitor = ip("192.0.2.41");
1261
1262 assert!(
1265 limiter
1266 .check_tiers(visitor, Some("/login"), None, None)
1267 .allowed()
1268 );
1269 let blocked = limiter.check_tiers(visitor, Some("/login"), None, None);
1270 assert!(!blocked.allowed());
1271 assert_eq!(blocked.tier(), RateLimitTier::Endpoint);
1272 assert_eq!(blocked.retry_after(), 60);
1273
1274 assert!(
1277 limiter
1278 .check_tiers(visitor, Some("/other"), None, None)
1279 .allowed()
1280 );
1281 assert_eq!(
1282 limiter.tracked_windows(),
1283 2,
1284 "the login endpoint window + the shared global window"
1285 );
1286 }
1287
1288 #[test]
1289 fn route_tier_uses_the_decorator_window_default_of_60() {
1290 let fake = FakeClock::default();
1291 let limiter = RateLimiter {
1292 config: enabled_config(100),
1293 windows: Arc::new(Mutex::new(LruCache::new(
1294 NonZeroUsize::new(16).expect("above zero"),
1295 ))),
1296 clock: fake.clock(),
1297 store: None,
1298 redis_prefix: String::from("guard_core:"),
1299 redis_fail_open: false,
1300 };
1301 let route = RouteRateLimits::new(Some(1), None, None).expect("valid route");
1302 let visitor = ip("192.0.2.42");
1303
1304 assert!(
1305 limiter
1306 .check_tiers(visitor, Some("/x"), Some(&route), None)
1307 .allowed()
1308 );
1309 let blocked = limiter.check_tiers(visitor, Some("/x"), Some(&route), None);
1310 assert!(!blocked.allowed());
1311 assert_eq!(blocked.tier(), RateLimitTier::Route);
1312 assert_eq!(
1313 blocked.retry_after(),
1314 60,
1315 "the reference default window (rate_limit_window or 60)"
1316 );
1317
1318 let route = RouteRateLimits::new(Some(1), Some(30), None).expect("valid route");
1320 let blocked = limiter.check_tiers(visitor, Some("/x"), Some(&route), None);
1321 assert!(!blocked.allowed());
1322 assert_eq!(blocked.retry_after(), 30);
1323 }
1324
1325 #[test]
1326 fn geo_tier_resolves_the_country_then_the_star_fallback() {
1327 let limiter = RateLimiter::new(RateLimitConfig {
1328 rate_limit: 100,
1329 ..RateLimitConfig::default()
1330 })
1331 .expect("valid config");
1332 let route = RouteRateLimits::new(
1333 None,
1334 None,
1335 Some(
1336 [
1337 ("RU".to_owned(), RateLimitEntry::new(1, 60).expect("valid")),
1338 ("*".to_owned(), RateLimitEntry::new(2, 45).expect("valid")),
1339 ]
1340 .into_iter()
1341 .collect(),
1342 ),
1343 )
1344 .expect("valid route");
1345
1346 let russian = ip("192.0.2.43");
1348 assert!(
1349 limiter
1350 .check_tiers(russian, Some("/x"), Some(&route), Some(&country_of("RU")))
1351 .allowed()
1352 );
1353 let blocked =
1354 limiter.check_tiers(russian, Some("/x"), Some(&route), Some(&country_of("RU")));
1355 assert!(!blocked.allowed());
1356 assert_eq!(blocked.tier(), RateLimitTier::Geo);
1357 assert_eq!(blocked.retry_after(), 60);
1358
1359 let other = ip("192.0.2.44");
1361 let resolver = country_of("DE");
1362 assert!(
1363 limiter
1364 .check_tiers(other, Some("/x"), Some(&route), Some(&resolver))
1365 .allowed()
1366 );
1367 assert!(
1368 limiter
1369 .check_tiers(other, Some("/x"), Some(&route), Some(&resolver))
1370 .allowed()
1371 );
1372 let blocked = limiter.check_tiers(other, Some("/x"), Some(&route), Some(&resolver));
1373 assert!(!blocked.allowed());
1374 assert_eq!(blocked.tier(), RateLimitTier::Geo);
1375 assert_eq!(blocked.retry_after(), 45);
1376
1377 let unknown: fn(IpAddr) -> Option<String> = |_ip| None;
1380 let third = ip("192.0.2.45");
1381 assert!(
1382 limiter
1383 .check_tiers(third, Some("/x"), Some(&route), Some(&unknown))
1384 .allowed()
1385 );
1386 assert!(
1387 limiter
1388 .check_tiers(third, Some("/x"), Some(&route), Some(&unknown))
1389 .allowed()
1390 );
1391 let blocked = limiter.check_tiers(third, Some("/x"), Some(&route), Some(&unknown));
1392 assert!(!blocked.allowed());
1393 assert_eq!(blocked.tier(), RateLimitTier::Geo);
1394 }
1395
1396 #[test]
1397 fn geo_tier_needs_a_configured_map_and_country_lookup() {
1398 let limiter = RateLimiter::new(RateLimitConfig {
1400 rate_limit: 1,
1401 ..RateLimitConfig::default()
1402 })
1403 .expect("valid config");
1404 let route = RouteRateLimits::new(None, None, Some(std::collections::HashMap::new()))
1405 .expect("valid route");
1406 let visitor = ip("192.0.2.46");
1407 assert!(
1408 limiter
1409 .check_tiers(visitor, Some("/x"), Some(&route), Some(&country_of("RU")))
1410 .allowed()
1411 );
1412 let blocked =
1413 limiter.check_tiers(visitor, Some("/x"), Some(&route), Some(&country_of("RU")));
1414 assert!(!blocked.allowed());
1415 assert_eq!(blocked.tier(), RateLimitTier::Global);
1416
1417 let route = RouteRateLimits::new(
1421 None,
1422 None,
1423 Some(
1424 std::iter::once(("*".to_owned(), RateLimitEntry::new(1, 60).expect("valid")))
1425 .collect(),
1426 ),
1427 )
1428 .expect("valid route");
1429 let fresh = ip("192.0.2.51");
1430 assert!(
1431 limiter
1432 .check_tiers(fresh, Some("/x"), Some(&route), None)
1433 .allowed()
1434 );
1435 let blocked = limiter.check_tiers(fresh, Some("/x"), Some(&route), None);
1436 assert!(!blocked.allowed());
1437 assert_eq!(
1438 blocked.tier(),
1439 RateLimitTier::Global,
1440 "no handler, no geo tier"
1441 );
1442 }
1443
1444 #[test]
1445 fn every_configured_tier_records_so_the_budgets_stay_independent() {
1446 let limiter = RateLimiter::new(RateLimitConfig {
1451 rate_limit: 2,
1452 ..RateLimitConfig::default()
1453 })
1454 .expect("valid config");
1455 let route = RouteRateLimits::new(Some(1_000), Some(60), None).expect("valid route");
1456 let visitor = ip("192.0.2.47");
1457
1458 assert!(
1459 limiter
1460 .check_tiers(visitor, Some("/x"), Some(&route), None)
1461 .allowed()
1462 );
1463 assert!(
1464 limiter
1465 .check_tiers(visitor, Some("/x"), Some(&route), None)
1466 .allowed()
1467 );
1468 let blocked = limiter.check_tiers(visitor, Some("/x"), Some(&route), None);
1469 assert!(!blocked.allowed());
1470 assert_eq!(
1471 blocked.tier(),
1472 RateLimitTier::Global,
1473 "the global tier crossed first"
1474 );
1475 assert_eq!(blocked.count(), 3);
1476 assert_eq!(limiter.tracked_windows(), 2, "route + global windows");
1477
1478 }
1483
1484 #[test]
1485 fn first_blocked_tier_wins_the_response_shape() {
1486 let limiter = RateLimiter::new(RateLimitConfig {
1490 rate_limit: 1,
1491 rate_limit_window: 120,
1492 endpoint_rate_limits: std::iter::once((
1493 "/login".to_owned(),
1494 RateLimitEntry::new(1, 30).expect("valid"),
1495 ))
1496 .collect(),
1497 ..RateLimitConfig::default()
1498 })
1499 .expect("valid config");
1500 let visitor = ip("192.0.2.48");
1501 assert!(
1502 limiter
1503 .check_tiers(visitor, Some("/login"), None, None)
1504 .allowed()
1505 );
1506 let blocked = limiter.check_tiers(visitor, Some("/login"), None, None);
1507 assert!(!blocked.allowed());
1508 assert_eq!(blocked.tier(), RateLimitTier::Endpoint);
1509 assert_eq!(blocked.retry_after(), 30, "the endpoint tier's window");
1510
1511 let blocked = limiter.check_tiers(visitor, Some("/other"), None, None);
1513 assert!(!blocked.allowed());
1514 assert_eq!(blocked.tier(), RateLimitTier::Global);
1515 assert_eq!(blocked.retry_after(), 120);
1516 }
1517
1518 #[test]
1519 fn tiers_without_a_request_path_skip_the_path_keyed_tiers() {
1520 let limiter = RateLimiter::new(RateLimitConfig {
1524 rate_limit: 1,
1525 endpoint_rate_limits: std::iter::once((
1526 "/login".to_owned(),
1527 RateLimitEntry::new(1, 60).expect("valid"),
1528 ))
1529 .collect(),
1530 ..RateLimitConfig::default()
1531 })
1532 .expect("valid config");
1533 let route = RouteRateLimits::new(Some(1), Some(60), None).expect("valid route");
1534 let visitor = ip("192.0.2.49");
1535 assert!(
1536 limiter
1537 .check_tiers(visitor, None, Some(&route), Some(&country_of("RU")))
1538 .allowed()
1539 );
1540 let blocked = limiter.check_tiers(visitor, None, Some(&route), Some(&country_of("RU")));
1541 assert!(!blocked.allowed());
1542 assert_eq!(blocked.tier(), RateLimitTier::Global);
1543 }
1544
1545 #[test]
1546 fn disabled_limiter_skips_the_tiers_without_recording() {
1547 let limiter = RateLimiter::new(RateLimitConfig {
1548 enable_rate_limiting: false,
1549 endpoint_rate_limits: std::iter::once((
1550 "/login".to_owned(),
1551 RateLimitEntry::new(1, 60).expect("valid"),
1552 ))
1553 .collect(),
1554 ..RateLimitConfig::default()
1555 })
1556 .expect("valid config");
1557 let route = RouteRateLimits::new(Some(1), Some(60), None).expect("valid route");
1558 let visitor = ip("192.0.2.50");
1559 for _ in 0..10 {
1560 let decision = limiter.check_tiers(visitor, Some("/login"), Some(&route), None);
1561 assert!(decision.allowed());
1562 }
1563 assert_eq!(limiter.tracked_windows(), 0, "nothing was recorded");
1564 }
1565}