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

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

    Always prefers the key injected into this process's env
    (`$BOT_BOTTLE_LAUNCH_BROKER_KEY`). `allow_host_file` decides the fallback when
    it is absent, and the distinction is a security boundary:

    - **Host-side** processes — the host controller and the host dev-harness — pass
      ``allow_host_file=True`` to read (minting on first use) the durable host key
      file (``bot_bottle_root()/launch-broker-key``) they legitimately own.
    - The **guest orchestrator** (``--broker http``) keeps the default ``False``.
      It runs inside a container/VM whose ``bot_bottle_root()`` is process-local,
      so minting a file there would silently create a key UNRELATED to the host
      controller's — startup would succeed but every launch would be rejected 401.
      It must instead be *given* the key by its launcher, and fail closed (None)
      if it wasn't, rather than diverge.

    None when no key is available (a guest with no injection, or an unwritable
    host root)."""
    key = LAUNCH_BROKER.key_from_env(environ)
    if not key and allow_host_file:
        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."""

    # 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 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. As the host-side owner of the key, it may mint/read the host
    key file (`allow_host_file=True`)."""
    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(allow_host_file=True)
    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())