Spaces:
Sleeping
Sleeping
File size: 2,090 Bytes
62516b8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 | from __future__ import annotations
from collections import defaultdict, deque
from threading import Lock
from time import monotonic
from src.config import settings
class RateLimitExceededError(Exception):
pass
class TokenEndpointRateLimiter:
def __init__(self) -> None:
self._ip_hits: dict[str, deque[float]] = defaultdict(deque)
self._client_hits: dict[str, deque[float]] = defaultdict(deque)
self._lock = Lock()
@staticmethod
def _consume(
*,
bucket: deque[float],
now: float,
window_seconds: int,
limit: int,
) -> bool:
while bucket and (now - bucket[0]) >= window_seconds:
bucket.popleft()
if len(bucket) >= limit:
return False
bucket.append(now)
return True
def enforce(self, *, ip: str, client_id: str) -> None:
window_seconds = max(1, settings.oauth_token_rate_limit_window_seconds)
per_ip_limit = max(1, settings.oauth_token_rate_limit_per_ip)
per_client_limit = max(1, settings.oauth_token_rate_limit_per_client)
key_ip = ip or "unknown"
key_client = client_id or "unknown"
with self._lock:
now = monotonic()
ip_allowed = self._consume(
bucket=self._ip_hits[key_ip],
now=now,
window_seconds=window_seconds,
limit=per_ip_limit,
)
if not ip_allowed:
raise RateLimitExceededError("Too many token requests from this IP.")
client_allowed = self._consume(
bucket=self._client_hits[key_client],
now=now,
window_seconds=window_seconds,
limit=per_client_limit,
)
if not client_allowed:
raise RateLimitExceededError("Too many token requests for this client.")
def reset(self) -> None:
with self._lock:
self._ip_hits.clear()
self._client_hits.clear()
oauth_token_rate_limiter = TokenEndpointRateLimiter()
|