diff --git a/backend/app/rate_limit.py b/backend/app/rate_limit.py index 28f69dc..6d8bd69 100644 --- a/backend/app/rate_limit.py +++ b/backend/app/rate_limit.py @@ -54,6 +54,12 @@ def _client_key(request: Request) -> str: # 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. @@ -68,6 +74,41 @@ async def auth_rate_limit(request: Request) -> None: _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.""" - _auth_limiter.reset() + for limiter in _registered_limiters: + limiter.reset() diff --git a/backend/tests/test_rate_limit.py b/backend/tests/test_rate_limit.py index 4080548..9b56c95 100644 --- a/backend/tests/test_rate_limit.py +++ b/backend/tests/test_rate_limit.py @@ -1,6 +1,7 @@ """The auth endpoints are rate-limited per client IP.""" import pytest +from fastapi import Depends, FastAPI from fastapi.testclient import TestClient from app import rate_limit @@ -54,3 +55,70 @@ def test_limiter_window_is_per_key(): with pytest.raises(Exception): limiter.check("a") limiter.check("b") # different key, unaffected + + +def _make_named_limit_app(): + """A tiny standalone app exercising two independent named limits, so + these tests don't depend on any sibling endpoint existing yet.""" + limit_a = rate_limit.rate_limit("test_named_a", max_hits=2, window_seconds=60) + limit_b = rate_limit.rate_limit("test_named_b", max_hits=2, window_seconds=60) + + app = FastAPI() + + @app.post("/a", dependencies=[Depends(limit_a)]) + def _a(): + return {"ok": True} + + @app.post("/b", dependencies=[Depends(limit_b)]) + def _b(): + return {"ok": True} + + return TestClient(app) + + +def test_named_limits_do_not_share_counters(): + client = _make_named_limit_app() + + # Burn out limit "a" entirely. + assert client.post("/a").status_code == 200 + assert client.post("/a").status_code == 200 + r = client.post("/a") + assert r.status_code == 429 + assert "Retry-After" in r.headers + + # Limit "b" is untouched -- same process, same client IP, different name. + assert client.post("/b").status_code == 200 + assert client.post("/b").status_code == 200 + assert client.post("/b").status_code == 429 + + +def test_named_limit_returns_429_with_retry_after_once_exhausted(): + client = _make_named_limit_app() + client.post("/a") + client.post("/a") + r = client.post("/a") + assert r.status_code == 429 + retry_after = r.headers["Retry-After"] + assert int(retry_after) >= 1 + + +def test_named_limit_is_a_noop_when_rate_limiting_disabled(): + prev = settings.RATE_LIMIT_ENABLED + settings.RATE_LIMIT_ENABLED = False + try: + client = _make_named_limit_app() + codes = [client.post("/a").status_code for _ in range(10)] + assert codes == [200] * 10 + finally: + settings.RATE_LIMIT_ENABLED = prev + + +def test_reset_all_clears_named_limiters_too(): + client = _make_named_limit_app() + client.post("/a") + client.post("/a") + assert client.post("/a").status_code == 429 + + rate_limit.reset_all() + + assert client.post("/a").status_code == 200