atlas-iac/services/hermes/scripts/execution_pool_protocol.py

#!/usr/bin/env python3
"""Authenticated, bounded, restart-safe Hermes execution-pool protocol."""

from __future__ import annotations

import hashlib
import hmac
import json
import os
import re
import stat
import threading
import time
import uuid
from http.server import ThreadingHTTPServer
from pathlib import Path
from typing import Any


MAX_WIRE_BYTES = 64 * 1024
MAX_ACTIVITY_BYTES = 12 * 1024
MAX_CLOCK_SKEW = 30
MAX_ENVELOPE_LIFETIME = 300
PROTOCOL_VERSION = 2
IDENTIFIER = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.-]{0,127}$")
KINDS = frozenset({"poll", "assignment", "heartbeat", "result", "ack"})


class ProtocolError(ValueError):
    """A request failed the authenticated pool contract."""


class BoundedHTTPServer(ThreadingHTTPServer):
    """Bound concurrent requests and slow clients on internal pool channels."""

    daemon_threads = True
    request_queue_size = 8

    def __init__(self, *args: Any, max_workers: int = 8, **kwargs: Any):
        self._slots = threading.BoundedSemaphore(max(1, min(max_workers, 16)))
        super().__init__(*args, **kwargs)

    def get_request(self) -> tuple[Any, Any]:
        request, address = super().get_request()
        request.settimeout(15)
        return request, address

    def process_request(self, request: Any, client_address: Any) -> None:
        if not self._slots.acquire(blocking=False):
            self.shutdown_request(request)
            return
        try:
            super().process_request(request, client_address)
        except Exception:
            self._slots.release()
            raise

    def process_request_thread(self, request: Any, client_address: Any) -> None:
        try:
            super().process_request_thread(request, client_address)
        finally:
            self._slots.release()


def canonical_json(value: Any) -> bytes:
    """Encode one value deterministically for digests and signatures."""
    return json.dumps(value, ensure_ascii=False, separators=(",", ":"), sort_keys=True).encode("utf-8")


def atomic_json(path: Path, value: dict[str, Any], mode: int = 0o600) -> None:
    """Durably replace one bounded pool document without following symlinks."""
    path.parent.mkdir(mode=0o700, parents=True, exist_ok=True)
    if path.parent.is_symlink():
        raise ProtocolError("pool state directory must not be a symlink")
    temporary = path.with_name(f".{path.name}.{uuid.uuid4().hex}.tmp")
    flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0)
    descriptor = os.open(temporary, flags, mode)
    try:
        remaining = memoryview(json.dumps(value, indent=2, sort_keys=True).encode() + b"\n")
        while remaining:
            remaining = remaining[os.write(descriptor, remaining) :]
        os.fsync(descriptor)
        os.fchmod(descriptor, mode)
    finally:
        os.close(descriptor)
    try:
        os.replace(temporary, path)
    finally:
        temporary.unlink(missing_ok=True)


def payload_digest(payload: Any) -> str:
    return hashlib.sha256(canonical_json(payload)).hexdigest()


def derive_ordinal_key(master: bytes, ordinal: int) -> bytes:
    """Derive one cryptographically isolated worker authority from the pool root."""
    if ordinal not in range(3) or not 32 <= len(master) <= 4096:
        raise ProtocolError("pool key derivation input is invalid")
    context = f"hermes-execution-pool-v2:worker:{ordinal}".encode()
    return hmac.new(master, context, hashlib.sha256).hexdigest().encode()


def envelope_ordinal(envelope: Any) -> int:
    """Read only the bounded ordinal needed to select a verification key."""
    if not isinstance(envelope, dict):
        raise ProtocolError("invalid pool message")
    value = envelope.get("worker_ordinal")
    if type(value) is not int or value not in range(3):
        raise ProtocolError("worker binding is outside the pool")
    return value


def read_key(path: Path) -> bytes:
    """Read a private regular file without following a final symlink."""
    flags = os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0)
    try:
        descriptor = os.open(path, flags)
    except OSError as error:
        raise ProtocolError(f"pool key is unavailable: {error}") from error
    try:
        info = os.fstat(descriptor)
        if not stat.S_ISREG(info.st_mode) or info.st_mode & 0o077:
            raise ProtocolError("pool key must be a private regular file")
        value = os.read(descriptor, 4097).strip()
    finally:
        os.close(descriptor)
    if len(value) < 32 or len(value) > 4096:
        raise ProtocolError("pool key length is outside the safe range")
    return value


def _identifier(name: str, value: Any, *, allow_empty: bool = False) -> str:
    if not isinstance(value, str):
        raise ProtocolError(f"invalid {name}")
    text = value
    if allow_empty and not text:
        return text
    if not IDENTIFIER.fullmatch(text):
        raise ProtocolError(f"invalid {name}")
    return text


