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 RollingQuota:
    """How many times a key may do something in a rolling window — as opposed to
    RateLimiter above, which punishes *failures* with a growing delay.

    B-58: registration was throttled with the failure limiter, and recorded a
    "failure" on every attempt, successful ones included. Five legitimate signups
    from one shared or NAT address locked the sixth real user out for up to 600s,
    with the backoff doubling from there — while an attacker sidestepped the whole
    thing through B-54. The intent (bound how many accounts one source can create)
    is right; failure backoff is the wrong instrument for it, since nothing here is
    a failed guess at a secret. A quota says exactly what is meant: this many
    accounts per source per window, and the answer to the one over it is "not yet",
    with an accurate wait rather than a punishment that grows.
    """

    def __init__(
        self,
        limit: int,
        window_seconds: float,
        max_keys: int = _MAX_BUCKETS,
        sweep_interval_seconds: float = _SWEEP_INTERVAL_SECONDS,
    ) -> None:
        self._limit = limit
        self._window_seconds = window_seconds
        self._max_keys = max_keys
        self._sweep_interval_seconds = sweep_interval_seconds
        self._events: dict[str, list[float]] = {}
        self._last_sweep_at = time.monotonic()

    def _live_events(self, key: str, now: float) -> list[float]:
        """The key's events still inside the window, pruned in place."""
        events = self._events.get(key)
        if events is None:
            return []
        cutoff = now - self._window_seconds
        while events and events[0] <= cutoff:
            events.pop(0)
        if not events:
            del self._events[key]
        return events

    def _prune(self, now: float) -> None:
        # Same bound as RateLimiter (B-56): the keys are caller-chosen, so the dict
        # needs both a sweep and a hard cap. Eviction order is likewise "closest to
        # leaving the window first" — dropping a key with room left in its quota
        # changes nothing, dropping a full one hands out free accounts.
        for key in list(self._events):
            self._live_events(key, now)
        self._last_sweep_at = now

        excess = len(self._events) - self._max_keys
        if excess > 0:
            by_oldest = sorted(self._events.items(), key=lambda item: item[1][-1])
            for key, _ in by_oldest[:excess]:
                del self._events[key]

    def retry_after(self, key: str) -> float:
        """Seconds until this key may act again — 0 while it is under quota."""
        now = time.monotonic()
        events = self._live_events(key, now)
        if len(events) < self._limit:
            return 0.0
        return events[0] + self._window_seconds - now

    def record(self, key: str) -> None:
        """Counts one *completed* action. Attempts that create nothing (a taken
        username, a validation error) deliberately don't consume the quota — the
        limit is on accounts that exist, not on requests."""
        now = time.monotonic()
        if now - self._last_sweep_at >= self._sweep_interval_seconds or len(self._events) > self._max_keys:
            self._prune(now)
        self._events.setdefault(key, []).append(now)


_REGISTRATIONS_PER_IP = 5
_REGISTRATION_WINDOW_SECONDS = 3600.0


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)
        # B-58: a quota, not failure backoff — 5 accounts per IP per hour. The one
        # over it waits only until the oldest of the five ages out, and a busy NAT
        # is slowed rather than locked out for progressively longer.
        self.register_ip = RollingQuota(
            limit=_REGISTRATIONS_PER_IP, window_seconds=_REGISTRATION_WINDOW_SECONDS
        )