Spaces:
Running
Running
File size: 1,650 Bytes
13b1a91 | 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 | """In-memory token bucket.
Clients are keyed by a salted hash of their address. The salt is random, lives
only in this process, and rotates hourly, so the key can't be reversed, joined
across restarts, or stored. Nothing here is logged or persisted.
"""
from __future__ import annotations
import hashlib
import secrets
import threading
import time
class RateLimiter:
def __init__(self, rate_per_minute: float = 60, burst: int = 30, rotate_seconds: float = 3600):
self.rate = rate_per_minute / 60.0
self.burst = burst
self.rotate_seconds = rotate_seconds
self._salt = secrets.token_bytes(16)
self._salt_at = time.time()
self._buckets: dict[str, tuple[float, float]] = {}
self._lock = threading.Lock()
def _key(self, client: str) -> str:
if time.time() - self._salt_at > self.rotate_seconds:
self._salt = secrets.token_bytes(16)
self._salt_at = time.time()
self._buckets.clear()
return hashlib.blake2b(client.encode(), key=self._salt, digest_size=12).hexdigest()
def allow(self, client: str, cost: float = 1.0) -> bool:
now = time.time()
with self._lock:
key = self._key(client)
tokens, last = self._buckets.get(key, (float(self.burst), now))
tokens = min(self.burst, tokens + (now - last) * self.rate)
if tokens < cost:
self._buckets[key] = (tokens, now)
return False
self._buckets[key] = (tokens - cost, now)
if len(self._buckets) > 50_000:
self._buckets.clear()
return True
|