| """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 |
| |
| 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, "" |
|
|
