From 41b1c3959e3764bba944ee3f0aca57c8a5981887 Mon Sep 17 00:00:00 2001 From: architect-bot Date: Tue, 18 Aug 2026 15:40:10 +0000 Subject: [PATCH] WIP checkpoint: Redis rate-limiting dependency --- backend/app/rate_limit.py | 197 +++++++++++++++++++++++++------------- 1 file changed, 131 insertions(+), 66 deletions(-) diff --git a/backend/app/rate_limit.py b/backend/app/rate_limit.py index 8613693..b3d8995 100644 --- a/backend/app/rate_limit.py +++ b/backend/app/rate_limit.py @@ -2,21 +2,36 @@ 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. +If Redis is unavailable the limiter can either allow the request through +(fail-open) or reject it with a ``503`` (fail-closed), depending on the +policy configured for that limit. Chat send uses fail-closed by default +because flooding/abuse is worse than a temporary outage. + +Example +------- + from app.rate_limit import rate_limit, named_limit + + chat_send_limit = rate_limit( + named_limit("chat_send", capacity=5, refill_rate=0.2, fail_open=False), + user_id_dependency=Depends(get_current_user_id), + ) + + @router.post("/messages") + async def send_message( + payload: MessageCreate, + current_user: CurrentUser, + db: DbDep, + _: None = Depends(chat_send_limit), + ): + ... """ from __future__ import annotations import time from dataclasses import dataclass -from typing import Any, Callable +from typing import Any, Awaitable, Callable from fastapi import Depends, HTTPException, Request, status from redis.asyncio import Redis @@ -66,8 +81,8 @@ else 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. +-- Keep the bucket state small and expire-able. Retaining the timestamp as a +-- float lets partial tokens accumulate smoothly between requests. redis.call("HSET", key, "tokens", tostring(tokens), "ts", tostring(now)) redis.call("PEXPIRE", key, 600000) @@ -76,7 +91,7 @@ return {allowed, tostring(tokens), retry_after} class RateLimitExceeded(HTTPException): - """Raised when a rate limit bucket has no tokens left.""" + """Raised when a rate-limit bucket has no tokens left.""" def __init__(self, retry_after: int) -> None: super().__init__( @@ -92,8 +107,8 @@ class RateLimit: ``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. + ``RateLimit(name="chat_send", capacity=10, refill_rate=1.0)`` allows a + burst of 10 requests and then one request per second on average. """ name: str @@ -104,36 +119,68 @@ class RateLimit: fail_open: bool = True def __post_init__(self) -> None: + if not self.name: + raise ValueError("name must be non-empty") if self.capacity <= 0: raise ValueError("capacity must be positive") if self.refill_rate <= 0: raise ValueError("refill_rate must be positive") +def default_user_scope(request: Request, user_id: int | None) -> str: + """Scope a bucket per authenticated user, or per client IP if anonymous.""" + if user_id is not None: + return f"user:{user_id}" + client = request.client + return f"ip:{client.host if client else 'unknown'}" + + @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] + # Builds the Redis key fragment for the request. The resolved user id is + # passed as the second argument (or ``None`` for anonymous limits). + scope: Callable[[Request, int | None], str] = default_user_scope @dataclass(frozen=True) class RateLimitResult: - """Result of a rate-limit check (used by tests).""" + """Outcome of a rate-limit check (used by tests).""" allowed: bool remaining: int retry_after: int +_redis_client: Redis | None = None + + +def get_redis() -> Redis: + """Return a shared async Redis client for rate limiting. + + The client is created lazily (``Redis.from_url`` does not connect until a + command is issued), so this dependency never blocks or fails just because + Redis is down. + """ + global _redis_client + if _redis_client is None: + _redis_client = Redis.from_url( + settings.REDIS_URL, + encoding="utf-8", + decode_responses=True, + socket_connect_timeout=settings.REDIS_TIMEOUT_SECONDS, + socket_timeout=settings.REDIS_TIMEOUT_SECONDS, + ) + return _redis_client + + class RateLimiter: - """Stateless dependency factory that enforces token buckets in Redis. + """Enforce a token bucket in Redis. Buckets are keyed as ``rate:{limit.name}:{scope_key}`` so multiple - endpoints can share a named limit while still being bucketed separately. + endpoints can share a named limit while remaining independently bucketed. """ def __init__(self, config: RateLimitConfig, redis: Redis | None = None) -> None: @@ -141,32 +188,25 @@ class RateLimiter: self._redis = redis def _redis_client(self) -> Redis: - return self._redis or Redis.from_url( - settings.REDIS_URL, - encoding="utf-8", - decode_responses=True, - ) + return self._redis or get_redis() - def _scope_key(self, request: Request, user_id: int | None) -> str: - return self._config.scope(request, user_id) + def _bucket_key(self, request: Request, user_id: int | None) -> str: + limit = self._config.limit + scope_key = self._config.scope(request, user_id) + return f"rate:{limit.name}:{scope_key}" - async def __call__( - self, - request: Request, - user_id: int | None = None, - ) -> None: - """Enforce the rate limit, raising 429 or 503 as appropriate. + async def check(self, request: Request, user_id: int | None = None) -> RateLimitResult: + """Atomically refill/take a token and return the outcome. - ``user_id`` is normally injected by another dependency (see - ``RateLimitDependency`` below); the limiter itself only needs the - resolved value. + Redis failures are mapped to ``allowed=True`` (fail-open) or + ``allowed=False`` (fail-closed, represented with ``retry_after=-1``) + according to the limit's policy, without raising. """ limit = self._config.limit - key = f"rate:{limit.name}:{self._scope_key(request, user_id)}" + key = self._bucket_key(request, user_id) try: - redis = self._redis_client() - result = await redis.eval( + result = await self._redis_client().eval( _TOKEN_BUCKET_LUA, 1, key, @@ -177,52 +217,77 @@ class RateLimiter: ) except RedisError: if limit.fail_open: - return + return RateLimitResult(allowed=True, remaining=0, retry_after=0) + # Fail-closed: no reliable retry estimate is available. + return RateLimitResult(allowed=False, remaining=0, retry_after=-1) + + # Redis may return ints or strings depending on decode_responses. + allowed = bool(int(result[0])) + remaining = int(result[1]) + retry_after = int(result[2]) + return RateLimitResult(allowed=allowed, remaining=remaining, retry_after=retry_after) + + async def __call__(self, request: Request, user_id: int | None = None) -> None: + """Enforce the rate limit, raising 429 or 503 as appropriate.""" + result = await self.check(request, user_id) + if result.allowed: + return + + if result.retry_after < 0: + # Redis was unavailable and this limit is fail-closed. 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) + ) + raise RateLimitExceeded(retry_after=result.retry_after) -def named_limit(name: str, **kwargs: Any) -> RateLimitConfig: - """Build a named rate-limit config with a sensible default scope. +def named_limit( + name: str, + *, + capacity: int, + refill_rate: float, + fail_open: bool = True, +) -> RateLimitConfig: + """Create a named rate-limit config with the default per-user/IP 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. + ``capacity`` is the burst size and ``refill_rate`` is tokens/second. """ return RateLimitConfig( - limit=RateLimit(name=name, **kwargs), - scope=default_user_scope, + limit=RateLimit( + name=name, + capacity=capacity, + refill_rate=refill_rate, + fail_open=fail_open, + ) ) -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, + *, + user_id_dependency: Any = None, +) -> Callable[[Request, Any], Awaitable[None]]: + """Build a FastAPI dependency enforcing ``config``. + ``user_id_dependency`` optionally injects the authenticated user id, e.g. + ``Depends(get_current_user_id)``. If omitted the limiter buckets by client + IP (the ``default_user_scope`` fallback), which is suitable for anonymous + endpoints but should be paired with an auth resolver for per-user limits. -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", ...))) + The returned callable is intended for ``Depends(...)``. """ + async def _anonymous_user_id() -> int | None: + return None + + resolver = user_id_dependency if user_id_dependency is not None else Depends(_anonymous_user_id) + async def dependency( request: Request, - user_id: int | None = None, - limiter: RateLimiter = Depends(lambda: RateLimiter(config)), + user_id: int | None = resolver, ) -> None: + limiter = RateLimiter(config) await limiter(request, user_id) return dependency