Add in-memory rate limiting middleware
Lightweight sliding-window per-IP rate limiter (100 req/min default) with no external dependency. Health endpoint is excluded. Returns 429 with Retry-After header when exceeded. Sufficient for single-process SQLite deployments; document the Redis upgrade path for scale.
This commit is contained in:
parent
c82fd66ffc
commit
d089bb8670
3 changed files with 184 additions and 0 deletions
89
document-parser/infra/rate_limiter.py
Normal file
89
document-parser/infra/rate_limiter.py
Normal file
|
|
@ -0,0 +1,89 @@
|
||||||
|
"""Lightweight in-memory rate limiter middleware for FastAPI.
|
||||||
|
|
||||||
|
Uses a sliding-window counter per client IP. No external dependency
|
||||||
|
required — suitable for single-process deployments with SQLite.
|
||||||
|
|
||||||
|
For multi-process or distributed setups, replace with a Redis-backed
|
||||||
|
solution (e.g. slowapi).
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import logging
|
||||||
|
import time
|
||||||
|
from collections import defaultdict
|
||||||
|
from dataclasses import dataclass, field
|
||||||
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
|
from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
|
||||||
|
from starlette.responses import JSONResponse, Response
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from starlette.requests import Request
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class _ClientBucket:
|
||||||
|
"""Sliding window of request timestamps for a single client."""
|
||||||
|
|
||||||
|
timestamps: list[float] = field(default_factory=list)
|
||||||
|
|
||||||
|
def count_recent(self, window: float, now: float) -> int:
|
||||||
|
"""Remove expired entries and return the count of recent requests."""
|
||||||
|
cutoff = now - window
|
||||||
|
self.timestamps = [t for t in self.timestamps if t > cutoff]
|
||||||
|
return len(self.timestamps)
|
||||||
|
|
||||||
|
def add(self, now: float) -> None:
|
||||||
|
self.timestamps.append(now)
|
||||||
|
|
||||||
|
|
||||||
|
class RateLimiterMiddleware(BaseHTTPMiddleware):
|
||||||
|
"""Per-IP rate limiter using in-memory sliding windows.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
app: The ASGI application.
|
||||||
|
requests_per_window: Max requests allowed per window.
|
||||||
|
window_seconds: Size of the sliding window in seconds.
|
||||||
|
exclude_paths: Paths exempt from rate limiting (e.g. health checks).
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
app,
|
||||||
|
*,
|
||||||
|
requests_per_window: int = 60,
|
||||||
|
window_seconds: float = 60.0,
|
||||||
|
exclude_paths: tuple[str, ...] = ("/api/health",),
|
||||||
|
):
|
||||||
|
super().__init__(app)
|
||||||
|
self._max_requests = requests_per_window
|
||||||
|
self._window = window_seconds
|
||||||
|
self._exclude = exclude_paths
|
||||||
|
self._buckets: dict[str, _ClientBucket] = defaultdict(_ClientBucket)
|
||||||
|
|
||||||
|
async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
|
||||||
|
if request.url.path in self._exclude:
|
||||||
|
return await call_next(request)
|
||||||
|
|
||||||
|
client_ip = request.client.host if request.client else "unknown"
|
||||||
|
now = time.monotonic()
|
||||||
|
|
||||||
|
bucket = self._buckets[client_ip]
|
||||||
|
recent = bucket.count_recent(self._window, now)
|
||||||
|
|
||||||
|
if recent >= self._max_requests:
|
||||||
|
retry_after = int(self._window)
|
||||||
|
logger.warning(
|
||||||
|
"Rate limit exceeded for %s (%d/%d)", client_ip, recent, self._max_requests
|
||||||
|
)
|
||||||
|
return JSONResponse(
|
||||||
|
status_code=429,
|
||||||
|
content={"detail": "Too many requests"},
|
||||||
|
headers={"Retry-After": str(retry_after)},
|
||||||
|
)
|
||||||
|
|
||||||
|
bucket.add(now)
|
||||||
|
return await call_next(request)
|
||||||
|
|
@ -20,6 +20,7 @@ from fastapi.middleware.cors import CORSMiddleware
|
||||||
|
|
||||||
from api.analyses import router as analyses_router
|
from api.analyses import router as analyses_router
|
||||||
from api.documents import router as documents_router
|
from api.documents import router as documents_router
|
||||||
|
from infra.rate_limiter import RateLimiterMiddleware
|
||||||
from infra.settings import settings
|
from infra.settings import settings
|
||||||
from persistence.database import get_connection, init_db
|
from persistence.database import get_connection, init_db
|
||||||
from services.analysis_service import AnalysisService
|
from services.analysis_service import AnalysisService
|
||||||
|
|
@ -94,6 +95,7 @@ app.add_middleware(
|
||||||
allow_methods=["GET", "POST", "DELETE", "OPTIONS"],
|
allow_methods=["GET", "POST", "DELETE", "OPTIONS"],
|
||||||
allow_headers=["Content-Type", "Authorization"],
|
allow_headers=["Content-Type", "Authorization"],
|
||||||
)
|
)
|
||||||
|
app.add_middleware(RateLimiterMiddleware, requests_per_window=100, window_seconds=60)
|
||||||
|
|
||||||
app.include_router(documents_router)
|
app.include_router(documents_router)
|
||||||
app.include_router(analyses_router)
|
app.include_router(analyses_router)
|
||||||
|
|
|
||||||
93
document-parser/tests/test_rate_limiter.py
Normal file
93
document-parser/tests/test_rate_limiter.py
Normal file
|
|
@ -0,0 +1,93 @@
|
||||||
|
"""Tests for the in-memory rate limiter middleware."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from unittest.mock import patch
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
from fastapi import FastAPI
|
||||||
|
from fastapi.testclient import TestClient
|
||||||
|
|
||||||
|
from infra.rate_limiter import RateLimiterMiddleware, _ClientBucket
|
||||||
|
|
||||||
|
|
||||||
|
class TestClientBucket:
|
||||||
|
def test_count_recent_filters_old_entries(self):
|
||||||
|
bucket = _ClientBucket(timestamps=[1.0, 2.0, 3.0, 10.0])
|
||||||
|
count = bucket.count_recent(window=5.0, now=12.0)
|
||||||
|
assert count == 1 # only 10.0 is within [7.0, 12.0]
|
||||||
|
|
||||||
|
def test_count_recent_keeps_all_when_within_window(self):
|
||||||
|
bucket = _ClientBucket(timestamps=[10.0, 11.0, 12.0])
|
||||||
|
count = bucket.count_recent(window=60.0, now=15.0)
|
||||||
|
assert count == 3
|
||||||
|
|
||||||
|
def test_add(self):
|
||||||
|
bucket = _ClientBucket()
|
||||||
|
bucket.add(1.0)
|
||||||
|
bucket.add(2.0)
|
||||||
|
assert len(bucket.timestamps) == 2
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def limited_app():
|
||||||
|
"""FastAPI app with a very low rate limit for testing."""
|
||||||
|
app = FastAPI()
|
||||||
|
app.add_middleware(
|
||||||
|
RateLimiterMiddleware,
|
||||||
|
requests_per_window=3,
|
||||||
|
window_seconds=60,
|
||||||
|
exclude_paths=("/health",),
|
||||||
|
)
|
||||||
|
|
||||||
|
@app.get("/test")
|
||||||
|
def test_endpoint():
|
||||||
|
return {"ok": True}
|
||||||
|
|
||||||
|
@app.get("/health")
|
||||||
|
def health():
|
||||||
|
return {"status": "ok"}
|
||||||
|
|
||||||
|
return app
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def client(limited_app):
|
||||||
|
return TestClient(limited_app)
|
||||||
|
|
||||||
|
|
||||||
|
class TestRateLimiterMiddleware:
|
||||||
|
def test_allows_requests_under_limit(self, client):
|
||||||
|
for _ in range(3):
|
||||||
|
resp = client.get("/test")
|
||||||
|
assert resp.status_code == 200
|
||||||
|
|
||||||
|
def test_blocks_requests_over_limit(self, client):
|
||||||
|
for _ in range(3):
|
||||||
|
client.get("/test")
|
||||||
|
|
||||||
|
resp = client.get("/test")
|
||||||
|
assert resp.status_code == 429
|
||||||
|
assert resp.json()["detail"] == "Too many requests"
|
||||||
|
assert "Retry-After" in resp.headers
|
||||||
|
|
||||||
|
def test_health_excluded_from_limit(self, client):
|
||||||
|
# Exhaust the limit
|
||||||
|
for _ in range(3):
|
||||||
|
client.get("/test")
|
||||||
|
|
||||||
|
# Health should still work
|
||||||
|
resp = client.get("/health")
|
||||||
|
assert resp.status_code == 200
|
||||||
|
|
||||||
|
def test_window_resets(self, client):
|
||||||
|
"""After the window expires, requests should be allowed again."""
|
||||||
|
for _ in range(3):
|
||||||
|
client.get("/test")
|
||||||
|
|
||||||
|
assert client.get("/test").status_code == 429
|
||||||
|
|
||||||
|
# Simulate time passing beyond the window
|
||||||
|
with patch("time.monotonic", return_value=1e12):
|
||||||
|
resp = client.get("/test")
|
||||||
|
assert resp.status_code == 200
|
||||||
Loading…
Reference in a new issue