Antoni Nuñez Romeu
54 lines
2.0 KiB
Python
54 lines
2.0 KiB
Python
import time
|
|
|
|
from fastapi import Request, Response
|
|
from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
|
|
|
|
from src.infrastructure.config.settings import settings
|
|
|
|
|
|
class RateLimitMiddleware(BaseHTTPMiddleware):
|
|
def __init__(self, app, requests_per_minute: int | None = None):
|
|
super().__init__(app)
|
|
self._rpm = requests_per_minute or settings.RATE_LIMIT_PER_MINUTE
|
|
self._clients: dict[str, list[float]] = {}
|
|
|
|
async def dispatch(
|
|
self, request: Request, call_next: RequestResponseEndpoint
|
|
) -> Response:
|
|
if request.method == "OPTIONS":
|
|
return await call_next(request)
|
|
|
|
path = request.url.path
|
|
if path.startswith("/docs") or path.startswith("/redoc") or path.startswith("/openapi") or path.startswith("/metrics"):
|
|
return await call_next(request)
|
|
|
|
client_key = self._get_client_key(request)
|
|
now = time.monotonic()
|
|
window = self._clients.setdefault(client_key, [])
|
|
window[:] = [t for t in window if now - t < 60]
|
|
|
|
if len(window) >= self._rpm:
|
|
return Response(
|
|
content='{"error":{"code":"RATE_LIMIT_EXCEEDED","message":"Rate limit exceeded. Try again later."}}',
|
|
status_code=429,
|
|
media_type="application/json",
|
|
headers={"Retry-After": "60"},
|
|
)
|
|
|
|
window.append(now)
|
|
response = await call_next(request)
|
|
remaining = self._rpm - len(window)
|
|
response.headers["X-RateLimit-Limit"] = str(self._rpm)
|
|
response.headers["X-RateLimit-Remaining"] = str(max(0, remaining))
|
|
response.headers["X-RateLimit-Reset"] = str(int(time.time()) + 60)
|
|
return response
|
|
|
|
@staticmethod
|
|
def _get_client_key(request: Request) -> str:
|
|
forwarded = request.headers.get("x-forwarded-for")
|
|
if forwarded:
|
|
return forwarded.split(",")[0].strip()
|
|
if request.client:
|
|
return request.client.host
|
|
return "unknown"
|