748 lines
31 KiB
Python
748 lines
31 KiB
Python
#!/usr/bin/env python3
|
|
"""Branch-focused tests for the lease-gated runtime executables."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import importlib.util
|
|
import io
|
|
import json
|
|
import os
|
|
import runpy
|
|
import socket
|
|
import stat
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import threading
|
|
import unittest
|
|
from contextlib import redirect_stderr
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from unittest.mock import patch
|
|
|
|
|
|
TOOLS_DIR = Path(__file__).parents[2] / "framework/tools/lease-broker"
|
|
if str(TOOLS_DIR) not in sys.path:
|
|
sys.path.insert(0, str(TOOLS_DIR))
|
|
|
|
|
|
def load_tool(module_name: str, filename: str):
|
|
spec = importlib.util.spec_from_file_location(module_name, TOOLS_DIR / filename)
|
|
if spec is None or spec.loader is None:
|
|
raise RuntimeError(f"unable to load {filename}")
|
|
module = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
LAUNCHER = load_tool("lease_runtime_launcher", "launch-runtime.py")
|
|
GATE = load_tool("lease_mutator_gate", "mutator-gate.py")
|
|
|
|
|
|
class FakeSocket:
|
|
def __init__(self, *chunks: bytes):
|
|
self.chunks = list(chunks)
|
|
self.timeout = None
|
|
self.connected = None
|
|
self.sent = b""
|
|
self.shutdown_how = None
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, *_args):
|
|
return False
|
|
|
|
def settimeout(self, value: float) -> None:
|
|
self.timeout = value
|
|
|
|
def connect(self, value: str) -> None:
|
|
self.connected = value
|
|
|
|
def sendall(self, value: bytes) -> None:
|
|
self.sent += value
|
|
|
|
def shutdown(self, how: int) -> None:
|
|
self.shutdown_how = how
|
|
|
|
def recv(self, _size: int) -> bytes:
|
|
return self.chunks.pop(0) if self.chunks else b""
|
|
|
|
|
|
class LaunchRuntimeTest(unittest.TestCase):
|
|
def test_success_registers_then_injects_session_before_exec(self) -> None:
|
|
calls: dict[str, object] = {}
|
|
session_id = "a" * 64
|
|
|
|
def request(path: Path, payload: dict[str, object]) -> dict[str, object]:
|
|
calls["path"] = path
|
|
calls["request"] = payload
|
|
return {"ok": True, "session_id": session_id}
|
|
|
|
def execute(command: str, argv: list[str], environment: dict[str, str]) -> None:
|
|
calls["execute"] = (command, argv, environment)
|
|
|
|
def initialize_generation(path: Path, generation: int) -> None:
|
|
calls["generation"] = (path, generation)
|
|
|
|
result = LAUNCHER.main(
|
|
["--runtime", "claude", "--", "claude", "--print", "hello"],
|
|
environ={
|
|
"MOSAIC_LEASE_BROKER_SOCKET": "/run/test/broker.sock",
|
|
"MOSAIC_RUNTIME_GENERATION": "7",
|
|
"PRESERVED": "yes",
|
|
},
|
|
request=request,
|
|
execute=execute,
|
|
initialize_generation=initialize_generation,
|
|
)
|
|
|
|
self.assertEqual(result, 0)
|
|
self.assertEqual(calls["path"], Path("/run/test/broker.sock"))
|
|
self.assertEqual(
|
|
calls["request"],
|
|
{"action": "register_anchor", "runtime_generation": 7},
|
|
)
|
|
command, argv, environment = calls["execute"]
|
|
self.assertEqual(command, "claude")
|
|
self.assertEqual(argv, ["claude", "--print", "hello"])
|
|
self.assertEqual(environment["MOSAIC_LEASE_SESSION_ID"], session_id)
|
|
self.assertEqual(environment["MOSAIC_RUNTIME_GENERATION"], "7")
|
|
self.assertEqual(environment["MOSAIC_LEASE_RUNTIME"], "claude")
|
|
self.assertEqual(
|
|
environment["MOSAIC_LEASE_GENERATION_FILE"],
|
|
f"/run/test/generation-{session_id}.state",
|
|
)
|
|
self.assertEqual(
|
|
calls["generation"],
|
|
(Path(f"/run/test/generation-{session_id}.state"), 7),
|
|
)
|
|
self.assertEqual(environment["PRESERVED"], "yes")
|
|
|
|
def test_dangerous_claude_mode_is_owned_and_injected_by_the_wrapper(self) -> None:
|
|
executed: list[tuple[str, list[str], dict[str, str]]] = []
|
|
result = LAUNCHER.main(
|
|
["--runtime", "claude", "--dangerous", "--", "claude", "-p", "hello"],
|
|
environ={"MOSAIC_LEASE_BROKER_SOCKET": "/broker"},
|
|
request=lambda *_args: {"ok": True, "session_id": "e" * 64},
|
|
execute=lambda *args: executed.append(args),
|
|
initialize_generation=lambda *_args: None,
|
|
)
|
|
self.assertEqual(result, 0)
|
|
self.assertEqual(
|
|
executed[0][1],
|
|
["claude", "--dangerously-skip-permissions", "-p", "hello"],
|
|
)
|
|
|
|
with redirect_stderr(io.StringIO()):
|
|
self.assertEqual(
|
|
LAUNCHER.main(
|
|
["--runtime", "pi", "--dangerous", "--", "pi"],
|
|
environ={"MOSAIC_LEASE_BROKER_SOCKET": "/broker"},
|
|
request=lambda *_args: {"ok": True, "session_id": "e" * 64},
|
|
execute=lambda *_args: self.fail("invalid dangerous runtime executed"),
|
|
),
|
|
64,
|
|
)
|
|
|
|
def test_command_without_separator_is_forwarded_unchanged(self) -> None:
|
|
executed: list[tuple[str, list[str], dict[str, str]]] = []
|
|
result = LAUNCHER.main(
|
|
["--runtime", "pi", "pi", "--print", "hello"],
|
|
environ={"MOSAIC_LEASE_BROKER_SOCKET": "/broker"},
|
|
request=lambda *_args: {"ok": True, "session_id": "f" * 64},
|
|
execute=lambda *args: executed.append(args),
|
|
initialize_generation=lambda *_args: None,
|
|
)
|
|
self.assertEqual(result, 0)
|
|
self.assertEqual(executed[0][0:2], ("pi", ["pi", "--print", "hello"]))
|
|
|
|
def test_missing_command_is_usage_error(self) -> None:
|
|
with redirect_stderr(io.StringIO()):
|
|
self.assertEqual(
|
|
LAUNCHER.main(
|
|
["--runtime", "pi", "--"],
|
|
environ={},
|
|
request=lambda *_args: {},
|
|
execute=lambda *_args: None,
|
|
),
|
|
64,
|
|
)
|
|
|
|
def test_registration_validation_and_environment_fail_closed(self) -> None:
|
|
good_session = "b" * 64
|
|
cases = [
|
|
({}, {"ok": True, "session_id": good_session}),
|
|
({"MOSAIC_LEASE_BROKER_SOCKET": "/x", "MOSAIC_RUNTIME_GENERATION": "bad"}, {}),
|
|
({"MOSAIC_LEASE_BROKER_SOCKET": "/x", "MOSAIC_RUNTIME_GENERATION": "-1"}, {}),
|
|
({"MOSAIC_LEASE_BROKER_SOCKET": "/x"}, {"ok": False, "session_id": good_session}),
|
|
({"MOSAIC_LEASE_BROKER_SOCKET": "/x"}, {"ok": True, "session_id": 4}),
|
|
({"MOSAIC_LEASE_BROKER_SOCKET": "/x"}, {"ok": True, "session_id": "b" * 63}),
|
|
({"MOSAIC_LEASE_BROKER_SOCKET": "/x"}, {"ok": True, "session_id": "z" * 64}),
|
|
]
|
|
for environment, reply in cases:
|
|
with self.subTest(environment=environment, reply=reply), redirect_stderr(io.StringIO()):
|
|
executed: list[object] = []
|
|
result = LAUNCHER.main(
|
|
["--runtime", "pi", "--", "pi"],
|
|
environ=environment,
|
|
request=lambda *_args, value=reply: value,
|
|
execute=lambda *args: executed.append(args),
|
|
)
|
|
self.assertEqual(result, 1)
|
|
self.assertEqual(executed, [])
|
|
|
|
def test_generation_initialization_failure_denies_before_exec(self) -> None:
|
|
with redirect_stderr(io.StringIO()):
|
|
self.assertEqual(
|
|
LAUNCHER.main(
|
|
["--runtime", "pi", "--", "pi"],
|
|
environ={"MOSAIC_LEASE_BROKER_SOCKET": "/broker"},
|
|
request=lambda *_args: {"ok": True, "session_id": "a" * 64},
|
|
execute=lambda *_args: self.fail("must not execute"),
|
|
initialize_generation=lambda *_args: (_ for _ in ()).throw(
|
|
OSError("unsafe state")
|
|
),
|
|
),
|
|
1,
|
|
)
|
|
|
|
def test_registration_exceptions_fail_closed(self) -> None:
|
|
failures = [ValueError("bad"), OSError("down"), json.JSONDecodeError("bad", "x", 0)]
|
|
for failure in failures:
|
|
with self.subTest(failure=type(failure).__name__), redirect_stderr(io.StringIO()):
|
|
def request(*_args, error=failure):
|
|
raise error
|
|
|
|
self.assertEqual(
|
|
LAUNCHER.main(
|
|
["--runtime", "claude", "--", "claude"],
|
|
environ={"MOSAIC_LEASE_BROKER_SOCKET": "/x"},
|
|
request=request,
|
|
execute=lambda *_args: self.fail("must not execute"),
|
|
),
|
|
1,
|
|
)
|
|
|
|
def test_exec_failure_is_fail_closed(self) -> None:
|
|
with redirect_stderr(io.StringIO()):
|
|
self.assertEqual(
|
|
LAUNCHER.main(
|
|
["--runtime", "pi", "--", "pi"],
|
|
environ={"MOSAIC_LEASE_BROKER_SOCKET": "/x"},
|
|
request=lambda *_args: {"ok": True, "session_id": "c" * 64},
|
|
execute=lambda *_args: (_ for _ in ()).throw(OSError("missing")),
|
|
initialize_generation=lambda *_args: None,
|
|
),
|
|
1,
|
|
)
|
|
|
|
def test_broker_reply_framing_and_shape_validation(self) -> None:
|
|
replies = [
|
|
(b'{"ok":true}\n', {"ok": True}),
|
|
(b'{"ok":true}', ValueError),
|
|
(b'[]\n', ValueError),
|
|
(b"x" * (LAUNCHER.MAX_FRAME + 1), ValueError),
|
|
]
|
|
for wire_reply, expected in replies:
|
|
with self.subTest(size=len(wire_reply)):
|
|
fake = FakeSocket(wire_reply)
|
|
with patch.object(LAUNCHER.socket, "socket", return_value=fake):
|
|
if isinstance(expected, type) and issubclass(expected, Exception):
|
|
with self.assertRaises(expected):
|
|
LAUNCHER.broker_request(Path("/broker"), {"action": "register_anchor"})
|
|
else:
|
|
self.assertEqual(
|
|
LAUNCHER.broker_request(Path("/broker"), {"action": "register_anchor"}),
|
|
expected,
|
|
)
|
|
self.assertEqual(fake.timeout, LAUNCHER.BROKER_TIMEOUT_SECONDS)
|
|
self.assertEqual(fake.connected, "/broker")
|
|
self.assertEqual(fake.shutdown_how, socket.SHUT_WR)
|
|
|
|
|
|
class ExecutableEntrypointTest(unittest.TestCase):
|
|
def test_real_claude_and_pi_gates_fail_closed_on_empty_or_truncated_reply(self) -> None:
|
|
for runtime in ("claude", "pi"):
|
|
for wire_reply in (b"", b'{"ok":true'):
|
|
with self.subTest(runtime=runtime, wire_reply=wire_reply), tempfile.TemporaryDirectory() as root:
|
|
socket_path = Path(root) / "broker.sock"
|
|
server = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
|
|
server.bind(str(socket_path))
|
|
server.listen(1)
|
|
|
|
def serve_reply() -> None:
|
|
with server:
|
|
connection, _ = server.accept()
|
|
with connection:
|
|
while connection.recv(4096):
|
|
pass
|
|
if wire_reply:
|
|
connection.sendall(wire_reply)
|
|
|
|
thread = threading.Thread(target=serve_reply, daemon=True)
|
|
thread.start()
|
|
environment = {
|
|
**os.environ,
|
|
"MOSAIC_LEASE_BROKER_SOCKET": str(socket_path),
|
|
"MOSAIC_LEASE_SESSION_ID": "d" * 64,
|
|
"MOSAIC_RUNTIME_GENERATION": "1",
|
|
}
|
|
try:
|
|
result = subprocess.run(
|
|
[sys.executable, str(TOOLS_DIR / "mutator-gate.py"), "--runtime", runtime],
|
|
input=b'{"tool_name":"Read"}\n',
|
|
capture_output=True,
|
|
env=environment,
|
|
check=False,
|
|
timeout=5,
|
|
)
|
|
except subprocess.TimeoutExpired as exc:
|
|
server.close()
|
|
thread.join(timeout=2)
|
|
self.fail(
|
|
f"{runtime} gate hung on wire reply {wire_reply!r}: {exc}"
|
|
)
|
|
thread.join(timeout=2)
|
|
|
|
self.assertFalse(thread.is_alive())
|
|
self.assertEqual(result.returncode, 2)
|
|
self.assertIn(b"GATE_UNAVAILABLE", result.stderr)
|
|
|
|
def test_launcher_entrypoint_returns_usage_without_a_command(self) -> None:
|
|
with patch.object(
|
|
sys,
|
|
"argv",
|
|
[str(TOOLS_DIR / "launch-runtime.py"), "--runtime", "claude"],
|
|
), redirect_stderr(io.StringIO()), self.assertRaises(SystemExit) as raised:
|
|
runpy.run_path(str(TOOLS_DIR / "launch-runtime.py"), run_name="__main__")
|
|
self.assertEqual(raised.exception.code, 64)
|
|
|
|
def test_revoker_entrypoint_denies_when_identity_environment_is_absent(self) -> None:
|
|
with patch.object(
|
|
sys,
|
|
"argv",
|
|
[
|
|
str(TOOLS_DIR / "revoke-lease.py"),
|
|
"--runtime",
|
|
"claude",
|
|
"--reason",
|
|
"pre-compact",
|
|
],
|
|
), patch.dict(os.environ, {}, clear=True), redirect_stderr(
|
|
io.StringIO()
|
|
), self.assertRaises(SystemExit) as raised:
|
|
runpy.run_path(str(TOOLS_DIR / "revoke-lease.py"), run_name="__main__")
|
|
self.assertEqual(raised.exception.code, 2)
|
|
|
|
def test_gate_entrypoint_denies_when_identity_environment_is_absent(self) -> None:
|
|
class Stdin:
|
|
buffer = io.BytesIO(b'{"tool_name":"Bash"}')
|
|
|
|
with patch.object(
|
|
sys,
|
|
"argv",
|
|
[str(TOOLS_DIR / "mutator-gate.py"), "--runtime", "claude"],
|
|
), patch.object(sys, "stdin", Stdin()), patch.dict(
|
|
os.environ, {}, clear=True
|
|
), redirect_stderr(io.StringIO()), self.assertRaises(SystemExit) as raised:
|
|
runpy.run_path(str(TOOLS_DIR / "mutator-gate.py"), run_name="__main__")
|
|
self.assertEqual(raised.exception.code, 2)
|
|
|
|
|
|
class MutatorGateTest(unittest.TestCase):
|
|
@staticmethod
|
|
def environment() -> dict[str, str]:
|
|
return {
|
|
"MOSAIC_LEASE_BROKER_SOCKET": "/run/test/broker.sock",
|
|
"MOSAIC_LEASE_SESSION_ID": "d" * 64,
|
|
"MOSAIC_RUNTIME_GENERATION": "2",
|
|
}
|
|
|
|
def run_main(self, *, tool: object = "Bash", reply: dict[str, object] | None = None):
|
|
calls: list[tuple[Path, dict[str, object]]] = []
|
|
|
|
def request(path: Path, payload: dict[str, object]) -> dict[str, object]:
|
|
calls.append((path, payload))
|
|
return reply if reply is not None else {"ok": True, "decision": "allow"}
|
|
|
|
stderr = io.StringIO()
|
|
with redirect_stderr(stderr):
|
|
result = GATE.main(
|
|
["--runtime", "claude"],
|
|
environ=self.environment(),
|
|
stream=io.BytesIO(json.dumps({"tool_name": tool}).encode()),
|
|
request=request,
|
|
)
|
|
return result, stderr.getvalue(), calls
|
|
|
|
def test_generation_file_is_the_effective_generation_authority(self) -> None:
|
|
calls: list[dict[str, object]] = []
|
|
result = GATE.main(
|
|
["--runtime", "pi"],
|
|
environ=self.environment(),
|
|
stream=io.BytesIO(b'{"tool_name":"bash"}'),
|
|
request=lambda _path, payload: calls.append(payload)
|
|
or {"ok": True, "decision": "allow"},
|
|
resolve_generation=lambda _environment: 9,
|
|
)
|
|
self.assertEqual(result, 0)
|
|
self.assertEqual(calls[0]["runtime_generation"], 9)
|
|
|
|
def test_allow_and_denial_decisions(self) -> None:
|
|
allowed, allowed_stderr, calls = self.run_main()
|
|
self.assertEqual(allowed, 0)
|
|
self.assertEqual(allowed_stderr, "")
|
|
self.assertEqual(calls[0][0], Path("/run/test/broker.sock"))
|
|
self.assertEqual(
|
|
calls[0][1],
|
|
{
|
|
"action": "authorize_tool",
|
|
"session_id": "d" * 64,
|
|
"runtime_generation": 2,
|
|
"runtime": "claude",
|
|
"tool_name": "Bash",
|
|
},
|
|
)
|
|
|
|
denied, denied_stderr, _ = self.run_main(reply={"ok": False, "code": "LEASE_EXPIRED"})
|
|
self.assertEqual(denied, 2)
|
|
self.assertIn("LEASE_EXPIRED", denied_stderr)
|
|
|
|
defaulted, defaulted_stderr, _ = self.run_main(reply={"ok": False, "code": 4})
|
|
self.assertEqual(defaulted, 2)
|
|
self.assertIn("MUTATOR_UNVERIFIED", defaulted_stderr)
|
|
|
|
def test_input_validation_fails_closed(self) -> None:
|
|
payloads = [
|
|
b"x" * (GATE.MAX_FRAME + 1),
|
|
b"[]",
|
|
b"{}",
|
|
json.dumps({"tool_name": ""}).encode(),
|
|
json.dumps({"tool_name": 4}).encode(),
|
|
json.dumps({"tool_name": "x" * 257}).encode(),
|
|
b"not-json",
|
|
]
|
|
for payload in payloads:
|
|
with self.subTest(size=len(payload)), redirect_stderr(io.StringIO()):
|
|
self.assertEqual(
|
|
GATE.main(
|
|
["--runtime", "pi"],
|
|
environ=self.environment(),
|
|
stream=io.BytesIO(payload),
|
|
request=lambda *_args: self.fail("invalid input reached broker"),
|
|
),
|
|
2,
|
|
)
|
|
|
|
def test_environment_generation_and_request_failures_deny(self) -> None:
|
|
environments = [
|
|
{},
|
|
{**self.environment(), "MOSAIC_RUNTIME_GENERATION": "bad"},
|
|
{**self.environment(), "MOSAIC_RUNTIME_GENERATION": "-1"},
|
|
]
|
|
for environment in environments:
|
|
with self.subTest(environment=environment), redirect_stderr(io.StringIO()):
|
|
self.assertEqual(
|
|
GATE.main(
|
|
["--runtime", "claude"],
|
|
environ=environment,
|
|
stream=io.BytesIO(b'{"tool_name":"Read"}'),
|
|
request=lambda *_args: {},
|
|
),
|
|
2,
|
|
)
|
|
|
|
failures = [ValueError("bad"), OSError("down"), json.JSONDecodeError("bad", "x", 0)]
|
|
for failure in failures:
|
|
with self.subTest(failure=type(failure).__name__), redirect_stderr(io.StringIO()):
|
|
def request(*_args, error=failure):
|
|
raise error
|
|
|
|
self.assertEqual(
|
|
GATE.main(
|
|
["--runtime", "claude"],
|
|
environ=self.environment(),
|
|
stream=io.BytesIO(b'{"tool_name":"Read"}'),
|
|
request=request,
|
|
),
|
|
2,
|
|
)
|
|
|
|
def test_broker_request_framing_payload_and_shape_validation(self) -> None:
|
|
with self.assertRaises(ValueError):
|
|
GATE.broker_request(Path("/broker"), {"session_id": "x" * GATE.MAX_FRAME})
|
|
|
|
replies = [
|
|
(b'{"ok":true,"decision":"allow"}\n', {"ok": True, "decision": "allow"}),
|
|
(b'{"ok":true}', ValueError),
|
|
(b'[]\n', ValueError),
|
|
(b"x" * (GATE.MAX_FRAME + 1), ValueError),
|
|
]
|
|
for wire_reply, expected in replies:
|
|
with self.subTest(size=len(wire_reply)):
|
|
fake = FakeSocket(wire_reply)
|
|
with patch.object(GATE.socket, "socket", return_value=fake):
|
|
if isinstance(expected, type) and issubclass(expected, Exception):
|
|
with self.assertRaises(expected):
|
|
GATE.broker_request(Path("/broker"), {"action": "authorize_tool"})
|
|
else:
|
|
self.assertEqual(
|
|
GATE.broker_request(Path("/broker"), {"action": "authorize_tool"}),
|
|
expected,
|
|
)
|
|
self.assertEqual(fake.timeout, GATE.BROKER_TIMEOUT_SECONDS)
|
|
self.assertEqual(fake.connected, "/broker")
|
|
self.assertEqual(fake.shutdown_how, socket.SHUT_WR)
|
|
|
|
|
|
class LeaseRevocationTest(unittest.TestCase):
|
|
def load_modules(self):
|
|
generation_path = TOOLS_DIR / "lease_generation.py"
|
|
revoker_path = TOOLS_DIR / "revoke-lease.py"
|
|
self.assertTrue(generation_path.is_file(), "lease_generation.py must be shipped")
|
|
self.assertTrue(revoker_path.is_file(), "revoke-lease.py must be shipped")
|
|
return (
|
|
load_tool("lease_generation_test", "lease_generation.py"),
|
|
load_tool("lease_revoker_test", "revoke-lease.py"),
|
|
)
|
|
|
|
def test_generation_parser_and_descriptor_security_rejections(self) -> None:
|
|
generation, _ = self.load_modules()
|
|
for value in (None, "", "-1", "no", "é", str(generation.MAX_GENERATION + 1)):
|
|
with self.subTest(value=value), self.assertRaises(ValueError):
|
|
generation.parse_generation(value)
|
|
for value in (-1, generation.MAX_GENERATION + 1):
|
|
with self.subTest(initialize=value), tempfile.TemporaryDirectory() as directory, self.assertRaises(ValueError):
|
|
generation.initialize_runtime_generation(Path(directory) / "state", value)
|
|
|
|
with patch.object(
|
|
generation.os,
|
|
"fstat",
|
|
return_value=SimpleNamespace(st_mode=stat.S_IFDIR | 0o700, st_uid=os.geteuid(), st_size=0),
|
|
), self.assertRaises(ValueError):
|
|
generation._validate_descriptor(4)
|
|
with patch.object(
|
|
generation.os,
|
|
"fstat",
|
|
return_value=SimpleNamespace(st_mode=stat.S_IFREG | 0o600, st_uid=os.geteuid() + 1, st_size=0),
|
|
), self.assertRaises(ValueError):
|
|
generation._validate_descriptor(4)
|
|
with patch.object(
|
|
generation.os,
|
|
"fstat",
|
|
return_value=SimpleNamespace(st_mode=stat.S_IFREG | 0o644, st_uid=os.geteuid(), st_size=0),
|
|
), self.assertRaises(ValueError):
|
|
generation._validate_descriptor(4)
|
|
with patch.object(
|
|
generation.os,
|
|
"fstat",
|
|
return_value=SimpleNamespace(
|
|
st_mode=stat.S_IFREG | 0o600,
|
|
st_uid=os.geteuid(),
|
|
st_size=generation.MAX_GENERATION_BYTES + 1,
|
|
),
|
|
), self.assertRaises(ValueError):
|
|
generation._validate_descriptor(4)
|
|
|
|
def test_generation_file_is_private_monotonic_and_rejects_unsafe_state(self) -> None:
|
|
generation, _ = self.load_modules()
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
path = Path(directory) / "runtime.generation"
|
|
generation.initialize_runtime_generation(path, 4)
|
|
self.assertEqual(generation.read_runtime_generation({
|
|
"MOSAIC_RUNTIME_GENERATION": "1",
|
|
"MOSAIC_LEASE_GENERATION_FILE": str(path),
|
|
}), 4)
|
|
self.assertEqual(generation.bump_runtime_generation({
|
|
"MOSAIC_RUNTIME_GENERATION": "1",
|
|
"MOSAIC_LEASE_GENERATION_FILE": str(path),
|
|
}), 5)
|
|
self.assertEqual(path.read_text(), "5\n")
|
|
self.assertEqual(stat.S_IMODE(path.stat().st_mode), 0o600)
|
|
|
|
path.write_text("bad\n")
|
|
with self.assertRaises(ValueError):
|
|
generation.read_runtime_generation({
|
|
"MOSAIC_RUNTIME_GENERATION": "1",
|
|
"MOSAIC_LEASE_GENERATION_FILE": str(path),
|
|
})
|
|
path.write_bytes(b"\xff\n")
|
|
with self.assertRaises(ValueError):
|
|
generation.read_runtime_generation({
|
|
"MOSAIC_LEASE_GENERATION_FILE": str(path),
|
|
})
|
|
path.write_text(f"{generation.MAX_GENERATION}\n")
|
|
with self.assertRaises(ValueError):
|
|
generation.bump_runtime_generation({
|
|
"MOSAIC_LEASE_GENERATION_FILE": str(path),
|
|
})
|
|
with self.assertRaises(ValueError):
|
|
generation.bump_runtime_generation({"MOSAIC_RUNTIME_GENERATION": "1"})
|
|
|
|
def test_generation_write_must_make_progress(self) -> None:
|
|
generation, _ = self.load_modules()
|
|
with tempfile.TemporaryDirectory() as directory, patch.object(
|
|
generation.os, "write", return_value=0
|
|
), self.assertRaises(OSError):
|
|
generation.initialize_runtime_generation(Path(directory) / "state", 1)
|
|
|
|
def test_revoker_reuses_broker_revoke_and_optional_generation_bump(self) -> None:
|
|
_, revoker = self.load_modules()
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
generation_file = Path(directory) / "runtime.generation"
|
|
generation_file.write_text("2\n")
|
|
generation_file.chmod(0o600)
|
|
environment = {
|
|
"MOSAIC_LEASE_BROKER_SOCKET": "/broker",
|
|
"MOSAIC_LEASE_SESSION_ID": "a" * 64,
|
|
"MOSAIC_RUNTIME_GENERATION": "2",
|
|
"MOSAIC_LEASE_GENERATION_FILE": str(generation_file),
|
|
}
|
|
calls: list[tuple[Path, dict[str, object]]] = []
|
|
result = revoker.main(
|
|
["--runtime", "pi", "--reason", "session-start-resume", "--bump-generation"],
|
|
environ=environment,
|
|
request=lambda path, payload: calls.append((path, payload))
|
|
or {"ok": True, "state": "UNVERIFIED"},
|
|
)
|
|
self.assertEqual(result, 0)
|
|
self.assertEqual(generation_file.read_text(), "3\n")
|
|
self.assertEqual(calls[0][0], Path("/broker"))
|
|
self.assertEqual(calls[0][1], {
|
|
"action": "revoke_lease",
|
|
"session_id": "a" * 64,
|
|
"runtime_generation": 3,
|
|
"reason": "session-start-resume",
|
|
"runtime": "pi",
|
|
})
|
|
|
|
calls.clear()
|
|
self.assertEqual(
|
|
revoker.main(
|
|
["--runtime", "pi", "--reason", "pi-context-after-compact"],
|
|
environ=environment,
|
|
request=lambda path, payload: calls.append((path, payload))
|
|
or {"ok": True, "state": "UNVERIFIED"},
|
|
),
|
|
0,
|
|
)
|
|
self.assertEqual(calls[0][1]["runtime_generation"], 3)
|
|
|
|
def test_revoker_broker_framing_and_shape_validation(self) -> None:
|
|
_, revoker = self.load_modules()
|
|
with self.assertRaises(ValueError):
|
|
revoker.broker_request(Path("/broker"), {"reason": "x" * revoker.MAX_FRAME})
|
|
replies = [
|
|
(b'{"ok":true,"state":"UNVERIFIED"}\n', {"ok": True, "state": "UNVERIFIED"}),
|
|
(b'{"ok":true}', ValueError),
|
|
(b'[]\n', ValueError),
|
|
(b'x' * (revoker.MAX_FRAME + 1), ValueError),
|
|
]
|
|
for wire_reply, expected in replies:
|
|
with self.subTest(size=len(wire_reply)):
|
|
fake = FakeSocket(wire_reply)
|
|
with patch.object(revoker.socket, "socket", return_value=fake):
|
|
if isinstance(expected, type) and issubclass(expected, Exception):
|
|
with self.assertRaises(expected):
|
|
revoker.broker_request(Path("/broker"), {"action": "revoke_lease"})
|
|
else:
|
|
self.assertEqual(
|
|
revoker.broker_request(Path("/broker"), {"action": "revoke_lease"}),
|
|
expected,
|
|
)
|
|
self.assertEqual(fake.timeout, revoker.BROKER_TIMEOUT_SECONDS)
|
|
self.assertEqual(fake.connected, "/broker")
|
|
self.assertEqual(fake.shutdown_how, socket.SHUT_WR)
|
|
|
|
def test_failed_observer_revocation_advances_the_local_generation_fence(self) -> None:
|
|
_, revoker = self.load_modules()
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
generation_file = Path(directory) / "runtime.generation"
|
|
generation_file.write_text("8\n")
|
|
generation_file.chmod(0o600)
|
|
environment = {
|
|
"MOSAIC_LEASE_BROKER_SOCKET": "/broker",
|
|
"MOSAIC_LEASE_SESSION_ID": "a" * 64,
|
|
"MOSAIC_RUNTIME_GENERATION": "8",
|
|
"MOSAIC_LEASE_GENERATION_FILE": str(generation_file),
|
|
}
|
|
with redirect_stderr(io.StringIO()):
|
|
result = revoker.main(
|
|
["--runtime", "claude", "--reason", "session-start-compact"],
|
|
environ=environment,
|
|
request=lambda *_args: (_ for _ in ()).throw(OSError("down")),
|
|
)
|
|
self.assertEqual(result, 2)
|
|
self.assertEqual(generation_file.read_text(), "9\n")
|
|
|
|
def test_revoker_fails_closed_on_identity_reply_and_transport_errors(self) -> None:
|
|
_, revoker = self.load_modules()
|
|
good = {
|
|
"MOSAIC_LEASE_BROKER_SOCKET": "/broker",
|
|
"MOSAIC_LEASE_SESSION_ID": "a" * 64,
|
|
"MOSAIC_RUNTIME_GENERATION": "1",
|
|
}
|
|
malformed_session = {**good, "MOSAIC_LEASE_SESSION_ID": "not-a-session"}
|
|
cases = [
|
|
({}, lambda *_args: {"ok": True, "state": "UNVERIFIED"}),
|
|
(malformed_session, lambda *_args: {"ok": True, "state": "UNVERIFIED"}),
|
|
(good, lambda *_args: {"ok": False, "state": "UNVERIFIED"}),
|
|
(good, lambda *_args: {"ok": True, "state": "VERIFIED"}),
|
|
(good, lambda *_args: (_ for _ in ()).throw(OSError("down"))),
|
|
(
|
|
good,
|
|
lambda *_args: (_ for _ in ()).throw(
|
|
json.JSONDecodeError("bad", "x", 0)
|
|
),
|
|
),
|
|
]
|
|
for environment, request in cases:
|
|
with self.subTest(environment=environment), redirect_stderr(io.StringIO()):
|
|
self.assertEqual(
|
|
revoker.main(
|
|
["--runtime", "claude", "--reason", "pre-compact"],
|
|
environ=environment,
|
|
request=request,
|
|
),
|
|
2,
|
|
)
|
|
|
|
with redirect_stderr(io.StringIO()):
|
|
self.assertEqual(
|
|
revoker.main(
|
|
["--runtime", "claude", "--reason", "x" * 129],
|
|
environ=good,
|
|
request=lambda *_args: self.fail("invalid reason reached broker"),
|
|
),
|
|
2,
|
|
)
|
|
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
generation_file = Path(directory) / "runtime.generation"
|
|
generation_file.write_text("1\n")
|
|
generation_file.chmod(0o600)
|
|
with redirect_stderr(io.StringIO()):
|
|
self.assertEqual(
|
|
revoker.main(
|
|
[
|
|
"--runtime",
|
|
"pi",
|
|
"--reason",
|
|
"session-start-resume",
|
|
"--bump-generation",
|
|
],
|
|
environ={
|
|
**good,
|
|
"MOSAIC_LEASE_GENERATION_FILE": str(generation_file),
|
|
},
|
|
request=lambda *_args: (_ for _ in ()).throw(OSError("down")),
|
|
),
|
|
2,
|
|
)
|
|
self.assertEqual(generation_file.read_text(), "2\n")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|