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}