de2d431571
Adds rate_limit(name, max_hits=, window_seconds=) as a factory for FastAPI dependencies, generalizing the auth_rate_limit pattern for future endpoints (chat send, message report) without touching auth_rate_limit's own behavior. Reuses SlidingWindowLimiter for counting and the existing 429 + Retry-After response, and respects RATE_LIMIT_ENABLED the same way auth_rate_limit does. reset_all() now clears every registered limiter (auth's plus any named ones), so the conftest autouse fixture keeps working for sibling tasks automatically. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Y2FESCaMXeJpC9gSctiEeD
115 lines
4.2 KiB
Python
115 lines
4.2 KiB
Python
"""Lightweight in-process rate limiting.
|
|
|
|
The backend runs single-worker (required by the game-server supervisor), so a
|
|
process-local sliding-window counter is sufficient and needs no Redis. Keyed
|
|
by client IP. Used to blunt credential brute-force on the auth endpoints.
|
|
|
|
If the backend is ever scaled to multiple workers/hosts this must move to a
|
|
shared store (Redis) -- flagged here rather than silently under-enforcing.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import time
|
|
from collections import defaultdict, deque
|
|
|
|
from fastapi import HTTPException, Request, status
|
|
|
|
from app.config import settings
|
|
|
|
|
|
class SlidingWindowLimiter:
|
|
def __init__(self, max_hits: int, window_seconds: float) -> None:
|
|
self.max_hits = max_hits
|
|
self.window = window_seconds
|
|
self._hits: dict[str, deque[float]] = defaultdict(deque)
|
|
|
|
def check(self, key: str) -> None:
|
|
now = time.monotonic()
|
|
q = self._hits[key]
|
|
cutoff = now - self.window
|
|
while q and q[0] < cutoff:
|
|
q.popleft()
|
|
if len(q) >= self.max_hits:
|
|
retry = max(1, int(q[0] + self.window - now))
|
|
raise HTTPException(
|
|
status_code=status.HTTP_429_TOO_MANY_REQUESTS,
|
|
detail="Too many attempts. Try again shortly.",
|
|
headers={"Retry-After": str(retry)},
|
|
)
|
|
q.append(now)
|
|
|
|
def reset(self) -> None:
|
|
self._hits.clear()
|
|
|
|
|
|
def _client_key(request: Request) -> str:
|
|
# X-Forwarded-For first hop when behind the reverse proxy, else the socket.
|
|
fwd = request.headers.get("x-forwarded-for", "")
|
|
if fwd:
|
|
return fwd.split(",")[0].strip()
|
|
return request.client.host if request.client else "unknown"
|
|
|
|
|
|
# 10 auth attempts / minute / IP -- generous for a human, painful for a
|
|
# brute-forcer.
|
|
_auth_limiter = SlidingWindowLimiter(max_hits=10, window_seconds=60)
|
|
|
|
# Every limiter in the process -- named limits built via `rate_limit()`
|
|
# below register themselves here too, so `reset_all()` (the test hook) and
|
|
# any future process-wide accounting stay correct without each caller
|
|
# having to remember to wire it up.
|
|
_registered_limiters: list[SlidingWindowLimiter] = [_auth_limiter]
|
|
|
|
|
|
async def auth_rate_limit(request: Request) -> None:
|
|
"""FastAPI dependency: rate-limit auth attempts per client IP.
|
|
|
|
A no-op when ``RATE_LIMIT_ENABLED`` is false -- the checkpoint smoke
|
|
suite runs every test against one backend from one client IP, so the
|
|
per-IP cap would trip mid-suite. The dedicated ``test_rate_limit``
|
|
module exercises the real behaviour regardless of this flag.
|
|
"""
|
|
if not settings.RATE_LIMIT_ENABLED:
|
|
return
|
|
_auth_limiter.check(f"auth:{_client_key(request)}")
|
|
|
|
|
|
def rate_limit(name: str, *, max_hits: int, window_seconds: float):
|
|
"""Build a named, per-client-IP FastAPI rate-limit dependency.
|
|
|
|
Generalizes the ``auth_rate_limit`` pattern above for other endpoints
|
|
(e.g. chat send, message report) that want their own independent cap:
|
|
|
|
report_rate_limit = rate_limit("message_report", max_hits=5, window_seconds=60)
|
|
|
|
@router.post("/report", dependencies=[Depends(report_rate_limit)])
|
|
...
|
|
|
|
Each call gets its own ``SlidingWindowLimiter`` -- counters for one
|
|
named limit never bleed into another, even for the same client IP --
|
|
and reuses ``SlidingWindowLimiter.check`` for the actual sliding-window
|
|
accounting and the 429 + ``Retry-After`` response. Like
|
|
``auth_rate_limit``, the returned dependency is a no-op when
|
|
``settings.RATE_LIMIT_ENABLED`` is false, so the checkpoint smoke suite
|
|
(one client IP for the whole run) doesn't trip on it.
|
|
|
|
``name`` only needs to be unique within the process; it namespaces the
|
|
per-IP key so it doesn't collide with ``auth_rate_limit`` or any other
|
|
named limit.
|
|
"""
|
|
limiter = SlidingWindowLimiter(max_hits=max_hits, window_seconds=window_seconds)
|
|
_registered_limiters.append(limiter)
|
|
|
|
async def _dependency(request: Request) -> None:
|
|
if not settings.RATE_LIMIT_ENABLED:
|
|
return
|
|
limiter.check(f"{name}:{_client_key(request)}")
|
|
|
|
return _dependency
|
|
|
|
|
|
def reset_all() -> None:
|
|
"""Test hook."""
|
|
for limiter in _registered_limiters:
|
|
limiter.reset()
|