src/api/middleware.py

"""API observability: per-request ID tagging, in-memory rate limiting, structured access logging."""

from __future__ import annotations

import json
import logging
import os
import stat
import time
import uuid
from collections import OrderedDict
from pathlib import Path
from typing import Optional

from fastapi import Request

log = logging.getLogger("gpp.access")

RATE_LIMIT_RPM = 60

# Upper bound on distinct client_keys tracked at once — without it, a client that can vary its apparent source IP (trivial over IPv6) has an unbounded memory-exhaustion vector; matches this codebase's `lru_cache(maxsize=...)` bounding pattern elsewhere.
_MAX_TRACKED_CLIENTS = 100_000

# Throttles state persistence to at most once per this many seconds, bounding disk I/O under many distinct clients; a crash between persists loses at most this much accounting, not the whole bucket set.
_PERSIST_INTERVAL_S = 1.0

# A rejection still forces a persist sooner than the interval above (see allow()), but not unconditionally: write cost scales with the *total* tracked-client count (every persist serializes the whole dict), not just the rejected client, so a client retrying immediately after a 429 — with no backoff, malicious or just buggy — would otherwise force a full-dict write on every single request. Measured: ~2.5ms/request at 5,000 tracked clients, ~66ms/request at the 100k cap, run synchronously on the single event-loop thread (allow() is called with no `await` before it in `async def observability`), so it stalls every other visitor too. This interval bounds that to a fixed worst-case rate regardless of how fast rejections arrive.
_FORCE_PERSIST_INTERVAL_S = 0.2

# Generously covers _MAX_TRACKED_CLIENTS entries even with long IPv6 keys; matches manifest.py/gpu_spec_db.py's convention of capping untrusted-file reads before parsing them.
_MAX_STATE_BYTES = 10 * 1024 * 1024  # 10 MB


def new_request_id() -> str:
    return uuid.uuid4().hex


