Spaces:
Running
Running
| """Bounded, per-process HTTP controls for ResearchBee. | |
| Uvicorn must trust ONLY the deployment's actual proxy addresses. This module | |
| uses the already-resolved ASGI client; it never trusts raw forwarding headers. | |
| Counters reset on restart and are not shared between workers. Use an edge/shared | |
| quota service and provider billing controls when a durable ceiling is required. | |
| """ | |
| from __future__ import annotations | |
| import ipaddress | |
| import json | |
| import os | |
| import threading | |
| import time | |
| from collections import OrderedDict, deque | |
| from dataclasses import dataclass | |
| from typing import Any, Callable | |
| def _env_int(name: str, default: int, low: int, high: int) -> int: | |
| try: | |
| return min(high, max(low, int(os.getenv(name, str(default))))) | |
| except (TypeError, ValueError): | |
| return default | |
| def validated_client_ip(value: Any) -> str: | |
| try: | |
| return str(ipaddress.ip_address(str(value))) | |
| except ValueError: | |
| return "unknown" | |
| def get_client_ip(request: Any) -> str: | |
| client = request.client | |
| return validated_client_ip(client.host) if client else "unknown" | |
| class SlidingWindowLimiter: | |
| """Fail closed for new identities when the bounded table is full.""" | |
| def __init__(self, max_keys: int = 10000, clock: Callable[[], float] = time.monotonic): | |
| self.max_keys = max_keys | |
| self.clock = clock | |
| self._records: OrderedDict[str, tuple[float, deque]] = OrderedDict() | |
| self._lock = threading.Lock() | |
| def allow(self, key: str, limit: int, window: float) -> bool: | |
| now = self.clock() | |
| with self._lock: | |
| # Entries are ordered by access; windows can differ, so inspect all | |
| # when capacity is reached. Routine calls remove expired head entries. | |
| while self._records: | |
| first, (expiry, _) = next(iter(self._records.items())) | |
| if expiry > now: | |
| break | |
| self._records.pop(first, None) | |
| if key not in self._records and len(self._records) >= self.max_keys: | |
| for old_key in [k for k, (expiry, _) in self._records.items() if expiry <= now]: | |
| self._records.pop(old_key, None) | |
| if len(self._records) >= self.max_keys: | |
| return False | |
| _, events = self._records.pop(key, (now + window, deque())) | |
| while events and events[0] <= now - window: | |
| events.popleft() | |
| allowed = len(events) < limit | |
| if allowed: | |
| events.append(now) | |
| self._records[key] = (now + window, events) | |
| return allowed | |
| RATE_LIMITS = { | |
| "global": (_env_int("RB_API_REQUESTS_PER_MINUTE", 120, 1, 10000), 60), | |
| "chat": (_env_int("RB_CHAT_REQUESTS_PER_MINUTE", 30, 1, 1000), 60), | |
| } | |
| _rate_limiter = SlidingWindowLimiter(_env_int("RB_RATE_MAX_IDENTITIES", 10000, 100, 100000)) | |
| def check_rate_limit(ip: str, bucket: str = "global") -> bool: | |
| limit, window = RATE_LIMITS[bucket] | |
| return _rate_limiter.allow(f"{validated_client_ip(ip)}:{bucket}", limit, window) | |
| PAID_ROUTES = frozenset({ | |
| ("POST", "/api/analyze-journal"), | |
| ("POST", "/api/browse-subject"), | |
| ("POST", "/api/check-license"), | |
| ("POST", "/api/check-journal-policies"), | |
| ("POST", "/api/find-repository"), | |
| ("POST", "/api/cover-letter"), | |
| ("POST", "/api/chat"), | |
| ("GET", "/api/lookup-journal-meta"), | |
| }) | |
| class SecurityConfig: | |
| paid_per_minute: int = 60 | |
| paid_daily_limit: int = 2000 | |
| max_concurrent: int = 4 | |
| max_body_bytes: int = 512 * 1024 | |
| max_identities: int = 10000 | |
| def from_env(cls) -> "SecurityConfig": | |
| return cls( | |
| paid_per_minute=_env_int("RB_PAID_REQUESTS_PER_MINUTE", 60, 1, 10000), | |
| paid_daily_limit=_env_int("RB_PAID_DAILY_REQUEST_LIMIT", 2000, 1, 1000000), | |
| max_concurrent=_env_int("RB_MAX_CONCURRENT_PAID_REQUESTS", 4, 1, 100), | |
| max_body_bytes=_env_int("RB_MAX_BODY_BYTES", 512 * 1024, 16384, 4 * 1024 * 1024), | |
| max_identities=_env_int("RB_RATE_MAX_IDENTITIES", 10000, 100, 100000), | |
| ) | |
| class PaidRequestGate: | |
| def __init__(self, config: SecurityConfig, clock: Callable[[], float] = time.time): | |
| self.config = config | |
| self.clock = clock | |
| self._limiter = SlidingWindowLimiter(config.max_identities) | |
| self._lock = threading.Lock() | |
| self._day = -1 | |
| self._daily = 0 | |
| self._active = 0 | |
| def enter(self, identity: str) -> tuple[bool, str, int]: | |
| with self._lock: | |
| now = self.clock() | |
| day = int(now // 86400) | |
| if day != self._day: | |
| self._day, self._daily = day, 0 | |
| if self._daily >= self.config.paid_daily_limit: | |
| return False, "Daily service request limit reached. Please try again tomorrow.", max(1, int((day + 1) * 86400 - now)) | |
| if self._active >= self.config.max_concurrent: | |
| return False, "ResearchBee is busy. Please try again shortly.", 10 | |
| if not self._limiter.allow(identity, self.config.paid_per_minute, 60): | |
| return False, "Too many requests. Please wait a minute and try again.", 60 | |
| self._daily += 1 | |
| self._active += 1 | |
| return True, "", 0 | |
| def leave(self) -> None: | |
| with self._lock: | |
| self._active = max(0, self._active - 1) | |
| class RuntimeSecurityMiddleware: | |
| """Pure ASGI: bounded body reads and concurrency slots held through SSE. | |
| Register before CORSMiddleware so CORS also wraps rejection responses. | |
| """ | |
| def __init__(self, app: Any, config: SecurityConfig | None = None): | |
| self.app = app | |
| self.config = config or SecurityConfig.from_env() | |
| self.gate = PaidRequestGate(self.config) | |
| async def __call__(self, scope: dict, receive: Any, send: Any) -> None: | |
| if scope["type"] != "http": | |
| await self.app(scope, receive, send) | |
| return | |
| path = scope.get("path", "") | |
| method = scope.get("method", "GET") | |
| is_admin = path == "/admin" or path.startswith("/admin/") | |
| protected = is_admin or path.startswith("/api/") | |
| async def secure_send(message: dict) -> None: | |
| if message["type"] == "http.response.start" and is_admin: | |
| headers = [(k, v) for k, v in message.get("headers", []) | |
| if k.lower() not in {b"cache-control", b"pragma", b"x-content-type-options", b"x-frame-options"}] | |
| headers += [(b"cache-control", b"no-store"), (b"pragma", b"no-cache"), | |
| (b"x-content-type-options", b"nosniff"), (b"x-frame-options", b"DENY")] | |
| message = {**message, "headers": headers} | |
| await send(message) | |
| async def reject(status: int, detail: str, retry_after: int = 0) -> None: | |
| body = json.dumps({"detail": detail}).encode() | |
| headers = [(b"content-type", b"application/json"), (b"content-length", str(len(body)).encode()), | |
| (b"cache-control", b"no-store")] | |
| if retry_after: | |
| headers.append((b"retry-after", str(retry_after).encode())) | |
| await secure_send({"type": "http.response.start", "status": status, "headers": headers}) | |
| await secure_send({"type": "http.response.body", "body": body}) | |
| # Buffer only the small bounded JSON body before invoking application | |
| # parsing. Check actual bytes too; Content-Length alone is not trusted. | |
| if protected and method in {"POST", "PUT", "PATCH"}: | |
| lengths = [v for k, v in scope.get("headers", []) if k.lower() == b"content-length"] | |
| try: | |
| if len(lengths) > 1 or (lengths and int(lengths[0]) < 0): | |
| raise ValueError | |
| if lengths and int(lengths[0]) > self.config.max_body_bytes: | |
| await reject(413, "Request is too large.") | |
| return | |
| except ValueError: | |
| await reject(400, "Invalid Content-Length.") | |
| return | |
| body = bytearray() | |
| while True: | |
| event = await receive() | |
| if event["type"] == "http.disconnect": | |
| return | |
| if event["type"] != "http.request": | |
| continue | |
| chunk = event.get("body", b"") | |
| if len(body) + len(chunk) > self.config.max_body_bytes: | |
| await reject(413, "Request is too large.") | |
| return | |
| body.extend(chunk) | |
| if not event.get("more_body", False): | |
| break | |
| original_receive = receive | |
| delivered = False | |
| async def replay_receive() -> dict: | |
| nonlocal delivered | |
| if not delivered: | |
| delivered = True | |
| return {"type": "http.request", "body": bytes(body), "more_body": False} | |
| return await original_receive() | |
| receive = replay_receive | |
| entered = False | |
| if (method, path.rstrip("/")) in PAID_ROUTES: | |
| client = scope.get("client") | |
| identity = validated_client_ip(client[0]) if client else "unknown" | |
| entered, detail, retry_after = self.gate.enter(identity) | |
| if not entered: | |
| await reject(429, detail, retry_after) | |
| return | |
| try: | |
| await self.app(scope, receive, secure_send) | |
| finally: | |
| if entered: | |
| self.gate.leave() | |