829 lines
34 KiB
Python
829 lines
34 KiB
Python
#!/usr/bin/env python3
|
|
"""Mosaic external lease broker for Linux SO_PEERCRED authenticated clients."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import copy
|
|
import errno
|
|
import hmac
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
import json
|
|
import os
|
|
import secrets
|
|
import signal
|
|
import socket
|
|
import stat
|
|
import struct
|
|
import sys
|
|
import threading
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Final
|
|
|
|
# This script is loaded both as an executable and through isolated stdlib tests.
|
|
# Keep its co-located receipt implementation importable in both modes.
|
|
_MODULE_DIRECTORY = str(Path(__file__).resolve().parent)
|
|
if _MODULE_DIRECTORY not in sys.path:
|
|
sys.path.insert(0, _MODULE_DIRECTORY)
|
|
from normative_fragments import build_payload_from_wire
|
|
from receipt_challenge import is_verbatim_receipt, latest_assistant_digest, receipt_for
|
|
from receipt_observer import FileTestReceiptObserver, ReceiptObserver, UnavailableReceiptObserver
|
|
|
|
MAX_FRAME: Final = 64 * 1024
|
|
MAX_STATE: Final = 4 * 1024 * 1024
|
|
MAX_PENDING_TOKENS: Final = 256
|
|
MAX_IN_FLIGHT_CONNECTIONS: Final = 16
|
|
MAX_LEASE_TTL_SECONDS: Final = 300
|
|
STATE_VERSION: Final = 1
|
|
READ_DEADLINE_SECONDS: Final = 1.0
|
|
HANDLE_QUEUE_TIMEOUT_SECONDS: Final = 1.0
|
|
SEND_TIMEOUT_SECONDS: Final = 1.0
|
|
HEX_256_LENGTH: Final = 64
|
|
LEASE_UNVERIFIED: Final = "UNVERIFIED"
|
|
LEASE_PENDING: Final = "PENDING_VERIFICATION"
|
|
LEASE_PENDING_PROMOTION: Final = "PENDING_PROMOTION"
|
|
LEASE_VERIFIED: Final = "VERIFIED"
|
|
READ_ONLY_TOOLS: Final = {
|
|
"claude": frozenset({"Read", "Grep", "Glob", "Ls", "Find"}),
|
|
"pi": frozenset({"read", "grep", "find", "ls"}),
|
|
}
|
|
RECOVERY_TOOL: Final = "mosaic_context_recover"
|
|
|
|
|
|
class BrokerFailure(Exception):
|
|
def __init__(self, code: str) -> None:
|
|
super().__init__(code)
|
|
self.code = code
|
|
|
|
|
|
class StateCommitUncertain(RuntimeError):
|
|
def __init__(self) -> None:
|
|
super().__init__("STATE_COMMIT_UNCERTAIN")
|
|
|
|
|
|
def is_non_negative_integer(value: object) -> bool:
|
|
return type(value) is int and value >= 0
|
|
|
|
|
|
def is_hex_256(value: object) -> bool:
|
|
return (
|
|
isinstance(value, str)
|
|
and len(value) == HEX_256_LENGTH
|
|
and all(character in "0123456789abcdef" for character in value)
|
|
)
|
|
|
|
|
|
def is_positive_decimal(value: object) -> bool:
|
|
return (
|
|
isinstance(value, str)
|
|
and len(value) > 0
|
|
and value[0] in "123456789"
|
|
and all(character in "0123456789" for character in value)
|
|
)
|
|
|
|
|
|
def valid_binding(binding: object) -> bool:
|
|
if not isinstance(binding, dict):
|
|
return False
|
|
required = {"compaction_epoch", "request_epoch", "h_source", "h_payload", "schema_version"}
|
|
if set(binding) != required:
|
|
return False
|
|
if not all(
|
|
is_non_negative_integer(binding[field])
|
|
for field in ("compaction_epoch", "request_epoch", "schema_version")
|
|
):
|
|
return False
|
|
return all(
|
|
is_hex_256(binding[field])
|
|
for field in ("h_source", "h_payload")
|
|
)
|
|
|
|
|
|
def valid_receipt_evidence(evidence: object) -> bool:
|
|
return (
|
|
isinstance(evidence, dict)
|
|
and set(evidence) == {"h_latest_assistant"}
|
|
and is_hex_256(evidence["h_latest_assistant"])
|
|
)
|
|
|
|
|
|
def validate_state(value: object) -> dict[str, object]:
|
|
if not isinstance(value, dict) or set(value) != {"version", "sessions", "tokens"}:
|
|
raise BrokerFailure("STATE_INTEGRITY")
|
|
if type(value["version"]) is not int or value["version"] != STATE_VERSION:
|
|
raise BrokerFailure("STATE_INTEGRITY")
|
|
sessions = value["sessions"]
|
|
tokens = value["tokens"]
|
|
if not isinstance(sessions, dict) or not isinstance(tokens, dict):
|
|
raise BrokerFailure("STATE_INTEGRITY")
|
|
if len(tokens) > MAX_PENDING_TOKENS:
|
|
raise BrokerFailure("STATE_INTEGRITY")
|
|
|
|
anchors: set[tuple[int, str]] = set()
|
|
for session_id, session in sessions.items():
|
|
if not is_hex_256(session_id) or not isinstance(session, dict):
|
|
raise BrokerFailure("STATE_INTEGRITY")
|
|
if set(session) != {"anchor_pid", "anchor_starttime", "runtime_generation"}:
|
|
raise BrokerFailure("STATE_INTEGRITY")
|
|
anchor_pid = session["anchor_pid"]
|
|
anchor_starttime = session["anchor_starttime"]
|
|
if type(anchor_pid) is not int or anchor_pid <= 0:
|
|
raise BrokerFailure("STATE_INTEGRITY")
|
|
if not is_positive_decimal(anchor_starttime):
|
|
raise BrokerFailure("STATE_INTEGRITY")
|
|
if not is_non_negative_integer(session["runtime_generation"]):
|
|
raise BrokerFailure("STATE_INTEGRITY")
|
|
anchor = (anchor_pid, anchor_starttime)
|
|
if anchor in anchors:
|
|
raise BrokerFailure("STATE_INTEGRITY")
|
|
anchors.add(anchor)
|
|
|
|
for token_value, token in tokens.items():
|
|
if not is_hex_256(token_value) or not isinstance(token, dict):
|
|
raise BrokerFailure("STATE_INTEGRITY")
|
|
token_fields = set(token)
|
|
if token_fields not in (
|
|
{"session_id", "runtime_generation", "binding", "consumed"},
|
|
{"session_id", "runtime_generation", "binding", "consumed", "evidence"},
|
|
):
|
|
raise BrokerFailure("STATE_INTEGRITY")
|
|
session_id = token["session_id"]
|
|
generation = token["runtime_generation"]
|
|
session = sessions.get(session_id) if isinstance(session_id, str) else None
|
|
if not is_hex_256(session_id) or not isinstance(session, dict):
|
|
raise BrokerFailure("STATE_INTEGRITY")
|
|
if not is_non_negative_integer(generation):
|
|
raise BrokerFailure("STATE_INTEGRITY")
|
|
if generation != session["runtime_generation"]:
|
|
raise BrokerFailure("STATE_INTEGRITY")
|
|
if not valid_binding(token["binding"]) or token["consumed"] is not False:
|
|
raise BrokerFailure("STATE_INTEGRITY")
|
|
evidence = token.get("evidence")
|
|
if evidence is not None and not valid_receipt_evidence(evidence):
|
|
raise BrokerFailure("STATE_INTEGRITY")
|
|
return value
|
|
|
|
|
|
def proc_node(pid: int) -> dict[str, int | str]:
|
|
try:
|
|
text = Path(f"/proc/{pid}/stat").read_text(encoding="utf-8")
|
|
except (FileNotFoundError, PermissionError, ProcessLookupError) as exc:
|
|
raise BrokerFailure("PID_UNAVAILABLE") from exc
|
|
close = text.rfind(")")
|
|
fields = text[close + 2 :].split()
|
|
if close < 0 or len(fields) < 20:
|
|
raise BrokerFailure("PROC_STAT_INVALID")
|
|
return {"pid": pid, "ppid": int(fields[1]), "starttime": fields[19]}
|
|
|
|
|
|
def verified_ancestry(peer_pid: int, anchor_pid: int, anchor_starttime: str) -> bool:
|
|
chain: list[dict[str, int | str]] = []
|
|
seen: set[int] = set()
|
|
current = peer_pid
|
|
while current > 0 and current not in seen:
|
|
seen.add(current)
|
|
node = proc_node(current)
|
|
chain.append(node)
|
|
if current == anchor_pid:
|
|
if node["starttime"] != anchor_starttime:
|
|
return False
|
|
break
|
|
current = int(node["ppid"])
|
|
else:
|
|
return False
|
|
if int(chain[-1]["pid"]) != anchor_pid:
|
|
return False
|
|
for original in chain:
|
|
repeated = proc_node(int(original["pid"]))
|
|
if repeated["starttime"] != original["starttime"]:
|
|
raise BrokerFailure("PID_STARTTIME_RACE")
|
|
return True
|
|
|
|
|
|
def secure_parent(path: Path) -> None:
|
|
parent = path.parent
|
|
if not parent.is_dir() or stat.S_IMODE(parent.stat().st_mode) != 0o700:
|
|
raise BrokerFailure("INSECURE_PARENT_MODE")
|
|
|
|
|
|
class StateStore:
|
|
def __init__(self, path: Path) -> None:
|
|
self.path = path
|
|
secure_parent(path)
|
|
self.poisoned = False
|
|
self.value: dict[str, object] = {"version": STATE_VERSION, "sessions": {}, "tokens": {}}
|
|
flags = os.O_RDONLY | getattr(os, "O_CLOEXEC", 0) | getattr(os, "O_NOFOLLOW", 0)
|
|
try:
|
|
descriptor = os.open(path, flags)
|
|
except FileNotFoundError:
|
|
return
|
|
except OSError as exc:
|
|
raise BrokerFailure("STATE_INTEGRITY") from exc
|
|
try:
|
|
metadata = os.fstat(descriptor)
|
|
if not stat.S_ISREG(metadata.st_mode):
|
|
raise BrokerFailure("STATE_INTEGRITY")
|
|
if stat.S_IMODE(metadata.st_mode) != 0o600:
|
|
raise BrokerFailure("INSECURE_STATE_MODE")
|
|
if metadata.st_size > MAX_STATE:
|
|
raise BrokerFailure("STATE_INTEGRITY")
|
|
chunks = bytearray()
|
|
while len(chunks) <= MAX_STATE:
|
|
chunk = os.read(descriptor, min(64 * 1024, MAX_STATE + 1 - len(chunks)))
|
|
if not chunk:
|
|
break
|
|
chunks.extend(chunk)
|
|
if len(chunks) > MAX_STATE:
|
|
raise BrokerFailure("STATE_INTEGRITY")
|
|
try:
|
|
loaded = json.loads(chunks)
|
|
except (json.JSONDecodeError, UnicodeDecodeError) as exc:
|
|
raise BrokerFailure("STATE_INTEGRITY") from exc
|
|
self.value = validate_state(loaded)
|
|
except OSError as exc:
|
|
raise BrokerFailure("STATE_INTEGRITY") from exc
|
|
finally:
|
|
os.close(descriptor)
|
|
|
|
def sessions(self) -> dict[str, dict[str, object]]:
|
|
sessions = self.value.get("sessions")
|
|
if not isinstance(sessions, dict):
|
|
raise BrokerFailure("STATE_INTEGRITY")
|
|
return sessions
|
|
|
|
def tokens(self) -> dict[str, dict[str, object]]:
|
|
tokens = self.value.get("tokens")
|
|
if not isinstance(tokens, dict):
|
|
raise BrokerFailure("STATE_INTEGRITY")
|
|
return tokens
|
|
|
|
def commit(self) -> None:
|
|
if self.poisoned:
|
|
raise StateCommitUncertain()
|
|
payload = (
|
|
json.dumps(self.value, sort_keys=True, separators=(",", ":")) + "\n"
|
|
).encode()
|
|
if len(payload) > MAX_STATE:
|
|
raise BrokerFailure("STATE_TOO_LARGE")
|
|
temporary = self.path.with_name(f".{self.path.name}.{os.getpid()}.tmp")
|
|
descriptor = os.open(temporary, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
|
|
replaced = False
|
|
try:
|
|
try:
|
|
remaining = memoryview(payload)
|
|
while remaining:
|
|
written = os.write(descriptor, remaining)
|
|
if written == 0:
|
|
raise OSError(errno.EIO, "state write made no progress")
|
|
remaining = remaining[written:]
|
|
os.fsync(descriptor)
|
|
finally:
|
|
os.close(descriptor)
|
|
os.replace(temporary, self.path)
|
|
replaced = True
|
|
try:
|
|
directory = os.open(self.path.parent, os.O_RDONLY | os.O_DIRECTORY)
|
|
try:
|
|
os.fsync(directory)
|
|
finally:
|
|
os.close(directory)
|
|
except OSError as exc:
|
|
self.poisoned = True
|
|
raise StateCommitUncertain() from exc
|
|
finally:
|
|
if not replaced:
|
|
try:
|
|
temporary.unlink()
|
|
except FileNotFoundError:
|
|
pass
|
|
|
|
|
|
class Broker:
|
|
def __init__(self, store: StateStore, observer: ReceiptObserver | None = None) -> None:
|
|
self.store = store
|
|
self.observer: ReceiptObserver = observer if observer is not None else UnavailableReceiptObserver()
|
|
# Set only by begin_verification after its mandatory revoke-first fence.
|
|
# It is preserved if later cycle admission is refused; all other broker
|
|
# actions retain the normal snapshot rollback behavior.
|
|
self._rejected_cycle_fence: tuple[dict[str, object], dict[str, dict[str, object]]] | None = None
|
|
# VERIFIED authority is deliberately volatile: broker restart revokes all
|
|
# leases while preserving WI-1 identity and pending-token integrity.
|
|
self.leases: dict[str, dict[str, object]] = {}
|
|
|
|
def authenticate(self, peer_pid: int, request: dict[str, object]) -> tuple[str, dict[str, object]]:
|
|
session_id = request.get("session_id")
|
|
generation = request.get("runtime_generation")
|
|
if not isinstance(session_id, str) or not is_non_negative_integer(generation):
|
|
raise BrokerFailure("INVALID_IDENTITY")
|
|
session = self.store.sessions().get(session_id)
|
|
if not isinstance(session, dict):
|
|
raise BrokerFailure("UNKNOWN_SESSION")
|
|
anchor_pid = session.get("anchor_pid")
|
|
anchor_starttime = session.get("anchor_starttime")
|
|
current_generation = session.get("runtime_generation")
|
|
if (
|
|
type(anchor_pid) is not int
|
|
or not isinstance(anchor_starttime, str)
|
|
or not is_non_negative_integer(current_generation)
|
|
):
|
|
raise BrokerFailure("STATE_INTEGRITY")
|
|
if not verified_ancestry(peer_pid, anchor_pid, anchor_starttime):
|
|
raise BrokerFailure("ANCESTRY_MISMATCH")
|
|
if generation < current_generation:
|
|
raise BrokerFailure("STALE_GENERATION")
|
|
if generation > current_generation:
|
|
session["runtime_generation"] = generation
|
|
self.revoke_session_authority(session_id)
|
|
return session_id, session
|
|
|
|
def session_for_anchor(
|
|
self, anchor_pid: int, anchor_starttime: str
|
|
) -> tuple[str, dict[str, object]] | None:
|
|
for session_id, session in self.store.sessions().items():
|
|
if (
|
|
session["anchor_pid"] == anchor_pid
|
|
and session["anchor_starttime"] == anchor_starttime
|
|
):
|
|
return session_id, session
|
|
return None
|
|
|
|
def revoke_session_tokens(self, session_id: str) -> None:
|
|
tokens = self.store.tokens()
|
|
for token_value in [
|
|
value for value, token in tokens.items() if token["session_id"] == session_id
|
|
]:
|
|
del tokens[token_value]
|
|
|
|
def revoke_session_authority(self, session_id: str) -> None:
|
|
self.revoke_session_tokens(session_id)
|
|
session = self.store.sessions().get(session_id)
|
|
generation = session.get("runtime_generation") if isinstance(session, dict) else None
|
|
self.leases[session_id] = {
|
|
"state": LEASE_UNVERIFIED,
|
|
"runtime_generation": generation,
|
|
}
|
|
|
|
def mint_token(
|
|
self,
|
|
session_id: str,
|
|
generation: int,
|
|
binding: dict[str, object],
|
|
) -> str:
|
|
if len(self.store.tokens()) >= MAX_PENDING_TOKENS:
|
|
raise BrokerFailure("TOKEN_CAPACITY")
|
|
token = secrets.token_hex(32)
|
|
self.store.tokens()[token] = {
|
|
"session_id": session_id,
|
|
"runtime_generation": generation,
|
|
"binding": copy.deepcopy(binding),
|
|
"consumed": False,
|
|
# Receipt evidence is durably committed before this challenge may
|
|
# be consumed and promotion made externally visible.
|
|
"evidence": None,
|
|
}
|
|
return token
|
|
|
|
def finish_promotion(self, session_id: str) -> None:
|
|
lease = self.leases.get(session_id)
|
|
if not isinstance(lease, dict) or lease.get("state") != LEASE_PENDING_PROMOTION:
|
|
raise BrokerFailure("STATE_INTEGRITY")
|
|
lease["state"] = LEASE_VERIFIED
|
|
|
|
def handle(self, peer: tuple[int, int, int], request: dict[str, object]) -> dict[str, object]:
|
|
if self.store.poisoned:
|
|
raise StateCommitUncertain()
|
|
previous = copy.deepcopy(self.store.value)
|
|
previous_leases = copy.deepcopy(self.leases)
|
|
self._rejected_cycle_fence = None
|
|
try:
|
|
response = self._handle(peer, request)
|
|
if self.store.value != previous:
|
|
self.store.commit()
|
|
if request.get("action") == "promote_lease":
|
|
session_id = request.get("session_id")
|
|
if not isinstance(session_id, str):
|
|
raise BrokerFailure("INVALID_IDENTITY")
|
|
self.finish_promotion(session_id)
|
|
response["state"] = LEASE_VERIFIED
|
|
return response
|
|
except StateCommitUncertain:
|
|
raise
|
|
except Exception:
|
|
fence = self._rejected_cycle_fence
|
|
if fence is None:
|
|
self.store.value = previous
|
|
self.leases = previous_leases
|
|
else:
|
|
# A refused re-verification must never resurrect the preceding
|
|
# VERIFIED authority. Preserve only this post-revoke fence;
|
|
# every unrelated partial-write failure still rolls back.
|
|
self.store.value, self.leases = fence
|
|
raise
|
|
finally:
|
|
self._rejected_cycle_fence = None
|
|
|
|
def _handle(self, peer: tuple[int, int, int], request: dict[str, object]) -> dict[str, object]:
|
|
peer_pid, peer_uid, peer_gid = peer
|
|
action = request.get("action")
|
|
if action == "register_anchor":
|
|
if "session_id" in request:
|
|
raise BrokerFailure("CALLER_SESSION_ID_REFUSED")
|
|
generation = request.get("runtime_generation")
|
|
if not is_non_negative_integer(generation):
|
|
raise BrokerFailure("INVALID_GENERATION")
|
|
anchor = proc_node(peer_pid)
|
|
anchor_starttime = str(anchor["starttime"])
|
|
existing = self.session_for_anchor(peer_pid, anchor_starttime)
|
|
if existing is None:
|
|
session_id = secrets.token_hex(32)
|
|
self.store.sessions()[session_id] = {
|
|
"anchor_pid": peer_pid,
|
|
"anchor_starttime": anchor_starttime,
|
|
"runtime_generation": generation,
|
|
}
|
|
else:
|
|
session_id, session = existing
|
|
current_generation = session["runtime_generation"]
|
|
if generation < current_generation:
|
|
raise BrokerFailure("STALE_GENERATION")
|
|
if generation > current_generation:
|
|
session["runtime_generation"] = generation
|
|
self.revoke_session_authority(session_id)
|
|
return {"ok": True, "session_id": session_id, "peer": {"pid": peer_pid, "uid": peer_uid, "gid": peer_gid, "starttime": anchor["starttime"]}}
|
|
if action == "authenticate":
|
|
self.authenticate(peer_pid, request)
|
|
return {"ok": True}
|
|
if action == "mint_token":
|
|
session_id, _ = self.authenticate(peer_pid, request)
|
|
binding = request.get("binding")
|
|
if not valid_binding(binding):
|
|
raise BrokerFailure("INVALID_BINDING")
|
|
token = self.mint_token(session_id, request["runtime_generation"], binding)
|
|
return {"ok": True, "token": token}
|
|
if action == "consume_token":
|
|
session_id, _ = self.authenticate(peer_pid, request)
|
|
token_value = request.get("token")
|
|
token = self.store.tokens().get(token_value) if isinstance(token_value, str) else None
|
|
if not isinstance(token, dict) or token.get("session_id") != session_id or token.get("runtime_generation") != request.get("runtime_generation") or token.get("consumed") is not False:
|
|
raise BrokerFailure("TOKEN_REPLAY")
|
|
del self.store.tokens()[token_value]
|
|
return {"ok": True}
|
|
if action == "begin_verification":
|
|
session_id, _ = self.authenticate(peer_pid, request)
|
|
runtime = request.get("runtime")
|
|
binding = request.get("binding")
|
|
construction = request.get("construction")
|
|
ttl_seconds = request.get("ttl_seconds", MAX_LEASE_TTL_SECONDS)
|
|
if runtime not in READ_ONLY_TOOLS:
|
|
raise BrokerFailure("INVALID_RUNTIME")
|
|
if not valid_binding(binding):
|
|
raise BrokerFailure("INVALID_BINDING")
|
|
if (
|
|
type(ttl_seconds) is not int
|
|
or ttl_seconds <= 0
|
|
or ttl_seconds > MAX_LEASE_TTL_SECONDS
|
|
):
|
|
raise BrokerFailure("INVALID_LEASE_TTL")
|
|
# Revoke-first is a broker operation, not advisory adapter order.
|
|
self.revoke_session_authority(session_id)
|
|
# If subsequent construction admission rejects, handle() restores
|
|
# this fence rather than the pre-cycle VERIFIED snapshot.
|
|
self._rejected_cycle_fence = (
|
|
copy.deepcopy(self.store.value), copy.deepcopy(self.leases)
|
|
)
|
|
try:
|
|
constructed = build_payload_from_wire(construction)
|
|
except ValueError as exc:
|
|
raise BrokerFailure("INVALID_CONSTRUCTION") from exc
|
|
if (
|
|
constructed.injectionDecision != "ACCEPTED"
|
|
or not constructed.promotion
|
|
or not isinstance(constructed.h_source, str)
|
|
or not isinstance(constructed.h_payload, str)
|
|
):
|
|
raise BrokerFailure("PAYLOAD_CONSTRUCTION_REFUSED")
|
|
if (
|
|
not hmac.compare_digest(binding["h_source"], constructed.h_source)
|
|
or not hmac.compare_digest(binding["h_payload"], constructed.h_payload)
|
|
):
|
|
raise BrokerFailure("PAYLOAD_BINDING_MISMATCH")
|
|
cycle_binding = copy.deepcopy(binding)
|
|
cycle_binding["runtime_generation"] = request["runtime_generation"]
|
|
challenge = self.mint_token(session_id, request["runtime_generation"], binding)
|
|
self.leases[session_id] = {
|
|
"state": LEASE_PENDING,
|
|
"runtime": runtime,
|
|
"runtime_generation": request["runtime_generation"],
|
|
"binding": copy.deepcopy(binding),
|
|
"receipt_challenge": challenge,
|
|
"ttl_seconds": ttl_seconds,
|
|
}
|
|
# The broker constructs both the bound challenge and the delivery
|
|
# text. The model's role is an exact copy, never hashing its output.
|
|
return {
|
|
"ok": True,
|
|
"state": LEASE_PENDING,
|
|
"receipt_challenge": challenge,
|
|
"receipt": receipt_for(challenge, cycle_binding),
|
|
"binding": cycle_binding,
|
|
}
|
|
if action == "observe_receipt":
|
|
session_id, _ = self.authenticate(peer_pid, request)
|
|
challenge = request.get("receipt_challenge")
|
|
if "latest_assistant_message" in request:
|
|
raise BrokerFailure("INVALID_RECEIPT")
|
|
if not isinstance(challenge, str):
|
|
raise BrokerFailure("INVALID_RECEIPT")
|
|
token = self.store.tokens().get(challenge)
|
|
if (
|
|
not isinstance(token, dict)
|
|
or token.get("session_id") != session_id
|
|
or token.get("runtime_generation") != request.get("runtime_generation")
|
|
or token.get("consumed") is not False
|
|
):
|
|
raise BrokerFailure("RECEIPT_REPLAY")
|
|
lease = self.leases.get(session_id)
|
|
if not isinstance(lease, dict) or lease.get("state") != LEASE_PENDING:
|
|
raise BrokerFailure("RECEIPT_REPLAY")
|
|
expected_challenge = lease.get("receipt_challenge")
|
|
binding = lease.get("binding")
|
|
if (
|
|
not isinstance(expected_challenge, str)
|
|
or not secrets.compare_digest(challenge, expected_challenge)
|
|
or not isinstance(binding, dict)
|
|
or token.get("binding") != binding
|
|
):
|
|
raise BrokerFailure("RECEIPT_REPLAY")
|
|
receipt_binding = copy.deepcopy(binding)
|
|
receipt_binding["runtime_generation"] = request["runtime_generation"]
|
|
message = self.observer.observe_latest_assistant_message(
|
|
session_id, str(lease["runtime"]), request["runtime_generation"], receipt_binding
|
|
)
|
|
if not isinstance(message, str):
|
|
raise BrokerFailure("RECEIPT_OBSERVATION_UNAVAILABLE")
|
|
# This accepts one exact current-cycle assistant entry only. It is
|
|
# deliberately not a transcript search and rejects quoted/extra text.
|
|
if not is_verbatim_receipt(message, challenge, receipt_binding):
|
|
raise BrokerFailure("RECEIPT_MISMATCH")
|
|
evidence = {"h_latest_assistant": latest_assistant_digest(message)}
|
|
token["evidence"] = evidence
|
|
lease["state"] = LEASE_PENDING_PROMOTION
|
|
lease["evidence"] = copy.deepcopy(evidence)
|
|
return {"ok": True, "state": LEASE_PENDING_PROMOTION}
|
|
if action == "promote_lease":
|
|
session_id, _ = self.authenticate(peer_pid, request)
|
|
challenge = request.get("receipt_challenge")
|
|
if not isinstance(challenge, str):
|
|
raise BrokerFailure("INVALID_RECEIPT")
|
|
lease = self.leases.get(session_id)
|
|
if not isinstance(lease, dict) or lease.get("state") == LEASE_UNVERIFIED:
|
|
raise BrokerFailure("INVALID_LEASE_TRANSITION")
|
|
token = self.store.tokens().get(challenge)
|
|
if (
|
|
not isinstance(token, dict)
|
|
or token.get("session_id") != session_id
|
|
or token.get("runtime_generation") != request.get("runtime_generation")
|
|
or token.get("consumed") is not False
|
|
):
|
|
raise BrokerFailure("RECEIPT_REPLAY")
|
|
if lease.get("state") != LEASE_PENDING_PROMOTION:
|
|
raise BrokerFailure("INVALID_LEASE_TRANSITION")
|
|
expected_challenge = lease.get("receipt_challenge")
|
|
if not isinstance(expected_challenge, str) or not secrets.compare_digest(challenge, expected_challenge):
|
|
raise BrokerFailure("RECEIPT_REPLAY")
|
|
if token.get("binding") != lease.get("binding") or not valid_receipt_evidence(token.get("evidence")):
|
|
raise BrokerFailure("PROMOTION_TOKEN_INVALID")
|
|
del self.store.tokens()[challenge]
|
|
lease["expires_at"] = time.monotonic() + int(lease["ttl_seconds"])
|
|
# handle() commits the evidence-backed consumption before
|
|
# finish_promotion() makes VERIFIED externally visible: promote-last.
|
|
return {"ok": True, "state": LEASE_PENDING_PROMOTION}
|
|
if action == "revoke_lease":
|
|
session_id, _ = self.authenticate(peer_pid, request)
|
|
self.revoke_session_authority(session_id)
|
|
return {"ok": True, "state": LEASE_UNVERIFIED}
|
|
if action == "authorize_tool":
|
|
session_id, _ = self.authenticate(peer_pid, request)
|
|
runtime = request.get("runtime")
|
|
tool_name = request.get("tool_name")
|
|
if runtime not in READ_ONLY_TOOLS:
|
|
raise BrokerFailure("INVALID_RUNTIME")
|
|
if not isinstance(tool_name, str) or not tool_name or len(tool_name) > 256:
|
|
raise BrokerFailure("INVALID_TOOL")
|
|
lease = self.leases.get(session_id)
|
|
state = lease.get("state") if isinstance(lease, dict) else LEASE_UNVERIFIED
|
|
if tool_name in READ_ONLY_TOOLS[runtime] or tool_name == RECOVERY_TOOL:
|
|
return {"ok": True, "decision": "allow", "state": state}
|
|
if not isinstance(lease, dict) or lease.get("state") != LEASE_VERIFIED:
|
|
return {
|
|
"ok": False,
|
|
"code": "MUTATOR_UNVERIFIED",
|
|
"decision": "deny",
|
|
"state": LEASE_UNVERIFIED,
|
|
}
|
|
if (
|
|
lease.get("runtime") != runtime
|
|
or lease.get("runtime_generation") != request.get("runtime_generation")
|
|
):
|
|
return {
|
|
"ok": False,
|
|
"code": "MUTATOR_UNVERIFIED",
|
|
"decision": "deny",
|
|
"state": LEASE_UNVERIFIED,
|
|
}
|
|
expires_at = lease.get("expires_at")
|
|
if not isinstance(expires_at, (int, float)) or time.monotonic() >= expires_at:
|
|
self.revoke_session_authority(session_id)
|
|
return {
|
|
"ok": False,
|
|
"code": "LEASE_EXPIRED",
|
|
"decision": "deny",
|
|
"state": LEASE_UNVERIFIED,
|
|
}
|
|
return {
|
|
"ok": True,
|
|
"decision": "allow",
|
|
"state": LEASE_VERIFIED,
|
|
}
|
|
raise BrokerFailure("UNKNOWN_ACTION")
|
|
|
|
|
|
def read_frame(connection: socket.socket, deadline: float) -> dict[str, object]:
|
|
data = bytearray()
|
|
while len(data) <= MAX_FRAME:
|
|
remaining = deadline - time.monotonic()
|
|
if remaining <= 0:
|
|
raise BrokerFailure("MALFORMED_REQUEST")
|
|
connection.settimeout(remaining)
|
|
chunk = connection.recv(min(4096, MAX_FRAME + 1 - len(data)))
|
|
if not chunk:
|
|
break
|
|
data.extend(chunk)
|
|
if len(data) > MAX_FRAME:
|
|
while True:
|
|
remaining = deadline - time.monotonic()
|
|
if remaining <= 0:
|
|
raise BrokerFailure("MALFORMED_REQUEST")
|
|
connection.settimeout(remaining)
|
|
if not connection.recv(4096):
|
|
break
|
|
raise BrokerFailure("MALFORMED_REQUEST")
|
|
if not data.endswith(b"\n") or data.count(b"\n") != 1:
|
|
raise BrokerFailure("MALFORMED_REQUEST")
|
|
try:
|
|
value = json.loads(data)
|
|
except (json.JSONDecodeError, UnicodeDecodeError) as exc:
|
|
raise BrokerFailure("MALFORMED_REQUEST") from exc
|
|
if not isinstance(value, dict):
|
|
raise BrokerFailure("MALFORMED_REQUEST")
|
|
return value
|
|
|
|
|
|
def handle_connection(
|
|
connection: socket.socket,
|
|
broker: Broker,
|
|
broker_lock: threading.Lock,
|
|
) -> None:
|
|
with connection:
|
|
read_deadline = time.monotonic() + READ_DEADLINE_SECONDS
|
|
try:
|
|
raw = connection.getsockopt(socket.SOL_SOCKET, socket.SO_PEERCRED, 12)
|
|
peer = struct.unpack("3i", raw)
|
|
request = read_frame(connection, read_deadline)
|
|
except BrokerFailure as exc:
|
|
reply = {"ok": False, "code": exc.code}
|
|
except OSError:
|
|
return
|
|
else:
|
|
# Queueing is bounded and fails closed before broker.handle mutates
|
|
# state. Once handling starts it must finish atomically; its reply
|
|
# then receives an independent send budget so a slow fsync cannot
|
|
# consume the write opportunity and create client/broker ambiguity.
|
|
acquired = broker_lock.acquire(timeout=HANDLE_QUEUE_TIMEOUT_SECONDS)
|
|
if not acquired:
|
|
reply = {"ok": False, "code": "BROKER_BUSY"}
|
|
else:
|
|
try:
|
|
try:
|
|
reply = broker.handle(peer, request)
|
|
except BrokerFailure as exc:
|
|
reply = {"ok": False, "code": exc.code}
|
|
finally:
|
|
broker_lock.release()
|
|
try:
|
|
connection.settimeout(SEND_TIMEOUT_SECONDS)
|
|
connection.sendall((json.dumps(reply, separators=(",", ":")) + "\n").encode())
|
|
except OSError:
|
|
return
|
|
|
|
|
|
def serve(
|
|
socket_path: Path,
|
|
state_path: Path,
|
|
observer: ReceiptObserver | None = None,
|
|
) -> None:
|
|
secure_parent(socket_path)
|
|
if socket_path.exists() or socket_path.is_symlink():
|
|
raise BrokerFailure("SOCKET_ALREADY_EXISTS")
|
|
store = StateStore(state_path)
|
|
broker = Broker(store, observer)
|
|
broker_lock = threading.Lock()
|
|
slots = threading.BoundedSemaphore(MAX_IN_FLIGHT_CONNECTIONS)
|
|
fatal_lock = threading.Lock()
|
|
fatal_errors: list[Exception] = []
|
|
executor = ThreadPoolExecutor(
|
|
max_workers=MAX_IN_FLIGHT_CONNECTIONS,
|
|
thread_name_prefix="mosaic-lease-broker",
|
|
)
|
|
server = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
|
|
server.bind(str(socket_path))
|
|
os.chmod(socket_path, 0o600)
|
|
owned = (socket_path.stat().st_dev, socket_path.stat().st_ino)
|
|
stopping = False
|
|
|
|
def stop(_signum: int, _frame: object) -> None:
|
|
nonlocal stopping
|
|
stopping = True
|
|
server.close()
|
|
|
|
def process_connection(connection: socket.socket) -> None:
|
|
try:
|
|
handle_connection(connection, broker, broker_lock)
|
|
except Exception as exc:
|
|
with fatal_lock:
|
|
if not fatal_errors:
|
|
fatal_errors.append(exc)
|
|
finally:
|
|
slots.release()
|
|
|
|
def fatal_error() -> Exception | None:
|
|
with fatal_lock:
|
|
return fatal_errors[0] if fatal_errors else None
|
|
|
|
signal.signal(signal.SIGTERM, stop)
|
|
signal.signal(signal.SIGINT, stop)
|
|
server.listen(MAX_IN_FLIGHT_CONNECTIONS)
|
|
server.settimeout(0.1)
|
|
print("READY", flush=True)
|
|
try:
|
|
while not stopping:
|
|
failure = fatal_error()
|
|
if failure is not None:
|
|
raise failure
|
|
if not slots.acquire(timeout=0.1):
|
|
continue
|
|
try:
|
|
connection, _ = server.accept()
|
|
except socket.timeout:
|
|
slots.release()
|
|
continue
|
|
except OSError:
|
|
slots.release()
|
|
failure = fatal_error()
|
|
if failure is not None:
|
|
raise failure
|
|
if stopping:
|
|
break
|
|
raise
|
|
try:
|
|
executor.submit(process_connection, connection)
|
|
except Exception:
|
|
slots.release()
|
|
connection.close()
|
|
raise
|
|
finally:
|
|
server.close()
|
|
executor.shutdown(wait=True)
|
|
try:
|
|
current = socket_path.stat()
|
|
if (current.st_dev, current.st_ino) == owned:
|
|
socket_path.unlink()
|
|
except FileNotFoundError:
|
|
pass
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--socket", required=True, type=Path)
|
|
parser.add_argument("--state", required=True, type=Path)
|
|
parser.add_argument("--test-observer-file", type=Path)
|
|
arguments = parser.parse_args()
|
|
observer = (
|
|
FileTestReceiptObserver(arguments.test_observer_file)
|
|
if arguments.test_observer_file is not None
|
|
else None
|
|
)
|
|
serve(arguments.socket, arguments.state, observer)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
try:
|
|
main()
|
|
except BrokerFailure as failure:
|
|
print(failure.code, file=sys.stderr)
|
|
raise SystemExit(1)
|
|
except StateCommitUncertain as failure:
|
|
print(str(failure), file=sys.stderr)
|
|
raise SystemExit(1)
|