From ed0f95f4457dd72355012ba35f129595e981330f Mon Sep 17 00:00:00 2001 From: claude Date: Sun, 26 Jul 2026 07:16:56 +0000 Subject: [PATCH 1/3] feat(orchestrator): host control server transport (#468) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Chunk 1 of the host-control-server stack: close the PRD's **transport** gap. Today LaunchBroker.submit(token) is an in-process method call from OrchestratorCore; this makes it a real out-of-process service reached over HTTP. - host_server.py: the host control server. A pure dispatch() (POST /broker verifies a signed token via the existing verify_request + _launch/_teardown path, GET /health) wrapped by a thin http.server adapter, mirroring orchestrator/server.py. Only the signed token crosses the wire; provenance/schema failures are fail-closed 401s that never touch the backend, a backend launch failure is a 502. - broker_client.py: BrokerClient — a drop-in submit(token) that POSTs the signed token to the host controller. A 401 re-raises as BrokerAuthError so the launch path's rollback is identical local or remote. - broker.py: SubmitBroker Protocol — the one method OrchestratorCore depends on, satisfied by both LaunchBroker and BrokerClient, so the core is unchanged (service.py annotation only). - __main__.py: wire `--broker http` behind the shared-secret env var (BOT_BOTTLE_BROKER_SECRET, hex) — a chunk-1 stopgap the durable TrustDomain key (chunk 2, #476) replaces. Tested: pure-dispatch cases, BrokerClient with HTTP mocked, and a real-socket sign -> POST -> verify -> act round-trip (incl. fail-closed forged token). pyright clean; pylint 9.86. Co-Authored-By: Claude Opus 4.8 --- bot_bottle/orchestrator/__main__.py | 37 ++- bot_bottle/orchestrator/broker.py | 12 + bot_bottle/orchestrator/broker_client.py | 119 +++++++++ bot_bottle/orchestrator/host_server.py | 243 ++++++++++++++++++ bot_bottle/orchestrator/service.py | 4 +- tests/unit/test_orchestrator_broker_client.py | 84 ++++++ tests/unit/test_orchestrator_host_server.py | 156 +++++++++++ 7 files changed, 645 insertions(+), 10 deletions(-) create mode 100644 bot_bottle/orchestrator/broker_client.py create mode 100644 bot_bottle/orchestrator/host_server.py create mode 100644 tests/unit/test_orchestrator_broker_client.py create mode 100644 tests/unit/test_orchestrator_host_server.py diff --git a/bot_bottle/orchestrator/__main__.py b/bot_bottle/orchestrator/__main__.py index 20fd4ec6..66ee895d 100644 --- a/bot_bottle/orchestrator/__main__.py +++ b/bot_bottle/orchestrator/__main__.py @@ -17,7 +17,9 @@ from pathlib import Path from .. import log from .store.store_manager import StoreManager -from .broker import LaunchBroker, StubBroker +from .broker import StubBroker, SubmitBroker +from .broker_client import BrokerClient +from .host_server import BROKER_SECRET_ENV, DEFAULT_PORT, broker_secret_from_env from .server import make_server from .docker_broker import DockerBroker from .store.registry_store import RegistryStore, default_db_path @@ -34,8 +36,13 @@ def main(argv: list[str] | None = None) -> int: help=f"registry DB path (default: {default_db_path()})", ) parser.add_argument( - "--broker", choices=("stub", "docker"), default="stub", - help="launch broker: 'stub' records requests; 'docker' runs containers", + "--broker", choices=("stub", "docker", "http"), default="stub", + help="launch broker: 'stub' records requests; 'docker' runs containers " + "in-process; 'http' relays signed requests to a host control server", + ) + parser.add_argument( + "--host-controller-url", default=f"http://127.0.0.1:{DEFAULT_PORT}", + help="host control server URL (used only with --broker http)", ) args = parser.parse_args(argv) @@ -47,11 +54,25 @@ def main(argv: list[str] | None = None) -> int: # operator reaches it over HTTP (never a second, disconnected DB). StoreManager(registry.db_path).migrate() - # An ephemeral signing secret ties the orchestrator (signer) to its - # broker (verifier). 'stub' records launches instead of starting - # anything; 'docker' runs real containers (firecracker drops in later). - secret = secrets.token_bytes(32) - broker: LaunchBroker = DockerBroker(secret) if args.broker == "docker" else StubBroker(secret) + # A signing secret ties the orchestrator (signer) to its broker (verifier). + # 'stub' records launches instead of starting anything; 'docker' runs real + # containers in-process; 'http' relays signed requests to a separate host + # control server, which verifies and launches. For 'stub'/'docker' the + # secret is ephemeral (signer and verifier share this process); for 'http' + # it must be the SAME secret the host controller holds, so it is read from + # the shared env var (the chunk-1 stand-in for out-of-band provisioning). + broker: SubmitBroker + if args.broker == "http": + secret = broker_secret_from_env() + if secret is None: + parser.error( + f"--broker http requires a shared signing secret in " + f"${BROKER_SECRET_ENV} (hex), matching the host control server" + ) + broker = BrokerClient(args.host_controller_url) + else: + secret = secrets.token_bytes(32) + broker = DockerBroker(secret) if args.broker == "docker" else StubBroker(secret) orchestrator = OrchestratorCore(registry, broker, secret) server = make_server(orchestrator, host=args.host, port=args.port) diff --git a/bot_bottle/orchestrator/broker.py b/bot_bottle/orchestrator/broker.py index 937cbc2c..1cd91687 100644 --- a/bot_bottle/orchestrator/broker.py +++ b/bot_bottle/orchestrator/broker.py @@ -29,6 +29,7 @@ import json import secrets import time from dataclasses import dataclass +from typing import Protocol _JWT_HEADER = {"alg": "HS256", "typ": "JWT"} _ALLOWED_OPS = ("launch", "teardown") @@ -123,6 +124,16 @@ def verify_request(token: str, secret: bytes) -> LaunchRequest: # --- the broker itself ------------------------------------------------------ +class SubmitBroker(Protocol): + """The single method `OrchestratorCore` depends on: verify a signed token and + perform its op, returning the verified request. Both the in-process + `LaunchBroker` and the out-of-process `BrokerClient` (which relays the token + to the host control server) satisfy it structurally, so the core is unchanged + whether the backend is local or a real host service.""" + + def submit(self, token: str) -> LaunchRequest: ... + + class LaunchBroker(abc.ABC): """Verifies a signed request came from the orchestrator, then performs the backend-native launch/teardown. Subclasses implement `_launch` / @@ -169,6 +180,7 @@ class StubBroker(LaunchBroker): __all__ = [ "BrokerAuthError", "LaunchRequest", + "SubmitBroker", "LaunchBroker", "StubBroker", "sign_request", diff --git a/bot_bottle/orchestrator/broker_client.py b/bot_bottle/orchestrator/broker_client.py new file mode 100644 index 00000000..79879def --- /dev/null +++ b/bot_bottle/orchestrator/broker_client.py @@ -0,0 +1,119 @@ +"""Orchestrator-side broker transport (issue #468, chunk 1). + +The signer's half of the launch-broker transport gap. `BrokerClient` satisfies +the exact `submit(token)` contract `OrchestratorCore` already depends on (see +`broker.SubmitBroker`), but instead of verifying and launching in-process it POSTs +the signed token to the host control server over HTTP (stdlib `urllib`, like +`orchestrator/client.py`). Because it is drop-in for that interface, wiring a real +out-of-process backend does not change the core: it still signs a request and +calls `submit()`; only the wire is new. + +A provenance/schema rejection from the host controller (HTTP 401) is re-raised as +the same `BrokerAuthError` the in-process broker raises, so the launch path's +rollback-on-failure (`OrchestratorCore.launch_bottle`) behaves identically whether +the broker is local or remote. +""" + +from __future__ import annotations + +import json +import urllib.error +import urllib.request + +from .broker import BrokerAuthError, LaunchRequest + +DEFAULT_TIMEOUT_SECONDS = 5.0 + + +class BrokerClientError(RuntimeError): + """A brokered launch/teardown could not be delivered to the host control + server: unreachable, or an unexpected status other than the fail-closed 401 + (which surfaces as `BrokerAuthError`, matching the in-process broker).""" + + +class BrokerClient: + """Drop-in `submit(token)` that relays a signed request to the host control + server. Holds no secret — provenance rides entirely in the signed token, so a + caller that can reach this client still cannot forge a launch.""" + + def __init__(self, base_url: str, *, timeout: float = DEFAULT_TIMEOUT_SECONDS) -> None: + self._base = base_url.rstrip("/") + self._timeout = timeout + + def submit(self, token: str) -> LaunchRequest: + """POST the signed token to the host controller and return the request it + verified and acted on. + + Raises `BrokerAuthError` on a fail-closed 401 (bad provenance/schema — + the same exception the in-process broker raises), or `BrokerClientError` + if the controller is unreachable, times out, or returns any other + non-success status.""" + data = json.dumps({"token": token}).encode() + req = urllib.request.Request( + f"{self._base}/broker", data=data, method="POST", + headers={"Content-Type": "application/json"}, + ) + try: + with urllib.request.urlopen(req, timeout=self._timeout) as resp: + return _request_from(_json_object(resp.read())) + except urllib.error.HTTPError as e: + detail = _error_detail(e) + if e.code == 401: + raise BrokerAuthError( + detail or "host controller rejected the request" + ) from e + raise BrokerClientError( + f"POST /broker: HTTP {e.code} {detail}".rstrip() + ) from e + except (urllib.error.URLError, TimeoutError, OSError, ValueError) as e: + raise BrokerClientError(f"POST /broker: {e}") from e + + +def _json_object(raw: bytes) -> dict[str, object]: + """Parse a JSON object, tolerating an empty or malformed body (→ {}), like + the orchestrator client — a bad body becomes a clean 'missing field' error + downstream rather than an opaque JSON crash.""" + if not raw: + return {} + try: + obj = json.loads(raw) + except ValueError: + return {} + return obj if isinstance(obj, dict) else {} + + +def _error_detail(e: urllib.error.HTTPError) -> str: + """The `error` string from a structured error response, best-effort — an + error body may be absent or unreadable, in which case there is no detail.""" + try: + detail = _json_object(e.read()).get("error", "") + except Exception: # noqa: BLE001 — the error body is advisory only + return "" + return detail if isinstance(detail, str) else "" + + +def _request_from(payload: dict[str, object]) -> LaunchRequest: + """Reconstruct the verified `LaunchRequest` the controller echoed, so the + returned value matches the in-process broker's (which returns the request it + acted on). A missing op/bottle_id means a malformed response.""" + op = payload.get("op") + bottle_id = payload.get("bottle_id") + if not isinstance(op, str) or not isinstance(bottle_id, str) or not bottle_id: + raise BrokerClientError("host controller response missing op/bottle_id") + source_ip = payload.get("source_ip") + image_ref = payload.get("image_ref") + slot = payload.get("slot") + return LaunchRequest( + op=op, + bottle_id=bottle_id, + source_ip=source_ip if isinstance(source_ip, str) else "", + image_ref=image_ref if isinstance(image_ref, str) else "", + slot=slot if isinstance(slot, int) and not isinstance(slot, bool) else None, + ) + + +__all__ = [ + "BrokerClient", + "BrokerClientError", + "DEFAULT_TIMEOUT_SECONDS", +] diff --git a/bot_bottle/orchestrator/host_server.py b/bot_bottle/orchestrator/host_server.py new file mode 100644 index 00000000..e974f887 --- /dev/null +++ b/bot_bottle/orchestrator/host_server.py @@ -0,0 +1,243 @@ +"""Host control server (issue #468) — the launch broker as a real host service. + +Chunk 1 of the host-control-server stack closes the **transport** gap the PRD +opens with: today `LaunchBroker.submit(token)` is an in-process method call from +`OrchestratorCore`, and a real host service needs it reachable over the wire. +This module is that service — the single privileged host component — reached over +**HTTP** (the universal transport 0070 chose), mirroring the orchestrator control +plane's shape (`orchestrator/server.py`): a pure `dispatch()` for socket-free +testing, wrapped by a thin stdlib `http.server` adapter. + + GET /health -> 200 {"status": "ok"} + POST /broker -> 200 {"op", "bottle_id", "source_ip", "image_ref", "slot"} + 400 (bad body) | 401 (bad provenance/schema) | 502 (backend) + body: {"token": ""} + +Only the **signed token** crosses the wire; the server holds the shared HS256 +secret and a real `LaunchBroker` (e.g. `DockerBroker`) and runs the existing +`verify_request` + `_launch`/`_teardown` path behind the endpoint, so nothing +free-form ever reaches it. Provenance/schema failures are fail-closed 401s that +never touch the backend (`LaunchBroker.submit` verifies before acting), and a +backend launch failure is a 502 the caller must surface — neither takes the +controller down. + +The signed launch token *is* the endpoint's authentication (its provenance is the +whole point of the JWS), so `/broker` needs no separate caller credential; the +host controller's own lifecycle endpoints, which do, arrive with the durable +`TrustDomain` key in a later chunk. + +The shared signing secret is read from `$BOT_BOTTLE_BROKER_SECRET` (hex). That is +a **chunk-1 stopgap**: it must be provisioned to signer and verifier out of band, +which is exactly what the durable `TrustDomain` key in chunk 2 (#476) replaces. +""" + +from __future__ import annotations + +import argparse +import http.server +import json +import os +import socketserver +import sys +import typing +from urllib.parse import urlsplit + +from .. import log +from .broker import BrokerAuthError, LaunchBroker +from .docker_broker import DockerBroker + +# JSON body payload type (parsed request / rendered response). +Json = dict[str, object] + +# The hex-encoded HS256 secret shared with the request signer (the orchestrator). +# Chunk-1 stopgap for the durable, out-of-band `TrustDomain` key of chunk 2. +BROKER_SECRET_ENV = "BOT_BOTTLE_BROKER_SECRET" + +# Default host-controller port. Distinct from the orchestrator control plane +# (8099) — a separate privileged component listening on its own socket. +DEFAULT_PORT = 8091 + + +def _parse_json_object(body: bytes) -> Json: + """Parse a JSON object body. Raises ValueError for non-objects / bad JSON.""" + if not body: + return {} + obj = json.loads(body) # raises json.JSONDecodeError (a ValueError) + if not isinstance(obj, dict): + raise ValueError("request body must be a JSON object") + return obj + + +def broker_secret_from_env(environ: typing.Mapping[str, str] | None = None) -> bytes | None: + """The shared HS256 secret from `$BOT_BOTTLE_BROKER_SECRET` (hex), or None + when unset or not valid hex. The signer (orchestrator, `--broker http`) and + the verifier (this server) read the same env var so both hold the same key — + the chunk-1 stand-in for out-of-band provisioning.""" + env = os.environ if environ is None else environ + raw = env.get(BROKER_SECRET_ENV, "").strip() + if not raw: + return None + try: + return bytes.fromhex(raw) + except ValueError: + return None + + +def dispatch( # pylint: disable=too-many-return-statements + broker: LaunchBroker, method: str, path: str, body: bytes, +) -> tuple[int, Json]: + """Route one host-control request to a (status, payload) pair. Pure — the + only side effect is the broker's own backend launch — so routing is testable + without a socket. + + Total by design: a provenance/schema failure becomes 401 and a backend launch + failure becomes 502 rather than raising, so one bad request can neither act + on the backend nor take the controller down for the next caller.""" + route = urlsplit(path).path.rstrip("/") or "/" + + if method == "GET" and route == "/health": + return 200, {"status": "ok"} + + if method == "POST" and route == "/broker": + try: + data = _parse_json_object(body) + except ValueError as e: + return 400, {"error": f"invalid JSON: {e}"} + token = data.get("token") + if not isinstance(token, str) or not token: + return 400, {"error": "token (string) is required"} + try: + req = broker.submit(token) + except BrokerAuthError as e: + # Fail-closed: bad signature, malformed token, or off-schema payload. + # `submit` verifies before acting, so nothing was launched. + return 401, {"error": f"broker auth failed: {e}"} + except Exception as e: # noqa: BLE001 — a backend launch failure (docker + # down, image gone) is operational, not a control-plane bug; the + # caller must see it as a distinct 502, and the server must stay up. + return 502, {"error": f"backend launch failed: {e}"} + return 200, { + "op": req.op, + "bottle_id": req.bottle_id, + "source_ip": req.source_ip, + "image_ref": req.image_ref, + "slot": req.slot, + } + + return 404, {"error": "not found"} + + +class Handler(http.server.BaseHTTPRequestHandler): + """Thin stdlib adapter: read the body, call `dispatch`, write JSON.""" + + # Quiet by default; opt back into stdlib access logging with + # BOT_BOTTLE_HOST_CONTROLLER_DEBUG (the controller has its own logging). + def log_message(self, format: str, *args: typing.Any) -> None: # noqa: A002 + if os.environ.get("BOT_BOTTLE_HOST_CONTROLLER_DEBUG"): + super().log_message(format, *args) + + def _serve(self, method: str) -> None: + """Read the request body, dispatch it, and write the JSON reply. A + dispatch that raises (it shouldn't — dispatch is total) still returns a + 500 rather than dropping the connection.""" + server = self.server + assert isinstance(server, HostControlServer) + length = int(self.headers.get("Content-Length") or 0) + body = self.rfile.read(length) if length > 0 else b"" + try: + status, payload = dispatch(server.broker, method, self.path, body) + except Exception as e: # noqa: BLE001 — the controller must stay up + sys.stderr.write(f"host controller: {method} {self.path} failed: {e!r}\n") + sys.stderr.flush() + status, payload = 500, {"error": f"internal error: {e}"} + data = json.dumps(payload).encode() + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(data))) + self.end_headers() + self.wfile.write(data) + + def do_GET(self) -> None: + self._serve("GET") + + def do_POST(self) -> None: + self._serve("POST") + + +class HostControlServer(socketserver.ThreadingMixIn, http.server.HTTPServer): + """Threading HTTP server that carries the launch broker for its handlers. + + The broker holds the shared signing secret and performs the backend-native + launch/teardown; the server itself keeps no secret of its own — provenance + rides entirely in each request's signed token.""" + + daemon_threads = True + allow_reuse_address = True + + def __init__(self, address: tuple[str, int], broker: LaunchBroker) -> None: + self.broker = broker + super().__init__(address, Handler) + + +def make_host_server( + broker: LaunchBroker, host: str = "127.0.0.1", port: int = DEFAULT_PORT +) -> HostControlServer: + """Build (but do not start) a host control server. `port=0` binds an + ephemeral port — read `server.server_address` for the actual one.""" + return HostControlServer((host, port), broker) + + +def main(argv: list[str] | None = None) -> int: + """Run the host control server as a plain process (dev-harness). + + python -m bot_bottle.orchestrator.host_server [--host H] [--port P] + + Fail-closed: without a shared `$BOT_BOTTLE_BROKER_SECRET` the server can + verify no request's provenance, so it refuses to start rather than run a + launcher that accepts unsigned input.""" + parser = argparse.ArgumentParser(prog="bot_bottle.orchestrator.host_server") + parser.add_argument("--host", default="127.0.0.1", help="bind address") + parser.add_argument("--port", type=int, default=DEFAULT_PORT, help="bind port (0 = ephemeral)") + args = parser.parse_args(argv) + + secret = broker_secret_from_env() + if secret is None: + sys.stderr.write( + f"host controller: refusing to start without a shared signing secret " + f"(${BROKER_SECRET_ENV}, hex) — it could verify no request's " + "provenance and would relay unsigned launches to the backend\n" + ) + sys.stderr.flush() + return 2 + + broker = DockerBroker(secret) + server = make_host_server(broker, host=args.host, port=args.port) + bound_host, bound_port = server.server_address[0], server.server_address[1] + log.info( + "host control server listening", + context={"host": bound_host, "port": bound_port}, + ) + try: + server.serve_forever() + except KeyboardInterrupt: + log.info("host controller shutting down") + finally: + server.server_close() + return 0 + + +__all__ = [ + "dispatch", + "Handler", + "HostControlServer", + "make_host_server", + "broker_secret_from_env", + "main", + "Json", + "BROKER_SECRET_ENV", + "DEFAULT_PORT", +] + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/bot_bottle/orchestrator/service.py b/bot_bottle/orchestrator/service.py index a43d9b14..96ca9f19 100644 --- a/bot_bottle/orchestrator/service.py +++ b/bot_bottle/orchestrator/service.py @@ -25,7 +25,7 @@ import json from collections.abc import Iterable from datetime import datetime, timezone -from .broker import LaunchBroker, LaunchRequest, sign_request +from .broker import LaunchRequest, SubmitBroker, sign_request from .store.registry_store import DEFAULT_REAP_GRACE_SECONDS, BottleRecord, RegistryStore from .supervisor import ( AuditEntry, @@ -62,7 +62,7 @@ class OrchestratorCore: def __init__( self, registry: RegistryStore, - broker: LaunchBroker, + broker: SubmitBroker, sign_secret: bytes, supervisor: Supervisor | None = None, ) -> None: diff --git a/tests/unit/test_orchestrator_broker_client.py b/tests/unit/test_orchestrator_broker_client.py new file mode 100644 index 00000000..e85d51d2 --- /dev/null +++ b/tests/unit/test_orchestrator_broker_client.py @@ -0,0 +1,84 @@ +"""Unit: orchestrator-side broker client (issue #468, chunk 1). HTTP mocked.""" + +from __future__ import annotations + +import io +import json +import unittest +import urllib.error +from unittest.mock import MagicMock, patch + +from bot_bottle.orchestrator.broker import BrokerAuthError, LaunchRequest +from bot_bottle.orchestrator.broker_client import BrokerClient, BrokerClientError + +_URLOPEN = "bot_bottle.orchestrator.broker_client.urllib.request.urlopen" + + +def _resp(payload: object) -> MagicMock: + m = MagicMock() + m.__enter__.return_value.read.return_value = json.dumps(payload).encode() + return m + + +def _http_error(code: int, payload: object = None) -> urllib.error.HTTPError: + body = json.dumps(payload).encode() if payload is not None else b"" + return urllib.error.HTTPError( + "http://host/broker", code, "err", {}, io.BytesIO(body)) # type: ignore[arg-type] + + +class TestSubmit(unittest.TestCase): + def setUp(self) -> None: + self.c = BrokerClient("http://host:8091") + + def test_returns_the_verified_request(self) -> None: + echo = { + "op": "launch", "bottle_id": "b1", "source_ip": "10.0.0.1", + "image_ref": "img", "slot": 3, + } + with patch(_URLOPEN, return_value=_resp(echo)): + got = self.c.submit("tok") + self.assertEqual( + LaunchRequest(op="launch", bottle_id="b1", source_ip="10.0.0.1", + image_ref="img", slot=3), + got, + ) + + def test_posts_token_to_broker_endpoint(self) -> None: + with patch(_URLOPEN, return_value=_resp({"op": "teardown", "bottle_id": "b1"})) as m: + self.c.submit("signed-token") + request = m.call_args.args[0] + self.assertEqual("POST", request.get_method()) + self.assertTrue(request.full_url.endswith("/broker")) + self.assertEqual({"token": "signed-token"}, json.loads(request.data)) + + def test_401_raises_broker_auth_error(self) -> None: + # A fail-closed provenance/schema rejection surfaces as the SAME exception + # the in-process broker raises, so the launch path's rollback is identical. + with patch(_URLOPEN, side_effect=_http_error(401, {"error": "bad signature"})): + with self.assertRaises(BrokerAuthError): + self.c.submit("forged") + + def test_502_raises_broker_client_error(self) -> None: + with patch(_URLOPEN, side_effect=_http_error(502, {"error": "docker down"})): + with self.assertRaises(BrokerClientError): + self.c.submit("tok") + + def test_unreachable_raises_broker_client_error(self) -> None: + with patch(_URLOPEN, side_effect=urllib.error.URLError("refused")): + with self.assertRaises(BrokerClientError): + self.c.submit("tok") + + def test_malformed_success_body_raises(self) -> None: + with patch(_URLOPEN, return_value=_resp({"op": "launch"})): # missing bottle_id + with self.assertRaises(BrokerClientError): + self.c.submit("tok") + + def test_empty_error_body_is_tolerated(self) -> None: + # An error with no readable JSON body still classifies by status code. + with patch(_URLOPEN, side_effect=_http_error(401)): + with self.assertRaises(BrokerAuthError): + self.c.submit("forged") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/test_orchestrator_host_server.py b/tests/unit/test_orchestrator_host_server.py new file mode 100644 index 00000000..3428280a --- /dev/null +++ b/tests/unit/test_orchestrator_host_server.py @@ -0,0 +1,156 @@ +"""Unit tests for the host control server (issue #468, chunk 1). + +Mostly exercises the pure `dispatch()` (socket-free, like the orchestrator +server tests), plus a real-socket round-trip through `BrokerClient` that proves +the full sign -> POST -> verify -> act seam over HTTP. +""" + +from __future__ import annotations + +import json +import secrets +import threading +import unittest + +from bot_bottle.orchestrator.broker import ( + BrokerAuthError, + LaunchBroker, + LaunchRequest, + StubBroker, + sign_request, +) +from bot_bottle.orchestrator.broker_client import BrokerClient +from bot_bottle.orchestrator.host_server import ( + broker_secret_from_env, + dispatch, + make_host_server, +) + + +def _body(obj: object) -> bytes: + return json.dumps(obj).encode() + + +class _RaisingBroker(LaunchBroker): + """A broker whose backend launch always fails — exercises the 502 path (an + operational backend failure, distinct from a fail-closed provenance 401).""" + + def _launch(self, req: LaunchRequest) -> None: + raise RuntimeError("docker down") + + def _teardown(self, req: LaunchRequest) -> None: + raise RuntimeError("docker down") + + +class TestDispatch(unittest.TestCase): + def setUp(self) -> None: + self.secret = secrets.token_bytes(16) + self.broker = StubBroker(self.secret) + + def _token(self, **kwargs: object) -> str: + return sign_request(LaunchRequest(**kwargs), self.secret) # type: ignore[arg-type] + + def test_health(self) -> None: + status, payload = dispatch(self.broker, "GET", "/health", b"") + self.assertEqual(200, status) + self.assertEqual("ok", payload["status"]) + + def test_broker_launch_verifies_and_acts(self) -> None: + token = self._token( + op="launch", bottle_id="b1", source_ip="10.243.0.1", + image_ref="img", slot=2, + ) + status, payload = dispatch(self.broker, "POST", "/broker", _body({"token": token})) + self.assertEqual(200, status) + self.assertEqual("launch", payload["op"]) + self.assertEqual("b1", payload["bottle_id"]) + self.assertEqual("img", payload["image_ref"]) + self.assertEqual(2, payload["slot"]) + self.assertEqual(["b1"], [r.bottle_id for r in self.broker.launched]) + + def test_broker_teardown_acts(self) -> None: + token = self._token(op="teardown", bottle_id="b1") + status, _ = dispatch(self.broker, "POST", "/broker", _body({"token": token})) + self.assertEqual(200, status) + self.assertEqual(["b1"], [r.bottle_id for r in self.broker.torn_down]) + + def test_forged_token_is_401_and_nothing_acted(self) -> None: + forged = sign_request( + LaunchRequest(op="launch", bottle_id="b1"), secrets.token_bytes(16)) + status, payload = dispatch(self.broker, "POST", "/broker", _body({"token": forged})) + self.assertEqual(401, status) + self.assertIn("broker auth failed", str(payload["error"])) + self.assertEqual([], self.broker.launched) # fail-closed: never launched + + def test_backend_failure_is_502(self) -> None: + broker = _RaisingBroker(self.secret) + token = self._token(op="launch", bottle_id="b1", image_ref="img") + status, payload = dispatch(broker, "POST", "/broker", _body({"token": token})) + self.assertEqual(502, status) + self.assertIn("backend launch failed", str(payload["error"])) + + def test_missing_token_is_400(self) -> None: + status, _ = dispatch(self.broker, "POST", "/broker", _body({})) + self.assertEqual(400, status) + + def test_bad_json_is_400(self) -> None: + status, _ = dispatch(self.broker, "POST", "/broker", b"{not json") + self.assertEqual(400, status) + + def test_unknown_route_404(self) -> None: + status, _ = dispatch(self.broker, "GET", "/nope", b"") + self.assertEqual(404, status) + + def test_trailing_slash_normalized(self) -> None: + status, _ = dispatch(self.broker, "GET", "/health/", b"") + self.assertEqual(200, status) + + +class TestBrokerSecretFromEnv(unittest.TestCase): + def test_reads_hex_secret(self) -> None: + s = secrets.token_bytes(16) + self.assertEqual(s, broker_secret_from_env({"BOT_BOTTLE_BROKER_SECRET": s.hex()})) + + def test_unset_is_none(self) -> None: + self.assertIsNone(broker_secret_from_env({})) + + def test_invalid_hex_is_none(self) -> None: + self.assertIsNone(broker_secret_from_env({"BOT_BOTTLE_BROKER_SECRET": "not-hex"})) + + +class TestSeamRoundTrip(unittest.TestCase): + """The whole point of chunk 1: a request signed by the orchestrator side is + POSTed to a real host control server, verified there, and acted on — over + HTTP, not an in-process call.""" + + def _serve(self, broker: LaunchBroker) -> BrokerClient: + server = make_host_server(broker, "127.0.0.1", 0) + self.addCleanup(server.server_close) + threading.Thread(target=server.serve_forever, daemon=True).start() + self.addCleanup(server.shutdown) + host, port = server.server_address[0], server.server_address[1] + return BrokerClient(f"http://{host}:{port}") + + def test_sign_post_verify_act_over_http(self) -> None: + secret = secrets.token_bytes(16) + broker = StubBroker(secret) + client = self._serve(broker) + req = LaunchRequest( + op="launch", bottle_id="b1", source_ip="10.0.0.1", image_ref="img", slot=1) + got = client.submit(sign_request(req, secret)) + self.assertEqual(req, got) # the controller echoes the verified request + self.assertEqual(["b1"], [r.bottle_id for r in broker.launched]) + + def test_forged_token_raises_broker_auth_error_over_http(self) -> None: + secret = secrets.token_bytes(16) + broker = StubBroker(secret) + client = self._serve(broker) + forged = sign_request( + LaunchRequest(op="launch", bottle_id="b1"), secrets.token_bytes(16)) + with self.assertRaises(BrokerAuthError): + client.submit(forged) + self.assertEqual([], broker.launched) # fail-closed across the wire + + +if __name__ == "__main__": + unittest.main() -- 2.52.0 From ec953ceda719b87445acb206313e5b5d6bdf8f9e Mon Sep 17 00:00:00 2001 From: claude Date: Sun, 26 Jul 2026 08:43:00 +0000 Subject: [PATCH 2/3] fix(orchestrator): address review on host control server transport (#468) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Codex review on #496: - **High — ambiguous delivery no longer orphans a launched bottle.** A timeout / dropped response from the host controller is now the ambiguous BrokerUnavailableError (distinct from the definite BrokerAuthError / BrokerClientError). OrchestratorCore.launch_bottle keeps the registry row on the ambiguous case instead of deregistering — deregistering would orphan a running container with no record (reconcile reaps rows, never containers). The row is left for reconcile to reap iff the bottle is not actually live. Definite failures still roll back, so a real failure leaves no orphan row. - **Medium — the privileged endpoint bounds request bodies.** The host server rejects an oversized Content-Length with 413 before reading it, and sets a per-request socket timeout, so a caller that can merely reach the socket (no signed token) can't exhaust memory or a handler thread. Tests: ambiguous-keep vs definite-rollback in the launch path; the BrokerUnavailableError/BrokerClientError split in BrokerClient; the 413 body cap + handler error paths (driven in-thread, since daemon request threads lose coverage) plus a deterministic real-socket check that declares an oversized Content-Length but sends a sliver (rejection on the header, no unread-body reset race); and the __main__ entrypoint broker selection. Diff-coverage 98%; pyright clean; pylint 9.8. Co-Authored-By: Claude Opus 4.8 --- bot_bottle/orchestrator/broker.py | 17 ++- bot_bottle/orchestrator/broker_client.py | 25 ++-- bot_bottle/orchestrator/host_server.py | 35 ++++- bot_bottle/orchestrator/service.py | 23 +++- tests/unit/test_orchestrator_broker_client.py | 41 +++++- tests/unit/test_orchestrator_host_server.py | 129 ++++++++++++++++++ tests/unit/test_orchestrator_main.py | 64 +++++++++ tests/unit/test_orchestrator_service.py | 36 ++++- 8 files changed, 340 insertions(+), 30 deletions(-) create mode 100644 tests/unit/test_orchestrator_main.py diff --git a/bot_bottle/orchestrator/broker.py b/bot_bottle/orchestrator/broker.py index 1cd91687..6d5c567f 100644 --- a/bot_bottle/orchestrator/broker.py +++ b/bot_bottle/orchestrator/broker.py @@ -38,7 +38,21 @@ _ALLOWED_OPS = ("launch", "teardown") class BrokerAuthError(Exception): """A broker request failed provenance or schema verification — bad/absent signature, malformed token, or a payload that doesn't match - the fixed launch-request shape. Fail-closed: the broker must not act.""" + the fixed launch-request shape. Fail-closed: the broker must not act. + + A **definite** negative: nothing was launched, so a caller may safely roll + back as if the op never happened.""" + + +class BrokerUnavailableError(Exception): + """A brokered request could not be carried to a verdict: the broker (or the + wire to it) was unreachable, timed out, or dropped the response. + + Crucially **ambiguous** — unlike `BrokerAuthError`, the op MAY already have + taken effect on the backend before the response was lost, so a caller must + NOT assume it did nothing (e.g. must not roll a registry row back as if no + launch happened, which would orphan a running container). Only the in-process + brokers never raise this; the out-of-process `BrokerClient` does.""" @dataclass(frozen=True) @@ -179,6 +193,7 @@ class StubBroker(LaunchBroker): __all__ = [ "BrokerAuthError", + "BrokerUnavailableError", "LaunchRequest", "SubmitBroker", "LaunchBroker", diff --git a/bot_bottle/orchestrator/broker_client.py b/bot_bottle/orchestrator/broker_client.py index 79879def..62783418 100644 --- a/bot_bottle/orchestrator/broker_client.py +++ b/bot_bottle/orchestrator/broker_client.py @@ -20,15 +20,17 @@ import json import urllib.error import urllib.request -from .broker import BrokerAuthError, LaunchRequest +from .broker import BrokerAuthError, BrokerUnavailableError, LaunchRequest DEFAULT_TIMEOUT_SECONDS = 5.0 class BrokerClientError(RuntimeError): - """A brokered launch/teardown could not be delivered to the host control - server: unreachable, or an unexpected status other than the fail-closed 401 - (which surfaces as `BrokerAuthError`, matching the in-process broker).""" + """The host control server *responded*, but with an unexpected status other + than the fail-closed 401 (which surfaces as `BrokerAuthError`) — e.g. a 502 + backend failure or a malformed body. A definite negative: the host processed + the request and it did not launch. (A *no-response* failure — unreachable / + timeout / dropped — is the ambiguous `BrokerUnavailableError` instead.)""" class BrokerClient: @@ -45,9 +47,11 @@ class BrokerClient: verified and acted on. Raises `BrokerAuthError` on a fail-closed 401 (bad provenance/schema — - the same exception the in-process broker raises), or `BrokerClientError` - if the controller is unreachable, times out, or returns any other - non-success status.""" + the same exception the in-process broker raises); `BrokerClientError` if + the host *responds* with any other non-success status or a malformed + body (a definite negative); or `BrokerUnavailableError` if no response is + obtained (unreachable / timeout / dropped) — the **ambiguous** case, where + the host may already have acted, so the caller must not roll back.""" data = json.dumps({"token": token}).encode() req = urllib.request.Request( f"{self._base}/broker", data=data, method="POST", @@ -65,8 +69,11 @@ class BrokerClient: raise BrokerClientError( f"POST /broker: HTTP {e.code} {detail}".rstrip() ) from e - except (urllib.error.URLError, TimeoutError, OSError, ValueError) as e: - raise BrokerClientError(f"POST /broker: {e}") from e + except (urllib.error.URLError, TimeoutError, OSError) as e: + # No usable response — unreachable, timed out, or the connection + # dropped mid-exchange. Ambiguous: the request may already have + # launched the bottle, so this is NOT a definite failure. + raise BrokerUnavailableError(f"POST /broker: {e}") from e def _json_object(raw: bytes) -> dict[str, object]: diff --git a/bot_bottle/orchestrator/host_server.py b/bot_bottle/orchestrator/host_server.py index e974f887..fe8130a0 100644 --- a/bot_bottle/orchestrator/host_server.py +++ b/bot_bottle/orchestrator/host_server.py @@ -57,6 +57,16 @@ BROKER_SECRET_ENV = "BOT_BOTTLE_BROKER_SECRET" # (8099) — a separate privileged component listening on its own socket. DEFAULT_PORT = 8091 +# Cap on the request body. A signed broker request is tiny, so rejecting anything +# larger *before reading it* keeps a caller that can merely reach the socket (no +# signed token needed) from exhausting memory or a handler thread with a huge +# Content-Length — the signed token, not mere reachability, is the authority. +MAX_BODY_BYTES = 64 * 1024 + +# Per-request socket timeout, bounding how long a stalled / slow-loris caller can +# hold a handler thread on this privileged listener. +REQUEST_TIMEOUT_SECONDS = 15 + def _parse_json_object(body: bytes) -> Json: """Parse a JSON object body. Raises ValueError for non-objects / bad JSON.""" @@ -130,6 +140,10 @@ def dispatch( # pylint: disable=too-many-return-statements class Handler(http.server.BaseHTTPRequestHandler): """Thin stdlib adapter: read the body, call `dispatch`, write JSON.""" + # Socket timeout per request (applied by StreamRequestHandler.setup) so a + # stalled caller can't pin a handler thread on this privileged listener. + timeout = REQUEST_TIMEOUT_SECONDS + # Quiet by default; opt back into stdlib access logging with # BOT_BOTTLE_HOST_CONTROLLER_DEBUG (the controller has its own logging). def log_message(self, format: str, *args: typing.Any) -> None: # noqa: A002 @@ -137,12 +151,21 @@ class Handler(http.server.BaseHTTPRequestHandler): super().log_message(format, *args) def _serve(self, method: str) -> None: - """Read the request body, dispatch it, and write the JSON reply. A - dispatch that raises (it shouldn't — dispatch is total) still returns a - 500 rather than dropping the connection.""" + """Read the request body (bounded), dispatch it, and write the JSON + reply. A dispatch that raises (it shouldn't — dispatch is total) still + returns a 500 rather than dropping the connection.""" server = self.server assert isinstance(server, HostControlServer) - length = int(self.headers.get("Content-Length") or 0) + try: + length = int(self.headers.get("Content-Length") or 0) + except ValueError: + self._reply(400, {"error": "invalid Content-Length"}) + return + if length < 0 or length > MAX_BODY_BYTES: + # Reject before reading: nothing legitimate is this big, so an + # oversized declared length is a bug or a resource-exhaustion attempt. + self._reply(413, {"error": "request body too large"}) + return body = self.rfile.read(length) if length > 0 else b"" try: status, payload = dispatch(server.broker, method, self.path, body) @@ -150,6 +173,10 @@ class Handler(http.server.BaseHTTPRequestHandler): sys.stderr.write(f"host controller: {method} {self.path} failed: {e!r}\n") sys.stderr.flush() status, payload = 500, {"error": f"internal error: {e}"} + self._reply(status, payload) + + def _reply(self, status: int, payload: typing.Mapping[str, object]) -> None: + """Write one JSON response with an explicit Content-Length.""" data = json.dumps(payload).encode() self.send_response(status) self.send_header("Content-Type", "application/json") diff --git a/bot_bottle/orchestrator/service.py b/bot_bottle/orchestrator/service.py index 96ca9f19..70de69d7 100644 --- a/bot_bottle/orchestrator/service.py +++ b/bot_bottle/orchestrator/service.py @@ -25,7 +25,7 @@ import json from collections.abc import Iterable from datetime import datetime, timezone -from .broker import LaunchRequest, SubmitBroker, sign_request +from .broker import BrokerUnavailableError, LaunchRequest, SubmitBroker, sign_request from .store.registry_store import DEFAULT_REAP_GRACE_SECONDS, BottleRecord, RegistryStore from .supervisor import ( AuditEntry, @@ -111,14 +111,23 @@ class OrchestratorCore: image_ref=image_ref, slot=slot, ) - launched = False try: self._broker.submit(sign_request(req, self._secret)) - launched = True - finally: - if not launched: - self.registry.deregister(rec.bottle_id) - self._tokens.pop(rec.bottle_id, None) + except BrokerUnavailableError: + # Ambiguous delivery failure (timeout / dropped response): the broker + # may already have launched the bottle before the response was lost. + # Do NOT deregister — that would orphan a running container with no + # registry row (reconcile reaps rows, never containers). Keep the row + # so reconcile reaps it iff the bottle is not actually live; surface + # the error so the caller knows the launch is unconfirmed. + raise + except Exception: + # A definite failure — a fail-closed rejection, a backend launch + # error, or the host reporting it did not launch: nothing is running, + # so roll the registry entry back to leave no orphan. + self.registry.deregister(rec.bottle_id) + self._tokens.pop(rec.bottle_id, None) + raise return rec def teardown_bottle(self, bottle_id: str) -> bool: diff --git a/tests/unit/test_orchestrator_broker_client.py b/tests/unit/test_orchestrator_broker_client.py index e85d51d2..e0de4747 100644 --- a/tests/unit/test_orchestrator_broker_client.py +++ b/tests/unit/test_orchestrator_broker_client.py @@ -8,7 +8,11 @@ import unittest import urllib.error from unittest.mock import MagicMock, patch -from bot_bottle.orchestrator.broker import BrokerAuthError, LaunchRequest +from bot_bottle.orchestrator.broker import ( + BrokerAuthError, + BrokerUnavailableError, + LaunchRequest, +) from bot_bottle.orchestrator.broker_client import BrokerClient, BrokerClientError _URLOPEN = "bot_bottle.orchestrator.broker_client.urllib.request.urlopen" @@ -58,14 +62,25 @@ class TestSubmit(unittest.TestCase): with self.assertRaises(BrokerAuthError): self.c.submit("forged") - def test_502_raises_broker_client_error(self) -> None: + def test_502_is_a_definite_client_error(self) -> None: + # The host responded — it processed the request and did not launch, so a + # definite BrokerClientError (the caller may safely roll back). with patch(_URLOPEN, side_effect=_http_error(502, {"error": "docker down"})): with self.assertRaises(BrokerClientError): self.c.submit("tok") - def test_unreachable_raises_broker_client_error(self) -> None: + def test_unreachable_is_ambiguous_unavailable(self) -> None: + # No response at all — the request may already have launched, so the + # AMBIGUOUS BrokerUnavailableError (the caller must NOT roll back). with patch(_URLOPEN, side_effect=urllib.error.URLError("refused")): - with self.assertRaises(BrokerClientError): + with self.assertRaises(BrokerUnavailableError): + self.c.submit("tok") + + def test_timeout_is_ambiguous_unavailable(self) -> None: + # A dropped/late response after the request was sent is the exact orphan + # risk: the host may have launched. Must be ambiguous, not a definite fail. + with patch(_URLOPEN, side_effect=TimeoutError("read timed out")): + with self.assertRaises(BrokerUnavailableError): self.c.submit("tok") def test_malformed_success_body_raises(self) -> None: @@ -79,6 +94,24 @@ class TestSubmit(unittest.TestCase): with self.assertRaises(BrokerAuthError): self.c.submit("forged") + def test_non_json_success_body_raises(self) -> None: + # A 200 whose body isn't JSON is tolerated into {} then fails the + # missing-field check — a definite client error, not a crash. + m = MagicMock() + m.__enter__.return_value.read.return_value = b"not json at all" + with patch(_URLOPEN, return_value=m): + with self.assertRaises(BrokerClientError): + self.c.submit("tok") + + def test_unreadable_error_body_is_tolerated(self) -> None: + # An HTTPError whose body can't be read (fp=None) still classifies by + # status — the error detail is best-effort. + err = urllib.error.HTTPError( + "http://host/broker", 502, "err", {}, None) # type: ignore[arg-type] + with patch(_URLOPEN, side_effect=err): + with self.assertRaises(BrokerClientError): + self.c.submit("tok") + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/test_orchestrator_host_server.py b/tests/unit/test_orchestrator_host_server.py index 3428280a..0217a2c8 100644 --- a/tests/unit/test_orchestrator_host_server.py +++ b/tests/unit/test_orchestrator_host_server.py @@ -7,10 +7,14 @@ the full sign -> POST -> verify -> act seam over HTTP. from __future__ import annotations +import http.client +import io import json import secrets import threading +import typing import unittest +from unittest.mock import MagicMock, patch from bot_bottle.orchestrator.broker import ( BrokerAuthError, @@ -21,8 +25,12 @@ from bot_bottle.orchestrator.broker import ( ) from bot_bottle.orchestrator.broker_client import BrokerClient from bot_bottle.orchestrator.host_server import ( + MAX_BODY_BYTES, + Handler, + HostControlServer, broker_secret_from_env, dispatch, + main, make_host_server, ) @@ -97,6 +105,16 @@ class TestDispatch(unittest.TestCase): status, _ = dispatch(self.broker, "POST", "/broker", b"{not json") self.assertEqual(400, status) + def test_empty_body_is_missing_token_400(self) -> None: + # Empty body parses to {} (no token) → 400, never reaching the broker. + status, _ = dispatch(self.broker, "POST", "/broker", b"") + self.assertEqual(400, status) + self.assertEqual([], self.broker.launched) + + def test_non_object_body_is_400(self) -> None: + status, _ = dispatch(self.broker, "POST", "/broker", b"[1, 2]") + self.assertEqual(400, status) + def test_unknown_route_404(self) -> None: status, _ = dispatch(self.broker, "GET", "/nope", b"") self.assertEqual(404, status) @@ -152,5 +170,116 @@ class TestSeamRoundTrip(unittest.TestCase): self.assertEqual([], broker.launched) # fail-closed across the wire +class TestRequestLimits(unittest.TestCase): + """The privileged listener must not let a caller that can merely reach the + socket (no signed token) exhaust it via an oversized declared body — and it + rejects on the Content-Length *header*, before reading the body.""" + + def _addr(self) -> tuple[str, int]: + self.broker = StubBroker(secrets.token_bytes(16)) + server = make_host_server(self.broker, "127.0.0.1", 0) + self.addCleanup(server.server_close) + threading.Thread(target=server.serve_forever, daemon=True).start() + self.addCleanup(server.shutdown) + host, port = server.server_address[:2] + return typing.cast(str, host), port + + def test_oversized_content_length_is_rejected_before_reading(self) -> None: + host, port = self._addr() + conn = http.client.HTTPConnection(host, port, timeout=5) + self.addCleanup(conn.close) + # Declare an oversized body but send only a sliver: the server must reject + # on the header before reading, so the caller gets a clean, deterministic + # 413 (no large unread body to race a connection reset). + conn.putrequest("POST", "/broker", skip_accept_encoding=True) + conn.putheader("Content-Type", "application/json") + conn.putheader("Content-Length", str(MAX_BODY_BYTES + 1)) + conn.endheaders() + conn.send(b"{}") # far short of the declared length; never read + resp = conn.getresponse() + self.assertEqual(413, resp.status) + self.assertEqual([], self.broker.launched) # never reached the broker + + +class TestServeUnit(unittest.TestCase): + """Drive `Handler._serve` directly (no socket). The real per-request handler + runs in a daemon thread whose coverage/trace data is lost, so the + bounded-body and error paths are exercised here in the main thread instead.""" + + def _handler(self, broker: LaunchBroker, headers: dict[str, str], + body: bytes = b"") -> tuple[Handler, MagicMock]: + server = HostControlServer.__new__(HostControlServer) + server.broker = broker + h = Handler.__new__(Handler) + h.server = server + h.headers = headers # type: ignore[assignment] — dict is a valid .get() stand-in + h.path = "/broker" + h.rfile = io.BytesIO(body) + h.wfile = io.BytesIO() + send_response = MagicMock() + h.send_response = send_response # type: ignore[method-assign] + h.send_header = MagicMock() # type: ignore[method-assign] + h.end_headers = MagicMock() # type: ignore[method-assign] + return h, send_response + + def test_oversized_content_length_is_413(self) -> None: + broker = StubBroker(secrets.token_bytes(16)) + h, send_response = self._handler(broker, {"Content-Length": str(MAX_BODY_BYTES + 1)}) + h.do_POST() # exercises do_POST -> _serve + send_response.assert_called_once_with(413) + self.assertEqual([], broker.launched) # rejected before the broker + + def test_invalid_content_length_is_400(self) -> None: + h, send_response = self._handler(StubBroker(secrets.token_bytes(16)), + {"Content-Length": "not-a-number"}) + h._serve("POST") + send_response.assert_called_once_with(400) + + def test_valid_request_dispatches_200(self) -> None: + secret = secrets.token_bytes(16) + broker = StubBroker(secret) + body = _body({"token": sign_request( + LaunchRequest(op="teardown", bottle_id="b1"), secret)}) + h, send_response = self._handler(broker, {"Content-Length": str(len(body))}, body) + h._serve("POST") + send_response.assert_called_once_with(200) + self.assertEqual(["b1"], [r.bottle_id for r in broker.torn_down]) + + def test_dispatch_exception_becomes_500(self) -> None: + # dispatch is total, but the handler still guards it: a raised dispatch + # returns 500 rather than dropping the connection. + h, send_response = self._handler( + StubBroker(secrets.token_bytes(16)), {"Content-Length": "0"}) + with patch("bot_bottle.orchestrator.host_server.dispatch", + side_effect=RuntimeError("boom")): + h._serve("POST") + send_response.assert_called_once_with(500) + + def test_health_over_do_get(self) -> None: + h, send_response = self._handler(StubBroker(secrets.token_bytes(16)), {}) + h.path = "/health" + h.do_GET() + send_response.assert_called_once_with(200) + + +class TestMain(unittest.TestCase): + def test_fail_closed_without_secret(self) -> None: + with patch("bot_bottle.orchestrator.host_server.broker_secret_from_env", + return_value=None): + self.assertEqual(2, main(["--port", "0"])) + + def test_serves_then_shuts_down_cleanly(self) -> None: + fake = MagicMock() + fake.server_address = ("127.0.0.1", 0) + fake.serve_forever.side_effect = KeyboardInterrupt + with patch("bot_bottle.orchestrator.host_server.broker_secret_from_env", + return_value=b"k"), \ + patch("bot_bottle.orchestrator.host_server.make_host_server", + return_value=fake): + self.assertEqual(0, main(["--port", "0"])) + fake.serve_forever.assert_called_once() + fake.server_close.assert_called_once() + + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/test_orchestrator_main.py b/tests/unit/test_orchestrator_main.py new file mode 100644 index 00000000..366eaf12 --- /dev/null +++ b/tests/unit/test_orchestrator_main.py @@ -0,0 +1,64 @@ +"""Unit: the orchestrator dev-harness entrypoint (`python -m bot_bottle.orchestrator`). + +Exercises broker selection (stub / docker / http) and the fail-closed http path, +patching `make_server` so the serve loop returns instead of blocking. +""" + +from __future__ import annotations + +import os +import secrets +import tempfile +import unittest +from pathlib import Path +from unittest.mock import MagicMock, patch + +from bot_bottle.orchestrator.__main__ import main + + +def _fake_server() -> MagicMock: + fake = MagicMock() + fake.server_address = ("127.0.0.1", 0) + # Break out of serve_forever immediately, exercising the try/finally. + fake.serve_forever.side_effect = KeyboardInterrupt + return fake + + +class TestMain(unittest.TestCase): + def _run(self, broker: str, env: dict[str, str] | None = None) -> tuple[int, MagicMock]: + fake = _fake_server() + with tempfile.TemporaryDirectory() as d: + argv = ["--db", str(Path(d) / "r.db"), "--port", "0", "--broker", broker] + with patch("bot_bottle.orchestrator.__main__.make_server", return_value=fake), \ + patch.dict("os.environ", env or {}, clear=False): + if env is None: + os.environ.pop("BOT_BOTTLE_BROKER_SECRET", None) + rc = main(argv) + return rc, fake + + def test_stub_broker_serves_and_closes(self) -> None: + rc, fake = self._run("stub") + self.assertEqual(0, rc) + fake.serve_forever.assert_called_once() + fake.server_close.assert_called_once() + + def test_docker_broker_serves(self) -> None: + rc, _ = self._run("docker") + self.assertEqual(0, rc) + + def test_http_broker_with_secret_serves(self) -> None: + rc, _ = self._run( + "http", env={"BOT_BOTTLE_BROKER_SECRET": secrets.token_bytes(16).hex()}) + self.assertEqual(0, rc) + + def test_http_broker_without_secret_exits(self) -> None: + # Fail-closed: --broker http with no shared secret is a usage error. + with tempfile.TemporaryDirectory() as d: + with patch.dict("os.environ", {}, clear=False): + os.environ.pop("BOT_BOTTLE_BROKER_SECRET", None) + with self.assertRaises(SystemExit): + main(["--db", str(Path(d) / "r.db"), "--broker", "http"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/test_orchestrator_service.py b/tests/unit/test_orchestrator_service.py index 36931075..76d662c0 100644 --- a/tests/unit/test_orchestrator_service.py +++ b/tests/unit/test_orchestrator_service.py @@ -11,7 +11,12 @@ from contextlib import closing from pathlib import Path from unittest.mock import patch -from bot_bottle.orchestrator.broker import LaunchBroker, LaunchRequest, StubBroker +from bot_bottle.orchestrator.broker import ( + BrokerUnavailableError, + LaunchBroker, + LaunchRequest, + StubBroker, +) from bot_bottle.orchestrator.store.registry_store import RegistryStore from bot_bottle.orchestrator.service import OrchestratorCore from bot_bottle.orchestrator.store.secret_store import new_env_var_secret @@ -25,8 +30,8 @@ from bot_bottle.orchestrator.supervisor import ( class _FailingBroker(LaunchBroker): - """Verifies the token like any broker, then fails the launch — to - exercise the orchestrator's registry rollback.""" + """Verifies the token like any broker, then fails the launch *definitely* — + to exercise the orchestrator's registry rollback.""" def _launch(self, req: LaunchRequest) -> None: raise RuntimeError("launch failed") @@ -35,6 +40,18 @@ class _FailingBroker(LaunchBroker): pass +class _UnavailableBroker(LaunchBroker): + """Verifies the token, then raises the *ambiguous* BrokerUnavailableError — + the host may already have launched — so the orchestrator must KEEP the + registry row rather than orphan a running container.""" + + def _launch(self, req: LaunchRequest) -> None: + raise BrokerUnavailableError("delivery dropped after send") + + def _teardown(self, req: LaunchRequest) -> None: + pass + + class TestOrchestrator(unittest.TestCase): def setUp(self) -> None: self._tmp = tempfile.TemporaryDirectory() @@ -144,11 +161,20 @@ class TestOrchestrator(unittest.TestCase): self.assertIsNotNone(self.orch.resolve("10.243.0.3", rec.identity_token)) self.assertIsNone(self.orch.resolve("10.243.0.3", "wrong-token")) - def test_launch_rolls_back_registry_on_broker_failure(self) -> None: + def test_launch_rolls_back_registry_on_definite_broker_failure(self) -> None: orch = OrchestratorCore(self.store, _FailingBroker(self.secret), self.secret) with self.assertRaises(RuntimeError): orch.launch_bottle("10.243.0.9") - self.assertEqual([], self.store.all()) # no orphan + self.assertEqual([], self.store.all()) # no orphan row + + def test_launch_keeps_registry_on_ambiguous_broker_failure(self) -> None: + # The host may already have launched the bottle before the response was + # lost, so deregistering would orphan a running container with no row. + # The row is kept for reconcile to reap iff the bottle is not live. + orch = OrchestratorCore(self.store, _UnavailableBroker(self.secret), self.secret) + with self.assertRaises(BrokerUnavailableError): + orch.launch_bottle("10.243.0.9") + self.assertEqual(1, len(self.store.all())) # row survives — no orphan container def test_gateway_status_reports_unconfigured(self) -> None: # The orchestrator no longer owns a standalone gateway lifecycle; the -- 2.52.0 From 7a48ea2b0cb8b2b3266f439856deb6e9b3d15e7c Mon Sep 17 00:00:00 2001 From: claude Date: Sun, 26 Jul 2026 17:11:44 +0000 Subject: [PATCH 3/3] fix(secret): authenticate bottled-secret encryption (#468) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The per-bottle egress-secret encryption (secret_store.py) was unauthenticated CTR/XOR: decrypting with the WRONG ENV_VAR_SECRET produced garbage that decrypt_value only rejected when it wasn't valid UTF-8. For short token values that garbage is coincidentally valid UTF-8 ~5% of the time, so `reprovision_from_secret` would occasionally "succeed" with a wrong key and inject a garbage egress credential — and test_reprovision_rejects_missing_rows_and_wrong_key failed ~5% of runs (flaky CI, surfaced by this stack's unit job). Switch to authenticated encrypt-then-MAC: append an HMAC-SHA256 tag over `nonce || ciphertext`, keyed by a domain-separated MAC subkey derived from the ENV_VAR_SECRET. decrypt_value verifies the tag (constant-time) before returning any plaintext, so a wrong key or tampered ciphertext is rejected deterministically. Blob format is now `nonce || ciphertext || tag` (the stored rows are transient — re-written every launch — so no migration is needed). Pre-existing bug on main, unrelated to the transport work, but it blocks this stack's CI. Tests: wrong key rejected 200/200; tampered ciphertext rejected; round-trips unchanged. Deterministic now (was ~5% flaky). Co-Authored-By: Claude Opus 4.8 --- bot_bottle/orchestrator/store/secret_store.py | 70 +++++++++++++------ tests/unit/test_orchestrator_secret_store.py | 31 ++++---- 2 files changed, 66 insertions(+), 35 deletions(-) diff --git a/bot_bottle/orchestrator/store/secret_store.py b/bot_bottle/orchestrator/store/secret_store.py index 7daf0767..a6271c24 100644 --- a/bot_bottle/orchestrator/store/secret_store.py +++ b/bot_bottle/orchestrator/store/secret_store.py @@ -12,12 +12,22 @@ reattachment path reads ENV_VAR_SECRET from the running agent container via ``POST /bottles//reprovision_gateway``; the orchestrator decrypts the stored rows and re-populates ``_tokens``. -Encryption scheme: HMAC-SHA256 used as a PRF in CTR mode (stdlib-only, -no external deps). Each value is encrypted independently. The output blob is -``nonce (16 bytes) || ciphertext`` encoded as URL-safe base64 (no padding). +Encryption scheme: HMAC-SHA256 used as a PRF in CTR mode, **authenticated** +encrypt-then-MAC (stdlib-only, no external deps). Each value is encrypted +independently. The output blob is ``nonce (16 bytes) || ciphertext || tag +(32 bytes)`` encoded as URL-safe base64 (no padding). keystream_block_i = HMAC-SHA256(key, nonce || i.to_bytes(4, "big")) ciphertext_i = plaintext_i XOR keystream_block_i[:len(plaintext_i)] + mac_key = HMAC-SHA256(key, "bottled-secret-mac-v1") + tag = HMAC-SHA256(mac_key, nonce || ciphertext) + +The tag is what makes a **wrong key deterministically detectable**: without it, +CTR decryption with the wrong key yields garbage that only fails when it isn't +valid UTF-8 (so ``reprovision`` would sometimes "succeed" with a wrong +ENV_VAR_SECRET and inject garbage egress tokens). The MAC key is derived from +the ENV_VAR_SECRET by a domain-separated HMAC so the same key never both +generates the keystream and signs the tag with the same message shape. """ from __future__ import annotations @@ -29,6 +39,7 @@ import secrets _KEY_BYTES = 32 # 256-bit key from ENV_VAR_SECRET _NONCE_BYTES = 16 # 128-bit random nonce per encrypt call +_TAG_BYTES = 32 # HMAC-SHA256 authentication tag _BLOCK = 32 # HMAC-SHA256 output width == one keystream block # Env-var name the agent container receives at startup. @@ -50,45 +61,58 @@ def _keystream(key: bytes, nonce: bytes, block_index: int) -> bytes: ).digest() +def _tag(key: bytes, nonce: bytes, ciphertext: bytes) -> bytes: + """The authentication tag over ``nonce || ciphertext``, keyed by a MAC + subkey domain-separated from the keystream key.""" + mac_key = hmac.new(key, b"bottled-secret-mac-v1", hashlib.sha256).digest() + return hmac.new(mac_key, nonce + ciphertext, hashlib.sha256).digest() + + +def _ctr(key: bytes, nonce: bytes, data: bytes) -> bytes: + """CTR keystream XOR — its own inverse, so it both encrypts and decrypts.""" + out = bytearray() + for i in range(0, len(data), _BLOCK): + chunk = data[i : i + _BLOCK] + ks = _keystream(key, nonce, i)[: len(chunk)] + out.extend(b ^ k for b, k in zip(chunk, ks)) + return bytes(out) + + def encrypt_value(secret_b64: str, plaintext: str) -> str: """Encrypt a single string value with *secret_b64* (the ENV_VAR_SECRET). - Returns a URL-safe base64 blob ``nonce || ciphertext`` suitable for + Returns a URL-safe base64 blob ``nonce || ciphertext || tag`` suitable for the ``bottled_agent_secrets.value`` column.""" key = _b64dec(secret_b64) - pt = plaintext.encode() nonce = secrets.token_bytes(_NONCE_BYTES) - ct = bytearray() - for i in range(0, len(pt), _BLOCK): - chunk = pt[i : i + _BLOCK] - ks = _keystream(key, nonce, i)[: len(chunk)] - ct.extend(p ^ k for p, k in zip(chunk, ks)) - return base64.urlsafe_b64encode(nonce + bytes(ct)).rstrip(b"=").decode() + ct = _ctr(key, nonce, plaintext.encode()) + tag = _tag(key, nonce, ct) + return base64.urlsafe_b64encode(nonce + ct + tag).rstrip(b"=").decode() def decrypt_value(secret_b64: str, blob_b64: str) -> str: """Decrypt a blob produced by :func:`encrypt_value`. Returns the original plaintext string. Raises ``ValueError`` for malformed - input or a key mismatch (wrong key produces garbage, not an error, unless - the plaintext is non-UTF-8 — treat all such failures as wrong key).""" + input, a **wrong key**, or a tampered ciphertext — all caught by the + authentication tag before any plaintext is returned, so a wrong + ENV_VAR_SECRET is rejected deterministically (never a garbage token).""" key = _b64dec(secret_b64) try: blob = _b64dec(blob_b64) except Exception as exc: raise ValueError(f"invalid ciphertext blob: {exc}") from exc - if len(blob) < _NONCE_BYTES: + if len(blob) < _NONCE_BYTES + _TAG_BYTES: raise ValueError("ciphertext blob too short") - nonce, ciphertext = blob[:_NONCE_BYTES], blob[_NONCE_BYTES:] - pt = bytearray() - for i in range(0, len(ciphertext), _BLOCK): - chunk = ciphertext[i : i + _BLOCK] - ks = _keystream(key, nonce, i)[: len(chunk)] - pt.extend(c ^ k for c, k in zip(chunk, ks)) + nonce = blob[:_NONCE_BYTES] + tag = blob[-_TAG_BYTES:] + ciphertext = blob[_NONCE_BYTES:-_TAG_BYTES] + if not hmac.compare_digest(tag, _tag(key, nonce, ciphertext)): + raise ValueError("ciphertext failed authentication (wrong key or tampered)") try: - return bytes(pt).decode() - except UnicodeDecodeError as exc: - raise ValueError(f"decryption produced non-UTF-8 output (wrong key?): {exc}") from exc + return _ctr(key, nonce, ciphertext).decode() + except UnicodeDecodeError as exc: # pragma: no cover - authenticated, so unreachable + raise ValueError(f"decryption produced non-UTF-8 output: {exc}") from exc __all__ = ["ENV_VAR_SECRET_NAME", "new_env_var_secret", "encrypt_value", "decrypt_value"] diff --git a/tests/unit/test_orchestrator_secret_store.py b/tests/unit/test_orchestrator_secret_store.py index 7301b126..22acd079 100644 --- a/tests/unit/test_orchestrator_secret_store.py +++ b/tests/unit/test_orchestrator_secret_store.py @@ -2,10 +2,12 @@ from __future__ import annotations +import base64 import unittest from bot_bottle.orchestrator.store.secret_store import ( ENV_VAR_SECRET_NAME, + _NONCE_BYTES, decrypt_value, encrypt_value, new_env_var_secret, @@ -65,22 +67,27 @@ class TestDecryptErrors(unittest.TestCase): def setUp(self) -> None: self.secret = new_env_var_secret() - def test_wrong_key_raises_value_error(self) -> None: + def test_wrong_key_always_raises_value_error(self) -> None: + # Deterministic: the authentication tag rejects a wrong key every time, + # so reprovision can never inject a garbage token. Repeat across many + # random keys (the old unauthenticated scheme let ~5% through when the + # garbage happened to decode as valid UTF-8). + for _ in range(200): + ct = encrypt_value(self.secret, "secret-token") + with self.assertRaises(ValueError): + decrypt_value(new_env_var_secret(), ct) + + def test_tampered_ciphertext_raises_value_error(self) -> None: ct = encrypt_value(self.secret, "secret-token") - other_key = new_env_var_secret() - # Wrong key produces garbage bytes; decrypt_value raises ValueError - # when the result is non-UTF-8 (which is very likely for 12-char data). - # We allow it to succeed only if garbage happens to be valid UTF-8, but - # the plaintext must not match. - try: - result = decrypt_value(other_key, ct) - self.assertNotEqual("secret-token", result) - except ValueError: - pass + raw = bytearray(base64.urlsafe_b64decode(ct + "=" * (-len(ct) % 4))) + raw[_NONCE_BYTES] ^= 0x01 # flip a bit in the ciphertext body → tag mismatch + tampered = base64.urlsafe_b64encode(bytes(raw)).rstrip(b"=").decode() + with self.assertRaises(ValueError): + decrypt_value(self.secret, tampered) def test_truncated_blob_raises_value_error(self) -> None: with self.assertRaises(ValueError): - decrypt_value(self.secret, "dG9vc2hvcnQ") # "tooshort" — under 16 nonce bytes + decrypt_value(self.secret, "dG9vc2hvcnQ") # "tooshort" — under nonce+tag def test_invalid_base64_raises_value_error(self) -> None: with self.assertRaises(ValueError): -- 2.52.0