Spaces:
Sleeping
Sleeping
| """Authenticate and bound machine requests before parsing private payloads.""" | |
| import hashlib | |
| import hmac | |
| from fastapi import Request | |
| from fastapi.exceptions import RequestValidationError | |
| from fastapi.responses import JSONResponse | |
| from fastapi.security import HTTPBearer | |
| from core.config import settings | |
| archive_bearer = HTTPBearer(auto_error=False) | |
| MAX_BODY_BYTES = 512000 | |
| class ArchiveSecurityMiddleware: | |
| def __init__(self, app, configuration=settings): | |
| self.app = app | |
| self.configuration = configuration | |
| async def __call__(self, scope, receive, send): | |
| path = scope.get("path", "") | |
| if scope["type"] != "http" or not path.startswith("/api/v1/archive/"): | |
| return await self.app(scope, receive, send) | |
| async def reject(status, detail): | |
| response = JSONResponse({"detail": detail}, status_code=status, | |
| headers={"Cache-Control": "no-store"}) | |
| await response(scope, receive, send) | |
| key = self.configuration.ARCHIVE_INFERENCE_API_KEY.get_secret_value() | |
| if len(key) < 32: | |
| return await reject(503, "Archive inference is not configured") | |
| headers = dict(scope.get("headers", [])) | |
| authorization = headers.get(b"authorization", b"") | |
| if not authorization.startswith(b"Bearer ") or len(authorization) > 512: | |
| return await reject(401, "Machine authentication required") | |
| supplied = authorization[7:] | |
| if not hmac.compare_digest(hashlib.sha256(supplied).digest(), | |
| hashlib.sha256(key.encode("utf-8")).digest()): | |
| return await reject(401, "Machine authentication required") | |
| content_length = headers.get(b"content-length") | |
| if content_length is not None: | |
| try: | |
| length = int(content_length) | |
| except ValueError: | |
| return await reject(400, "Invalid request length") | |
| if length < 0 or length > MAX_BODY_BYTES: | |
| return await reject(413, "Request body is too large") | |
| body = bytearray() | |
| while True: | |
| message = await receive() | |
| if message["type"] == "http.disconnect": | |
| return | |
| body.extend(message.get("body", b"")) | |
| if len(body) > MAX_BODY_BYTES: | |
| return await reject(413, "Request body is too large") | |
| if not message.get("more_body", False): | |
| break | |
| delivered = False | |
| async def bounded_receive(): | |
| nonlocal delivered | |
| if not delivered: | |
| delivered = True | |
| return {"type": "http.request", "body": bytes(body), "more_body": False} | |
| return await receive() | |
| async def no_store_send(message): | |
| if message["type"] == "http.response.start": | |
| message["headers"] = [pair for pair in message.get("headers", []) | |
| if pair[0].lower() != b"cache-control"] | |
| message["headers"].append((b"cache-control", b"no-store")) | |
| await send(message) | |
| await self.app(scope, bounded_receive, no_store_send) | |
| async def safe_validation_errors(request: Request, exc: RequestValidationError): | |
| # Pydantic's default error response contains the rejected input. Do not echo | |
| # unpublished documents, tokens, or malformed private input back in errors. | |
| errors = [{"loc": error["loc"], "type": error["type"], "msg": "Invalid request"} | |
| for error in exc.errors()] | |
| return JSONResponse({"detail": errors}, status_code=422, | |
| headers={"Cache-Control": "no-store"}) | |