From d4abe654512b8183fb96535456fd1d3216c74c3f Mon Sep 17 00:00:00 2001 From: architect-bot Date: Tue, 18 Aug 2026 15:37:58 +0000 Subject: [PATCH] WIP checkpoint: Redis rate-limiting dependency --- backend/app/rate_limit.py | 228 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 228 insertions(+) create mode 100644 backend/app/rate_limit.py diff --git a/backend/app/rate_limit.py b/backend/app/rate_limit.py new file mode 100644 index 0000000..8613693 --- /dev/null +++ b/backend/app/rate_limit.py @@ -0,0 +1,228 @@ +"""Redis-backed token-bucket rate limiting. + +This module provides a reusable FastAPI dependency that rate-limits requests +per endpoint scope (optionally per-user) using an atomic Lua script on Redis. +If Redis is unavailable the dependency can either allow the request through +(fail-open) or reject it with a ``503`` (fail-closed), depending on the policy +for the limit. + +Rate limits are intentionally dependency-only: they do not wrap a whole router +as middleware because per-scope capacity/refill and per-user bucketing make a +dependency (with explicit ``Request`` access) the most predictable place to +implement them. +""" + +from __future__ import annotations + +import time +from dataclasses import dataclass +from typing import Any, Callable + +from fastapi import Depends, HTTPException, Request, status +from redis.asyncio import Redis +from redis.exceptions import RedisError + +from app.config import settings + +# Atomic token-bucket implementation. +# +# KEYS[1] is the bucket key. ARGV[1] is the bucket capacity, ARGV[2] is the +# refill rate in tokens per second, ARGV[3] is the current Unix time in +# seconds (float), and ARGV[4] is the request cost (tokens to take). +# +# The bucket state is stored as a hash with "tokens" and "ts". The script +# refills the bucket up to capacity, then takes the requested tokens. It +# returns an array: +# [0] allowed (1 or 0) +# [1] tokens remaining after the take +# [2] retry-after seconds (0 when allowed) +_TOKEN_BUCKET_LUA = """ +local key = KEYS[1] +local capacity = tonumber(ARGV[1]) +local refill_rate = tonumber(ARGV[2]) +local now = tonumber(ARGV[3]) +local cost = tonumber(ARGV[4]) + +local state = redis.call("HMGET", key, "tokens", "ts") +local tokens = capacity +local ts = now + +if state[1] ~= false then + tokens = tonumber(state[1]) + ts = tonumber(state[2]) or now +end + +local elapsed = math.max(0, now - ts) +tokens = math.min(capacity, tokens + (elapsed * refill_rate)) + +local allowed = 0 +local retry_after = 0 + +if tokens >= cost then + tokens = tokens - cost + allowed = 1 +else + -- Time until enough tokens for a single request have accrued. + retry_after = math.ceil((cost - tokens) / refill_rate) +end + +-- Keep bucket state small/expire-able. Retain the stamp as a float so +-- partial tokens accumulate smoothly between requests. +redis.call("HSET", key, "tokens", tostring(tokens), "ts", tostring(now)) +redis.call("PEXPIRE", key, 600000) + +return {allowed, tostring(tokens), retry_after} +""" + + +class RateLimitExceeded(HTTPException): + """Raised when a rate limit bucket has no tokens left.""" + + def __init__(self, retry_after: int) -> None: + super().__init__( + status_code=status.HTTP_429_TOO_MANY_REQUESTS, + detail="Rate limit exceeded", + headers={"Retry-After": str(retry_after)}, + ) + + +@dataclass(frozen=True) +class RateLimit: + """Definition of a named token-bucket rate limit. + + ``capacity`` is the burst size and ``refill_rate`` is the steady-state + token accumulation rate (tokens/second). For example, + ``RateLimit(10, 1.0)`` allows a burst of 10 requests and then one request + per second on average. + """ + + name: str + capacity: int + refill_rate: float + # When Redis is unreachable: ``True`` allows the request (fail-open), + # ``False`` rejects it with 503 (fail-closed). + fail_open: bool = True + + def __post_init__(self) -> None: + if self.capacity <= 0: + raise ValueError("capacity must be positive") + if self.refill_rate <= 0: + raise ValueError("refill_rate must be positive") + + +@dataclass(frozen=True) +class RateLimitConfig: + """A named rate limit plus how it is applied to a request.""" + + limit: RateLimit + # Builds the Redis key prefix for the request. The authenticated user id + # is passed as the first argument (or ``None`` for anonymous limits). + scope: Callable[[Request, int | None], str] + + +@dataclass(frozen=True) +class RateLimitResult: + """Result of a rate-limit check (used by tests).""" + + allowed: bool + remaining: int + retry_after: int + + +class RateLimiter: + """Stateless dependency factory that enforces token buckets in Redis. + + Buckets are keyed as ``rate:{limit.name}:{scope_key}`` so multiple + endpoints can share a named limit while still being bucketed separately. + """ + + def __init__(self, config: RateLimitConfig, redis: Redis | None = None) -> None: + self._config = config + self._redis = redis + + def _redis_client(self) -> Redis: + return self._redis or Redis.from_url( + settings.REDIS_URL, + encoding="utf-8", + decode_responses=True, + ) + + def _scope_key(self, request: Request, user_id: int | None) -> str: + return self._config.scope(request, user_id) + + async def __call__( + self, + request: Request, + user_id: int | None = None, + ) -> None: + """Enforce the rate limit, raising 429 or 503 as appropriate. + + ``user_id`` is normally injected by another dependency (see + ``RateLimitDependency`` below); the limiter itself only needs the + resolved value. + """ + limit = self._config.limit + key = f"rate:{limit.name}:{self._scope_key(request, user_id)}" + + try: + redis = self._redis_client() + result = await redis.eval( + _TOKEN_BUCKET_LUA, + 1, + key, + str(limit.capacity), + str(limit.refill_rate), + str(time.time()), + "1", # each request consumes one token + ) + except RedisError: + if limit.fail_open: + return + raise HTTPException( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + detail="Rate limiting service unavailable", + ) from None + + # Redis may return ints or bytes depending on decode_responses. + allowed = bool(int(result[0])) + retry_after = int(result[2]) + if not allowed: + raise RateLimitExceeded(retry_after=retry_after) + + +def named_limit(name: str, **kwargs: Any) -> RateLimitConfig: + """Build a named rate-limit config with a sensible default scope. + + The default scope buckets per authenticated user (falling back to the + client IP for unauthenticated requests), which matches the chat-send and + report-message use cases. + """ + return RateLimitConfig( + limit=RateLimit(name=name, **kwargs), + scope=default_user_scope, + ) + + +def default_user_scope(request: Request, user_id: int | None) -> str: + """Return the authenticated user id, or client IP for anonymous callers.""" + if user_id is not None: + return f"user:{user_id}" + return f"ip:{request.client.host if request.client else 'unknown'}" + + +def rate_limit(config: RateLimitConfig) -> Callable[..., Any]: + """FastAPI dependency that resolves the current user and enforces a limit. + + This keeps call sites declarative, e.g.:: + + ChatSendDep = Depends(rate_limit(named_limit("chat_send", ...))) + """ + + async def dependency( + request: Request, + user_id: int | None = None, + limiter: RateLimiter = Depends(lambda: RateLimiter(config)), + ) -> None: + await limiter(request, user_id) + + return dependency