fix(orchestrator): bound streamed request bodies
test / image-input-builds (pull_request) Successful in 49s
test / integration-docker (pull_request) Has been cancelled
test / unit (pull_request) Successful in 1m4s
test / coverage (pull_request) Has been skipped
tracker-policy-pr / check-pr (pull_request) Failing after 12m11s
test / image-input-builds (pull_request) Successful in 49s
test / integration-docker (pull_request) Has been cancelled
test / unit (pull_request) Successful in 1m4s
test / coverage (pull_request) Has been skipped
tracker-policy-pr / check-pr (pull_request) Failing after 12m11s
This commit is contained in:
@@ -44,7 +44,7 @@ if TYPE_CHECKING:
|
||||
from ..gateway import Gateway, GatewayError
|
||||
from .lifecycle import Orchestrator
|
||||
from .service import OrchestratorCore
|
||||
from .server import OrchestratorServer, create_app, dispatch, make_server
|
||||
from .server import OrchestratorServer, create_app, make_server
|
||||
|
||||
|
||||
# Facade name -> submodule that defines it. Lazy so importing a leaf (or the
|
||||
@@ -68,7 +68,6 @@ _LAZY: dict[str, str] = {
|
||||
"Orchestrator": ".lifecycle",
|
||||
"OrchestratorCore": ".service",
|
||||
"create_app": ".server",
|
||||
"dispatch": ".server",
|
||||
"OrchestratorServer": ".server",
|
||||
"make_server": ".server",
|
||||
}
|
||||
@@ -102,7 +101,6 @@ __all__ = [
|
||||
"Orchestrator",
|
||||
"OrchestratorCore",
|
||||
"create_app",
|
||||
"dispatch",
|
||||
"OrchestratorServer",
|
||||
"make_server",
|
||||
]
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import math
|
||||
import sys
|
||||
from fastapi import FastAPI, HTTPException
|
||||
@@ -13,7 +14,11 @@ from starlette.types import ASGIApp, Message, Receive, Scope, Send
|
||||
from ..orchestrator_auth import ROLE_CLI, ROLES
|
||||
from ..supervisor.types import TOOLS
|
||||
from ..trust_domain import CONTROL_PLANE
|
||||
from .http_contract import MAX_BODY_BYTES, ORCHESTRATOR_AUTH_HEADER
|
||||
from .http_contract import (
|
||||
MAX_BODY_BYTES,
|
||||
ORCHESTRATOR_AUTH_HEADER,
|
||||
REQUEST_BODY_TIMEOUT_SECONDS,
|
||||
)
|
||||
from .service import OrchestratorCore
|
||||
|
||||
_GATEWAY_ROUTES = frozenset({
|
||||
@@ -105,7 +110,7 @@ class ControlPlaneBoundary:
|
||||
if role not in allowed:
|
||||
await self._reject(scope, send, 403, "insufficient role for this route")
|
||||
return
|
||||
scope["state"]["role"] = role
|
||||
scope.setdefault("state", {})["role"] = role
|
||||
raw_length = dict(scope["headers"]).get(b"content-length")
|
||||
if raw_length is not None:
|
||||
try:
|
||||
@@ -120,7 +125,15 @@ class ControlPlaneBoundary:
|
||||
await self._reject(scope, send, 413, "request body too large")
|
||||
return
|
||||
try:
|
||||
await self.app(scope, self._bounded_receive(receive), send)
|
||||
body = await self._read_body(receive)
|
||||
except _BodyTooLarge:
|
||||
await self._reject(scope, send, 413, "request body too large")
|
||||
return
|
||||
except TimeoutError:
|
||||
await self._reject(scope, send, 408, "request body read timed out")
|
||||
return
|
||||
try:
|
||||
await self.app(scope, self._replay_body(body), send)
|
||||
except Exception as exc: # noqa: BLE001 - redact control-plane failures
|
||||
sys.stderr.write(
|
||||
f"orchestrator: {method} {route} failed "
|
||||
@@ -141,19 +154,36 @@ class ControlPlaneBoundary:
|
||||
return {"type": "http.disconnect"}
|
||||
|
||||
@staticmethod
|
||||
def _bounded_receive(receive: Receive) -> Receive:
|
||||
consumed = 0
|
||||
async def _read_body(receive: Receive) -> bytes:
|
||||
body = bytearray()
|
||||
async with asyncio.timeout(REQUEST_BODY_TIMEOUT_SECONDS):
|
||||
while True:
|
||||
message = await receive()
|
||||
if message["type"] != "http.request":
|
||||
break
|
||||
body.extend(message.get("body", b""))
|
||||
if len(body) > MAX_BODY_BYTES:
|
||||
raise _BodyTooLarge
|
||||
if not message.get("more_body", False):
|
||||
break
|
||||
return bytes(body)
|
||||
|
||||
async def bounded() -> Message:
|
||||
nonlocal consumed
|
||||
message = await receive()
|
||||
if message["type"] == "http.request":
|
||||
consumed += len(message.get("body", b""))
|
||||
if consumed > MAX_BODY_BYTES:
|
||||
raise HTTPException(413, "request body too large")
|
||||
return message
|
||||
@staticmethod
|
||||
def _replay_body(body: bytes) -> Receive:
|
||||
sent = False
|
||||
|
||||
return bounded
|
||||
async def replay() -> Message:
|
||||
nonlocal sent
|
||||
if sent:
|
||||
return {"type": "http.disconnect"}
|
||||
sent = True
|
||||
return {"type": "http.request", "body": body, "more_body": False}
|
||||
|
||||
return replay
|
||||
|
||||
|
||||
class _BodyTooLarge(Exception):
|
||||
"""The streamed request exceeded the control-plane body limit."""
|
||||
|
||||
|
||||
def _required(value: str, name: str) -> str:
|
||||
|
||||
@@ -2,5 +2,10 @@
|
||||
|
||||
ORCHESTRATOR_AUTH_HEADER = "x-bot-bottle-orchestrator-auth"
|
||||
MAX_BODY_BYTES = 1 * 1024 * 1024
|
||||
REQUEST_BODY_TIMEOUT_SECONDS = 10.0
|
||||
|
||||
__all__ = ["MAX_BODY_BYTES", "ORCHESTRATOR_AUTH_HEADER"]
|
||||
__all__ = [
|
||||
"MAX_BODY_BYTES",
|
||||
"ORCHESTRATOR_AUTH_HEADER",
|
||||
"REQUEST_BODY_TIMEOUT_SECONDS",
|
||||
]
|
||||
|
||||
@@ -8,7 +8,6 @@ import threading
|
||||
|
||||
import uvicorn
|
||||
|
||||
from ..orchestrator_auth import ROLE_CLI, mint
|
||||
from ..trust_domain import CONTROL_PLANE
|
||||
from .api import create_app
|
||||
from .http_contract import MAX_BODY_BYTES, ORCHESTRATOR_AUTH_HEADER
|
||||
@@ -18,42 +17,6 @@ MAX_REQUESTS = 32
|
||||
KEEP_ALIVE_TIMEOUT_SECONDS = 10
|
||||
|
||||
|
||||
def dispatch(
|
||||
orchestrator: OrchestratorCore,
|
||||
method: str,
|
||||
path: str,
|
||||
body: bytes,
|
||||
*,
|
||||
role: str | None = ROLE_CLI,
|
||||
) -> tuple[int, dict[str, object]]:
|
||||
"""Socket-free compatibility adapter for route unit tests.
|
||||
|
||||
Production requests always enter through the FastAPI ASGI application.
|
||||
"""
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
key = "in-process-dispatch-key"
|
||||
headers: dict[str, str] = {"content-type": "application/json"}
|
||||
if role is not None:
|
||||
headers[ORCHESTRATOR_AUTH_HEADER] = mint(role, key)
|
||||
response = TestClient(create_app(orchestrator, signing_key=key)).request(
|
||||
method, path, content=body, headers=headers,
|
||||
)
|
||||
payload = response.json()
|
||||
if response.status_code == 422:
|
||||
detail = payload.get("detail", []) if isinstance(payload, dict) else []
|
||||
field = ""
|
||||
if isinstance(detail, list) and detail and isinstance(detail[0], dict):
|
||||
location = detail[0].get("loc", ())
|
||||
if isinstance(location, (list, tuple)) and len(location) > 1:
|
||||
field = str(location[1])
|
||||
suffix = f": {field}" if field else ""
|
||||
return 400, {"error": f"invalid request body{suffix}"}
|
||||
if isinstance(payload, dict) and "detail" in payload and "error" not in payload:
|
||||
payload = {"error": payload["detail"]}
|
||||
return response.status_code, payload
|
||||
|
||||
|
||||
class OrchestratorServer:
|
||||
"""Small lifecycle wrapper around Uvicorn with an eagerly bound socket."""
|
||||
|
||||
@@ -113,6 +76,5 @@ __all__ = [
|
||||
"ORCHESTRATOR_AUTH_HEADER",
|
||||
"OrchestratorServer",
|
||||
"create_app",
|
||||
"dispatch",
|
||||
"make_server",
|
||||
]
|
||||
|
||||
Reference in New Issue
Block a user