grishka/lemmy
1
0
Fork 0
Code Issues Pull requests Projects Releases Wiki Activity

Fix rate limiter (#4560)

Browse source
This commit is contained in:
Sander Saarend 2024-03-25 13:56:03 +02:00 committed by GitHub
parent ef4bb3cc40
commit 99d01e186a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 19 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

20
crates/utils/src/rate_limit/rate_limiter.rs
View file

@ -158,7 +158,7 @@ impl<K: Eq + Hash, C: MapLevel> MapLevel for Map<K, C> {
// Evaluated if `some_children_remaining` is false // Evaluated if `some_children_remaining` is false
let total_has_refill_in_future = || { let total_has_refill_in_future = || {
group.total.into_iter().all(|(action_type, bucket)| { group.total.into_iter().any(|(action_type, bucket)| {
#[allow(clippy::indexing_slicing)] #[allow(clippy::indexing_slicing)]
let config = configs[action_type]; let config = configs[action_type];
bucket.update(now, config).tokens != config.capacity bucket.update(now, config).tokens != config.capacity
@ -416,5 +416,23 @@ mod tests {
rate_limiter.remove_full_buckets(now); rate_limiter.remove_full_buckets(now);
assert!(rate_limiter.ipv4_buckets.is_empty()); assert!(rate_limiter.ipv4_buckets.is_empty());
assert!(rate_limiter.ipv6_buckets.is_empty()); assert!(rate_limiter.ipv6_buckets.is_empty());
// `remove full buckets` should not remove empty buckets
let ip = "1.1.1.1".parse().unwrap();
// empty the bucket with 2 requests
assert!(rate_limiter.check(ActionType::Post, ip, now));
assert!(rate_limiter.check(ActionType::Post, ip, now));
rate_limiter.remove_full_buckets(now);
assert!(!rate_limiter.ipv4_buckets.is_empty());
// `remove full buckets` should not remove partial buckets
now.secs += 2;
let ip = "1.1.1.1".parse().unwrap();
// Only make one request, so bucket still has 1 token
assert!(rate_limiter.check(ActionType::Post, ip, now));
rate_limiter.remove_full_buckets(now);
assert!(!rate_limiter.ipv4_buckets.is_empty());
} }
} }