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 )