"""Launch broker contract (PRD 0070). A VM/container can't spawn its own host-network siblings, so the orchestrator brokers agent launches through a small privileged component. The request it sends is: * **structured** — static flags + ids only (bottle id, pool slot, a content-addressed image ref), never a free-form path or argv, so a request can't be coerced into launching an arbitrary payload; and * **signed** — wrapped as a compact JWS/JWT the broker verifies before acting, so a *compromised co-located component* (an agent-facing gateway, say) can't forge a launch. This is the concrete form of the "structured requests only" rule in the PRD security review. Signing is **HS256** over a secret shared by the orchestrator (signer) and the broker (verifier) — appropriate here because both are trusted host components provisioned together; the secret is exactly what an agent-facing component does not have. (Stdlib only — the project takes no runtime deps; a cross-host deployment could swap in an asymmetric alg.) """ from __future__ import annotations import abc import base64 import hashlib import hmac import json import secrets import time from dataclasses import dataclass from typing import Protocol _JWT_HEADER = {"alg": "HS256", "typ": "JWT"} # The closed op vocabulary, split by kind (PRD "ids + static flags" rule): # * mutation ops act on ONE bottle — they carry its id (+ static launch flags); # * query ops enumerate host state — they carry NO ids or flags at all. # Each new host-privileged op is added here deliberately; `verify_request` # enforces the per-kind shape so the privileged surface can't drift. _MUTATION_OPS = ("launch", "teardown") _QUERY_OPS = ("list_live",) _ALLOWED_OPS = _MUTATION_OPS + _QUERY_OPS # Every claim key a signed request may carry — the request fields plus the two # per-signature envelope claims. A strict allow-list (open question 1, resolved # "yes"): an unknown key is rejected outright, so the schema can't silently widen # as ops are added. _ALLOWED_CLAIMS = frozenset( {"op", "bottle_id", "source_ip", "image_ref", "slot", "jti", "iat"} ) 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. 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) class BrokerRequest: """A structured, un-coercible broker request. Ids + static flags only — `image_ref` is a content-addressed id, never a path/argv. A **mutation** op (`launch`/`teardown`) names the bottle it acts on in `bottle_id` (plus launch's static flags). A **query** op (`list_live`) acts on no single bottle, so it carries nothing but its `op` name — `bottle_id` and every flag stay empty, and `verify_request` rejects a query that smuggles any in. `LaunchRequest` is kept as an alias for the launch/teardown callers.""" op: str # one of _ALLOWED_OPS bottle_id: str = "" source_ip: str = "" image_ref: str = "" slot: int | None = None # Historical name — the request type predates the query ops. Kept so the launch # path (`OrchestratorCore`, `DockerBroker`) reads unchanged. LaunchRequest = BrokerRequest # --- minimal JWS/JWT (HS256), stdlib only ---------------------------------- def _b64url(data: bytes) -> str: return base64.urlsafe_b64encode(data).rstrip(b"=").decode("ascii") def _b64url_decode(text: str) -> bytes: return base64.urlsafe_b64decode(text + "=" * (-len(text) % 4)) def _mac(signing_input: str, secret: bytes) -> str: digest = hmac.new(secret, signing_input.encode("ascii"), hashlib.sha256).digest() return _b64url(digest) def sign_request(req: LaunchRequest, secret: bytes) -> str: """Serialize `req` to signed-JWT form. Adds a per-signature `jti` (replay id) and `iat` (issued-at).""" claims: dict[str, object] = { "op": req.op, "bottle_id": req.bottle_id, "source_ip": req.source_ip, "image_ref": req.image_ref, "slot": req.slot, "jti": secrets.token_hex(8), "iat": int(time.time()), } header = _b64url(json.dumps(_JWT_HEADER, sort_keys=True, separators=(",", ":")).encode()) payload = _b64url(json.dumps(claims, sort_keys=True, separators=(",", ":")).encode()) signing_input = f"{header}.{payload}" return f"{signing_input}.{_mac(signing_input, secret)}" def verify_request(token: str, secret: bytes) -> LaunchRequest: """Verify the signature and the request shape; return the parsed request. Raises `BrokerAuthError` on any failure (fail-closed).""" parts = token.split(".") if len(parts) != 3: raise BrokerAuthError("malformed token") header_b, payload_b, sig = parts if not hmac.compare_digest(_mac(f"{header_b}.{payload_b}", secret), sig): raise BrokerAuthError("bad signature") try: header = json.loads(_b64url_decode(header_b)) claims = json.loads(_b64url_decode(payload_b)) except (ValueError, json.JSONDecodeError) as e: raise BrokerAuthError("malformed token payload") from e if not isinstance(header, dict) or header.get("alg") != "HS256": raise BrokerAuthError("unexpected header/alg") if not isinstance(claims, dict): raise BrokerAuthError("claims are not an object") # Strict schema (open question 1): a signed request may carry ONLY the known # claim keys, so the privileged surface can't widen by smuggling an extra # claim past the fixed field checks below. unknown = set(claims) - _ALLOWED_CLAIMS if unknown: raise BrokerAuthError(f"unknown claim(s): {', '.join(sorted(unknown))}") op = claims.get("op") if not isinstance(op, str) or op not in _ALLOWED_OPS: raise BrokerAuthError("request does not match the broker-request schema") bottle_id = claims.get("bottle_id", "") source_ip = claims.get("source_ip", "") image_ref = claims.get("image_ref", "") slot = claims.get("slot") if ( not isinstance(bottle_id, str) or not isinstance(source_ip, str) or not isinstance(image_ref, str) or not (slot is None or isinstance(slot, int)) ): raise BrokerAuthError("request does not match the broker-request schema") # Per-kind shape: a mutation names exactly one bottle; a query names none and # carries no flags (the "no arguments" rule for `list_live`). Enforcing both # halves keeps `bottle_id`/flags from being a coercible field on a query op. if op in _MUTATION_OPS: if not bottle_id: raise BrokerAuthError(f"{op} requires a bottle_id") elif bottle_id or source_ip or image_ref or slot is not None: raise BrokerAuthError(f"{op} is a query op and takes no ids or flags") return BrokerRequest( op=op, bottle_id=bottle_id, source_ip=source_ip, image_ref=image_ref, slot=slot ) # --- the broker itself ------------------------------------------------------ class SubmitBroker(Protocol): """The broker surface `OrchestratorCore` depends on. `submit` verifies a signed *mutation* token and performs its op, returning the verified request; `list_live` verifies a signed `list_live` token and returns the source IPs of the bottles the backend currently has running (reconcile's live set). 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: ... def list_live(self, token: str) -> list[str]: ... class LaunchBroker(abc.ABC): """Verifies a signed request came from the orchestrator, then performs the backend-native launch/teardown. Subclasses implement `_launch` / `_teardown`; verification is shared and fail-closed.""" def __init__(self, secret: bytes) -> None: self._secret = secret def submit(self, token: str) -> LaunchRequest: """Verify `token` and perform its mutation op. Returns the verified request; raises `BrokerAuthError` if provenance/shape fails, or if the token is a non-mutation op (a `list_live` token routed here by mistake).""" req = verify_request(token, self._secret) if req.op == "launch": self._launch(req) elif req.op == "teardown": self._teardown(req) else: raise BrokerAuthError(f"{req.op} is not a submit op") return req def list_live(self, token: str) -> list[str]: """Verify a `list_live` `token` and return the source IPs of the bottles the backend currently has running. Raises `BrokerAuthError` on bad provenance/shape or a non-`list_live` op; a backend enumeration failure is converted to `BrokerUnavailableError` — the single "live set could not be determined" signal reconcile catches to skip the sweep rather than reaping healthy rows against a partial set (the out-of-process `BrokerClient` raises the same on an enumeration 502 / no response).""" req = verify_request(token, self._secret) if req.op != "list_live": raise BrokerAuthError(f"{req.op} is not a list_live op") try: return list(self._list_live()) except Exception as e: # noqa: BLE001 — any enumeration failure means the # live set is unknown; reconcile must skip, never reap against it. raise BrokerUnavailableError(f"list_live enumeration failed: {e}") from e @abc.abstractmethod def _launch(self, req: LaunchRequest) -> None: ... @abc.abstractmethod def _teardown(self, req: LaunchRequest) -> None: ... @abc.abstractmethod def _list_live(self) -> list[str]: """Every running bottle's source IP, backend-native. Must raise (not return a partial list) if the live set can't be determined authoritatively, so reconcile skips rather than reaping healthy rows.""" ... class StubBroker(LaunchBroker): """Dev-harness broker: records verified requests without launching anything. Exercises the full sign -> verify -> act contract in-process. Its live set (`_list_live`) is, by default, every bottle it recorded a launch for and no teardown since — the in-process analogue of enumerating the backend. Tests that need to simulate a bottle dying out from under the registry (the case reconcile exists for) set `live_source_ips` explicitly to override that derived set.""" def __init__(self, secret: bytes) -> None: super().__init__(secret) self.launched: list[LaunchRequest] = [] self.torn_down: list[LaunchRequest] = [] # When not None, the exact live set `_list_live` reports — lets a test # say "only these IPs are still up" regardless of what was launched. self.live_source_ips: list[str] | None = None def _launch(self, req: LaunchRequest) -> None: self.launched.append(req) def _teardown(self, req: LaunchRequest) -> None: self.torn_down.append(req) def _list_live(self) -> list[str]: if self.live_source_ips is not None: return list(self.live_source_ips) torn = {r.bottle_id for r in self.torn_down} return [r.source_ip for r in self.launched if r.bottle_id not in torn and r.source_ip] __all__ = [ "BrokerAuthError", "BrokerUnavailableError", "BrokerRequest", "LaunchRequest", "SubmitBroker", "LaunchBroker", "StubBroker", "sign_request", "verify_request", ]