"""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"), }) @dataclass(frozen=True) 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 @classmethod 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()