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..6d5c567f 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") @@ -37,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) @@ -123,6 +138,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` / @@ -168,7 +193,9 @@ class StubBroker(LaunchBroker): __all__ = [ "BrokerAuthError", + "BrokerUnavailableError", "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..62783418 --- /dev/null +++ b/bot_bottle/orchestrator/broker_client.py @@ -0,0 +1,126 @@ +"""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, BrokerUnavailableError, LaunchRequest + +DEFAULT_TIMEOUT_SECONDS = 5.0 + + +class BrokerClientError(RuntimeError): + """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: + """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); `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", + 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) 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]: + """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..fe8130a0 --- /dev/null +++ b/bot_bottle/orchestrator/host_server.py @@ -0,0 +1,270 @@ +"""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 + +# 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.""" + 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.""" + + # 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 + if os.environ.get("BOT_BOTTLE_HOST_CONTROLLER_DEBUG"): + super().log_message(format, *args) + + def _serve(self, method: str) -> None: + """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) + 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) + 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}"} + 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") + 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..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 LaunchBroker, LaunchRequest, 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, @@ -62,7 +62,7 @@ class OrchestratorCore: def __init__( self, registry: RegistryStore, - broker: LaunchBroker, + broker: SubmitBroker, sign_secret: bytes, supervisor: Supervisor | None = None, ) -> None: @@ -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/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_broker_client.py b/tests/unit/test_orchestrator_broker_client.py new file mode 100644 index 00000000..e0de4747 --- /dev/null +++ b/tests/unit/test_orchestrator_broker_client.py @@ -0,0 +1,117 @@ +"""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, + BrokerUnavailableError, + 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_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_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(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: + 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") + + 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 new file mode 100644 index 00000000..0217a2c8 --- /dev/null +++ b/tests/unit/test_orchestrator_host_server.py @@ -0,0 +1,285 @@ +"""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 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, + LaunchBroker, + LaunchRequest, + StubBroker, + sign_request, +) +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, +) + + +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_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) + + 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 + + +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_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): 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