Bound the rate limiter's bucket dict (B-56)
_buckets is keyed by strings the caller chooses — any username, and (before B-54) any IP — and only ever grew: decay_seconds aged a bucket's counter but never removed the entry, so hammering login with random usernames was an unbounded memory leak. A bucket is "spent" once its lockout has expired *and* its failure count would decay to zero on the next failure anyway — at which point keeping it and dropping it are indistinguishable, which is what makes eviction safe. Those are swept on record_failure (at most once every 60s) and on the read path, so a key that's merely being probed never leaves an entry behind. That alone holds the dict at the size of the genuinely active attack surface. _MAX_BUCKETS = 50_000 is the backstop for a burst faster than the sweep interval, when nothing has had time to expire. Over it, the entries closest to expiry go first: what an attacker gets from a successful flood is the loss of the shallowest, nearly-over lockouts, never the deep ones actually holding an attack back. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
+53
-1
@@ -1,6 +1,15 @@
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
|
||||
# B-56: bounds on _buckets, which is keyed by strings the caller chooses.
|
||||
# 50k entries is a few MB at ~100 bytes each — far more than any real deployment's
|
||||
# active attacker set, and small enough that filling it isn't a memory attack.
|
||||
_MAX_BUCKETS = 50_000
|
||||
# How often record_failure sweeps out spent entries. Cheap (one pass over a dict
|
||||
# that the sweep itself keeps small) and off the request's critical path in the
|
||||
# normal case, since a successful login records no failure at all.
|
||||
_SWEEP_INTERVAL_SECONDS = 60.0
|
||||
|
||||
|
||||
@dataclass
|
||||
class _Bucket:
|
||||
@@ -29,22 +38,65 @@ class RateLimiter:
|
||||
base_delay: float = 2.0,
|
||||
max_delay: float = 300.0,
|
||||
decay_seconds: float = 900.0,
|
||||
max_buckets: int = _MAX_BUCKETS,
|
||||
sweep_interval_seconds: float = _SWEEP_INTERVAL_SECONDS,
|
||||
) -> None:
|
||||
self._threshold = threshold
|
||||
self._base_delay = base_delay
|
||||
self._max_delay = max_delay
|
||||
self._decay_seconds = decay_seconds
|
||||
self._max_buckets = max_buckets
|
||||
self._sweep_interval_seconds = sweep_interval_seconds
|
||||
self._buckets: dict[str, _Bucket] = {}
|
||||
self._last_sweep_at = time.monotonic()
|
||||
|
||||
def _is_spent(self, bucket: _Bucket, now: float) -> bool:
|
||||
"""Nothing left to remember: the lockout has expired *and* the failure count
|
||||
would decay to zero on the next failure anyway. Dropping such a bucket is
|
||||
indistinguishable from keeping it — which is what makes eviction safe."""
|
||||
return bucket.locked_until <= now and now - bucket.last_failure_at > self._decay_seconds
|
||||
|
||||
def _prune(self, now: float) -> None:
|
||||
"""B-56: the dict was keyed by attacker-chosen strings (any username, and via
|
||||
B-54 any IP) and only ever grew — `decay_seconds` aged a bucket's counter but
|
||||
never removed the entry, so hammering login with random usernames was an
|
||||
unbounded memory leak.
|
||||
|
||||
Spent buckets go first, and they carry no information, so that alone keeps
|
||||
the dict at the size of the genuinely active attack surface. The hard cap
|
||||
below is the backstop for a burst faster than the sweep interval: it evicts
|
||||
the entries closest to expiry, i.e. the ones whose loss buys an attacker the
|
||||
least — never the freshest lockouts, which are the ones actually holding an
|
||||
attack back."""
|
||||
for key in [k for k, b in self._buckets.items() if self._is_spent(b, now)]:
|
||||
del self._buckets[key]
|
||||
self._last_sweep_at = now
|
||||
|
||||
excess = len(self._buckets) - self._max_buckets
|
||||
if excess > 0:
|
||||
by_expiry = sorted(
|
||||
self._buckets.items(), key=lambda item: (item[1].locked_until, item[1].last_failure_at)
|
||||
)
|
||||
for key, _ in by_expiry[:excess]:
|
||||
del self._buckets[key]
|
||||
|
||||
def retry_after(self, key: str) -> float:
|
||||
bucket = self._buckets.get(key)
|
||||
if bucket is None:
|
||||
return 0.0
|
||||
remaining = bucket.locked_until - time.monotonic()
|
||||
now = time.monotonic()
|
||||
if self._is_spent(bucket, now):
|
||||
# Self-cleaning read path: a key that's merely being probed never
|
||||
# accumulates an entry that outlives its own usefulness.
|
||||
del self._buckets[key]
|
||||
return 0.0
|
||||
remaining = bucket.locked_until - now
|
||||
return remaining if remaining > 0 else 0.0
|
||||
|
||||
def record_failure(self, key: str) -> None:
|
||||
now = time.monotonic()
|
||||
if now - self._last_sweep_at >= self._sweep_interval_seconds or len(self._buckets) > self._max_buckets:
|
||||
self._prune(now)
|
||||
bucket = self._buckets.setdefault(key, _Bucket())
|
||||
if bucket.failures and now - bucket.last_failure_at > self._decay_seconds:
|
||||
bucket.failures = 0
|
||||
|
||||
Reference in New Issue
Block a user