class RateLimiter:
    """In-memory token bucket, one bucket per client IP, LRU-bounded to `_MAX_TRACKED_CLIENTS` entries. Soft abuse prevention, not a hard SLA guarantee. State is per-process by default (`state_path=None`, same as the module-level spec/model caches elsewhere in this codebase); passing `state_path` opts into surviving a process restart by persisting to a symlink-guarded, atomically-written JSON file, matching `train_final.py`'s convention — `time.monotonic()` values stay valid to compare across the old and new process since CLOCK_MONOTONIC is boot-relative, not process-relative."""

    def __init__(
        self,
        rpm: int = RATE_LIMIT_RPM,
        max_clients: int = _MAX_TRACKED_CLIENTS,
        state_path: Path | None = None,
    ) -> None:
        self._capacity = float(rpm)
        self._refill_per_sec = rpm / 60.0
        self._max_clients = max_clients
        self._state_path = state_path
        self._last_persist = 0.0
        # client_key -> (tokens, last_refill_monotonic); insertion order = LRU order, moved to the end on every access so the evicted entry is genuinely the least-recently-seen client.
        self._buckets: OrderedDict[str, tuple[float, float]] = OrderedDict()
        if self._state_path is not None:
            self._load_state()

    def allow(self, client_key: str) -> bool:
        now = time.monotonic()
        # A never-seen client starts with a full bucket and zero elapsed time — seeding via `defaultdict(time.monotonic)` instead would call time.monotonic() *after* `now` is already captured, making `elapsed` spuriously negative and shaving a fraction of a token off every client's first request.
        tokens, last = self._buckets.get(client_key, (self._capacity, now))
        elapsed = now - last
        tokens = min(self._capacity, tokens + elapsed * self._refill_per_sec)
        allowed = tokens >= 1.0
        if allowed:
            tokens -= 1.0

        self._buckets[client_key] = (tokens, now)
        self._buckets.move_to_end(client_key)
        if len(self._buckets) > self._max_clients:
            self._buckets.popitem(last=False)

        if self._state_path is not None:
            # A rejection is forced through the throttle sooner (see _FORCE_PERSIST_INTERVAL_S), not unconditionally: a rapid burst can otherwise exhaust the bucket in-memory well before the next throttled write reaches disk, so a crash right after the burst would recover a stale, under-depleted snapshot and grant free requests — the one state a crash most needs to preserve is "this client is already being blocked." Bounded, not immediate, so a client retrying rejections in a tight loop can't force a full-dict write on every single request.
            self._persist_state(now, force=not allowed)
        return allowed

    def _load_state(self) -> None:
        # A single lstat() call refuses both a symlinked path and an oversized file before this reads/parses it — matching manifest.py's _refuse_symlink_and_oversize convention for any file this codebase treats as untrusted external state, not just MLPerf sources.
        try:
            st = os.lstat(self._state_path)
        except FileNotFoundError:
            return
        except OSError:
            log.warning("Could not stat persisted rate-limiter state at %s; starting fresh", self._state_path)
            return
        if stat.S_ISLNK(st.st_mode):
            log.warning("Refusing to load rate-limiter state through a symlink: %s", self._state_path)
            return
        if st.st_size > _MAX_STATE_BYTES:
            log.warning("Persisted rate-limiter state at %s is too large (%d bytes); starting fresh", self._state_path, st.st_size)
            return
        # The whole parse-and-shape-check lives in one try/except, not just json.loads(): a file that's valid JSON but the wrong shape (a list, or values that aren't a [tokens, last] pair of numbers) must degrade to "start fresh," never raise past this constructor — this runs at src/api/main.py's module import time, so an uncaught exception here would crash the app on every single startup, turning a supervisord crash-restart into a permanent crash loop instead of recovering it. Demonstrated exploitable before this fix: a state file containing a JSON list, or a dict with a non-numeric or wrong-length value, each crashed __init__ with an unhandled AttributeError/ValueError.
        try:
            raw = json.loads(self._state_path.read_text())
            if not isinstance(raw, dict):
                raise ValueError(f"expected a JSON object at the top level, got {type(raw).__name__}")
            loaded = {str(k): (float(tokens), float(last)) for k, (tokens, last) in raw.items()}
        except (OSError, ValueError, TypeError):
            log.warning("Could not parse persisted rate-limiter state from %s; starting fresh", self._state_path)
            return
        self._buckets.update(loaded)
        while len(self._buckets) > self._max_clients:
            self._buckets.popitem(last=False)

    def _persist_state(self, now: float, force: bool = False) -> None:
        interval = _FORCE_PERSIST_INTERVAL_S if force else _PERSIST_INTERVAL_S
        if now - self._last_persist < interval:
            return
        self._last_persist = now
        try:
            # The tmp-file + rename() below already can't write *through* a symlink at this path — POSIX rename() replaces the destination directory entry itself, it never dereferences it, unlike a bare open("w") (the risk train_final.py/03_model_training.ipynb's atomic-write helpers guard against for that reason). What this check protects against instead: silently clobbering an unexpected pre-existing symlink at this path into a plain file. Refusing and warning is more honest than overwriting something that isn't supposed to be here.
            if self._state_path.is_symlink():
                log.warning("Refusing to persist rate-limiter state through a symlink: %s", self._state_path)
                return
            tmp_path = self._state_path.with_name(self._state_path.name + ".tmp")
            tmp_path.write_text(json.dumps({k: list(v) for k, v in self._buckets.items()}))
            tmp_path.replace(self._state_path)
        except OSError:
            log.warning("Could not persist rate-limiter state to %s", self._state_path)


def log_access(
    request_id: str,
    request: Request,
    status_code: int,
    latency_ms: float,
) -> None:
    """One structured JSON line per request: request_id, ts, method, route, status, latency_ms, client_ip."""
    client_ip: Optional[str] = request.client.host if request.client else None
    log.info(json.dumps({
        "request_id": request_id,
        "ts": time.time(),
        "method": request.method,
        "route": request.url.path,
        "status": status_code,
        "latency_ms": round(latency_ms, 2),
        "client_ip": client_ip,
    }))