"""A tiny in-process sliding-window rate limiter. Generic groundwork: keyed by an arbitrary string (e.g. a client — `client_key`), so it can throttle any endpoint. Single-worker uvicorn → in-memory state is sufficient; it resets on restart, which is fine for abuse throttling (not for anything that must survive a deploy). Not shared across processes — if we ever run multiple workers, swap the backing store for Redis behind the same ``allow()`` interface. """ import ipaddress import threading import time from fastapi import Request from app.config import settings class RateLimiter: """At most `max_events` attempts per `window_seconds` per key — or, when each attempt passes a `cost`, at most that much total cost (a byte budget counts bytes, not calls).""" def __init__(self, max_events: float, window_seconds: float): self.max_events = max_events self.window = window_seconds # key → (monotonic time, cost) of every counted attempt still inside the window self._hits: dict[str, list[tuple[float, float]]] = {} self._lock = threading.Lock() def reset(self) -> None: """Forget every recorded attempt. Each limiter is a module-level singleton, so a test suite sharing one process would otherwise pool all of its attempts into a single window and eventually trip an unrelated test with a 429.""" with self._lock: self._hits.clear() def exhausted(self, key: str) -> bool: """Whether ``key``'s window is already full (``allow`` would refuse it), WITHOUT counting an attempt — for a limit that counts only some outcomes: ask it first, ``allow`` the ones that count.""" cutoff = time.monotonic() - self.window with self._lock: return sum(c for t, c in self._hits.get(key, []) if t > cutoff) >= self.max_events def room(self, key: str) -> float: """How much cost ``key`` may still spend in the current window, WITHOUT counting anything — for work counted only once it has succeeded: ask it first, ``allow`` what was done.""" cutoff = time.monotonic() - self.window with self._lock: used = sum(c for t, c in self._hits.get(key, []) if t > cutoff) return max(0.0, self.max_events - used) def allow(self, key: str, cost: float = 1) -> bool: """Record an attempt costing ``cost`` for ``key`` and return whether it is within the limit. True -> within the cap (the attempt is counted). False -> the attempt would pass the cap for the current window (attempt NOT counted, so a blocked caller can't keep pushing the window forward).""" now = time.monotonic() cutoff = now - self.window with self._lock: hits = [(t, c) for t, c in self._hits.get(key, []) if t > cutoff] if sum(c for _, c in hits) + cost > self.max_events: self._hits[key] = hits return False hits.append((now, cost)) self._hits[key] = hits # Opportunistic cleanup so abandoned keys don't accumulate unboundedly. if len(self._hits) > 4096: for k in list(self._hits): fresh = [(t, c) for t, c in self._hits[k] if t > cutoff] if fresh: self._hits[k] = fresh else: del self._hits[k] return True def _norm_ip(s: str) -> ipaddress.IPv4Address | ipaddress.IPv6Address | None: """Parse an IP for comparison, unwrapping IPv4-mapped IPv6 (`::ffff:10.0.0.1` → `10.0.0.1`) so a plain-IPv4 allowlist still matches a peer that surfaces in mapped form. None for a non-IP string.""" try: ip = ipaddress.ip_address(s.strip()) # silent-ok: parsing an untrusted header value, where "not an IP" is a documented answer this # function returns — and a line per request is what an attacker would aim for. except ValueError: return None return ip.ipv4_mapped if getattr(ip, "ipv4_mapped", None) else ip def client_key(request: Request) -> str: """Best-effort client identity for rate limiting: its IP, or for IPv6 the /64 the address sits in — a subscriber is handed a whole /64, so keyed on the full address one client would have endless identities to spread its attempts over. X-Forwarded-For is client-controlled, so we trust it ONLY when the request actually arrived from a configured reverse proxy (settings.trusted_proxy_ips — the address the proxy connects FROM, e.g. the VPS Caddy's WireGuard peer IP). Our proxy APPENDS the client it saw to XFF, so the RIGHTMOST entry is the real client and can't be forged by a client pre-seeding the header. Any other peer (someone hitting the app port directly, bypassing the proxy) is untrusted — we use its real socket IP, so it can't spoof its rate-limit identity. SAFETY: this relies on `request.client` being the TRUE socket peer. uvicorn is started WITHOUT `--proxy-headers` (see entrypoint.sh) precisely so it never rewrites `request.client` from XFF — do NOT add that flag, or the trust check below would compare against a forgeable value.""" peer = request.client.host if request.client else "unknown" client = peer trusted = {_norm_ip(p) for p in settings.trusted_proxy_ip_set} - {None} if _norm_ip(peer) in trusted: xff = request.headers.get("x-forwarded-for") parts = [p.strip() for p in (xff or "").split(",") if p.strip()] if parts: client = parts[-1] ip = _norm_ip(client) if isinstance(ip, ipaddress.IPv6Address): return str(ipaddress.IPv6Network((ip, 64), strict=False)) return client