47 lines
1.0 KiB
Python
47 lines
1.0 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
from collections import deque
|
||
|
|
from threading import Lock
|
||
|
|
from time import monotonic
|
||
|
|
|
||
|
|
from fastapi import HTTPException
|
||
|
|
|
||
|
|
_BUCKETS: dict[str, deque[float]] = {}
|
||
|
|
_LOCK = Lock()
|
||
|
|
|
||
|
|
|
||
|
|
async def enforce_rate_limit(
|
||
|
|
*,
|
||
|
|
scope: str,
|
||
|
|
identifier: str,
|
||
|
|
limit: int,
|
||
|
|
window_seconds: int,
|
||
|
|
) -> None:
|
||
|
|
_enforce_rate_limit_in_memory(
|
||
|
|
key=f"auth:rate_limit:{scope}:{identifier.lower()}",
|
||
|
|
limit=limit,
|
||
|
|
window_seconds=window_seconds,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def _enforce_rate_limit_in_memory(
|
||
|
|
*,
|
||
|
|
key: str,
|
||
|
|
limit: int,
|
||
|
|
window_seconds: int,
|
||
|
|
) -> None:
|
||
|
|
now = monotonic()
|
||
|
|
with _LOCK:
|
||
|
|
bucket = _BUCKETS.setdefault(key, deque())
|
||
|
|
cutoff = now - float(window_seconds)
|
||
|
|
while bucket and bucket[0] <= cutoff:
|
||
|
|
bucket.popleft()
|
||
|
|
if len(bucket) >= limit:
|
||
|
|
raise HTTPException(status_code=429, detail="Too many requests")
|
||
|
|
bucket.append(now)
|
||
|
|
|
||
|
|
|
||
|
|
def reset_rate_limit_state() -> None:
|
||
|
|
with _LOCK:
|
||
|
|
_BUCKETS.clear()
|