40 lines
1.2 KiB
Python
40 lines
1.2 KiB
Python
from __future__ import annotations
|
|
|
|
import time
|
|
from collections import deque
|
|
from dataclasses import dataclass
|
|
from threading import Lock
|
|
|
|
from fastapi import HTTPException, Request, status
|
|
|
|
|
|
@dataclass
|
|
class SlidingWindowRateLimiter:
|
|
limit: int
|
|
window_seconds: int
|
|
|
|
def __post_init__(self) -> None:
|
|
self._events: dict[str, deque[float]] = {}
|
|
self._lock = Lock()
|
|
|
|
def hit(self, key: str) -> None:
|
|
now = time.time()
|
|
floor = now - self.window_seconds
|
|
|
|
with self._lock:
|
|
bucket = self._events.setdefault(key, deque())
|
|
while bucket and bucket[0] <= floor:
|
|
bucket.popleft()
|
|
if len(bucket) >= self.limit:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_429_TOO_MANY_REQUESTS,
|
|
detail="Too many requests. Please try again later.",
|
|
)
|
|
bucket.append(now)
|
|
|
|
|
|
def request_client_key(request: Request, *, suffix: str = "") -> str:
|
|
forwarded_for = request.headers.get("x-forwarded-for", "")
|
|
client_ip = forwarded_for.split(",", 1)[0].strip() if forwarded_for else (request.client.host if request.client else "unknown")
|
|
return f"{client_ip}:{suffix}" if suffix else client_ip
|