fix(orchestrator): bound streamed request bodies
test / integration-docker (pull_request) Has been cancelled
test / image-input-builds (pull_request) Successful in 42s
test / unit (pull_request) Failing after 14m6s
test / coverage (pull_request) Has been skipped
tracker-policy-pr / check-pr (pull_request) Failing after 13s
test / integration-docker (pull_request) Has been cancelled
test / image-input-builds (pull_request) Successful in 42s
test / unit (pull_request) Failing after 14m6s
test / coverage (pull_request) Has been skipped
tracker-policy-pr / check-pr (pull_request) Failing after 13s
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",
|
||||
]
|
||||
|
||||
@@ -3,6 +3,8 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import subprocess
|
||||
import sys
|
||||
import unittest
|
||||
import urllib.error
|
||||
from unittest.mock import MagicMock, patch
|
||||
@@ -20,6 +22,30 @@ from bot_bottle.orchestrator.client import (
|
||||
_URLOPEN = "bot_bottle.orchestrator.client.urllib.request.urlopen"
|
||||
|
||||
|
||||
class TestImportBoundary(unittest.TestCase):
|
||||
def test_host_client_does_not_import_server_dependencies(self) -> None:
|
||||
script = """
|
||||
import importlib.abc
|
||||
import sys
|
||||
|
||||
class BlockServerDependencies(importlib.abc.MetaPathFinder):
|
||||
def find_spec(self, fullname, path=None, target=None):
|
||||
if fullname.split(".", 1)[0] in {"fastapi", "uvicorn"}:
|
||||
raise ImportError(f"host import reached {fullname}")
|
||||
return None
|
||||
|
||||
sys.meta_path.insert(0, BlockServerDependencies())
|
||||
import bot_bottle.orchestrator.client
|
||||
"""
|
||||
result = subprocess.run(
|
||||
[sys.executable, "-c", script],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=False,
|
||||
)
|
||||
self.assertEqual(0, result.returncode, result.stderr)
|
||||
|
||||
|
||||
class TestHostAuthToken(unittest.TestCase):
|
||||
def test_mints_a_cli_token_from_the_host_key(self) -> None:
|
||||
# The CLI mints its `cli` token from the control-plane trust domain's
|
||||
|
||||
@@ -6,6 +6,7 @@ server tests), plus one real-socket round-trip to prove the handler wiring.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import http.client
|
||||
import io
|
||||
@@ -17,13 +18,17 @@ import threading
|
||||
import unittest
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
import httpx
|
||||
from contextlib import closing
|
||||
from collections.abc import Iterator
|
||||
from pathlib import Path
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from bot_bottle.orchestrator_auth import ROLE_CLI, ROLE_GATEWAY, mint
|
||||
from bot_bottle.orchestrator import api as orchestrator_api
|
||||
from bot_bottle.orchestrator.broker import StubBroker
|
||||
from bot_bottle.orchestrator.server import MAX_BODY_BYTES, dispatch, make_server
|
||||
from bot_bottle.orchestrator.http_contract import ORCHESTRATOR_AUTH_HEADER
|
||||
from bot_bottle.orchestrator.server import MAX_BODY_BYTES, create_app, make_server
|
||||
from bot_bottle.orchestrator.store.registry_store import BottleRecord, RegistryStore
|
||||
from bot_bottle.orchestrator.service import OrchestratorCore
|
||||
from bot_bottle.orchestrator.store.store_manager import StoreManager
|
||||
@@ -38,6 +43,47 @@ def _body(obj: object) -> bytes:
|
||||
return json.dumps(obj).encode()
|
||||
|
||||
|
||||
def dispatch(
|
||||
orchestrator: OrchestratorCore,
|
||||
method: str,
|
||||
path: str,
|
||||
body: bytes,
|
||||
*,
|
||||
role: str | None = ROLE_CLI,
|
||||
) -> tuple[int, dict[str, object]]:
|
||||
"""Exercise the real ASGI application without a network socket."""
|
||||
key = "in-process-dispatch-key"
|
||||
headers = {"content-type": "application/json"}
|
||||
if role is not None:
|
||||
headers[ORCHESTRATOR_AUTH_HEADER] = mint(role, key)
|
||||
|
||||
async def request() -> httpx.Response:
|
||||
transport = httpx.ASGITransport(app=create_app(orchestrator, signing_key=key))
|
||||
async with httpx.AsyncClient(
|
||||
transport=transport,
|
||||
base_url="http://orchestrator",
|
||||
follow_redirects=True,
|
||||
) as client:
|
||||
return await client.request(
|
||||
method, path, content=body, headers=headers,
|
||||
)
|
||||
|
||||
response = asyncio.run(request())
|
||||
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
|
||||
|
||||
|
||||
def _orchestrator(db_path: Path) -> OrchestratorCore:
|
||||
store = RegistryStore(db_path)
|
||||
store.migrate()
|
||||
@@ -361,6 +407,76 @@ class TestServerRoundTrip(unittest.TestCase):
|
||||
self.assertNotIn("SENSITIVE", output)
|
||||
|
||||
|
||||
class TestControlPlaneBoundary(unittest.IsolatedAsyncioTestCase):
|
||||
@staticmethod
|
||||
def _scope(key: str) -> dict[str, object]:
|
||||
return {
|
||||
"type": "http",
|
||||
"asgi": {"version": "3.0"},
|
||||
"http_version": "1.1",
|
||||
"method": "POST",
|
||||
"scheme": "http",
|
||||
"path": "/bottles",
|
||||
"raw_path": b"/bottles",
|
||||
"query_string": b"",
|
||||
"headers": [(
|
||||
ORCHESTRATOR_AUTH_HEADER.encode(),
|
||||
mint(ROLE_CLI, key).encode(),
|
||||
)],
|
||||
"client": ("127.0.0.1", 1),
|
||||
"server": ("127.0.0.1", 80),
|
||||
"state": {},
|
||||
}
|
||||
|
||||
async def test_chunked_oversized_body_returns_413(self) -> None:
|
||||
key = "stream-limit-key"
|
||||
called = False
|
||||
first: dict[str, object] = {
|
||||
"type": "http.request",
|
||||
"body": b"x" * MAX_BODY_BYTES,
|
||||
"more_body": True,
|
||||
}
|
||||
last: dict[str, object] = {
|
||||
"type": "http.request", "body": b"x", "more_body": False,
|
||||
}
|
||||
chunks: Iterator[dict[str, object]] = iter([first, last])
|
||||
sent: list[dict[str, object]] = []
|
||||
|
||||
async def receive() -> dict[str, object]:
|
||||
return next(chunks)
|
||||
|
||||
async def send(message: dict[str, object]) -> None:
|
||||
sent.append(message)
|
||||
|
||||
async def inner(*_args: object) -> None:
|
||||
nonlocal called
|
||||
called = True
|
||||
|
||||
boundary = orchestrator_api.ControlPlaneBoundary(inner, key)
|
||||
await boundary(self._scope(key), receive, send) # type: ignore[arg-type]
|
||||
self.assertFalse(called)
|
||||
self.assertEqual(413, sent[0]["status"])
|
||||
|
||||
async def test_slow_stream_returns_408(self) -> None:
|
||||
key = "stream-timeout-key"
|
||||
sent: list[dict[str, object]] = []
|
||||
|
||||
async def receive() -> dict[str, object]:
|
||||
await asyncio.sleep(1)
|
||||
return {"type": "http.request", "body": b"", "more_body": False}
|
||||
|
||||
async def send(message: dict[str, object]) -> None:
|
||||
sent.append(message)
|
||||
|
||||
async def inner(*_args: object) -> None:
|
||||
self.fail("timed-out body reached the application")
|
||||
|
||||
boundary = orchestrator_api.ControlPlaneBoundary(inner, key)
|
||||
with patch.object(orchestrator_api, "REQUEST_BODY_TIMEOUT_SECONDS", 0.01):
|
||||
await boundary(self._scope(key), receive, send) # type: ignore[arg-type]
|
||||
self.assertEqual(408, sent[0]["status"])
|
||||
|
||||
|
||||
class TestOrchestratorAuth(unittest.TestCase):
|
||||
"""Role-scoped control-plane tokens (issue #400 / #469 review): every route
|
||||
but /health needs a valid token, and the token's role gates which routes it
|
||||
|
||||
Reference in New Issue
Block a user