Files
stack/docs/compaction-refresh/probes/p3_generation_broker.py
2026-07-18 16:15:50 -05:00

232 lines
10 KiB
Python

#!/usr/bin/env python3
"""P3 broker prototype: peercred-keyed runtime_generation and lease revocation."""
from __future__ import annotations
import argparse
import importlib.util
import json
import os
import secrets
import socket
import struct
from collections.abc import Callable, Mapping
from pathlib import Path
from typing import Any
def proc_starttime(pid: int) -> int:
text = Path(f"/proc/{pid}/stat").read_text()
close = text.rfind(")")
return int(text[close + 2 :].split()[19])
def emit(log: Path, value: dict[str, Any]) -> None:
with log.open("a", encoding="utf-8") as out:
out.write(json.dumps(value, sort_keys=True) + "\n")
def load_generation_functions(
path: Path,
) -> tuple[Callable[[Mapping[str, str]], int], Callable[[Mapping[str, str]], int]]:
spec = importlib.util.spec_from_file_location("d4_lease_generation", path)
if spec is None or spec.loader is None:
raise ValueError("generation module is unavailable")
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
reader = getattr(module, "read_runtime_generation", None)
bumper = getattr(module, "bump_runtime_generation", None)
if not callable(reader) or not callable(bumper):
raise ValueError("generation module has no read/bump functions")
return reader, bumper
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--socket", required=True)
parser.add_argument("--log", required=True)
parser.add_argument("--generation-module", required=True, type=Path)
ns = parser.parse_args()
socket_path = Path(ns.socket)
log_path = Path(ns.log)
read_runtime_generation, bump_runtime_generation = load_generation_functions(
ns.generation_module
)
socket_path.parent.mkdir(parents=True, exist_ok=True)
os.chmod(socket_path.parent, 0o700)
socket_path.unlink(missing_ok=True)
log_path.unlink(missing_ok=True)
server = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
server.bind(str(socket_path))
os.chmod(socket_path, 0o600)
server.listen(8)
generations: dict[tuple[int, int], int] = {}
lease_state: dict[tuple[int, int], str] = {}
# The gated launcher registers its own exec-preserved PID here. This is
# deliberately volatile fixture state; nothing is written outside root.
launcher_sessions: dict[tuple[int, int], str] = {}
generation_files: dict[tuple[int, int], Path] = {}
def generation_environment(identity: tuple[int, int]) -> dict[str, str]:
state_path = generation_files.get(identity)
if state_path is None or state_path.parent != socket_path.parent:
raise ValueError("generation file is outside the fixture root")
return {"MOSAIC_LEASE_GENERATION_FILE": str(state_path)}
def file_generation(identity: tuple[int, int]) -> int:
return read_runtime_generation(generation_environment(identity))
emit(log_path, {"event": "listen", "pid": os.getpid(), "socket": str(socket_path)})
while True:
conn, _ = server.accept()
with conn:
raw = conn.getsockopt(socket.SOL_SOCKET, socket.SO_PEERCRED, 12)
pid, uid, gid = struct.unpack("3i", raw)
starttime = proc_starttime(pid)
request = json.loads(conn.makefile("r", encoding="utf-8").readline())
if request.get("action") == "shutdown-broker":
conn.sendall(b'{"ok":true}\n')
break
identity = (pid, starttime)
if request.get("action") == "register_anchor":
generation = request.get("runtime_generation")
if type(generation) is not int or generation < 0:
conn.sendall(b'{"ok":false,"code":"INVALID_GENERATION"}\n')
continue
session_id = launcher_sessions.setdefault(identity, secrets.token_hex(32))
generation_file = socket_path.parent / f"generation-{session_id}.state"
generation_files[identity] = generation_file
record = {
"event": "lease_anchor_registered",
"peercred": {"pid": pid, "uid": uid, "gid": gid},
"starttime_ticks": starttime,
"runtime_generation": generation,
"session_id_shape": "hex-256",
"generation_file": str(generation_file),
}
emit(log_path, record)
reply = {
"ok": True,
"session_id": session_id,
"peer": {"pid": pid, "uid": uid, "gid": gid, "starttime": str(starttime)},
}
conn.sendall((json.dumps(reply, sort_keys=True) + "\n").encode())
continue
# The D4 extension requests this at each post-start lifecycle
# boundary; the exact WI-3 helper mutates the launcher-created file.
if request.get("action") == "bump-generation":
generation = bump_runtime_generation(generation_environment(identity))
record = {
"event": "generation_state_bumped",
"peercred": {"pid": pid, "uid": uid, "gid": gid},
"starttime_ticks": starttime,
"generation": generation,
"generation_file": str(generation_files[identity]),
"generation_source": "state-file",
}
emit(log_path, record)
conn.sendall((json.dumps(record, sort_keys=True) + "\n").encode())
continue
if request.get("action") == "promote-probe":
generation = file_generation(identity)
lease_state[identity] = "VERIFIED"
record = {
"event": "probe_lease_promoted",
"peercred": {"pid": pid, "uid": uid, "gid": gid},
"starttime_ticks": starttime,
"generation": generation,
"generation_file": str(generation_files[identity]),
"generation_source": "state-file",
"new_lease_state": "VERIFIED",
}
emit(log_path, record)
conn.sendall((json.dumps(record, sort_keys=True) + "\n").encode())
continue
# D4 fixture-only authorization observation. It exposes the broker's
# current versus superseded generation disposition without changing it.
if request.get("action") == "authorize-probe":
generation = request.get("generation")
if type(generation) is not int or generation < 0:
conn.sendall(b'{"ok":false,"code":"INVALID_GENERATION"}\n')
continue
current_generation = file_generation(identity)
current_lease = lease_state.get(identity, "NONE")
if generation < current_generation:
code = "STALE_GENERATION"
elif generation > current_generation:
code = "FUTURE_GENERATION"
elif current_lease != "VERIFIED":
code = "MUTATOR_UNVERIFIED"
else:
code = "ALLOW"
record = {
"event": "generation_authorization",
"peercred": {"pid": pid, "uid": uid, "gid": gid},
"starttime_ticks": starttime,
"requested_generation": generation,
"current_generation": current_generation,
"generation_file": str(generation_files[identity]),
"generation_source": "state-file",
"lease_state": current_lease,
"ok": code == "ALLOW",
"code": code,
}
emit(log_path, record)
conn.sendall((json.dumps(record, sort_keys=True) + "\n").encode())
continue
if request.get("action") == "source-invalid":
old_lease = lease_state.get(identity, "NONE")
lease_state[identity] = "REVOKED"
record = {
"event": "source_invalidation_revoke",
"peercred": {"pid": pid, "uid": uid, "gid": gid},
"starttime_ticks": starttime,
"generation": generations.get(identity, 0),
"source_reason": request.get("reason"),
"prior_lease": old_lease,
"new_lease_state": "REVOKED",
"promotion": False,
}
emit(log_path, record)
conn.sendall((json.dumps(record, sort_keys=True) + "\n").encode())
continue
if request.get("action") != "lifecycle":
conn.sendall(b'{"ok":false,"reason":"invalid-action"}\n')
continue
old_generation = generations.get(identity, 0)
old_lease = lease_state.get(identity, "NONE")
new_generation = file_generation(identity)
if new_generation <= old_generation:
conn.sendall(b'{"ok":false,"code":"NON_MONOTONIC_STATE_FILE"}\n')
continue
generations[identity] = new_generation
# Every lifecycle boundary revokes first. A start establishes a new
# UNVERIFIED incarnation; it never inherits prior VERIFIED state.
lease_state[identity] = "UNVERIFIED" if request.get("phase") == "start" else "REVOKED"
record = {
"event": "runtime_generation_bump",
"peercred": {"pid": pid, "uid": uid, "gid": gid},
"starttime_ticks": starttime,
"phase": request.get("phase"),
"reason": request.get("reason"),
"old_generation": old_generation,
"new_generation": new_generation,
"generation_file": str(generation_files[identity]),
"generation_source": "state-file",
"prior_lease": old_lease,
"prior_lease_revoked": True,
"new_lease_state": lease_state[identity],
}
emit(log_path, record)
conn.sendall((json.dumps(record, sort_keys=True) + "\n").encode())
server.close()
socket_path.unlink(missing_ok=True)
if __name__ == "__main__":
main()