From ded9cb3bf39aebc223a0e202680114c4220261b6 Mon Sep 17 00:00:00 2001 From: claude Date: Sun, 26 Jul 2026 07:16:56 +0000 Subject: [PATCH] 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()