Skip to main content

rmcp_server_kit/
bounded_limiter.rs

1//! Memory-bounded keyed rate limiter.
2//!
3//! [`crate::bounded_limiter::BoundedKeyedLimiter`] wraps a map of per-key
4//! [`governor::DefaultDirectRateLimiter`] instances behind a hard cap on the
5//! number of tracked keys, with an idle-eviction policy and configurable
6//! full-table behaviour when the cap is reached.
7//!
8//! # Why
9//!
10//! The `governor` crate ships a [`governor::RateLimiter::keyed`] state store
11//! whose memory grows monotonically with the number of distinct keys
12//! observed. For server use cases keyed by source IP this is a
13//! denial-of-service vector: an attacker spraying packets from spoofed or
14//! distinct source addresses can exhaust process memory regardless of the
15//! per-key quota.
16//!
17//! [`crate::bounded_limiter::BoundedKeyedLimiter`] addresses this by:
18//!
19//! 1. Holding a [`std::collections::HashMap`] of `K -> Entry` where each
20//!    `Entry` carries its own direct (per-key) limiter and a `last_seen`
21//!    timestamp.
22//! 2. Capping the map at `max_tracked_keys` entries.
23//! 3. On insert when the map is full, first pruning entries whose
24//!    `last_seen` is older than `idle_eviction`, then applying
25//!    [`KeyEvictionPolicy`](crate::bounded_limiter::KeyEvictionPolicy). The default policy evicts the entry with the
26//!    oldest `last_seen` ("LRU eviction") so the new key is inserted.
27//! 4. Updating `last_seen` on **every** check (including rate-limit
28//!    rejections) so an actively-firing attacker cannot dodge eviction by
29//!    appearing idle.
30//! 5. Optionally spawning a best-effort background prune task. Cap
31//!    enforcement does **not** depend on this task running -- it is
32//!    purely an optimization that reclaims memory between admission
33//!    events.
34//!
35//! # Trade-offs
36//!
37//! - When a previously-evicted key reappears it gets a **fresh** quota.
38//!   This is documented behaviour: a key under sustained load keeps its
39//!   `last_seen` updated and therefore is never evicted; eviction only
40//!   targets idle keys.
41//! - The map uses [`std::sync::Mutex`] (not [`tokio::sync::Mutex`]) since
42//!   admission checks must be synchronous and never `.await`.
43//! - We do not log inside the critical section.
44
45use std::{
46    collections::HashMap,
47    hash::Hash,
48    num::{NonZeroU32, NonZeroUsize},
49    str::FromStr,
50    sync::{Arc, Mutex, PoisonError, Weak},
51    time::{Duration, Instant},
52};
53
54use governor::{
55    DefaultDirectRateLimiter, Quota, RateLimiter,
56    clock::{Clock as _, DefaultClock},
57};
58
59/// Reason a [`BoundedKeyedLimiter::check_key`] call rejected a request.
60///
61/// Currently only carries a single variant; modelled as an enum (rather
62/// than a unit struct) so callers can `match` exhaustively and to leave
63/// room for future reasons (e.g. burst-debt or distinct quota classes).
64#[non_exhaustive]
65#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
66pub enum BoundedLimiterError {
67    /// The key has exceeded its per-key quota for the current window.
68    #[error("rate limit exceeded for key")]
69    RateLimited,
70}
71
72/// Reason a detailed bounded-limiter check denied a request.
73#[non_exhaustive]
74#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
75pub enum BoundedLimiterDeny {
76    /// The key has exceeded its per-key quota for the current window.
77    #[error("rate limit exceeded; retry after {0:?}")]
78    RateLimited(Duration),
79    /// The limiter is at its tracked-key capacity and the configured policy
80    /// rejects unseen keys instead of evicting an existing bucket.
81    #[error("tracked-key capacity is full")]
82    CapacityFull,
83}
84
85/// Behaviour when a new key arrives after the tracked-key table reaches capacity.
86#[non_exhaustive]
87#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, serde::Deserialize)]
88#[serde(rename_all = "snake_case")]
89pub enum KeyEvictionPolicy {
90    /// Evict the least-recently-seen key and admit the new key.
91    #[default]
92    EvictLru,
93    /// Reject unseen keys while preserving buckets for already-tracked keys.
94    RejectNew,
95}
96
97impl FromStr for KeyEvictionPolicy {
98    type Err = ();
99
100    fn from_str(value: &str) -> Result<Self, Self::Err> {
101        match value {
102            "evict_lru" => Ok(Self::EvictLru),
103            "reject_new" => Ok(Self::RejectNew),
104            _ => Err(()),
105        }
106    }
107}
108
109/// Per-key limiter entry: the underlying direct limiter plus the wall-clock
110/// timestamp of the most recent admission attempt for this key.
111struct Entry {
112    limiter: DefaultDirectRateLimiter,
113    last_seen: Instant,
114}
115
116/// Inner shared state. Held behind an [`Arc`] in [`BoundedKeyedLimiter`]
117/// and a [`Weak`] inside the optional background prune task so the task
118/// self-terminates once the limiter is dropped.
119struct Inner<K: Eq + Hash + Clone> {
120    map: Mutex<HashMap<K, Entry>>,
121    quota: Quota,
122    max_tracked_keys: usize,
123    idle_eviction: Duration,
124    key_eviction_policy: KeyEvictionPolicy,
125}
126
127/// Memory-bounded keyed rate limiter.
128///
129/// Cheaply cloneable; clones share state.
130#[allow(
131    missing_debug_implementations,
132    reason = "wraps governor RateLimiter which has no Debug impl"
133)]
134pub struct BoundedKeyedLimiter<K: Eq + Hash + Clone> {
135    inner: Arc<Inner<K>>,
136}
137
138impl<K: Eq + Hash + Clone> Clone for BoundedKeyedLimiter<K> {
139    fn clone(&self) -> Self {
140        Self {
141            inner: Arc::clone(&self.inner),
142        }
143    }
144}
145
146impl<K: Eq + Hash + Clone + Send + Sync + 'static> BoundedKeyedLimiter<K> {
147    /// Create a new bounded keyed limiter.
148    ///
149    /// * `quota` -- the per-key rate-limit quota applied to every entry.
150    /// * `max_tracked_keys` -- hard cap on the number of simultaneously
151    ///   tracked keys. When reached, an insert first prunes idle entries
152    ///   then falls back to LRU eviction.
153    /// * `idle_eviction` -- entries whose `last_seen` is older than this
154    ///   are eligible for opportunistic pruning.
155    ///
156    /// # Background prune task
157    ///
158    /// If a Tokio runtime is available at construction time, a best-effort
159    /// background task is spawned that periodically prunes idle entries.
160    /// Cap enforcement does **not** depend on this task; it is purely an
161    /// optimisation that reclaims memory between admission events. The
162    /// task self-terminates when the last [`BoundedKeyedLimiter`] clone is
163    /// dropped (it holds only a [`Weak`] reference to the inner state).
164    ///
165    /// If no Tokio runtime is available (e.g. unit tests using
166    /// `#[test]` rather than `#[tokio::test]`), no task is spawned and
167    /// pruning happens lazily on every full-table insert. Both behaviours
168    /// are correct.
169    #[must_use]
170    pub(crate) fn new(
171        quota: Quota,
172        max_tracked_keys: NonZeroUsize,
173        idle_eviction: Duration,
174    ) -> Self {
175        Self::new_with_policy(
176            quota,
177            max_tracked_keys,
178            idle_eviction,
179            KeyEvictionPolicy::default(),
180        )
181    }
182
183    /// Create a new bounded keyed limiter with explicit full-table behaviour.
184    #[must_use]
185    pub(crate) fn new_with_policy(
186        quota: Quota,
187        max_tracked_keys: NonZeroUsize,
188        idle_eviction: Duration,
189        key_eviction_policy: KeyEvictionPolicy,
190    ) -> Self {
191        let inner = Arc::new(Inner {
192            map: Mutex::new(HashMap::new()),
193            quota,
194            max_tracked_keys: max_tracked_keys.get(),
195            idle_eviction,
196            key_eviction_policy,
197        });
198        Self::spawn_prune_task(&inner);
199        Self { inner }
200    }
201
202    /// Construct a [`BoundedKeyedLimiter`] with a per-minute quota.
203    ///
204    /// Convenience constructor that builds a per-minute [`Quota`] from
205    /// `requests_per_minute`. The rate is clamped to a minimum of `1`
206    /// request/min so a misconfigured `0` does not panic at startup.
207    ///
208    /// * `requests_per_minute` -- per-key rate, clamped to `>= 1`.
209    /// * `max_tracked_keys` -- hard cap on simultaneously tracked keys,
210    ///   clamped to `>= 1`. `McpServerConfig` validation rejects `0`
211    ///   upstream, so the clamp is defense-in-depth for direct callers.
212    ///   When reached, an insert first prunes idle entries then falls
213    ///   back to LRU eviction.
214    /// * `idle_eviction` -- entries whose `last_seen` is older than this
215    ///   are eligible for opportunistic pruning.
216    #[must_use]
217    pub fn with_per_minute(
218        requests_per_minute: u32,
219        max_tracked_keys: usize,
220        idle_eviction: Duration,
221    ) -> Self {
222        let rate = NonZeroU32::new(requests_per_minute.max(1)).unwrap_or(NonZeroU32::MIN);
223        let max_tracked_keys = NonZeroUsize::new(max_tracked_keys).unwrap_or(NonZeroUsize::MIN);
224        Self::new(Quota::per_minute(rate), max_tracked_keys, idle_eviction)
225    }
226
227    /// Construct a [`BoundedKeyedLimiter`] with a per-minute quota and policy.
228    #[must_use]
229    pub fn with_per_minute_and_policy(
230        requests_per_minute: u32,
231        max_tracked_keys: usize,
232        idle_eviction: Duration,
233        key_eviction_policy: KeyEvictionPolicy,
234    ) -> Self {
235        let rate = NonZeroU32::new(requests_per_minute.max(1)).unwrap_or(NonZeroU32::MIN);
236        let max_tracked_keys = NonZeroUsize::new(max_tracked_keys).unwrap_or(NonZeroUsize::MIN);
237        Self::new_with_policy(
238            Quota::per_minute(rate),
239            max_tracked_keys,
240            idle_eviction,
241            key_eviction_policy,
242        )
243    }
244
245    /// Construct a [`BoundedKeyedLimiter`] with a per-second quota.
246    ///
247    /// Convenience constructor that builds a per-second [`Quota`] from
248    /// `requests_per_second`. The rate is clamped to a minimum of `1`
249    /// request/sec so a misconfigured `0` does not panic at startup.
250    ///
251    /// * `requests_per_second` -- per-key rate, clamped to `>= 1`.
252    /// * `max_tracked_keys` -- hard cap on simultaneously tracked keys,
253    ///   clamped to `>= 1`. `McpServerConfig` validation rejects `0`
254    ///   upstream, so the clamp is defense-in-depth for direct callers.
255    ///   When reached, an insert first prunes idle entries then falls
256    ///   back to LRU eviction.
257    /// * `idle_eviction` -- entries whose `last_seen` is older than this
258    ///   are eligible for opportunistic pruning.
259    #[must_use]
260    pub fn with_per_second(
261        requests_per_second: u32,
262        max_tracked_keys: usize,
263        idle_eviction: Duration,
264    ) -> Self {
265        let rate = NonZeroU32::new(requests_per_second.max(1)).unwrap_or(NonZeroU32::MIN);
266        let max_tracked_keys = NonZeroUsize::new(max_tracked_keys).unwrap_or(NonZeroUsize::MIN);
267        Self::new(Quota::per_second(rate), max_tracked_keys, idle_eviction)
268    }
269
270    /// Construct a [`BoundedKeyedLimiter`] with a per-second quota and policy.
271    #[must_use]
272    pub fn with_per_second_and_policy(
273        requests_per_second: u32,
274        max_tracked_keys: usize,
275        idle_eviction: Duration,
276        key_eviction_policy: KeyEvictionPolicy,
277    ) -> Self {
278        let rate = NonZeroU32::new(requests_per_second.max(1)).unwrap_or(NonZeroU32::MIN);
279        let max_tracked_keys = NonZeroUsize::new(max_tracked_keys).unwrap_or(NonZeroUsize::MIN);
280        Self::new_with_policy(
281            Quota::per_second(rate),
282            max_tracked_keys,
283            idle_eviction,
284            key_eviction_policy,
285        )
286    }
287
288    /// Spawn the optional background prune task. No-op if there is no
289    /// current Tokio runtime.
290    fn spawn_prune_task(inner: &Arc<Inner<K>>) {
291        let Ok(handle) = tokio::runtime::Handle::try_current() else {
292            return;
293        };
294        let weak: Weak<Inner<K>> = Arc::downgrade(inner);
295        // Prune at most once every quarter of `idle_eviction`, but never
296        // less than once per minute (to avoid waking up too often when
297        // operators configure a very long eviction window).
298        let interval = (inner.idle_eviction / 4).max(Duration::from_mins(1));
299        handle.spawn(async move {
300            let mut ticker = tokio::time::interval(interval);
301            // We just woke up from `Handle::spawn`; don't burn the first tick.
302            ticker.tick().await;
303            loop {
304                ticker.tick().await;
305                let Some(inner) = weak.upgrade() else {
306                    return;
307                };
308                Self::prune_idle(&inner);
309            }
310        });
311    }
312
313    /// Drop entries whose `last_seen` is older than `idle_eviction`.
314    fn prune_idle(inner: &Inner<K>) {
315        let mut guard = inner.map.lock().unwrap_or_else(PoisonError::into_inner);
316        let cutoff = Instant::now()
317            .checked_sub(inner.idle_eviction)
318            .unwrap_or_else(Instant::now);
319        guard.retain(|_, entry| entry.last_seen >= cutoff);
320    }
321
322    /// Evict the single entry with the oldest `last_seen`. Caller must hold
323    /// the map lock. Used only when the table is full *after* idle pruning.
324    fn evict_lru(map: &mut HashMap<K, Entry>) {
325        let oldest_key = map
326            .iter()
327            .min_by_key(|(_, entry)| entry.last_seen)
328            .map(|(k, _)| k.clone());
329        if let Some(key) = oldest_key {
330            map.remove(&key);
331        }
332    }
333
334    /// Test the per-key quota for `key`.
335    ///
336    /// Returns `Ok(())` if the request is allowed. The `last_seen`
337    /// timestamp is updated on **every** call -- including rate-limit
338    /// rejections -- so an actively firing attacker cannot age out into
339    /// a fresh quota by appearing idle.
340    ///
341    /// When inserting a new key into a full table, idle entries are pruned
342    /// first; if the table is still full, the entry with the oldest
343    /// `last_seen` is evicted (LRU). The new key is always inserted --
344    /// honest new clients are never rejected because the table is full.
345    ///
346    /// # Errors
347    ///
348    /// Returns [`BoundedLimiterError::RateLimited`] when `key` has
349    /// exceeded its per-key quota for the current window.
350    pub fn check_key(&self, key: &K) -> Result<(), BoundedLimiterError> {
351        self.check_key_wait(key)
352            .map_err(|_| BoundedLimiterError::RateLimited)
353    }
354
355    /// Test the per-key quota for `key`, returning the wait time on deny.
356    ///
357    /// Identical admission semantics to [`check_key`](Self::check_key)
358    /// (same `last_seen` refresh, idle-prune, and LRU-eviction behavior);
359    /// the two methods share one code path.
360    ///
361    /// # Errors
362    ///
363    /// On deny, returns the **best-effort current wait** until the next
364    /// request for this key could be admitted, measured against
365    /// governor's default clock at the moment of the failed check. The
366    /// value is a raw [`Duration`]; rounding (e.g. ceiling to whole
367    /// seconds for a `Retry-After` header) is the caller's concern.
368    pub fn check_key_wait(&self, key: &K) -> Result<(), Duration> {
369        let mut guard = self
370            .inner
371            .map
372            .lock()
373            .unwrap_or_else(PoisonError::into_inner);
374        let now = Instant::now();
375        if let Some(entry) = guard.get_mut(key) {
376            entry.last_seen = now;
377            return entry
378                .limiter
379                .check()
380                .map_err(|not_until| not_until.wait_time_from(DefaultClock::default().now()));
381        }
382        // New key: make room if necessary, then insert.
383        if guard.len() >= self.inner.max_tracked_keys {
384            // Prune idle first.
385            let cutoff = now
386                .checked_sub(self.inner.idle_eviction)
387                .unwrap_or_else(Instant::now);
388            guard.retain(|_, entry| entry.last_seen >= cutoff);
389            // If still full, evict LRU.
390            if guard.len() >= self.inner.max_tracked_keys {
391                Self::evict_lru(&mut guard);
392            }
393        }
394        let limiter = RateLimiter::direct(self.inner.quota);
395        let result = limiter
396            .check()
397            .map_err(|not_until| not_until.wait_time_from(DefaultClock::default().now()));
398        guard.insert(
399            key.clone(),
400            Entry {
401                limiter,
402                last_seen: now,
403            },
404        );
405        result
406    }
407
408    /// Test the per-key quota for `key`, preserving capacity-denial details.
409    ///
410    /// Unlike [`check_key_wait`](Self::check_key_wait), this method honors the
411    /// configured [`KeyEvictionPolicy`] and can report full-table rejection via
412    /// [`BoundedLimiterDeny::CapacityFull`]. Existing callers that need legacy
413    /// always-evict behaviour can keep using [`check_key`](Self::check_key) or
414    /// [`check_key_wait`](Self::check_key_wait).
415    ///
416    /// # Errors
417    ///
418    /// Returns [`BoundedLimiterDeny::RateLimited`] when an established bucket is
419    /// over quota, or [`BoundedLimiterDeny::CapacityFull`] when an unseen key is
420    /// rejected by [`KeyEvictionPolicy::RejectNew`].
421    pub fn check_key_detailed(&self, key: &K) -> Result<(), BoundedLimiterDeny> {
422        let mut guard = self
423            .inner
424            .map
425            .lock()
426            .unwrap_or_else(PoisonError::into_inner);
427        let now = Instant::now();
428        if let Some(entry) = guard.get_mut(key) {
429            entry.last_seen = now;
430            return entry.limiter.check().map_err(|not_until| {
431                BoundedLimiterDeny::RateLimited(
432                    not_until.wait_time_from(DefaultClock::default().now()),
433                )
434            });
435        }
436        if guard.len() >= self.inner.max_tracked_keys {
437            let cutoff = now
438                .checked_sub(self.inner.idle_eviction)
439                .unwrap_or_else(Instant::now);
440            guard.retain(|_, entry| entry.last_seen >= cutoff);
441            if guard.len() >= self.inner.max_tracked_keys {
442                match self.inner.key_eviction_policy {
443                    KeyEvictionPolicy::EvictLru => Self::evict_lru(&mut guard),
444                    KeyEvictionPolicy::RejectNew => return Err(BoundedLimiterDeny::CapacityFull),
445                }
446            }
447        }
448        let limiter = RateLimiter::direct(self.inner.quota);
449        let result = limiter.check().map_err(|not_until| {
450            BoundedLimiterDeny::RateLimited(not_until.wait_time_from(DefaultClock::default().now()))
451        });
452        guard.insert(
453            key.clone(),
454            Entry {
455                limiter,
456                last_seen: now,
457            },
458        );
459        result
460    }
461
462    /// Number of currently tracked keys. Used by tests and admin endpoints.
463    #[must_use]
464    pub fn len(&self) -> usize {
465        self.inner
466            .map
467            .lock()
468            .unwrap_or_else(PoisonError::into_inner)
469            .len()
470    }
471
472    /// `true` when no keys are currently tracked.
473    #[must_use]
474    pub fn is_empty(&self) -> bool {
475        self.len() == 0
476    }
477}
478
479#[cfg(test)]
480mod tests {
481    use std::{
482        net::IpAddr,
483        num::{NonZeroU32, NonZeroUsize},
484        time::Duration,
485    };
486
487    use governor::Quota;
488
489    use super::{BoundedKeyedLimiter, BoundedLimiterDeny, BoundedLimiterError, KeyEvictionPolicy};
490
491    fn ip(n: u32) -> IpAddr {
492        IpAddr::from(n.to_be_bytes())
493    }
494
495    fn cap(n: usize) -> NonZeroUsize {
496        NonZeroUsize::new(n).unwrap_or(NonZeroUsize::MIN)
497    }
498
499    /// Deny on the existing-key branch must report a positive,
500    /// quota-bounded wait time.
501    #[test]
502    fn check_key_wait_existing_key_deny_returns_bounded_wait() {
503        let quota = Quota::per_minute(NonZeroU32::new(1).unwrap());
504        let limiter: BoundedKeyedLimiter<IpAddr> =
505            BoundedKeyedLimiter::new(quota, cap(10), Duration::from_hours(1));
506        assert!(limiter.check_key_wait(&ip(1)).is_ok(), "burst admits first");
507        let wait = limiter
508            .check_key_wait(&ip(1))
509            .expect_err("second call within the window must deny");
510        assert!(wait > Duration::ZERO, "wait must be positive, got {wait:?}");
511        assert!(
512            wait <= Duration::from_secs(60),
513            "per-minute quota wait must be <= 60s, got {wait:?}"
514        );
515    }
516
517    /// The new-key branch always admits the first check: a freshly
518    /// constructed governor limiter starts with a full bucket and burst
519    /// capacity is `NonZeroU32` (>= 1). The deny arm on that branch is
520    /// defensive symmetry, not a reachable path.
521    #[test]
522    fn check_key_wait_new_key_first_check_admits() {
523        let quota = Quota::per_minute(NonZeroU32::new(1).unwrap());
524        let limiter: BoundedKeyedLimiter<IpAddr> =
525            BoundedKeyedLimiter::new(quota, cap(10), Duration::from_hours(1));
526        for i in 0..5_u32 {
527            assert!(
528                limiter.check_key_wait(&ip(i)).is_ok(),
529                "first check for new key {i} must admit"
530            );
531        }
532    }
533
534    /// `check_key` delegates to `check_key_wait`: identical admission
535    /// decisions, error mapped to the reason-only enum.
536    #[test]
537    fn check_key_delegates_to_wait_path() {
538        let quota = Quota::per_minute(NonZeroU32::new(1).unwrap());
539        let limiter: BoundedKeyedLimiter<IpAddr> =
540            BoundedKeyedLimiter::new(quota, cap(10), Duration::from_hours(1));
541        assert!(limiter.check_key(&ip(7)).is_ok());
542        assert_eq!(
543            limiter.check_key(&ip(7)),
544            Err(BoundedLimiterError::RateLimited)
545        );
546    }
547
548    #[test]
549    fn check_key_detailed_reports_rate_limit_wait_under_default_policy() {
550        let quota = Quota::per_minute(NonZeroU32::new(1).unwrap());
551        let limiter: BoundedKeyedLimiter<IpAddr> =
552            BoundedKeyedLimiter::new(quota, cap(10), Duration::from_hours(1));
553        assert!(limiter.check_key_detailed(&ip(7)).is_ok());
554        let deny = limiter
555            .check_key_detailed(&ip(7))
556            .expect_err("second call within the window must deny");
557        match deny {
558            BoundedLimiterDeny::RateLimited(wait) => {
559                assert!(wait > Duration::ZERO, "wait must be positive, got {wait:?}");
560                assert!(
561                    wait <= Duration::from_secs(60),
562                    "per-minute quota wait must be <= 60s, got {wait:?}"
563                );
564            }
565            BoundedLimiterDeny::CapacityFull => panic!("default policy must not reject capacity"),
566        }
567    }
568
569    /// The hard cap on tracked keys must never be exceeded, even under a
570    /// stream of distinct keys far larger than the cap.
571    #[test]
572    fn never_exceeds_max_tracked_keys() {
573        let quota = Quota::per_minute(NonZeroU32::new(10).unwrap());
574        let limiter: BoundedKeyedLimiter<IpAddr> =
575            BoundedKeyedLimiter::new(quota, cap(100), Duration::from_hours(1));
576        for i in 0..10_000_u32 {
577            let _ = limiter.check_key(&ip(i));
578            assert!(
579                limiter.len() <= 100,
580                "tracked keys exceeded cap at iteration {i}: {} > 100",
581                limiter.len()
582            );
583        }
584        assert_eq!(limiter.len(), 100, "table should be full at the cap");
585    }
586
587    #[test]
588    fn reject_new_at_cap_denies_unseen_key_but_keeps_established_key() {
589        let quota = Quota::per_minute(NonZeroU32::new(2).unwrap());
590        let limiter: BoundedKeyedLimiter<IpAddr> = BoundedKeyedLimiter::new_with_policy(
591            quota,
592            cap(1),
593            Duration::from_hours(1),
594            KeyEvictionPolicy::RejectNew,
595        );
596        let established = ip(10);
597        assert!(limiter.check_key_detailed(&established).is_ok());
598        assert_eq!(limiter.len(), 1);
599
600        let unseen = ip(11);
601        assert_eq!(
602            limiter.check_key_detailed(&unseen),
603            Err(BoundedLimiterDeny::CapacityFull)
604        );
605        assert_eq!(limiter.len(), 1);
606        assert!(
607            limiter.check_key_detailed(&established).is_ok(),
608            "established key keeps its existing bucket and remaining quota"
609        );
610    }
611
612    #[test]
613    fn evict_lru_policy_at_cap_admits_new_key_and_evicts_lru() {
614        let quota = Quota::per_minute(NonZeroU32::new(2).unwrap());
615        let limiter: BoundedKeyedLimiter<IpAddr> = BoundedKeyedLimiter::new_with_policy(
616            quota,
617            cap(1),
618            Duration::from_hours(1),
619            KeyEvictionPolicy::EvictLru,
620        );
621        let first = ip(20);
622        assert!(limiter.check_key_detailed(&first).is_ok());
623        assert!(limiter.check_key_detailed(&first).is_ok());
624        assert!(limiter.check_key_detailed(&first).is_err());
625
626        std::thread::sleep(Duration::from_millis(5));
627        assert!(limiter.check_key_detailed(&ip(21)).is_ok());
628        assert_eq!(limiter.len(), 1);
629        assert!(
630            limiter.check_key_detailed(&first).is_ok(),
631            "LRU-evicted key returns with fresh quota under EvictLru"
632        );
633    }
634
635    /// When a previously-evicted key reappears, it must get a fresh quota.
636    /// This is *documented* behaviour, not a bug: keys under sustained
637    /// load keep their `last_seen` updated and therefore are not evicted.
638    #[test]
639    fn evicted_keys_get_fresh_quota() {
640        let quota = Quota::per_minute(NonZeroU32::new(2).unwrap());
641        let limiter: BoundedKeyedLimiter<IpAddr> =
642            BoundedKeyedLimiter::new(quota, cap(2), Duration::from_hours(1));
643
644        let target = ip(1);
645        // Burn the quota for `target`.
646        assert!(limiter.check_key(&target).is_ok(), "first ok");
647        assert!(limiter.check_key(&target).is_ok(), "second ok");
648        assert!(limiter.check_key(&target).is_err(), "third blocked");
649
650        // Force eviction by inserting two unrelated keys (cap = 2). The
651        // attacker (`target`) is rate-limited -- it has a *recent*
652        // `last_seen` because of the failed check above. So inserting
653        // two new keys must NOT evict the attacker; instead one of the
654        // *other* unrelated keys gets evicted via LRU. We therefore
655        // need three unrelated keys to push `target` out by LRU.
656        //
657        // Sleep a tiny amount so unrelated keys have strictly newer
658        // last_seen than `target`'s last write.
659        std::thread::sleep(Duration::from_millis(5));
660        let _ = limiter.check_key(&ip(2));
661        std::thread::sleep(Duration::from_millis(5));
662        let _ = limiter.check_key(&ip(3));
663        // `target` is now the oldest entry; cap is 2. ip(3) eviction LRU'd
664        // either ip(2) or `target`. Inserting ip(4) again forces another
665        // eviction. After enough fresh inserts, `target` is gone.
666        std::thread::sleep(Duration::from_millis(5));
667        let _ = limiter.check_key(&ip(4));
668        std::thread::sleep(Duration::from_millis(5));
669        let _ = limiter.check_key(&ip(5));
670
671        // `target` should have been evicted by now -- a fresh check_key
672        // re-inserts with a fresh quota.
673        assert!(
674            limiter.check_key(&target).is_ok(),
675            "evicted key gets a fresh quota on reappearance"
676        );
677    }
678
679    /// An actively over-quota key must NOT be evicted just because new
680    /// keys are knocking. `last_seen` is updated on every check including
681    /// rate-limit rejections, so the attacker stays at the front of the
682    /// LRU queue. Other (older) entries are evicted instead.
683    #[test]
684    fn active_over_quota_key_not_evicted() {
685        let quota = Quota::per_minute(NonZeroU32::new(2).unwrap());
686        let limiter: BoundedKeyedLimiter<IpAddr> =
687            BoundedKeyedLimiter::new(quota, cap(3), Duration::from_hours(1));
688
689        // Seed the table with three idle entries so cap is reached.
690        for i in 100..103_u32 {
691            let _ = limiter.check_key(&ip(i));
692        }
693        assert_eq!(limiter.len(), 3);
694
695        // The attacker now starts firing. First two are allowed
696        // (fills quota), then we expect refusals -- but each refusal
697        // updates last_seen so the attacker stays "current".
698        std::thread::sleep(Duration::from_millis(5));
699        let attacker = ip(200);
700        // Inserting attacker evicts one of the older keys (cap=3).
701        let _ = limiter.check_key(&attacker);
702        let _ = limiter.check_key(&attacker);
703
704        // Interleave attacker hits with new-key knocks. The attacker
705        // keeps firing (last_seen always current), so when new keys
706        // arrive and force eviction, the LRU victim must be one of the
707        // *other* (older) entries, not the attacker.
708        for new_key in 300..310_u32 {
709            std::thread::sleep(Duration::from_millis(2));
710            let _ = limiter.check_key(&attacker); // attacker stays current
711            std::thread::sleep(Duration::from_millis(2));
712            let _ = limiter.check_key(&ip(new_key)); // forces eviction
713        }
714
715        // One final attacker hit immediately before the assertion to
716        // ensure no other key has been touched more recently.
717        let _ = limiter.check_key(&attacker);
718
719        // Attacker must STILL be rate-limited (quota exhausted, not a
720        // freshly-allocated entry). The check returns Err because the
721        // existing entry with exhausted quota is still there.
722        assert!(
723            limiter.check_key(&attacker).is_err(),
724            "actively over-quota attacker must not be evicted into a fresh quota"
725        );
726    }
727}