""" Rate-limit control plane for the HF proxy. The default backend is deliberately process-local and is an abuse gate only. Operators that need one quota decision shared by multiple proxy replicas may select the optional ``redis`` backend. Redis mode uses one atomic server-side Lua operation per request and HMACs the client identity before it leaves the process, so raw IP-like identifiers are never stored as Redis keys. The Redis guarantee is scoped to one Redis consistency domain. This module does not claim billing/accounting correctness across independent Redis deployments, Active-Active conflict domains, or a gateway that bypasses this service. """ from __future__ import annotations import asyncio import hashlib import hmac from typing import Any from ._redis_security import RedisSecurityError, redis_connection_kwargs class RateLimitBackendError(RuntimeError): """Stable, non-sensitive rate-limit backend error.""" def __init__(self, code: str) -> None: super().__init__(code) self.code = code _REDIS_FIXED_WINDOW_LUA = r""" local key = KEYS[1] local window_seconds = tonumber(ARGV[1]) local current = redis.call('INCR', key) if current == 1 then redis.call('EXPIRE', key, window_seconds) end local ttl = redis.call('TTL', key) return { current, ttl } """.strip() def _safe_component(value: str, fallback: str = "generic") -> str: out = "".join(ch for ch in str(value or "").lower() if ch.isalnum() or ch in "_-:") return out[:64] or fallback class RedisRateLimiter: """Shared fixed-window limiter backed by Redis atomic scripting.""" backend = "redis" shared = True authoritative = True consistency_scope = "single_redis_consistency_domain" def __init__( self, url: str, *, identity_secret: str, key_prefix: str = "sphinx-ai-assistant", socket_timeout_seconds: float = 2.0, client: Any | None = None, require_tls: bool = False, ) -> None: if not str(url or "").strip(): raise RateLimitBackendError("REDIS_URL_REQUIRED") if len(str(identity_secret or "").encode("utf-8")) < ( 32 # ruff: ignore[magic-value-comparison] ): raise RateLimitBackendError("IDENTITY_SECRET_TOO_SHORT") self.url = str(url).strip() self.require_tls = bool(require_tls) try: self._transport, self._connection_kwargs = redis_connection_kwargs( self.url, require_tls=self.require_tls, socket_timeout_seconds=socket_timeout_seconds, ) except RedisSecurityError as exc: raise RateLimitBackendError(exc.code) from exc self._secret = str(identity_secret).encode("utf-8") self.key_prefix = _safe_component(key_prefix, "sphinx-ai-assistant") self.socket_timeout_seconds = max( 0.25, min(float(socket_timeout_seconds), 10.0) ) self._client = client self._owns_client = client is None self._init_lock = asyncio.Lock() def manifest(self) -> dict[str, Any]: return { "backend": self.backend, "shared": self.shared, "authoritative": self.authoritative, "consistency_scope": self.consistency_scope, "identity_externalized": "hmac_sha256", **self._transport.manifest(), } async def initialize(self) -> None: async with self._init_lock: if self._client is None: try: import redis.asyncio as redis_async # type: ignore[import-not-found] # ruff: ignore[import-outside-top-level] except Exception as exc: # pragma: no cover - deployment dependency raise RateLimitBackendError("REDIS_DEPENDENCY_UNAVAILABLE") from exc self._client = redis_async.from_url(self.url, **self._connection_kwargs) try: await self._client.ping() except Exception as exc: raise RateLimitBackendError("REDIS_UNAVAILABLE") from exc async def close(self) -> None: if self._client is None or not self._owns_client: return closer = getattr(self._client, "aclose", None) if closer is None: closer = getattr(self._client, "close", None) if closer is not None: result = closer() if hasattr(result, "__await__"): await result self._client = None def _identity_key(self, identity: str, scope: str) -> str: digest = hmac.new( self._secret, str(identity or "unknown").encode("utf-8"), hashlib.sha256 ).hexdigest() return f"{self.key_prefix}:rl:{_safe_component(scope)}:{digest}" async def consume( self, identity: str, *, scope: str, limit: int, window_seconds: int = 3600, ) -> tuple[bool, int, int]: if self._client is None: raise RateLimitBackendError("REDIS_NOT_INITIALIZED") bounded_limit = max(1, min(int(limit), 1_000_000)) bounded_window = max(1, min(int(window_seconds), 86_400)) key = self._identity_key(identity, scope) try: result = await self._client.eval( _REDIS_FIXED_WINDOW_LUA, 1, key, bounded_window ) count = int(result[0]) ttl = int(result[1]) except Exception as exc: raise RateLimitBackendError("REDIS_CONSUME_FAILED") from exc retry_after = max(1, ttl if ttl > 0 else bounded_window) return count <= bounded_limit, count, retry_after __all__ = ["_REDIS_FIXED_WINDOW_LUA", "RateLimitBackendError", "RedisRateLimiter"]