From 78ce33896012076a0a452a712f7f7ee806b8157e Mon Sep 17 00:00:00 2001 From: Andreas Fuchs Date: Sun, 19 Mar 2023 16:17:38 -0400 Subject: [PATCH 1/3] Implement contains_key for rate limiters This allows systems to check whether they have already added a given key to a state store, which might help fight off DDoS attacks that vary the key (e.g. the origin IP address). --- governor/src/state.rs | 2 ++ governor/src/state/direct.rs | 6 ++++++ governor/src/state/in_memory.rs | 4 ++++ governor/src/state/keyed.rs | 16 ++++++++++++++++ governor/src/state/keyed/dashmap.rs | 4 ++++ governor/src/state/keyed/hashmap.rs | 5 +++++ governor/tests/keyed_dashmap.rs | 2 ++ governor/tests/keyed_hashmap.rs | 2 ++ 8 files changed, 41 insertions(+) diff --git a/governor/src/state.rs b/governor/src/state.rs index 2ea0e7b3..40a66717 100644 --- a/governor/src/state.rs +++ b/governor/src/state.rs @@ -46,6 +46,8 @@ pub trait StateStore { fn measure_and_replace(&self, key: &Self::Key, f: F) -> Result where F: Fn(Option) -> Result<(T, Nanos), E>; + + fn contains_key(&self, key: &Self::Key) -> bool; } /// A rate limiter. diff --git a/governor/src/state/direct.rs b/governor/src/state/direct.rs index 9310a6c8..7803060c 100644 --- a/governor/src/state/direct.rs +++ b/governor/src/state/direct.rs @@ -135,4 +135,10 @@ mod test { fn not_keyed_impls_coverage() { assert_eq!(NotKeyed::NonKey, NotKeyed::NonKey); } + + #[test] + fn not_keyed_contains_key_impls_coverage() { + let state = InMemoryState::default(); + assert!(state.contains_key(&NotKeyed::NonKey)); + } } diff --git a/governor/src/state/in_memory.rs b/governor/src/state/in_memory.rs index 100b0569..e78eccfa 100644 --- a/governor/src/state/in_memory.rs +++ b/governor/src/state/in_memory.rs @@ -59,6 +59,10 @@ impl StateStore for InMemoryState { { self.measure_and_replace_one(f) } + + fn contains_key(&self, _key: &Self::Key) -> bool { + true + } } impl Debug for InMemoryState { diff --git a/governor/src/state/keyed.rs b/governor/src/state/keyed.rs index 1ff92e77..84a06d7b 100644 --- a/governor/src/state/keyed.rs +++ b/governor/src/state/keyed.rs @@ -86,6 +86,17 @@ where C: clock::Clock, MW: RateLimitingMiddleware, { + /// Check whether a given key already exists in the state store. + /// + /// This method returns true if the key is already occupying + /// memory (that is, if an underlying hash map state store would + /// return true on its `contains_key` method). Particularly, it will + /// also return true on keys that would be dropped by + /// [`ShrinkableKeyedStateStore::retain_recent`]. + pub fn contains_key(&self, key: &K) -> bool { + self.state.contains_key(key) + } + /// Allow a single cell through the rate limiter for the given key. /// /// If the rate limit is reached, `check_key` returns information about the earliest @@ -252,6 +263,10 @@ mod test { impl StateStore for NaiveKeyedStateStore { type Key = K; + fn contains_key(&self, _key: &Self::Key) -> bool { + false + } + fn measure_and_replace(&self, _key: &Self::Key, f: F) -> Result where F: Fn(Option) -> Result<(T, Nanos), E>, @@ -283,6 +298,7 @@ mod test { NaiveKeyedStateStore::default(), &FakeRelativeClock::default(), ); + assert_eq!(false, lim.contains_key(&1u32)); assert_eq!(lim.check_key(&1u32), Ok(())); assert!(lim.is_empty()); assert_eq!(lim.len(), 0); diff --git a/governor/src/state/keyed/dashmap.rs b/governor/src/state/keyed/dashmap.rs index 43535110..820c73e6 100755 --- a/governor/src/state/keyed/dashmap.rs +++ b/governor/src/state/keyed/dashmap.rs @@ -27,6 +27,10 @@ impl StateStore for DashMapStateStore { let entry = self.entry(key.clone()).or_default(); (*entry).measure_and_replace_one(f) } + + fn contains_key(&self, key: &Self::Key) -> bool { + self.contains_key(key) + } } /// # Keyed rate limiters - [`DashMap`]-backed diff --git a/governor/src/state/keyed/hashmap.rs b/governor/src/state/keyed/hashmap.rs index ca34e128..c3f546ae 100644 --- a/governor/src/state/keyed/hashmap.rs +++ b/governor/src/state/keyed/hashmap.rs @@ -37,6 +37,11 @@ impl StateStore for HashMapStateStore { .or_insert_with(InMemoryState::default); entry.measure_and_replace_one(f) } + + fn contains_key(&self, key: &Self::Key) -> bool { + let map = self.lock(); + (*map).contains_key(key) + } } impl ShrinkableKeyedStateStore for HashMapStateStore { diff --git a/governor/tests/keyed_dashmap.rs b/governor/tests/keyed_dashmap.rs index 2c427f31..d8b69a9a 100755 --- a/governor/tests/keyed_dashmap.rs +++ b/governor/tests/keyed_dashmap.rs @@ -16,7 +16,9 @@ fn accepts_first_cell() { let clock = FakeRelativeClock::default(); let lb = RateLimiter::dashmap_with_clock(Quota::per_second(nonzero!(5u32)), &clock); for key in KEYS { + assert!(!lb.contains_key(&key)); assert_eq!(Ok(()), lb.check_key(&key), "key {:?}", key); + assert!(lb.contains_key(&key)); } } diff --git a/governor/tests/keyed_hashmap.rs b/governor/tests/keyed_hashmap.rs index f7eafd02..0cd1a759 100644 --- a/governor/tests/keyed_hashmap.rs +++ b/governor/tests/keyed_hashmap.rs @@ -14,7 +14,9 @@ fn accepts_first_cell() { let clock = FakeRelativeClock::default(); let lb = RateLimiter::hashmap_with_clock(Quota::per_second(nonzero!(5u32)), &clock); for key in KEYS { + assert!(!lb.contains_key(&key)); assert_eq!(Ok(()), lb.check_key(&key), "key {:?}", key); + assert!(lb.contains_key(&key)); } } From 8d138bd18b161568aab4cffc55e7eacf29731c8b Mon Sep 17 00:00:00 2001 From: Andreas Fuchs Date: Sun, 19 Mar 2023 19:48:16 -0400 Subject: [PATCH 2/3] Changelog --- governor/CHANGELOG.md | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/governor/CHANGELOG.md b/governor/CHANGELOG.md index 8487899c..cebf7f47 100644 --- a/governor/CHANGELOG.md +++ b/governor/CHANGELOG.md @@ -9,6 +9,10 @@ `DefaultKeyedRateLimiter` to cut down on type-typing of typical rate limiters in struct and function definitions. Requested in [#85](https://github.com/antifuchs/governor/issues/85). +* New method `contains_key` that can be used to check if a + `RateLimiter` has a given key (which can be used to prevent + arbitrary attacker-controlled memory usage in case of a DDoS + attack). ### Changed * The API for `.check_n` and `.until_n` (and their keyed counterpart) From dbfb73a10659816bb3b71f4aa5bf0a2d15b5d5a3 Mon Sep 17 00:00:00 2001 From: Andreas Fuchs Date: Sun, 19 Mar 2023 17:38:43 -0400 Subject: [PATCH 3/3] empty commit to trigger codecov