File size: 3,059 Bytes
7e2c39a | 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 | """Abuse controls for AI-powered Spaces.
Layers (see docs/rate-limiting.md):
1. Per-IP sliding window (server-side, survives incognito/new sessions).
2. Global daily cap — cannot be bypassed by identity rotation; bounds
worst-case daily spend to DAILY_CAP x cost-per-generation.
3. Queue concurrency (configured on the Blocks) bounds burn rate.
4. api_name=False on generate events removes the scripted-API surface.
State is in-memory and resets on Space restart; that is acceptable for
abuse control (not accounting).
"""
import os
import threading
import time
from collections import defaultdict
class RateLimiter:
def __init__(
self,
per_ip_limit: int = int(os.environ.get("LF_IP_LIMIT", 8)),
window_seconds: int = int(os.environ.get("LF_IP_WINDOW_SECONDS", 3600)),
daily_cap: int = int(os.environ.get("LF_DAILY_CAP", 500)),
):
self.per_ip_limit = per_ip_limit
self.window = window_seconds
self.daily_cap = daily_cap
self._hits: dict[str, list[float]] = defaultdict(list)
self._day: str = ""
self._day_count = 0
self._lock = threading.Lock()
@staticmethod
def client_ip(request) -> str:
"""Real client IP. On HF Spaces the app sits behind a proxy, so the
client is the first entry of x-forwarded-for, not request.client."""
if request is None:
return "unknown"
fwd = dict(request.headers).get("x-forwarded-for", "")
if fwd:
return fwd.split(",")[0].strip()
try:
return request.client.host
except Exception:
return "unknown"
def check(self, request) -> tuple[bool, str]:
"""Returns (allowed, user_safe_message)."""
ip = self.client_ip(request)
now = time.time()
today = time.strftime("%Y-%m-%d", time.gmtime(now))
with self._lock:
if today != self._day:
self._day, self._day_count = today, 0
if self._day_count >= self.daily_cap:
return False, (
"The forge is cooling down for today — daily generation "
"limit reached. Come back tomorrow!"
)
window_start = now - self.window
self._hits[ip] = [t for t in self._hits[ip] if t > window_start]
if len(self._hits[ip]) >= self.per_ip_limit:
minutes = max(1, int((self._hits[ip][0] + self.window - now) / 60))
return False, (
f"You've hit the hourly generation limit. Try again in about "
f"{minutes} minute{'s' if minutes != 1 else ''}."
)
self._hits[ip].append(now)
self._day_count += 1
# Opportunistic cleanup so the dict doesn't grow unbounded.
if len(self._hits) > 10000:
stale = [k for k, v in self._hits.items() if not v or v[-1] < window_start]
for k in stale:
del self._hits[k]
return True, ""
|