From c472deafaf622fabbf83c09c67ba7dc2a4456a14 Mon Sep 17 00:00:00 2001 From: codex Date: Mon, 27 Jul 2026 02:24:06 +0000 Subject: [PATCH] fix(orchestrator): bound streamed request bodies --- bot_bottle/orchestrator/__init__.py | 4 +- bot_bottle/orchestrator/api.py | 58 ++++++++--- bot_bottle/orchestrator/http_contract.py | 7 +- bot_bottle/orchestrator/server.py | 38 -------- tests/unit/test_orchestrator_client.py | 26 +++++ tests/unit/test_orchestrator_server.py | 118 ++++++++++++++++++++++- 6 files changed, 194 insertions(+), 57 deletions(-) diff --git a/bot_bottle/orchestrator/__init__.py b/bot_bottle/orchestrator/__init__.py index 3a2c775d..4ff9c1cd 100644 --- a/bot_bottle/orchestrator/__init__.py +++ b/bot_bottle/orchestrator/__init__.py @@ -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", ] diff --git a/bot_bottle/orchestrator/api.py b/bot_bottle/orchestrator/api.py index a97eb108..f2aefa05 100644 --- a/bot_bottle/orchestrator/api.py +++ b/bot_bottle/orchestrator/api.py @@ -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: diff --git a/bot_bottle/orchestrator/http_contract.py b/bot_bottle/orchestrator/http_contract.py index 4d2aeece..07480ec5 100644 --- a/bot_bottle/orchestrator/http_contract.py +++ b/bot_bottle/orchestrator/http_contract.py @@ -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", +] diff --git a/bot_bottle/orchestrator/server.py b/bot_bottle/orchestrator/server.py index bbd5189b..0378ceb6 100644 --- a/bot_bottle/orchestrator/server.py +++ b/bot_bottle/orchestrator/server.py @@ -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", ] diff --git a/tests/unit/test_orchestrator_client.py b/tests/unit/test_orchestrator_client.py index 5ee3a84c..9b2aac83 100644 --- a/tests/unit/test_orchestrator_client.py +++ b/tests/unit/test_orchestrator_client.py @@ -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 diff --git a/tests/unit/test_orchestrator_server.py b/tests/unit/test_orchestrator_server.py index e159c84d..f5b28f33 100644 --- a/tests/unit/test_orchestrator_server.py +++ b/tests/unit/test_orchestrator_server.py @@ -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