server/ratelimit.py

"""Visitor identity + sliding-window rate limiting (in-memory, single process).

Lifted from godseed/server/ratelimit.py (the proven June-12 pattern) with
PAREIDOLIA's names and numbers: publishing to the Menagerie is capped at
6/hour per identity (awaken itself needs no cap here — the caller's own
ZeroGPU quota is the natural limit, per ARCHITECTURE.md §6).

Identity: a signed cookie token hashed together with the client IP — stable for
a visitor, cheap to verify, no accounts. The HMAC secret comes from
``PAREIDOLIA_SECRET`` (set it on the Space so identities survive restarts) or
is generated per-process. Denied attempts are NOT recorded, so hammering the
endpoint never extends the lockout. The Space runs a single uvicorn worker, so
in-memory state is authoritative.
"""

from __future__ import annotations

import hashlib
import hmac
import os
import secrets
import time
from collections import OrderedDict, deque
from typing import Callable, Optional, Tuple

COOKIE_NAME = "pareidolia_id"
DEFAULT_LIMIT = 6  # publishes per window per identity (§6)
DEFAULT_WINDOW_SECONDS = 3600.0
# Defense-in-depth: a per-IP ceiling (several visitors can share a NAT, so it is
# intentionally looser than the per-identity limit). Cookie-clearing can mint
# fresh identities; this cap bounds that abuse.
DEFAULT_IP_LIMIT = 24
# Hard ceiling on distinct limiter keys (security review #2b): keys derive from
# client-influenced strings, so without a bound a flood of fabricated
# identities grows the map until memory dies. 50k keys of a few floats each is
# a few MB — roomy for real traffic, fatal for the attack.
DEFAULT_MAX_KEYS = 50_000


def client_ip(request) -> str:
    """Best client IP we can get.

    On HF Spaces the app sits behind a trusted router that APPENDS the real
    peer address to ``X-Forwarded-For`` — so the RIGHTMOST hop is the only one
    the client cannot forge (security review #2a: the leftmost entries are
    attacker-supplied and were being used to mint fresh per-IP rate-limit
    buckets on demand). Locally, with no XFF, it's the socket peer."""
    forwarded = request.headers.get("x-forwarded-for", "")
    for hop in reversed(forwarded.split(",")):
        hop = hop.strip()
        if hop:
            return hop
    return request.client.host if request.client else "unknown"


class ClientIdentity:
    """Mints and verifies the signed visitor cookie; derives the rate-limit id."""

    def __init__(self, secret: Optional[str] = None) -> None:
        self.secret = (
            secret or os.environ.get("PAREIDOLIA_SECRET") or secrets.token_hex(32)
        )

    def _sign(self, token: str) -> str:
        return hmac.new(
            self.secret.encode("utf-8"), token.encode("utf-8"), hashlib.sha256
        ).hexdigest()[:32]

    def mint(self) -> str:
        token = secrets.token_hex(8)
        return f"{token}.{self._sign(token)}"

    def verify(self, cookie_value: Optional[str]) -> Optional[str]:
        """Return the token if the cookie is authentic, else None."""
        if not cookie_value or "." not in cookie_value:
            return None
        token, _, signature = cookie_value.rpartition(".")
        if not token or not hmac.compare_digest(signature, self._sign(token)):
            return None
        return token

    def resolve(self, request) -> Tuple[str, Optional[str]]:
        """Return ``(client_id, cookie_to_set)``.

        ``cookie_to_set`` is a freshly minted signed value when the request
        carried no valid cookie (the caller sets it on the response), else None.
        ``client_id = sha256(token | ip)`` — signed cookie + IP hash.
        """
        token = self.verify(request.cookies.get(COOKIE_NAME))
        new_cookie: Optional[str] = None
        if token is None:
            new_cookie = self.mint()
            token = new_cookie.split(".", 1)[0]
        ip = client_ip(request)
        cid = hashlib.sha256(f"{token}|{ip}".encode("utf-8")).hexdigest()[:16]
        return cid, new_cookie


class RateLimiter:
    """Sliding-window counter: at most ``limit`` recorded hits per ``window``
    seconds per key. ``clock`` is injectable for tests (monotonic seconds).

    Distinct keys are capped at ``max_keys`` with least-recently-hit eviction
    (an OrderedDict LRU), so client-minted key floods cannot grow memory
    without bound between hourly sweeps (security review #2b)."""

    def __init__(
        self,
        limit: int = DEFAULT_LIMIT,
        window: float = DEFAULT_WINDOW_SECONDS,
        clock: Callable[[], float] = time.monotonic,
        max_keys: int = DEFAULT_MAX_KEYS,
    ) -> None:
        self.limit = limit
        self.window = window
        self.clock = clock
        self.max_keys = max(1, max_keys)
        self._hits: OrderedDict[str, deque[float]] = OrderedDict()
        self._last_sweep = clock()

    def hit(self, key: str) -> Tuple[bool, float]:
        """Try to record one hit for ``key``.

        Returns ``(allowed, retry_after_seconds)``. Denied attempts are not
        recorded, so the window slides only on real (allowed) publishes.
        """
        now = self.clock()
        hits = self._hits.get(key)
        if hits is None:
            # New key: enforce the distinct-key ceiling first — evict the
            # least-recently-hit key(s) so the map is bounded by construction.
            while len(self._hits) >= self.max_keys:
                self._hits.popitem(last=False)
            hits = self._hits[key] = deque()
        else:
            self._hits.move_to_end(key)  # mark recently hit (LRU freshness)
        while hits and now - hits[0] >= self.window:
            hits.popleft()
        if len(hits) >= self.limit:
            return False, max(0.0, self.window - (now - hits[0]))
        hits.append(now)
        self._sweep_if_due(now)
        return True, 0.0

    def remaining(self, key: str) -> int:
        now = self.clock()
        hits = self._hits.get(key)
        if not hits:
            return self.limit
        live = sum(1 for t in hits if now - t < self.window)
        return max(0, self.limit - live)

    def _sweep_if_due(self, now: float) -> None:
        """Drop keys whose every hit has expired (bounds memory over long uptimes)."""
        if now - self._last_sweep < self.window:
            return
        self._last_sweep = now
        stale = [
            key
            for key, hits in self._hits.items()
            if not hits or now - hits[-1] >= self.window
        ]
        for key in stale:
            del self._hits[key]