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 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())