test(#829): regress claudex bypass and per-tool coverage gaps
This commit is contained in:
ms-lead-reviewer
322
packages/mosaic/src/mutator-gate/runtime_tools_unittest.py
Normal file
322
packages/mosaic/src/mutator-gate/runtime_tools_unittest.py
Normal file
@@ -0,0 +1,322 @@
|
||||
#!/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 socket
|
||||
import unittest
|
||||
from contextlib import redirect_stderr
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
|
||||
TOOLS_DIR = Path(__file__).parents[2] / "framework/tools/lease-broker"
|
||||
|
||||
|
||||
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)
|
||||
|
||||
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,
|
||||
)
|
||||
|
||||
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["PRESERVED"], "yes")
|
||||
|
||||
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_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")),
|
||||
),
|
||||
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 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_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)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user