bot-bottle/tests/unit/test_orchestrator_host_server.py

"""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 os
import secrets
import tempfile
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,
    dispatch,
    main,
    make_host_server,
)
from bot_bottle.paths import LAUNCH_BROKER_KEY_ENV


def _body(obj: object) -> bytes:
    return json.dumps(obj).encode()


class _RaisingBroker(LaunchBroker):
    """A broker whose backend launch/enumeration 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")

    def _list_live(self) -> list[str]:
        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_broker_live_returns_source_ips(self) -> None:
        # Two launched bottles -> the stub reports both as live.
        self.broker.submit(self._token(op="launch", bottle_id="b1", source_ip="10.0.0.1"))
        self.broker.submit(self._token(op="launch", bottle_id="b2", source_ip="10.0.0.2"))
        token = self._token(op="list_live")
        status, payload = dispatch(self.broker, "POST", "/broker/live", _body({"token": token}))
        self.assertEqual(200, status)
        self.assertEqual(
            ["10.0.0.1", "10.0.0.2"],
            sorted(typing.cast("list[str]", payload["source_ips"])),
        )

    def test_broker_live_forged_token_is_401(self) -> None:
        forged = sign_request(LaunchRequest(op="list_live"), secrets.token_bytes(16))
        status, payload = dispatch(self.broker, "POST", "/broker/live", _body({"token": forged}))
        self.assertEqual(401, status)
        self.assertIn("broker auth failed", str(payload["error"]))

    def test_broker_live_rejects_a_mutation_token(self) -> None:
        # A launch token routed to the query endpoint is a fail-closed 401 — the
        # endpoints don't share a schema even though they share a secret.
        token = self._token(op="launch", bottle_id="b1", image_ref="img")
        status, _ = dispatch(self.broker, "POST", "/broker/live", _body({"token": token}))
        self.assertEqual(401, status)

    def test_broker_live_backend_failure_is_502(self) -> None:
        broker = _RaisingBroker(self.secret)
        token = self._token(op="list_live")
        status, payload = dispatch(broker, "POST", "/broker/live", _body({"token": token}))
        self.assertEqual(502, status)
        self.assertIn("backend enumeration failed", str(payload["error"]))

    def test_broker_live_missing_token_is_400(self) -> None:
        status, _ = dispatch(self.broker, "POST", "/broker/live", _body({}))
        self.assertEqual(400, status)

    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 TestBrokerSecret(unittest.TestCase):
    """The durable launch-broker key (#468/#476): prefer the env-injected key,
    else the durable host key file, so signer and verifier resolve the same one."""

    def test_reads_injected_key_from_env(self) -> None:
        # The injected key is honoured regardless of allow_host_file — both the
        # host controller and the guest orchestrator take an injected key.
        self.assertEqual(
            b"injected-key", broker_secret({LAUNCH_BROKER_KEY_ENV: "injected-key"}))
        self.assertEqual(
            b"injected-key",
            broker_secret({LAUNCH_BROKER_KEY_ENV: "injected-key"}, allow_host_file=True))

    def test_guest_without_injection_fails_closed(self) -> None:
        # The default (guest orchestrator): no env key and NO host-file fallback,
        # so it returns None rather than mint a divergent process-local key.
        self.assertIsNone(broker_secret({}))

    def test_host_side_falls_back_to_the_durable_key_file(self) -> None:
        # allow_host_file=True (host controller / dev-harness): mint/read the
        # durable host key file, the same key on every call (restart re-adoption).
        with tempfile.TemporaryDirectory() as root:
            with patch.dict("os.environ", {"BOT_BOTTLE_ROOT": root}, clear=False):
                os.environ.pop(LAUNCH_BROKER_KEY_ENV, None)
                first = broker_secret(allow_host_file=True)
                second = broker_secret(allow_host_file=True)
        self.assertTrue(first)
        self.assertEqual(first, second)


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

    def test_list_live_enumerates_over_http(self) -> None:
        secret = secrets.token_bytes(16)
        broker = StubBroker(secret)
        broker.live_source_ips = ["10.0.0.7", "10.0.0.8"]
        client = self._serve(broker)
        got = client.list_live(sign_request(LaunchRequest(op="list_live"), secret))
        self.assertEqual(["10.0.0.7", "10.0.0.8"], got)


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