Files
plm-lottery/tests/unit/test_rate_limit.py
T
davideandClaude Opus 5 ab65728bdc 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>
2026-08-03 22:22:51 +02:00

92 lines
3.6 KiB
Python

"""B-56: RateLimiter._buckets is keyed by strings the caller chooses — any
username, and (before B-54) any IP — and used to only ever grow. decay_seconds
aged a bucket's *counter* but never removed the entry, so hammering login with
random usernames was an unbounded memory leak. These tests pin both halves of
the bound: spent entries are swept, and the dict has a hard cap.
The throttling behaviour itself is exercised end-to-end in test_auth.py; what
matters here is that pruning never hands an attacker a free pass.
"""
import time
from app.auth.rate_limit import RateLimiter
def _limiter(**kwargs) -> RateLimiter:
# sweep_interval_seconds=0 makes every record_failure sweep, so the tests are
# deterministic instead of depending on wall-clock timing.
kwargs.setdefault("sweep_interval_seconds", 0.0)
return RateLimiter(**kwargs)
def test_spent_buckets_are_swept_on_the_next_failure():
limiter = _limiter(decay_seconds=0.05)
for i in range(50):
limiter.record_failure(f"user:{i}")
assert len(limiter._buckets) == 50
time.sleep(0.06) # every bucket is now past decay_seconds and unlocked
limiter.record_failure("user:fresh")
assert list(limiter._buckets) == ["user:fresh"]
def test_a_bucket_still_locking_someone_out_is_never_swept():
"""The whole point of the entry: evicting it would reset the backoff and let the
attacker start over from a free attempt."""
limiter = _limiter(threshold=1, base_delay=300.0, decay_seconds=0.05)
limiter.record_failure("user:victim")
assert limiter.retry_after("user:victim") > 0
time.sleep(0.06) # past decay_seconds, but the lockout is still running
limiter.record_failure("user:someone-else")
assert "user:victim" in limiter._buckets
assert limiter.retry_after("user:victim") > 0
def test_the_dict_is_capped_even_within_one_sweep_interval():
"""The cap is the backstop for a burst faster than the sweep interval, where
nothing has had time to expire yet."""
limiter = _limiter(max_buckets=10, sweep_interval_seconds=3600.0, decay_seconds=3600.0)
for i in range(200):
limiter.record_failure(f"ip:{i}")
assert len(limiter._buckets) <= 11 # the cap, plus the entry recorded after the last prune
def test_the_cap_evicts_the_entries_closest_to_expiry_first():
"""What gets dropped under pressure must buy an attacker the least. The deepest
lockout — the one built up over the most failures, and so the one actually
holding an attack back — has to be the last thing evicted, not collateral of a
flood of one-failure keys."""
limiter = _limiter(max_buckets=3, threshold=1, base_delay=1.0, max_delay=600.0)
for _ in range(10):
limiter.record_failure("ip:persistent") # backoff doubles: locked for ~512s
for i in range(20):
limiter.record_failure(f"ip:filler{i}") # one failure each: locked for ~1s
assert "ip:persistent" in limiter._buckets
assert limiter.retry_after("ip:persistent") > 100
def test_retry_after_drops_a_spent_bucket_it_looks_at():
limiter = _limiter(decay_seconds=0.05)
limiter.record_failure("user:probe")
time.sleep(0.06)
assert limiter.retry_after("user:probe") == 0.0
assert limiter._buckets == {}
def test_pruning_does_not_reset_a_live_failure_count():
"""A bucket below the lockout threshold still carries state worth keeping: the
next failure must count as the second, not the first."""
limiter = _limiter(threshold=2, base_delay=300.0, decay_seconds=3600.0)
limiter.record_failure("user:a")
limiter.record_failure("user:b") # triggers a sweep
limiter.record_failure("user:a")
assert limiter.retry_after("user:a") > 0