def sign_envelope(
    key: bytes,
    kind: str,
    binding: dict[str, Any],
    payload: Any,
    *,
    now: int | None = None,
    lifetime: int = 120,
    delivery_id: str | None = None,
) -> dict[str, Any]:
    """Bind a message to the exact Kanban run, ordinal, and attempt."""
    current = int(time.time()) if now is None else int(now)
    lifetime = max(1, min(int(lifetime), MAX_ENVELOPE_LIFETIME))
    envelope = {
        "version": PROTOCOL_VERSION,
        "kind": kind,
        "board": str(binding.get("board") or ""),
        "task_id": str(binding.get("task_id") or ""),
        "run_id": str(binding.get("run_id") or ""),
        "worker_ordinal": int(binding.get("worker_ordinal", -1)),
        "attempt": int(binding.get("attempt", 0)),
        "delivery_id": delivery_id or str(uuid.uuid4()),
        "issued_at": current,
        "expires_at": current + lifetime,
        "payload_digest": payload_digest(payload),
        "payload": payload,
    }
    if kind not in KINDS:
        raise ProtocolError("unsupported message kind")
    unsigned = canonical_json(envelope)
    if len(unsigned) > MAX_WIRE_BYTES:
        raise ProtocolError("pool message exceeds the wire limit")
    envelope["signature"] = hmac.new(key, unsigned, hashlib.sha256).hexdigest()
    return envelope


def verify_envelope(
    key: bytes,
    envelope: Any,
    *,
    expected_kind: str | None = None,
    now: int | None = None,
) -> dict[str, Any]:
    """Verify structure, lifetime, digest, and HMAC before using a message."""
    if not isinstance(envelope, dict) or len(canonical_json(envelope)) > MAX_WIRE_BYTES:
        raise ProtocolError("invalid or oversized pool message")
    required = {"version", "kind", "board", "task_id", "run_id", "worker_ordinal", "attempt", "delivery_id", "issued_at", "expires_at", "payload_digest", "payload", "signature"}
    if set(envelope) != required or envelope.get("version") != PROTOCOL_VERSION:
        raise ProtocolError(
            f"pool message fields do not match version {PROTOCOL_VERSION}"
        )
    kind = envelope["kind"]
    if not isinstance(kind, str):
        raise ProtocolError("unexpected message kind")
    if kind not in KINDS or (expected_kind and kind != expected_kind):
        raise ProtocolError("unexpected message kind")
    empty_binding = kind in {"poll", "ack"}
    _identifier("board", envelope["board"], allow_empty=empty_binding)
    _identifier("task_id", envelope["task_id"], allow_empty=empty_binding)
    _identifier("run_id", envelope["run_id"], allow_empty=empty_binding)
    _identifier("delivery_id", envelope["delivery_id"])
    numeric = tuple(envelope[name] for name in ("worker_ordinal", "attempt", "issued_at", "expires_at"))
    if any(type(value) is not int for value in numeric):
        raise ProtocolError("invalid numeric binding")
    ordinal, attempt, issued, expires = numeric
    if ordinal not in range(3) or attempt < 0:
        raise ProtocolError("worker binding is outside the pool")
    current = int(time.time()) if now is None else int(now)
    if issued > current + MAX_CLOCK_SKEW or expires < current - MAX_CLOCK_SKEW:
        raise ProtocolError("pool message is outside its validity window")
    if expires <= issued or expires - issued > MAX_ENVELOPE_LIFETIME:
        raise ProtocolError("pool message lifetime is invalid")
    if envelope["payload_digest"] != payload_digest(envelope["payload"]):
        raise ProtocolError("pool payload digest mismatch")
    signature = str(envelope["signature"])
    unsigned = dict(envelope)
    unsigned.pop("signature")
    expected = hmac.new(key, canonical_json(unsigned), hashlib.sha256).hexdigest()
    if not hmac.compare_digest(signature, expected):
        raise ProtocolError("pool message authentication failed")
    return envelope


def parse_wire(body: bytes) -> dict[str, Any]:
    if not body or len(body) > MAX_WIRE_BYTES:
        raise ProtocolError("empty or oversized request")
    try:
        value = json.loads(body)
    except (UnicodeDecodeError, json.JSONDecodeError) as error:
        raise ProtocolError("malformed JSON request") from error
    if not isinstance(value, dict):
        raise ProtocolError("request must be a JSON object")
    return value


def __getattr__(name: str) -> Any:
    """Expose the durable store on the stable protocol surface.

    ``PoolStore`` moved to :mod:`execution_pool_store` to keep both modules under
    the managed line ceiling. It is resolved lazily here so the store can keep
    importing this module's primitives without an import cycle.
    """
    if name == "PoolStore":
        from execution_pool_store import PoolStore

        return PoolStore
    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")