bot-bottle/bot_bottle/orchestrator/host_server.py

"""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": "<signed launch/teardown JWT>"}

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 `host`-role
tokens of the separate `HOST_CONTROLLER` trust domain in a later chunk.

The shared signing secret is the durable **launch-broker `TrustDomain` key**
(#468/#476): a host-canonical key file minted 0600 on first use, provisioned to
the orchestrator (signer) and this server (verifier). A backend launcher injects
it via `$BOT_BOTTLE_LAUNCH_BROKER_KEY`; a host-side dev-harness process reads the
key file directly. Durability is the point — a restarted orchestrator re-verifies
against the same key, so re-adoption works.
"""

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 ..paths import LAUNCH_BROKER_KEY_ENV
from ..trust_domain import LAUNCH_BROKER
from .broker import BrokerAuthError, LaunchBroker
from .docker_broker import DockerBroker

# JSON body payload type (parsed request / rendered response).
Json = dict[str, object]

# 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(environ: typing.Mapping[str, str] | None = None) -> bytes | None:
    """The shared launch-broker HS256 secret, as this process should use it.

    Prefers the key injected into this process's env (`$BOT_BOTTLE_LAUNCH_BROKER_KEY`)
    — the path a containerized backend launcher uses — and falls back to the
    durable host key file (`bot_bottle_root()/launch-broker-key`, minted 0600 on
    first use) for a host-side process like the dev-harness. The signer
    (orchestrator, `--broker http`) resolves it the same way, so both hold the
    same key. None only when neither is available (an unwritable host root)."""
    key = LAUNCH_BROKER.key_from_env(environ)
    if not key:
        try:
            key = LAUNCH_BROKER.signing_key()  # host-canonical, minted on first use
        except OSError:
            return None
    return key.encode("utf-8") if key else 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 the launch-broker key 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()
    if secret is None:
        sys.stderr.write(
            f"host controller: refusing to start without the launch-broker key "
            f"(${LAUNCH_BROKER_KEY_ENV}, or a writable host root to mint it) — 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",
    "main",
    "Json",
    "DEFAULT_PORT",
]


if __name__ == "__main__":
    raise SystemExit(main())