Spaces:
Runtime error
Runtime error
| """Phase S - hand-rolled sliding-window rate limiter, backed by redis""" | |
| import time | |
| import uuid | |
| import redis.asyncio as aioredis | |
| from fastapi import HTTPException, Request | |
| from server.limits import REDIS_URL, RATE_LIMIT_ENABLED | |
| from server.user_auth import decode_access_token | |
| # Atomic sliding-window-log in one round trip | |
| # KEYS[1] = bucket key | |
| # ARGV[1] = now (ms) ARGV[2] = window (ms) ARGV[3] = limit ARGV[4] = member | |
| # Returns: {allowed(0/1), remaining, retry_after_ms} | |
| _SLIDING_WINDOW = """ | |
| redis.call('ZREMRANGEBYSCORE', KEYS[1], 0, tonumber(ARGV[1]) - tonumber(ARGV[2])) | |
| local count = redis.call('ZCARD', KEYS[1]) | |
| local limit = tonumber(ARGV[3]) | |
| if count < limit then | |
| redis.call('ZADD', KEYS[1], ARGV[1], ARGV[4]) | |
| redis.call('PEXPIRE', KEYS[1], ARGV[2]) | |
| return {1, limit - count - 1, 0} | |
| end | |
| -- Blocked: count >= limit >= 1, so the set is non-empty and the | |
| -- oldest member's score always exists (the script runs atomically). | |
| local oldest = redis.call('ZRANGE', KEYS[1], 0, 0) | |
| local score = redis.call('ZSCORE', KEYS[1], oldest[1]) | |
| local retry = (tonumber(score) + tonumber(ARGV[2])) - tonumber(ARGV[1]) | |
| if retry < 0 then retry = 0 end | |
| return {0, 0, retry} | |
| """ | |
| _redis = None | |
| _script = None | |
| async def init_redis() -> None: | |
| """Connect to Redis and register the Lua script. Call on startup.""" | |
| global _redis, _script | |
| if not RATE_LIMIT_ENABLED: | |
| return | |
| try: | |
| _redis = aioredis.from_url(REDIS_URL, encoding="utf-8", decode_responses=True) | |
| _script = _redis.register_script(_SLIDING_WINDOW) | |
| await _redis.ping() | |
| except Exception: | |
| _redis = None # fail open - limiter allows all until redis returns | |
| async def close_redis() -> None: | |
| """Close the Redis connection. Call on shutdown.""" | |
| global _redis | |
| if _redis is not None: | |
| await _redis.aclose() | |
| _redis = None | |
| class RateResult: | |
| def __init__(self, allowed, remaining, retry_after, limit): | |
| self.allowed = allowed | |
| self.remaining = remaining | |
| self.retry_after = retry_after | |
| self.limit = limit | |
| # ---- In-memory fallback (used when Redis isn't configured / reachable) ---- | |
| # Per-process sliding-window log, keyed exactly like the Redis path. Good enough | |
| # for a single-container deploy such as a free Hugging Face Space: it isn't | |
| # shared across replicas and resets on restart, but it keeps real limits in | |
| # force instead of failing open on a public endpoint. | |
| _mem_buckets: dict = {} | |
| def _hit_in_memory(key, limit, window) -> RateResult: | |
| now_ms = time.time() * 1000 | |
| win_ms = window * 1000 | |
| cutoff = now_ms - win_ms | |
| # Opportunistic cleanup so keys that went quiet don't leak memory. | |
| if len(_mem_buckets) > 10000: | |
| for k in [k for k, v in _mem_buckets.items() if not v]: | |
| _mem_buckets.pop(k, None) | |
| bucket = _mem_buckets.setdefault(key, []) | |
| drop = 0 | |
| for ts in bucket: | |
| if ts <= cutoff: | |
| drop += 1 | |
| else: | |
| break | |
| if drop: | |
| del bucket[:drop] | |
| if len(bucket) < limit: | |
| bucket.append(now_ms) | |
| return RateResult(True, limit - len(bucket), 0, limit) | |
| retry_ms = (bucket[0] + win_ms) - now_ms if bucket else win_ms | |
| if retry_ms < 0: | |
| retry_ms = 0 | |
| return RateResult(False, 0, (int(retry_ms) + 999) // 1000, limit) | |
| async def _hit(key, limit, window) -> RateResult: | |
| if not RATE_LIMIT_ENABLED: | |
| return RateResult(True, limit, 0, limit) | |
| if _redis is None: | |
| # No Redis (e.g. single-container HF Space) -> per-process limiter. | |
| return _hit_in_memory(key, limit, window) | |
| now_ms = int(time.time() * 1000) | |
| member = f"{now_ms}-{uuid.uuid4().hex[:8]}" # unique even within the same ms | |
| try: | |
| allowed, remaining, retry_ms = await _script( | |
| keys=[key], args=[now_ms, window * 1000, limit, member] | |
| ) | |
| except Exception: | |
| return _hit_in_memory(key, limit, window) # Redis hiccup -> fall back | |
| retry_after = (int(retry_ms) + 999) // 1000 # ms -> ceil seconds | |
| return RateResult(bool(allowed), int(remaining), retry_after, limit) | |
| def identity_from_request(request: Request) -> str: | |
| """Key by logged-in user if a valid JWT is present, else by client IP.""" | |
| auth = request.headers.get("authorization", "") | |
| if auth.lower().startswith("bearer "): | |
| uid = decode_access_token(auth[7:].strip()) | |
| if uid: | |
| return f"user:{uid}" | |
| fwd = request.headers.get("x-forwarded-for") # set by nginx | |
| if fwd: | |
| ip = fwd.split(",")[0].strip() | |
| elif request.client: | |
| ip = request.client.host | |
| else: | |
| ip = "unknown" | |
| return f"ip:{ip}" | |
| def rate_limit(limit: int, window: int, scope: str): | |
| """Build a FastAPI dependency enforcing `limit` requests / `window` seconds.""" | |
| async def _dep(request: Request) -> None: | |
| key = f"rl:{scope}:{identity_from_request(request)}" | |
| res = await _hit(key, limit, window) | |
| request.state.rate_headers = { # middleware copies these onto the response | |
| "X-RateLimit-Limit": str(res.limit), | |
| "X-RateLimit-Remaining": str(max(res.remaining, 0)), | |
| } | |
| if not res.allowed: | |
| raise HTTPException( | |
| status_code=429, | |
| detail="Rate limit exceeded. Please slow down.", | |
| headers={ | |
| "Retry-After": str(res.retry_after), | |
| "X-RateLimit-Limit": str(res.limit), | |
| "X-RateLimit-Remaining": "0", | |
| }, | |
| ) | |
| return _dep | |