"""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]