aster_forge_middleware/shared/
rate_limit.rs

1//! Framework-neutral governor quotas and normalized keyed rate limiting.
2use governor::clock::{Clock, DefaultClock, QuantaInstant};
3use governor::middleware::NoOpMiddleware;
4use governor::state::keyed::DefaultKeyedStateStore;
5use governor::{NotUntil, Quota, RateLimiter};
6use std::num::{NonZeroU32, NonZeroU64};
7use std::sync::Arc;
8use std::time::Duration;
9type StringKeyedLimiter =
10    RateLimiter<String, DefaultKeyedStateStore<String>, DefaultClock, NoOpMiddleware>;
11/// Returns the retry delay in whole seconds for a governor rejection.
12///
13/// Sub-second waits round up to one second so clients never see a zero delay
14/// that invites an immediate retry.
15#[must_use]
16pub fn retry_after_seconds(not_until: &NotUntil<QuantaInstant>) -> u64 {
17    let delay = not_until.wait_time_from(DefaultClock::default().now());
18    ceil_retry_after(delay)
19}
20
21fn ceil_retry_after(delay: Duration) -> u64 {
22    delay
23        .as_secs()
24        .saturating_add(u64::from(delay.subsec_nanos() > 0))
25        .max(1)
26}
27
28/// A keyed string rate limiter with product-neutral key normalization.
29///
30/// The limiter trims surrounding whitespace and lowercases keys before checking
31/// the quota. This suits usernames, email addresses, provider IDs, and similar
32/// business-unique identifiers where accidental case differences should not
33/// bypass a rate limit.
34#[derive(Clone)]
35pub struct NormalizedStringRateLimiter {
36    enabled: bool,
37    limiter: Arc<StringKeyedLimiter>,
38}
39
40impl NormalizedStringRateLimiter {
41    /// Whether checks are enabled. Disabled transport adapters bypass key extraction too.
42    #[must_use]
43    pub const fn enabled(&self) -> bool {
44        self.enabled
45    }
46
47    /// Removes expired buckets; schedule maintenance in the owning product runtime.
48    pub fn retain_recent(&self) {
49        self.limiter.retain_recent();
50    }
51
52    /// Builds a limiter from a non-zero quota and enabled flag.
53    #[must_use]
54    pub fn new(enabled: bool, seconds_per_request: NonZeroU64, burst_size: NonZeroU32) -> Self {
55        Self {
56            enabled,
57            limiter: Arc::new(RateLimiter::keyed(rate_limit_quota(
58                seconds_per_request,
59                burst_size,
60            ))),
61        }
62    }
63
64    /// Checks a raw key after trimming whitespace and lowercasing it.
65    #[must_use]
66    pub fn check(&self, raw_key: &str) -> Option<RateLimitRejection> {
67        if !self.enabled {
68            return None;
69        }
70
71        let key = raw_key.trim().to_ascii_lowercase();
72        self.limiter
73            .check_key(&key)
74            .err()
75            .map(|not_until| RateLimitRejection::from_not_until(&not_until))
76    }
77}
78
79/// Product-neutral rate-limit rejection metadata.
80#[derive(Debug, Clone, Copy, PartialEq, Eq)]
81pub struct RateLimitRejection {
82    retry_after_seconds: u64,
83}
84
85impl RateLimitRejection {
86    pub(crate) fn from_not_until(not_until: &NotUntil<QuantaInstant>) -> Self {
87        Self {
88            retry_after_seconds: retry_after_seconds(not_until),
89        }
90    }
91
92    /// Returns how many seconds clients should wait before retrying.
93    #[must_use]
94    pub const fn retry_after_seconds(self) -> u64 {
95        self.retry_after_seconds
96    }
97}
98
99/// Builds a governor quota from non-zero seconds and burst size.
100///
101/// # Panics
102/// Panics only if governor rejects a non-zero period.
103#[expect(clippy::expect_used, reason = "NonZeroU64 creates a non-zero duration")]
104#[must_use]
105pub fn rate_limit_quota(seconds_per_request: NonZeroU64, burst_size: NonZeroU32) -> Quota {
106    Quota::with_period(Duration::from_secs(seconds_per_request.get()))
107        .expect("non-zero period")
108        .allow_burst(burst_size)
109}
110
111#[cfg(test)]
112mod tests {
113    use super::*;
114    use std::num::{NonZeroU32, NonZeroU64};
115    #[test]
116    fn quota_refills_at_the_exact_boundary_without_sleep() {
117        use governor::clock::FakeRelativeClock;
118        let clock = FakeRelativeClock::default();
119        let quota = rate_limit_quota(NonZeroU64::new(2).unwrap(), NonZeroU32::new(2).unwrap());
120        let limiter = RateLimiter::dashmap_with_clock(quota, clock.clone());
121        assert!(limiter.check_key(&"one").is_ok());
122        assert!(limiter.check_key(&"one").is_ok());
123        assert!(limiter.check_key(&"one").is_err());
124        assert!(limiter.check_key(&"two").is_ok());
125        clock.advance(Duration::from_millis(1999));
126        assert!(limiter.check_key(&"one").is_err());
127        clock.advance(Duration::from_millis(1));
128        assert!(limiter.check_key(&"one").is_ok());
129        assert!(limiter.check_key(&"one").is_err());
130    }
131
132    #[test]
133    fn multi_second_retry_delay_rounds_up_instead_of_inviting_early_retry() {
134        assert_eq!(ceil_retry_after(Duration::from_millis(1999)), 2);
135        assert_eq!(ceil_retry_after(Duration::from_millis(59_999)), 60);
136        assert_eq!(ceil_retry_after(Duration::from_secs(2)), 2);
137        assert_eq!(ceil_retry_after(Duration::ZERO), 1);
138    }
139    #[test]
140    fn retry_after_seconds_rounds_sub_second_waits_up_to_one() {
141        let quota = governor::Quota::with_period(std::time::Duration::from_secs(1))
142            .unwrap()
143            .allow_burst(NonZeroU32::new(1).unwrap());
144        let limiter = governor::RateLimiter::keyed(quota);
145
146        assert!(limiter.check_key(&"key").is_ok());
147        let not_until = limiter
148            .check_key(&"key")
149            .expect_err("second immediate check should be rate limited");
150
151        // The remaining wait is strictly below one second (some nanoseconds have
152        // elapsed since the first check), so truncating whole seconds would
153        // report 0 and tell the client to retry immediately.
154        assert_eq!(retry_after_seconds(&not_until), 1);
155    }
156    #[test]
157    fn normalized_string_limiter_can_be_disabled() {
158        let limiter = NormalizedStringRateLimiter::new(
159            false,
160            NonZeroU64::new(60).unwrap(),
161            NonZeroU32::new(1).unwrap(),
162        );
163
164        assert!(limiter.check("admin@example.com").is_none());
165        assert!(limiter.check("admin@example.com").is_none());
166    }
167    #[test]
168    fn normalized_string_limiter_trims_and_lowercases_keys() {
169        let limiter = NormalizedStringRateLimiter::new(
170            true,
171            NonZeroU64::new(60).unwrap(),
172            NonZeroU32::new(1).unwrap(),
173        );
174
175        assert!(limiter.check("Admin@Example.com").is_none());
176        assert!(limiter.check("other@example.com").is_none());
177        assert!(limiter.check(" admin@example.com ").is_some());
178    }
179}
aster_forge_middleware/shared/rate_limit.rs

aster_forge_middleware/shared/
rate_limit.rs
