Spaces:
Sleeping
Sleeping
File size: 3,687 Bytes
f16b727 | 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 | """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"})
|