plm-lottery/app/auth/rate_limit.py

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:
    failures: int = 0
    locked_until: float = 0.0
    last_failure_at: float = 0.0


class RateLimiter:
    """In-process failed-attempt throttle with exponential backoff, keyed by an
    arbitrary string (username, IP...). Single-process-only, like UserLocks
    (app/tx/locks.py) — an accepted MVP constraint; a multi-worker deployment
    would need a shared store (Redis) instead (B-33).

    Brute-forcing a login here isn't a spammy client to be capped at N req/s —
    it's an attempt to withdraw someone else's funds — so failures are
    penalized with a delay that doubles each time past `threshold` free
    attempts, rather than a flat rate cap. `decay_seconds` ages a bucket back
    to zero once failures stop, so a shared/NAT IP isn't punished forever for
    someone else's earlier mistakes.
    """

    def __init__(
        self,
        threshold: int = 5,
        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
        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
        bucket.failures += 1
        bucket.last_failure_at = now
        if bucket.failures >= self._threshold:
            delay = min(self._max_delay, self._base_delay * 2 ** (bucket.failures - self._threshold))
            bucket.locked_until = now + delay

    def record_success(self, key: str) -> None:
        self._buckets.pop(key, None)


class AuthRateLimiters:
    """The three throttles B-33 needs, bundled so they can live on `app.state`
    (like `UserLocks`, see app/tx/locks.py) rather than as module globals.

    A module global would persist for the lifetime of the process — fine in
    production (one app instance), but wrong in the test suite, where every
    test builds its own FastAPI app against a fresh in-memory DB and expects a
    clean slate; a shared global would leak failure counts between unrelated
    tests. Per-`app.state` state gets a fresh instance per app automatically.
    """

    def __init__(self) -> None:
        self.login = RateLimiter(threshold=5, base_delay=2.0, max_delay=300.0)
        self.login_ip = RateLimiter(threshold=20, base_delay=2.0, max_delay=300.0)
        self.register_ip = RateLimiter(threshold=5, base_delay=5.0, max_delay=600.0)