Blame for server/ratelimit.py

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