Spaces:
Sleeping
Sleeping
File size: 6,423 Bytes
d623240 011e5b6 d623240 011e5b6 d623240 011e5b6 d623240 011e5b6 d623240 011e5b6 d623240 011e5b6 d623240 | 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 | from __future__ import annotations
import logging
import time
from collections import defaultdict, deque
from threading import Lock
from typing import Deque
from uuid import uuid4
from fastapi import Request
from fastapi.responses import JSONResponse
from starlette.middleware.base import BaseHTTPMiddleware
from app.core.config import Settings
from app.core.error_handlers import build_error_payload
logger = logging.getLogger("oraculo_api.middleware")
def resolve_client_ip(request: Request) -> str:
forwarded_for = request.headers.get("x-forwarded-for")
if forwarded_for:
return forwarded_for.split(",")[0].strip()
if request.client:
return request.client.host
return "unknown"
class RequestContextMiddleware(BaseHTTPMiddleware):
async def dispatch(self, request: Request, call_next):
request_id = request.headers.get("x-request-id", str(uuid4()))
request.state.request_id = request_id
request.state.started_at = time.perf_counter()
request.state.client_ip = resolve_client_ip(request)
response = await call_next(request)
duration_ms = (time.perf_counter() - request.state.started_at) * 1000
response.headers["X-Request-ID"] = request_id
response.headers["X-Process-Time-MS"] = f"{duration_ms:.2f}"
return response
class SecurityHeadersMiddleware(BaseHTTPMiddleware):
def __init__(self, app, settings: Settings):
super().__init__(app)
self.settings = settings
@staticmethod
def _is_huggingface_space_request(request: Request) -> bool:
host = (request.url.hostname or "").lower()
return host.endswith(".hf.space")
def _frame_ancestors_for_request(self, request: Request) -> str:
if self._is_huggingface_space_request(request):
return "https://huggingface.co https://*.huggingface.co"
return "'none'"
def _content_security_policy_for_request(self, request: Request) -> str:
path = request.url.path
frame_ancestors = self._frame_ancestors_for_request(request)
docs_paths = {
self.settings.docs_url,
self.settings.redoc_url,
self.settings.openapi_url,
}
if path in {value for value in docs_paths if value}:
return (
"default-src 'self'; "
"script-src 'self' 'unsafe-inline' https://cdn.jsdelivr.net; "
"style-src 'self' 'unsafe-inline' https://cdn.jsdelivr.net; "
"img-src 'self' data: https://fastapi.tiangolo.com https://cdn.jsdelivr.net; "
"font-src 'self' https://cdn.jsdelivr.net; "
"connect-src 'self'; "
f"frame-ancestors {frame_ancestors}; "
"base-uri 'self';"
)
return f"default-src 'none'; frame-ancestors {frame_ancestors}; base-uri 'none';"
async def dispatch(self, request: Request, call_next):
response = await call_next(request)
if not self.settings.security_headers_enabled:
return response
frame_ancestors = self._frame_ancestors_for_request(request)
response.headers["X-Content-Type-Options"] = "nosniff"
if frame_ancestors == "'none'":
response.headers["X-Frame-Options"] = "DENY"
response.headers["Referrer-Policy"] = "no-referrer"
response.headers["Permissions-Policy"] = "camera=(), microphone=(), geolocation=()"
response.headers["Cache-Control"] = "no-store"
response.headers["Pragma"] = "no-cache"
response.headers["Content-Security-Policy"] = self._content_security_policy_for_request(request)
return response
class MaxRequestSizeMiddleware(BaseHTTPMiddleware):
def __init__(self, app, max_request_size_bytes: int):
super().__init__(app)
self.max_request_size_bytes = max_request_size_bytes
async def dispatch(self, request: Request, call_next):
content_length = request.headers.get("content-length")
if content_length and int(content_length) > self.max_request_size_bytes:
payload = build_error_payload(
request,
code="payload_too_large",
message="Payload exceeds the maximum allowed size.",
detail={"max_request_size_bytes": self.max_request_size_bytes},
)
return JSONResponse(status_code=413, content=payload)
return await call_next(request)
class SimpleInMemoryRateLimiter:
def __init__(self, max_requests: int, window_seconds: int):
self.max_requests = max_requests
self.window_seconds = window_seconds
self._storage: dict[str, Deque[float]] = defaultdict(deque)
self._lock = Lock()
def is_allowed(self, client_key: str) -> tuple[bool, int]:
now = time.time()
with self._lock:
bucket = self._storage[client_key]
while bucket and now - bucket[0] > self.window_seconds:
bucket.popleft()
if len(bucket) >= self.max_requests:
retry_after = max(1, int(self.window_seconds - (now - bucket[0])))
return False, retry_after
bucket.append(now)
return True, 0
class RateLimitMiddleware(BaseHTTPMiddleware):
def __init__(self, app, settings: Settings):
super().__init__(app)
self.settings = settings
self.limiter = SimpleInMemoryRateLimiter(
max_requests=settings.rate_limit_requests,
window_seconds=settings.rate_limit_window_seconds,
)
async def dispatch(self, request: Request, call_next):
if not self.settings.rate_limit_enabled or request.url.path in self.settings.rate_limit_exempt_paths:
return await call_next(request)
client_key = resolve_client_ip(request)
is_allowed, retry_after = self.limiter.is_allowed(client_key)
if not is_allowed:
payload = build_error_payload(
request,
code="rate_limit_exceeded",
message="Rate limit exceeded. Please retry later.",
detail={"retry_after_seconds": retry_after},
)
return JSONResponse(
status_code=429,
content=payload,
headers={"Retry-After": str(retry_after)},
)
return await call_next(request)
|