Spaces:
Running
Running
File size: 9,853 Bytes
4b2991e | 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 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 | """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()
|