server/ratelimit.py

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

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 ``GODSEED_SECRET``
(set it on the Space so identities survive restarts) or is generated per-process.

Limit: 3 wishes / hour / client, true sliding window. 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 deque
from typing import Callable, Optional, Tuple

COOKIE_NAME = "godseed_id"
DEFAULT_LIMIT = 3
DEFAULT_WINDOW_SECONDS = 3600.0
# Defense-in-depth: a per-IP ceiling (several clients can share a NAT, so it is
# intentionally looser than the per-client limit). Cookie-clearing can mint fresh
# client ids; this cap bounds that abuse.
DEFAULT_IP_LIMIT = 12


def client_ip(request) -> str:
    """Best client IP we can get. SECURITY: the FIRST X-Forwarded-For hop is
    client-controlled (a forged header mints a fresh rate-limit bucket per
    request — verify-fleet finding). The trustworthy entry is the one the
    closest trusted proxy appended: the LAST hop. On HF Spaces the edge proxy
    appends the real client IP, so take the last entry; locally use the peer."""
    forwarded = request.headers.get("x-forwarded-for", "")
    if forwarded:
        hops = [h.strip() for h in forwarded.split(",") if h.strip()]
        if hops:
            return hops[-1]
    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("GODSEED_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 per contract.
        """
        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)."""

    def __init__(
        self,
        limit: int = DEFAULT_LIMIT,
        window: float = DEFAULT_WINDOW_SECONDS,
        clock: Callable[[], float] = time.monotonic,
    ) -> None:
        self.limit = limit
        self.window = window
        self.clock = clock
        self._hits: dict[str, deque[float]] = {}
        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) wishes.
        """
        now = self.clock()
        hits = self._hits.setdefault(key, deque())
        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]