Import source: axon/middleware/rate_limit
Import the existing local project snapshot. This commit organizes existing files; it does not represent newly developed functionality.
This commit is contained in:
@@ -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)
|
||||
Reference in New Issue
Block a user