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 pub fn with_config(&self, config: RateLimitConfig) -> Result<Self, RateLimitConfigError> {
570 let fresh = Self::with_config_and_clock(config, Arc::clone(&self.clock))?;
571 Ok(Self {
572 windows: Arc::clone(&self.windows),
573 store: self.store.clone(),
574 redis_prefix: self.redis_prefix.clone(),
575 redis_fail_open: self.redis_fail_open,
576 ..fresh
577 })
578 }
579
580 #[must_use]
596 pub fn with_distributed_store(
597 mut self,
598 store: Arc<dyn crate::distributed::SlidingWindowStore>,
599 redis_prefix: &str,
600 redis_fail_open: bool,
601 ) -> Self {
602 self.store = Some(store);
603 redis_prefix.clone_into(&mut self.redis_prefix);
604 self.redis_fail_open = redis_fail_open;
605 self
606 }
607
608 pub fn check_distributed(
616 &self,
617 ip: IpAddr,
618 endpoint: Option<&str>,
619 ) -> Result<RateLimitDecision, crate::distributed::StoreError> {
620 if !self.config.enable_rate_limiting {
621 return Ok(RateLimitDecision {
622 allowed: true,
623 count: 0,
624 window: self.config.rate_limit_window,
625 });
626 }
627 let (allowed, count, window) = self.record_and_decide_distributed(
628 ip,
629 endpoint,
630 self.config.rate_limit,
631 self.config.rate_limit_window,
632 )?;
633 Ok(RateLimitDecision {
634 allowed,
635 count,
636 window,
637 })
638 }
639
640 pub fn check_tiers_distributed(
649 &self,
650 ip: IpAddr,
651 url_path: Option<&str>,
652 route: Option<&RouteRateLimits>,
653 country_of_ip: Option<&dyn Fn(IpAddr) -> Option<String>>,
654 ) -> Result<TierDecision, crate::distributed::StoreError> {
655 let inert = TierDecision {
656 allowed: true,
657 count: 0,
658 tier: RateLimitTier::Global,
659 window: self.config.rate_limit_window,
660 };
661 if !self.config.enable_rate_limiting {
662 return Ok(inert);
663 }
664 if let Some(path) = url_path {
665 if let Some(entry) = self.config.endpoint_rate_limits.get(path) {
666 let (allowed, count, window) = self.record_and_decide_distributed(
667 ip,
668 Some(path),
669 entry.requests,
670 entry.window,
671 )?;
672 if !allowed {
673 return Ok(TierDecision {
674 allowed,
675 count,
676 tier: RateLimitTier::Endpoint,
677 window,
678 });
679 }
680 }
681 if let Some(route) = route {
682 if let Some(limit) = route.rate_limit() {
683 let window = route
684 .rate_limit_window()
685 .unwrap_or(DEFAULT_ROUTE_RATE_LIMIT_WINDOW);
686 let (allowed, count, window) =
687 self.record_and_decide_distributed(ip, Some(path), limit, window)?;
688 if !allowed {
689 return Ok(TierDecision {
690 allowed,
691 count,
692 tier: RateLimitTier::Route,
693 window,
694 });
695 }
696 }
697 if let (Some(limits), Some(country_of_ip)) =
698 (route.geo_rate_limits(), country_of_ip)
699 {
700 let country = country_of_ip(ip);
701 let entry = country
702 .as_ref()
703 .and_then(|code| limits.get(code))
704 .or_else(|| limits.get("*"));
705 if let Some(entry) = entry {
706 let (allowed, count, window) = self.record_and_decide_distributed(
707 ip,
708 Some(path),
709 entry.requests,
710 entry.window,
711 )?;
712 if !allowed {
713 return Ok(TierDecision {
714 allowed,
715 count,
716 tier: RateLimitTier::Geo,
717 window,
718 });
719 }
720 }
721 }
722 }
723 }
724 let (allowed, count, window) = self.record_and_decide_distributed(
725 ip,
726 None,
727 self.config.rate_limit,
728 self.config.rate_limit_window,
729 )?;
730 Ok(TierDecision {
731 allowed,
732 count,
733 tier: RateLimitTier::Global,
734 window,
735 })
736 }
737
738 pub fn with_clock(clock: Clock) -> Self {
739 Self {
740 config: RateLimitConfig::default(),
741 windows: new_window_store(),
742 store: None,
743 redis_prefix: String::from("guard_core:"),
744 redis_fail_open: false,
745 clock,
746 }
747 }
748
749 #[must_use]
751 pub const fn config(&self) -> &RateLimitConfig {
752 &self.config
753 }
754
755 #[allow(clippy::significant_drop_tightening)]
763 fn record_and_decide_distributed(
768 &self,
769 ip: IpAddr,
770 endpoint: Option<&str>,
771 limit: u32,
772 window_seconds: u64,
773 ) -> Result<(bool, u64, u64), crate::distributed::StoreError> {
774 let Some(store) = &self.store else {
775 return Ok(self.record_and_decide(ip, endpoint, limit, window_seconds));
776 };
777 let now = (self.clock)();
778 let key = endpoint.map_or_else(
779 || crate::distributed::rate_window_key(&self.redis_prefix, ip),
780 |path| crate::distributed::rate_window_key_endpoint(&self.redis_prefix, ip, path),
781 );
782 match store.record_hit(&key, now, window_seconds) {
783 Ok(count) => {
784 let allowed = count <= u64::from(limit);
785 Ok((allowed, count, window_seconds))
786 }
787 Err(error) if self.redis_fail_open => {
788 let _ = error;
793 Ok(self.record_and_decide(ip, endpoint, limit, window_seconds))
794 }
795 Err(error) => Err(error),
796 }
797 }
798
799 fn record_and_decide(
800 &self,
801 ip: IpAddr,
802 endpoint: Option<&str>,
803 limit: u32,
804 window_seconds: u64,
805 ) -> (bool, u64, u64) {
806 #[allow(clippy::cast_precision_loss)]
809 let window = window_seconds as f64;
810 let now = (self.clock)();
811 let window_start = now - window;
812 let key = WindowKey {
813 ip: canonical(ip),
814 endpoint: endpoint.map(str::to_owned),
815 };
816 let inside = {
818 let mut windows = self.windows.lock().expect("rate window store");
819 let timestamps = windows
820 .try_get_or_insert_mut(key, || Ok::<_, core::convert::Infallible>(VecDeque::new()))
821 .expect("key capacity just reserved");
822 while timestamps
823 .front()
824 .is_some_and(|&recorded| recorded <= window_start)
825 {
826 timestamps.pop_front();
827 }
828 let inside = timestamps.len();
829 timestamps.push_back(now);
830 drop(windows);
831 inside
832 };
833 let allowed = inside < usize::try_from(limit).unwrap_or(usize::MAX);
834 let count = u64::try_from(inside + 1).unwrap_or(u64::MAX);
835 (allowed, count, window_seconds)
836 }
837
838 #[must_use]
846 pub fn check(&self, ip: IpAddr, endpoint: Option<&str>) -> RateLimitDecision {
847 if !self.config.enable_rate_limiting {
848 return RateLimitDecision {
849 allowed: true,
850 count: 0,
851 window: self.config.rate_limit_window,
852 };
853 }
854 let (allowed, count, window) = self.record_and_decide(
855 ip,
856 endpoint,
857 self.config.rate_limit,
858 self.config.rate_limit_window,
859 );
860 RateLimitDecision {
861 allowed,
862 count,
863 window,
864 }
865 }
866
867 #[must_use]
885 pub fn check_tiers(
886 &self,
887 ip: IpAddr,
888 url_path: Option<&str>,
889 route: Option<&RouteRateLimits>,
890 country_of_ip: Option<&dyn Fn(IpAddr) -> Option<String>>,
891 ) -> TierDecision {
892 let inert = TierDecision {
893 allowed: true,
894 count: 0,
895 tier: RateLimitTier::Global,
896 window: self.config.rate_limit_window,
897 };
898 if !self.config.enable_rate_limiting {
899 return inert;
900 }
901
902 if let Some(path) = url_path {
905 if let Some(entry) = self.config.endpoint_rate_limits.get(path) {
906 let (allowed, count, window) =
907 self.record_and_decide(ip, Some(path), entry.requests, entry.window);
908 if !allowed {
909 return TierDecision {
910 allowed,
911 count,
912 tier: RateLimitTier::Endpoint,
913 window,
914 };
915 }
916 }
917
918 if let Some(route) = route {
921 if let Some(limit) = route.rate_limit() {
922 let window = route
923 .rate_limit_window()
924 .unwrap_or(DEFAULT_ROUTE_RATE_LIMIT_WINDOW);
925 let (allowed, count, window) =
926 self.record_and_decide(ip, Some(path), limit, window);
927 if !allowed {
928 return TierDecision {
929 allowed,
930 count,
931 tier: RateLimitTier::Route,
932 window,
933 };
934 }
935 }
936
937 if let (Some(limits), Some(country_of_ip)) =
941 (route.geo_rate_limits(), country_of_ip)
942 {
943 let country = country_of_ip(ip);
944 let entry = country
945 .as_ref()
946 .and_then(|code| limits.get(code))
947 .or_else(|| limits.get("*"));
948 if let Some(entry) = entry {
949 let (allowed, count, window) =
950 self.record_and_decide(ip, Some(path), entry.requests, entry.window);
951 if !allowed {
952 return TierDecision {
953 allowed,
954 count,
955 tier: RateLimitTier::Geo,
956 window,
957 };
958 }
959 }
960 }
961 }
962 }
963
964 let (allowed, count, window) = self.record_and_decide(
967 ip,
968 None,
969 self.config.rate_limit,
970 self.config.rate_limit_window,
971 );
972 TierDecision {
973 allowed,
974 count,
975 tier: RateLimitTier::Global,
976 window,
977 }
978 }
979
980 pub fn reset(&self) {
983 self.windows.lock().expect("rate window store").clear();
984 }
985
986 #[must_use]
989 pub fn tracked_windows(&self) -> usize {
990 self.windows.lock().expect("rate window store").len()
991 }
992}
993
994#[cfg(test)]
995mod tests {
996 use super::*;
997 use std::collections::HashMap;
998 use std::str::FromStr;
999 use std::sync::Arc;
1000 use std::sync::atomic::{AtomicU64, Ordering};
1001
1002 use crate::distributed::MemoryStore;
1003
1004 #[derive(Clone, Default)]
1007 struct FakeClock(Arc<AtomicU64>);
1008
1009 impl FakeClock {
1010 fn advance(&self, seconds: u64) {
1011 self.0.fetch_add(seconds, Ordering::Relaxed);
1012 }
1013
1014 fn clock(&self) -> Clock {
1015 let state = self.0.clone();
1016 #[allow(clippy::cast_precision_loss)]
1017 Arc::new(move || state.load(Ordering::Relaxed) as f64)
1018 }
1019 }
1020
1021 fn ip(text: &str) -> IpAddr {
1022 IpAddr::from_str(text).expect("test address")
1023 }
1024
1025 fn enabled_config(rate_limit: u32) -> RateLimitConfig {
1026 RateLimitConfig {
1027 enable_rate_limiting: true,
1028 rate_limit,
1029 ..RateLimitConfig::default()
1030 }
1031 }
1032
1033 #[test]
1034 fn default_config_matches_the_reference() {
1035 let config = RateLimitConfig::default();
1036 assert!(config.enable_rate_limiting);
1037 assert_eq!(config.rate_limit, 10);
1038 assert_eq!(config.rate_limit_window, 60);
1039 assert!(!config.enable_rate_limit_auto_ban);
1040 }
1041
1042 #[test]
1043 fn new_fails_closed_on_zero_limit_or_window() {
1044 let error = RateLimiter::new(RateLimitConfig {
1045 rate_limit: 0,
1046 ..RateLimitConfig::default()
1047 })
1048 .unwrap_err();
1049 assert_eq!(error.field, "rate_limit");
1050 assert_eq!(error.field, std::borrow::Cow::Borrowed("rate_limit"));
1052 assert_eq!(
1053 error.to_string(),
1054 "invalid rate_limit: must be at least 1 request per window"
1055 );
1056
1057 let error = RateLimiter::new(RateLimitConfig {
1058 enable_rate_limiting: true,
1059 rate_limit: 5,
1060 rate_limit_window: 0,
1061 enable_rate_limit_auto_ban: false,
1062 endpoint_rate_limits: std::collections::HashMap::new(),
1063 })
1064 .unwrap_err();
1065 assert_eq!(error.field, "rate_limit_window");
1066 }
1067
1068 #[test]
1069 fn new_fails_closed_on_invalid_endpoint_entries() {
1070 let error = RateLimiter::new(RateLimitConfig {
1071 endpoint_rate_limits: std::iter::once((
1072 "/login".to_owned(),
1073 RateLimitEntry {
1074 requests: 0,
1075 window: 60,
1076 },
1077 ))
1078 .collect(),
1079 ..RateLimitConfig::default()
1080 })
1081 .unwrap_err();
1082 assert_eq!(error.field, "endpoint_rate_limits[/login]");
1083 assert_eq!(
1084 error.to_string(),
1085 "invalid endpoint_rate_limits[/login]: every entry needs at least 1 request \
1086 per at least 1 second"
1087 );
1088
1089 let error = RateLimitEntry::new(1, 0).unwrap_err();
1090 assert_eq!(error.field, "window");
1091
1092 let error = RouteRateLimits::new(Some(0), Some(60), None).unwrap_err();
1093 assert_eq!(error.field, "rate_limit");
1094
1095 let error = RouteRateLimits::new(Some(10), Some(0), None).unwrap_err();
1096 assert_eq!(error.field, "rate_limit_window");
1097
1098 let error = RouteRateLimits::new(
1099 None,
1100 None,
1101 Some(
1102 std::iter::once((
1103 "RU".to_owned(),
1104 RateLimitEntry {
1105 requests: 1,
1106 window: 0,
1107 },
1108 ))
1109 .collect(),
1110 ),
1111 )
1112 .unwrap_err();
1113 assert_eq!(error.field, "geo_rate_limits[RU]");
1114
1115 assert!(RouteRateLimits::new(None, None, None).is_ok());
1118 assert!(RouteRateLimits::new(None, Some(30), None).is_ok());
1119 }
1120
1121 #[test]
1122 fn disabled_limiter_allows_and_records_nothing() {
1123 let limiter = RateLimiter::new(RateLimitConfig {
1124 enable_rate_limiting: false,
1125 ..RateLimitConfig::default()
1126 })
1127 .expect("valid config");
1128 for _ in 0..50 {
1129 assert!(limiter.check(ip("192.0.2.1"), None).allowed);
1130 }
1131 assert_eq!(limiter.tracked_windows(), 0, "no window was recorded");
1132 }
1133
1134 #[test]
1135 fn blocks_at_the_crossing_and_reports_the_crossing_count() {
1136 let limiter = RateLimiter::new(enabled_config(3)).expect("valid config");
1137 for expected in 1..=3 {
1138 let decision = limiter.check(ip("192.0.2.1"), None);
1139 assert!(decision.allowed, "request {expected} must pass");
1140 assert_eq!(decision.count, expected);
1141 }
1142 let decision = limiter.check(ip("192.0.2.1"), None);
1143 assert!(
1144 !decision.allowed,
1145 "the 4th request inside the window blocks"
1146 );
1147 assert_eq!(decision.count, 4);
1148 assert_eq!(decision.retry_after(), 60, "Retry-After is the window");
1149 }
1150
1151 #[test]
1152 fn window_slide_restores_the_budget() {
1153 let fake = FakeClock::default();
1154 let limiter = RateLimiter {
1155 config: enabled_config(2),
1156 windows: Arc::new(Mutex::new(LruCache::new(
1157 NonZeroUsize::new(16).expect("above zero"),
1158 ))),
1159 clock: fake.clock(),
1160 store: None,
1161 redis_prefix: String::from("guard_core:"),
1162 redis_fail_open: false,
1163 };
1164 assert!(limiter.check(ip("192.0.2.1"), None).allowed);
1165 assert!(limiter.check(ip("192.0.2.1"), None).allowed);
1166 assert!(!limiter.check(ip("192.0.2.1"), None).allowed);
1167
1168 fake.advance(30);
1170 assert!(!limiter.check(ip("192.0.2.1"), None).allowed);
1171 fake.advance(31);
1174 assert!(limiter.check(ip("192.0.2.1"), None).allowed);
1175 }
1176
1177 #[test]
1178 fn windows_are_per_ip() {
1179 let limiter = RateLimiter::new(enabled_config(1)).expect("valid config");
1180 assert!(limiter.check(ip("192.0.2.1"), None).allowed);
1181 assert!(!limiter.check(ip("192.0.2.1"), None).allowed);
1182 assert!(
1183 limiter.check(ip("192.0.2.2"), None).allowed,
1184 "another IP has its own budget"
1185 );
1186 }
1187
1188 #[test]
1189 fn endpoint_scope_is_isolated_from_the_global_scope() {
1190 let limiter = RateLimiter::new(enabled_config(1)).expect("valid config");
1191 assert!(limiter.check(ip("192.0.2.1"), None).allowed);
1192 assert!(
1193 limiter.check(ip("192.0.2.1"), Some("/login")).allowed,
1194 "the endpoint window is a separate budget"
1195 );
1196 assert!(!limiter.check(ip("192.0.2.1"), Some("/login")).allowed);
1197 assert!(
1198 limiter.check(ip("192.0.2.1"), Some("/signup")).allowed,
1199 "another endpoint has its own budget"
1200 );
1201 assert_eq!(limiter.tracked_windows(), 3);
1202 }
1203
1204 #[test]
1205 fn ipv4_mapped_requests_share_the_ipv4_budget() {
1206 let limiter = RateLimiter::new(enabled_config(1)).expect("valid config");
1207 assert!(limiter.check(ip("::ffff:192.0.2.1"), None).allowed);
1208 assert!(
1209 !limiter.check(ip("192.0.2.1"), None).allowed,
1210 "the mapped form must count toward the same window"
1211 );
1212 }
1213
1214 #[test]
1215 fn reset_drops_every_window() {
1216 let limiter = RateLimiter::new(enabled_config(1)).expect("valid config");
1217 assert!(limiter.check(ip("192.0.2.1"), None).allowed);
1218 assert!(!limiter.check(ip("192.0.2.1"), None).allowed);
1219 limiter.reset();
1220 assert_eq!(limiter.tracked_windows(), 0);
1221 assert!(limiter.check(ip("192.0.2.1"), None).allowed);
1222 }
1223
1224 #[test]
1225 fn large_windows_keep_their_exact_length() {
1226 let limiter = RateLimiter::new(RateLimitConfig {
1229 enable_rate_limiting: true,
1230 rate_limit: 1,
1231 rate_limit_window: u64::from(u32::MAX) + 1,
1232 ..RateLimitConfig::default()
1233 })
1234 .expect("valid config");
1235 assert!(limiter.check(ip("192.0.2.1"), None).allowed);
1236 let decision = limiter.check(ip("192.0.2.1"), None);
1237 assert!(!decision.allowed);
1238 assert_eq!(decision.retry_after(), u64::from(u32::MAX) + 1);
1239 }
1240
1241 fn country_of(code: &'static str) -> impl Fn(IpAddr) -> Option<String> {
1243 move |_ip| Some(code.to_owned())
1244 }
1245
1246 #[test]
1247 fn tierless_check_tiers_collapses_to_the_global_tier() {
1248 let limiter = RateLimiter::new(enabled_config(2)).expect("valid config");
1252 let visitor = ip("192.0.2.40");
1253 assert!(
1254 limiter
1255 .check_tiers(visitor, Some("/x"), None, None)
1256 .allowed()
1257 );
1258 assert!(
1259 limiter
1260 .check_tiers(visitor, Some("/x"), None, None)
1261 .allowed()
1262 );
1263 let blocked = limiter.check_tiers(visitor, Some("/x"), None, None);
1264 assert!(!blocked.allowed());
1265 assert_eq!(blocked.tier(), RateLimitTier::Global);
1266 assert_eq!(blocked.tier_name(), "global");
1267 assert_eq!(blocked.retry_after(), 60);
1268 assert_eq!(limiter.tracked_windows(), 1, "one global window");
1269 }
1270
1271 #[test]
1272 fn endpoint_tier_matches_the_exact_path_and_keys_its_own_window() {
1273 let limiter = RateLimiter::new(RateLimitConfig {
1274 endpoint_rate_limits: std::iter::once((
1275 "/login".to_owned(),
1276 RateLimitEntry::new(1, 60).expect("valid entry"),
1277 ))
1278 .collect(),
1279 ..RateLimitConfig::default()
1280 })
1281 .expect("valid config");
1282 let visitor = ip("192.0.2.41");
1283
1284 assert!(
1287 limiter
1288 .check_tiers(visitor, Some("/login"), None, None)
1289 .allowed()
1290 );
1291 let blocked = limiter.check_tiers(visitor, Some("/login"), None, None);
1292 assert!(!blocked.allowed());
1293 assert_eq!(blocked.tier(), RateLimitTier::Endpoint);
1294 assert_eq!(blocked.retry_after(), 60);
1295
1296 assert!(
1299 limiter
1300 .check_tiers(visitor, Some("/other"), None, None)
1301 .allowed()
1302 );
1303 assert_eq!(
1304 limiter.tracked_windows(),
1305 2,
1306 "the login endpoint window + the shared global window"
1307 );
1308 }
1309
1310 #[test]
1311 fn route_tier_uses_the_decorator_window_default_of_60() {
1312 let fake = FakeClock::default();
1313 let limiter = RateLimiter {
1314 config: enabled_config(100),
1315 windows: Arc::new(Mutex::new(LruCache::new(
1316 NonZeroUsize::new(16).expect("above zero"),
1317 ))),
1318 clock: fake.clock(),
1319 store: None,
1320 redis_prefix: String::from("guard_core:"),
1321 redis_fail_open: false,
1322 };
1323 let route = RouteRateLimits::new(Some(1), None, None).expect("valid route");
1324 let visitor = ip("192.0.2.42");
1325
1326 assert!(
1327 limiter
1328 .check_tiers(visitor, Some("/x"), Some(&route), None)
1329 .allowed()
1330 );
1331 let blocked = limiter.check_tiers(visitor, Some("/x"), Some(&route), None);
1332 assert!(!blocked.allowed());
1333 assert_eq!(blocked.tier(), RateLimitTier::Route);
1334 assert_eq!(
1335 blocked.retry_after(),
1336 60,
1337 "the reference default window (rate_limit_window or 60)"
1338 );
1339
1340 let route = RouteRateLimits::new(Some(1), Some(30), None).expect("valid route");
1342 let blocked = limiter.check_tiers(visitor, Some("/x"), Some(&route), None);
1343 assert!(!blocked.allowed());
1344 assert_eq!(blocked.retry_after(), 30);
1345 }
1346
1347 #[test]
1348 fn geo_tier_resolves_the_country_then_the_star_fallback() {
1349 let limiter = RateLimiter::new(RateLimitConfig {
1350 rate_limit: 100,
1351 ..RateLimitConfig::default()
1352 })
1353 .expect("valid config");
1354 let route = RouteRateLimits::new(
1355 None,
1356 None,
1357 Some(
1358 [
1359 ("RU".to_owned(), RateLimitEntry::new(1, 60).expect("valid")),
1360 ("*".to_owned(), RateLimitEntry::new(2, 45).expect("valid")),
1361 ]
1362 .into_iter()
1363 .collect(),
1364 ),
1365 )
1366 .expect("valid route");
1367
1368 let russian = ip("192.0.2.43");
1370 assert!(
1371 limiter
1372 .check_tiers(russian, Some("/x"), Some(&route), Some(&country_of("RU")))
1373 .allowed()
1374 );
1375 let blocked =
1376 limiter.check_tiers(russian, Some("/x"), Some(&route), Some(&country_of("RU")));
1377 assert!(!blocked.allowed());
1378 assert_eq!(blocked.tier(), RateLimitTier::Geo);
1379 assert_eq!(blocked.retry_after(), 60);
1380
1381 let other = ip("192.0.2.44");
1383 let resolver = country_of("DE");
1384 assert!(
1385 limiter
1386 .check_tiers(other, Some("/x"), Some(&route), Some(&resolver))
1387 .allowed()
1388 );
1389 assert!(
1390 limiter
1391 .check_tiers(other, Some("/x"), Some(&route), Some(&resolver))
1392 .allowed()
1393 );
1394 let blocked = limiter.check_tiers(other, Some("/x"), Some(&route), Some(&resolver));
1395 assert!(!blocked.allowed());
1396 assert_eq!(blocked.tier(), RateLimitTier::Geo);
1397 assert_eq!(blocked.retry_after(), 45);
1398
1399 let unknown: fn(IpAddr) -> Option<String> = |_ip| None;
1402 let third = ip("192.0.2.45");
1403 assert!(
1404 limiter
1405 .check_tiers(third, Some("/x"), Some(&route), Some(&unknown))
1406 .allowed()
1407 );
1408 assert!(
1409 limiter
1410 .check_tiers(third, Some("/x"), Some(&route), Some(&unknown))
1411 .allowed()
1412 );
1413 let blocked = limiter.check_tiers(third, Some("/x"), Some(&route), Some(&unknown));
1414 assert!(!blocked.allowed());
1415 assert_eq!(blocked.tier(), RateLimitTier::Geo);
1416 }
1417
1418 #[test]
1419 fn geo_tier_needs_a_configured_map_and_country_lookup() {
1420 let limiter = RateLimiter::new(RateLimitConfig {
1422 rate_limit: 1,
1423 ..RateLimitConfig::default()
1424 })
1425 .expect("valid config");
1426 let route = RouteRateLimits::new(None, None, Some(std::collections::HashMap::new()))
1427 .expect("valid route");
1428 let visitor = ip("192.0.2.46");
1429 assert!(
1430 limiter
1431 .check_tiers(visitor, Some("/x"), Some(&route), Some(&country_of("RU")))
1432 .allowed()
1433 );
1434 let blocked =
1435 limiter.check_tiers(visitor, Some("/x"), Some(&route), Some(&country_of("RU")));
1436 assert!(!blocked.allowed());
1437 assert_eq!(blocked.tier(), RateLimitTier::Global);
1438
1439 let route = RouteRateLimits::new(
1443 None,
1444 None,
1445 Some(
1446 std::iter::once(("*".to_owned(), RateLimitEntry::new(1, 60).expect("valid")))
1447 .collect(),
1448 ),
1449 )
1450 .expect("valid route");
1451 let fresh = ip("192.0.2.51");
1452 assert!(
1453 limiter
1454 .check_tiers(fresh, Some("/x"), Some(&route), None)
1455 .allowed()
1456 );
1457 let blocked = limiter.check_tiers(fresh, Some("/x"), Some(&route), None);
1458 assert!(!blocked.allowed());
1459 assert_eq!(
1460 blocked.tier(),
1461 RateLimitTier::Global,
1462 "no handler, no geo tier"
1463 );
1464 }
1465
1466 #[test]
1467 fn every_configured_tier_records_so_the_budgets_stay_independent() {
1468 let limiter = RateLimiter::new(RateLimitConfig {
1473 rate_limit: 2,
1474 ..RateLimitConfig::default()
1475 })
1476 .expect("valid config");
1477 let route = RouteRateLimits::new(Some(1_000), Some(60), None).expect("valid route");
1478 let visitor = ip("192.0.2.47");
1479
1480 assert!(
1481 limiter
1482 .check_tiers(visitor, Some("/x"), Some(&route), None)
1483 .allowed()
1484 );
1485 assert!(
1486 limiter
1487 .check_tiers(visitor, Some("/x"), Some(&route), None)
1488 .allowed()
1489 );
1490 let blocked = limiter.check_tiers(visitor, Some("/x"), Some(&route), None);
1491 assert!(!blocked.allowed());
1492 assert_eq!(
1493 blocked.tier(),
1494 RateLimitTier::Global,
1495 "the global tier crossed first"
1496 );
1497 assert_eq!(blocked.count(), 3);
1498 assert_eq!(limiter.tracked_windows(), 2, "route + global windows");
1499
1500 }
1505
1506 #[test]
1507 fn first_blocked_tier_wins_the_response_shape() {
1508 let limiter = RateLimiter::new(RateLimitConfig {
1512 rate_limit: 1,
1513 rate_limit_window: 120,
1514 endpoint_rate_limits: std::iter::once((
1515 "/login".to_owned(),
1516 RateLimitEntry::new(1, 30).expect("valid"),
1517 ))
1518 .collect(),
1519 ..RateLimitConfig::default()
1520 })
1521 .expect("valid config");
1522 let visitor = ip("192.0.2.48");
1523 assert!(
1524 limiter
1525 .check_tiers(visitor, Some("/login"), None, None)
1526 .allowed()
1527 );
1528 let blocked = limiter.check_tiers(visitor, Some("/login"), None, None);
1529 assert!(!blocked.allowed());
1530 assert_eq!(blocked.tier(), RateLimitTier::Endpoint);
1531 assert_eq!(blocked.retry_after(), 30, "the endpoint tier's window");
1532
1533 let blocked = limiter.check_tiers(visitor, Some("/other"), None, None);
1535 assert!(!blocked.allowed());
1536 assert_eq!(blocked.tier(), RateLimitTier::Global);
1537 assert_eq!(blocked.retry_after(), 120);
1538 }
1539
1540 #[test]
1541 fn tiers_without_a_request_path_skip_the_path_keyed_tiers() {
1542 let limiter = RateLimiter::new(RateLimitConfig {
1546 rate_limit: 1,
1547 endpoint_rate_limits: std::iter::once((
1548 "/login".to_owned(),
1549 RateLimitEntry::new(1, 60).expect("valid"),
1550 ))
1551 .collect(),
1552 ..RateLimitConfig::default()
1553 })
1554 .expect("valid config");
1555 let route = RouteRateLimits::new(Some(1), Some(60), None).expect("valid route");
1556 let visitor = ip("192.0.2.49");
1557 assert!(
1558 limiter
1559 .check_tiers(visitor, None, Some(&route), Some(&country_of("RU")))
1560 .allowed()
1561 );
1562 let blocked = limiter.check_tiers(visitor, None, Some(&route), Some(&country_of("RU")));
1563 assert!(!blocked.allowed());
1564 assert_eq!(blocked.tier(), RateLimitTier::Global);
1565 }
1566
1567 #[test]
1568 fn disabled_limiter_skips_the_tiers_without_recording() {
1569 let limiter = RateLimiter::new(RateLimitConfig {
1570 enable_rate_limiting: false,
1571 endpoint_rate_limits: std::iter::once((
1572 "/login".to_owned(),
1573 RateLimitEntry::new(1, 60).expect("valid"),
1574 ))
1575 .collect(),
1576 ..RateLimitConfig::default()
1577 })
1578 .expect("valid config");
1579 let route = RouteRateLimits::new(Some(1), Some(60), None).expect("valid route");
1580 let visitor = ip("192.0.2.50");
1581 for _ in 0..10 {
1582 let decision = limiter.check_tiers(visitor, Some("/login"), Some(&route), None);
1583 assert!(decision.allowed());
1584 }
1585 assert_eq!(limiter.tracked_windows(), 0, "nothing was recorded");
1586 }
1587
1588 #[test]
1589 fn zero_request_entries_are_rejected() {
1590 assert_eq!(RateLimitEntry::new(0, 60).err().unwrap().field, "requests");
1591 }
1592
1593 #[test]
1594 fn tier_names_read_like_the_reference_reasons() {
1595 assert_eq!(RateLimitTier::Endpoint.name(), "endpoint");
1596 assert_eq!(RateLimitTier::Route.name(), "route");
1597 assert_eq!(RateLimitTier::Geo.name(), "geo");
1598 assert_eq!(RateLimitTier::Global.name(), "global");
1599 }
1600
1601 #[test]
1602 fn limiter_debug_carries_the_config() {
1603 let limiter = RateLimiter::with_clock(Arc::new(system_clock));
1604 assert!(format!("{limiter:?}").contains("RateLimiter"));
1605 assert_eq!(limiter.config().rate_limit_window, 60);
1606 }
1607
1608 #[test]
1609 fn distributed_checks_short_circuit_when_disabled() {
1610 let limiter = RateLimiter::with_config_and_clock(
1611 RateLimitConfig {
1612 enable_rate_limiting: false,
1613 ..RateLimitConfig::default()
1614 },
1615 Arc::new(system_clock),
1616 )
1617 .expect("config");
1618 let decision = limiter
1619 .check_distributed(ip("192.0.2.51"), None)
1620 .expect("inert");
1621 assert!(decision.allowed);
1622
1623 let tiers = limiter
1624 .check_tiers_distributed(ip("192.0.2.51"), Some("/x"), None, None)
1625 .expect("inert");
1626 assert!(tiers.allowed());
1627 }
1628
1629 #[test]
1630 fn distributed_tiers_cover_endpoint_route_and_geo_windows() {
1631 let fake = FakeClock::default();
1632 let store = Arc::new(MemoryStore::default());
1633 let geo_limits = HashMap::from([
1634 ("RU".to_owned(), RateLimitEntry::new(2, 60).unwrap()),
1635 ("*".to_owned(), RateLimitEntry::new(5, 60).unwrap()),
1636 ]);
1637 let route = RouteRateLimits::new(Some(3), Some(30), Some(geo_limits)).expect("valid route");
1638 let limiter = RateLimiter::with_config_and_clock(
1639 RateLimitConfig {
1640 enable_rate_limiting: true,
1641 rate_limit: 100,
1642 endpoint_rate_limits: std::iter::once((
1643 "/api".to_owned(),
1644 RateLimitEntry::new(2, 60).unwrap(),
1645 ))
1646 .collect(),
1647 ..RateLimitConfig::default()
1648 },
1649 fake.clock(),
1650 )
1651 .expect("config")
1652 .with_distributed_store(
1653 Arc::clone(&store) as Arc<dyn crate::distributed::SlidingWindowStore>,
1654 "gc:",
1655 false,
1656 );
1657
1658 let country_ru = |probe: IpAddr| -> Option<String> {
1661 if probe == ip("192.0.2.52") {
1662 Some("RU".to_owned())
1663 } else {
1664 None
1665 }
1666 };
1667 let first = limiter
1671 .check_tiers_distributed(
1672 ip("192.0.2.52"),
1673 Some("/api"),
1674 Some(&route),
1675 Some(&country_ru),
1676 )
1677 .expect("store");
1678 assert!(!first.allowed());
1679 assert_eq!(first.tier(), RateLimitTier::Geo);
1680 assert_eq!(first.count(), 3);
1681
1682 let second = limiter
1685 .check_tiers_distributed(
1686 ip("192.0.2.53"),
1687 Some("/other"),
1688 Some(&route),
1689 Some(&country_ru),
1690 )
1691 .expect("store");
1692 assert!(second.allowed());
1693
1694 let third = limiter
1697 .check_tiers_distributed(
1698 ip("192.0.2.52"),
1699 Some("/other"),
1700 Some(&route),
1701 Some(&country_ru),
1702 )
1703 .expect("store");
1704 assert!(third.allowed());
1705 }
1706
1707 #[test]
1708 fn a_store_error_propagates_from_the_endpoint_tier() {
1709 let store = Arc::new(MemoryStore::default());
1710 let limiter = RateLimiter::with_config_and_clock(
1711 RateLimitConfig {
1712 enable_rate_limiting: true,
1713 rate_limit: 100,
1714 endpoint_rate_limits: std::iter::once((
1715 "/api".to_owned(),
1716 RateLimitEntry::new(2, 60).unwrap(),
1717 ))
1718 .collect(),
1719 ..RateLimitConfig::default()
1720 },
1721 FakeClock::default().clock(),
1722 )
1723 .expect("config")
1724 .with_distributed_store(
1725 Arc::clone(&store) as Arc<dyn crate::distributed::SlidingWindowStore>,
1726 "gc:",
1727 false,
1728 );
1729 store.fail.store(true, std::sync::atomic::Ordering::Relaxed);
1730 assert!(
1731 limiter
1732 .check_tiers_distributed(ip("192.0.2.60"), Some("/api"), None, None)
1733 .is_err()
1734 );
1735 }
1736
1737 #[test]
1738 fn a_store_error_propagates_from_the_route_tier() {
1739 let store = Arc::new(MemoryStore::default());
1740 let route = RouteRateLimits::new(Some(3), Some(60), None).expect("valid route");
1741 let limiter = RateLimiter::with_config_and_clock(
1742 RateLimitConfig {
1743 enable_rate_limiting: true,
1744 ..RateLimitConfig::default()
1745 },
1746 FakeClock::default().clock(),
1747 )
1748 .expect("config")
1749 .with_distributed_store(
1750 Arc::clone(&store) as Arc<dyn crate::distributed::SlidingWindowStore>,
1751 "gc:",
1752 false,
1753 );
1754 store.fail.store(true, std::sync::atomic::Ordering::Relaxed);
1755 assert!(
1758 limiter
1759 .check_tiers_distributed(ip("192.0.2.59"), Some("/r"), Some(&route), None)
1760 .is_err()
1761 );
1762 }
1763
1764 #[test]
1765 fn a_store_error_propagates_from_the_geo_tier() {
1766 let store = Arc::new(MemoryStore::default());
1767 let geo_limits = HashMap::from([("RU".to_owned(), RateLimitEntry::new(2, 60).unwrap())]);
1768 let route = RouteRateLimits::new(None, None, Some(geo_limits)).expect("valid route");
1769 let limiter = RateLimiter::with_config_and_clock(
1770 RateLimitConfig {
1771 enable_rate_limiting: true,
1772 ..RateLimitConfig::default()
1773 },
1774 FakeClock::default().clock(),
1775 )
1776 .expect("config")
1777 .with_distributed_store(
1778 Arc::clone(&store) as Arc<dyn crate::distributed::SlidingWindowStore>,
1779 "gc:",
1780 false,
1781 );
1782 store.fail.store(true, std::sync::atomic::Ordering::Relaxed);
1783 assert!(
1786 limiter
1787 .check_tiers_distributed(
1788 ip("192.0.2.61"),
1789 Some("/other"),
1790 Some(&route),
1791 Some(&|_| Some("RU".to_owned())),
1792 )
1793 .is_err()
1794 );
1795 }
1796
1797 #[test]
1798 fn a_geo_limit_miss_falls_through_to_the_global_tier() {
1799 let store = Arc::new(MemoryStore::default());
1800 let geo_limits = HashMap::from([("CN".to_owned(), RateLimitEntry::new(2, 60).unwrap())]);
1801 let route = RouteRateLimits::new(None, None, Some(geo_limits)).expect("valid route");
1802 let limiter = RateLimiter::with_config_and_clock(
1803 RateLimitConfig {
1804 enable_rate_limiting: true,
1805 ..RateLimitConfig::default()
1806 },
1807 FakeClock::default().clock(),
1808 )
1809 .expect("config")
1810 .with_distributed_store(
1811 Arc::clone(&store) as Arc<dyn crate::distributed::SlidingWindowStore>,
1812 "gc:",
1813 false,
1814 );
1815 let decision = limiter
1817 .check_tiers_distributed(
1818 ip("192.0.2.62"),
1819 Some("/other"),
1820 Some(&route),
1821 Some(&|_| Some("US".to_owned())),
1822 )
1823 .expect("store");
1824 assert!(decision.allowed());
1825 assert_eq!(decision.tier(), RateLimitTier::Global);
1826 }
1827}
1828
1829#[cfg(test)]
1830mod unit_twins {
1831 use super::*;
1832 use crate::distributed::MemoryStore;
1833 use std::str::FromStr;
1834
1835 fn ip(text: &str) -> IpAddr {
1836 IpAddr::from_str(text).expect("test address")
1837 }
1838
1839 fn failing_store_limiter(endpoint: Option<(&str, RateLimitEntry)>) -> RateLimiter {
1842 let store = MemoryStore::default();
1843 store.fail.store(true, std::sync::atomic::Ordering::Relaxed);
1844 let mut endpoint_rate_limits = HashMap::new();
1845 if let Some((path, entry)) = endpoint {
1846 endpoint_rate_limits.insert(path.to_owned(), entry);
1847 }
1848 RateLimiter::with_config_and_clock(
1849 RateLimitConfig {
1850 enable_rate_limiting: true,
1851 rate_limit: 100,
1852 endpoint_rate_limits,
1853 ..RateLimitConfig::default()
1854 },
1855 Arc::new(system_clock),
1856 )
1857 .expect("config")
1858 .with_distributed_store(
1859 Arc::new(store) as Arc<dyn crate::distributed::SlidingWindowStore>,
1860 "gc:",
1861 false,
1862 )
1863 }
1864
1865 #[test]
1866 fn the_endpoint_tier_surfaces_a_failing_backend() {
1867 let limiter = failing_store_limiter(Some((
1868 "/api",
1869 RateLimitEntry::new(2, 60).expect("valid entry"),
1870 )));
1871 assert!(
1872 limiter
1873 .check_tiers_distributed(ip("192.0.2.60"), Some("/api"), None, None)
1874 .is_err()
1875 );
1876 }
1877
1878 #[test]
1879 fn the_geo_tier_surfaces_a_failing_backend() {
1880 let geo_limits = HashMap::from([(
1881 "RU".to_owned(),
1882 RateLimitEntry::new(2, 60).expect("valid entry"),
1883 )]);
1884 let route = RouteRateLimits::new(None, None, Some(geo_limits)).expect("valid route");
1886 let limiter = failing_store_limiter(None);
1887 let country = |probe: IpAddr| -> Option<String> {
1888 if probe == ip("192.0.2.61") {
1889 Some("RU".to_owned())
1890 } else {
1891 None
1892 }
1893 };
1894 assert!(
1897 limiter
1898 .check_tiers_distributed(
1899 ip("192.0.2.61"),
1900 Some("/other"),
1901 Some(&route),
1902 Some(&country),
1903 )
1904 .is_err()
1905 );
1906 assert!(
1909 limiter
1910 .check_tiers_distributed(
1911 ip("192.0.2.62"),
1912 Some("/other"),
1913 Some(&route),
1914 Some(&country),
1915 )
1916 .is_err()
1917 );
1918 }
1919
1920 #[test]
1921 fn exhausted_tiers_answer_their_denials() {
1922 let store = Arc::new(MemoryStore::default());
1924 let geo_limits = HashMap::from([(
1925 "RU".to_owned(),
1926 RateLimitEntry::new(1, 60).expect("valid entry"),
1927 )]);
1928 let limiter = RateLimiter::with_config_and_clock(
1929 RateLimitConfig {
1930 enable_rate_limiting: true,
1931 rate_limit: 1,
1932 endpoint_rate_limits: std::iter::once((
1933 "/api".to_owned(),
1934 RateLimitEntry::new(1, 60).expect("valid entry"),
1935 ))
1936 .collect(),
1937 ..RateLimitConfig::default()
1938 },
1939 Arc::new(system_clock),
1940 )
1941 .expect("config")
1942 .with_distributed_store(
1943 Arc::clone(&store) as Arc<dyn crate::distributed::SlidingWindowStore>,
1944 "gc:",
1945 false,
1946 );
1947 let country_ru = |probe: IpAddr| -> Option<String> {
1948 if probe == ip("192.0.2.63") || probe == ip("192.0.2.67") {
1949 Some("RU".to_owned())
1950 } else {
1951 None
1952 }
1953 };
1954 let route = RouteRateLimits::new(Some(1), Some(30), None).expect("valid route");
1955 let geo_route = RouteRateLimits::new(None, None, Some(geo_limits)).expect("valid route");
1958
1959 let endpoint = limiter
1961 .check_tiers_distributed(ip("192.0.2.63"), Some("/api"), None, None)
1962 .expect("store");
1963 assert!(endpoint.allowed());
1964 let endpoint = limiter
1965 .check_tiers_distributed(ip("192.0.2.63"), Some("/api"), None, None)
1966 .expect("store");
1967 assert!(!endpoint.allowed());
1968
1969 let route_hit = limiter
1971 .check_tiers_distributed(ip("192.0.2.64"), Some("/route"), Some(&route), None)
1972 .expect("store");
1973 assert!(route_hit.allowed());
1974 let route_hit = limiter
1975 .check_tiers_distributed(ip("192.0.2.64"), Some("/route"), Some(&route), None)
1976 .expect("store");
1977 assert!(!route_hit.allowed());
1978
1979 let geo = limiter
1981 .check_tiers_distributed(
1982 ip("192.0.2.67"),
1983 Some("/geo"),
1984 Some(&geo_route),
1985 Some(&country_ru),
1986 )
1987 .expect("store");
1988 assert!(geo.allowed(), "first geo: {geo:?}");
1989 let skipped = limiter
1991 .check_tiers_distributed(
1992 ip("192.0.2.68"),
1993 Some("/geo"),
1994 Some(&geo_route),
1995 Some(&country_ru),
1996 )
1997 .expect("store");
1998 assert!(skipped.allowed(), "skipped geo: {skipped:?}");
1999 let geo = limiter
2000 .check_tiers_distributed(
2001 ip("192.0.2.67"),
2002 Some("/geo"),
2003 Some(&geo_route),
2004 Some(&country_ru),
2005 )
2006 .expect("store");
2007 assert!(!geo.allowed(), "second geo: {geo:?}");
2008 assert_eq!(geo.tier(), RateLimitTier::Geo);
2009
2010 let global = limiter
2012 .check_tiers_distributed(ip("192.0.2.65"), Some("/other"), None, None)
2013 .expect("store");
2014 assert!(global.allowed());
2015 let global = limiter
2016 .check_tiers_distributed(ip("192.0.2.65"), Some("/other"), None, None)
2017 .expect("store");
2018 assert!(!global.allowed());
2019 assert_eq!(global.tier(), RateLimitTier::Global);
2020 }
2021
2022 #[test]
2023 fn the_global_tier_surfaces_a_failing_backend() {
2024 let limiter = failing_store_limiter(None);
2027 assert!(
2028 limiter
2029 .check_tiers_distributed(ip("192.0.2.66"), Some("/other"), None, None)
2030 .is_err()
2031 );
2032 assert!(
2033 limiter
2034 .check_tiers_distributed(ip("192.0.2.66"), None, None, None)
2035 .is_err()
2036 );
2037 }
2038
2039 #[test]
2040 fn a_geo_tier_without_a_matching_country_falls_through_to_the_global_tier() {
2041 let store = Arc::new(MemoryStore::default());
2044 let geo_limits = HashMap::from([(
2045 "RU".to_owned(),
2046 RateLimitEntry::new(2, 60).expect("valid entry"),
2047 )]);
2048 let route = RouteRateLimits::new(Some(3), Some(30), Some(geo_limits)).expect("valid route");
2049 let limiter = RateLimiter::with_config_and_clock(
2050 RateLimitConfig {
2051 enable_rate_limiting: true,
2052 rate_limit: 100,
2053 ..RateLimitConfig::default()
2054 },
2055 Arc::new(system_clock),
2056 )
2057 .expect("config")
2058 .with_distributed_store(
2059 Arc::clone(&store) as Arc<dyn crate::distributed::SlidingWindowStore>,
2060 "gc:",
2061 false,
2062 );
2063 let no_country = |_: IpAddr| -> Option<String> { None };
2064 let decision = limiter
2065 .check_tiers_distributed(
2066 ip("192.0.2.62"),
2067 Some("/other"),
2068 Some(&route),
2069 Some(&no_country),
2070 )
2071 .expect("store");
2072 assert!(decision.allowed());
2073 assert_eq!(decision.tier(), RateLimitTier::Global);
2074 }
2075}