diff --git a/axon/middleware/rate_limit.py b/axon/middleware/rate_limit.py new file mode 100644 index 0000000..d43e277 --- /dev/null +++ b/axon/middleware/rate_limit.py @@ -0,0 +1,39 @@ +import time +from collections import defaultdict, deque + +from fastapi.responses import JSONResponse +from starlette.middleware.base import BaseHTTPMiddleware +from starlette.requests import Request + +from axon.config import settings + +_EXEMPT = {"/health", "/metrics", "/docs", "/openapi.json", "/redoc"} + + +class RateLimitMiddleware(BaseHTTPMiddleware): + def __init__(self, app, limit_per_minute: int | None = None) -> None: + super().__init__(app) + self.limit_per_minute = limit_per_minute or settings.rate_limit_per_minute + self._hits: dict[str, deque[float]] = defaultdict(deque) + + async def dispatch(self, request: Request, call_next): + if self.limit_per_minute <= 0 or request.url.path in _EXEMPT: + return await call_next(request) + + key = request.headers.get("X-API-Key") or request.client.host if request.client else "unknown" + now = time.monotonic() + window_start = now - 60.0 + hits = self._hits[key] + + while hits and hits[0] < window_start: + hits.popleft() + + if len(hits) >= self.limit_per_minute: + return JSONResponse( + status_code=429, + content={"detail": "Rate limit exceeded"}, + headers={"Retry-After": "60"}, + ) + + hits.append(now) + return await call_next(request)