File size: 3,427 Bytes
30e6218 | 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 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 | """In-process sliding-window rate limiter for ``/api`` routes.
Two ceilings: one per client and one global backstop. Both count requests in
a sliding 60 second window. State is process-local, which matches the
single-process deployment boundary.
Client identity: the socket address, unless the request carries a valid API
key *and* the configured client-IP header, in which case the header value is
used. A trusted proxy (the demo frontend) holds the key and forwards the real
visitor address in that header; anonymous callers cannot spoof it.
"""
from __future__ import annotations
import hmac
import math
import threading
import time
from collections import deque
from typing import Deque, Dict, Optional
from fastapi import Request
from ..core.config import Settings
WINDOW_SECONDS = 60.0
_MAX_TRACKED_CLIENTS = 10_000
class SlidingWindowLimiter:
def __init__(self, per_client: int, global_limit: int, window: float = WINDOW_SECONDS):
self.per_client = max(0, per_client)
self.global_limit = max(0, global_limit)
self.window = window
self._clients: Dict[str, Deque[float]] = {}
self._global: Deque[float] = deque()
self._lock = threading.Lock()
@property
def enabled(self) -> bool:
return self.per_client > 0 or self.global_limit > 0
def _prune(self, bucket: Deque[float], now: float) -> None:
cutoff = now - self.window
while bucket and bucket[0] <= cutoff:
bucket.popleft()
def check(self, client_key: str, now: Optional[float] = None) -> Optional[int]:
"""Record one request. Return ``None`` if allowed, else seconds to wait."""
if not self.enabled:
return None
now = time.monotonic() if now is None else now
with self._lock:
if self.global_limit:
self._prune(self._global, now)
if len(self._global) >= self.global_limit:
return max(1, math.ceil(self._global[0] + self.window - now))
if self.per_client:
bucket = self._clients.get(client_key)
if bucket is None:
if len(self._clients) >= _MAX_TRACKED_CLIENTS:
self._evict_idle(now)
bucket = self._clients[client_key] = deque()
self._prune(bucket, now)
if len(bucket) >= self.per_client:
return max(1, math.ceil(bucket[0] + self.window - now))
bucket.append(now)
if self.global_limit:
self._global.append(now)
return None
def _evict_idle(self, now: float) -> None:
for key in list(self._clients):
self._prune(self._clients[key], now)
if not self._clients[key]:
del self._clients[key]
def client_identity(request: Request, settings: Settings) -> str:
"""Socket address, or the forwarded address when the caller proves it holds the key."""
socket_host = request.client.host if request.client else "unknown"
if settings.api_key is None:
return socket_host
provided = request.headers.get("X-API-Key") or ""
if not hmac.compare_digest(provided.encode(), settings.api_key.encode()):
return socket_host
forwarded = (request.headers.get(settings.client_ip_header) or "").split(",")[0].strip()
return forwarded[:128] if forwarded else socket_host
|