258 lines
9.8 KiB
Python
258 lines
9.8 KiB
Python
#!/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 truncate_utf8(value: str, maximum: int) -> str:
|
|
"""Bound display text by UTF-8 bytes without splitting a character."""
|
|
if maximum < 0:
|
|
raise ProtocolError("text byte limit is invalid")
|
|
return value.encode("utf-8", "replace")[:maximum].decode("utf-8", "ignore")
|
|
|
|
|
|
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}")
|