Generalize rate_limit.py into a reusable named-limit dependency
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
This commit is contained in:
@@ -54,6 +54,12 @@ def _client_key(request: Request) -> str:
|
|||||||
# brute-forcer.
|
# brute-forcer.
|
||||||
_auth_limiter = SlidingWindowLimiter(max_hits=10, window_seconds=60)
|
_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:
|
async def auth_rate_limit(request: Request) -> None:
|
||||||
"""FastAPI dependency: rate-limit auth attempts per client IP.
|
"""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)}")
|
_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:
|
def reset_all() -> None:
|
||||||
"""Test hook."""
|
"""Test hook."""
|
||||||
_auth_limiter.reset()
|
for limiter in _registered_limiters:
|
||||||
|
limiter.reset()
|
||||||
|
|||||||
@@ -1,6 +1,7 @@
|
|||||||
"""The auth endpoints are rate-limited per client IP."""
|
"""The auth endpoints are rate-limited per client IP."""
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
|
from fastapi import Depends, FastAPI
|
||||||
from fastapi.testclient import TestClient
|
from fastapi.testclient import TestClient
|
||||||
|
|
||||||
from app import rate_limit
|
from app import rate_limit
|
||||||
@@ -54,3 +55,70 @@ def test_limiter_window_is_per_key():
|
|||||||
with pytest.raises(Exception):
|
with pytest.raises(Exception):
|
||||||
limiter.check("a")
|
limiter.check("a")
|
||||||
limiter.check("b") # different key, unaffected
|
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
|
||||||
|
|||||||
Reference in New Issue
Block a user