228 lines
9.3 KiB
Python
228 lines
9.3 KiB
Python
#!/usr/bin/env python3
|
|
"""RED-first contracts for the shipped receipt challenge and observer seam."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import base64
|
|
import copy
|
|
import hashlib
|
|
import importlib.util
|
|
import os
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
|
|
TOOLS = Path(__file__).parents[2] / "framework/tools/lease-broker"
|
|
DAEMON_PATH = TOOLS / "daemon.py"
|
|
FRAGMENTS_PATH = TOOLS / "normative_fragments.py"
|
|
OBSERVER_PATH = TOOLS / "receipt_observer.py"
|
|
|
|
|
|
def load_module(name: str, path: Path):
|
|
assert path.is_file(), f"shipped module is missing: {path}"
|
|
spec = importlib.util.spec_from_file_location(name, path)
|
|
if spec is None or spec.loader is None:
|
|
raise RuntimeError(f"unable to load {name}")
|
|
module = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
DAEMON = load_module("lease_broker_receipt_daemon", DAEMON_PATH)
|
|
FRAGMENTS = load_module("lease_broker_normative_fragments", FRAGMENTS_PATH)
|
|
|
|
|
|
class BrokerFixture(unittest.TestCase):
|
|
def setUp(self) -> None:
|
|
self.temporary = tempfile.TemporaryDirectory()
|
|
root = Path(self.temporary.name)
|
|
os.chmod(root, 0o700)
|
|
self.peer = (os.getpid(), os.getuid(), os.getgid())
|
|
self.broker = DAEMON.Broker(DAEMON.StateStore(root / "state.json"))
|
|
registered = self.broker.handle(self.peer, {
|
|
"action": "register_anchor",
|
|
"runtime_generation": 7,
|
|
})
|
|
self.session_id = registered["session_id"]
|
|
self.assertIsInstance(self.session_id, str)
|
|
|
|
def tearDown(self) -> None:
|
|
self.temporary.cleanup()
|
|
|
|
def construction(self) -> tuple[dict[str, object], dict[str, object]]:
|
|
content = b"Constitution\n"
|
|
expected_sha256 = hashlib.sha256(content).hexdigest()
|
|
construction = {
|
|
"manifest_version": 1,
|
|
"generator_version": "wi5-receipt-test",
|
|
"fragments": [{
|
|
"source_id": "authority/constitution",
|
|
"content_base64": base64.b64encode(content).decode("ascii"),
|
|
"expected_sha256": expected_sha256,
|
|
}],
|
|
}
|
|
result = FRAGMENTS.build_payload(
|
|
manifest_version=construction["manifest_version"],
|
|
generator_version=construction["generator_version"],
|
|
fragments=[FRAGMENTS.NormativeFragment("authority/constitution", content, expected_sha256)],
|
|
)
|
|
self.assertEqual(result.injectionDecision, "ACCEPTED")
|
|
self.assertTrue(result.promotion)
|
|
return construction, {
|
|
"compaction_epoch": 3,
|
|
"request_epoch": 8,
|
|
"h_source": result.h_source,
|
|
"h_payload": result.h_payload,
|
|
"schema_version": 1,
|
|
}
|
|
|
|
def begin(self, binding: dict[str, object], construction: dict[str, object]) -> dict[str, object]:
|
|
response = self.broker.handle(self.peer, {
|
|
"action": "begin_verification",
|
|
"session_id": self.session_id,
|
|
"runtime_generation": 7,
|
|
"runtime": "pi",
|
|
"binding": binding,
|
|
"construction": construction,
|
|
})
|
|
self.assertEqual(response["state"], DAEMON.LEASE_PENDING)
|
|
self.assertIsInstance(response.get("receipt_challenge"), str)
|
|
self.assertIsInstance(response.get("receipt"), str)
|
|
return response
|
|
|
|
|
|
class BuildPayloadAdmissionTest(BrokerFixture):
|
|
def test_b3_forged_h_source_or_h_payload_is_refused_against_shipped_build_payload(self) -> None:
|
|
construction, trusted = self.construction()
|
|
for field in ("h_source", "h_payload"):
|
|
with self.subTest(field=field):
|
|
forged = dict(trusted)
|
|
forged[field] = "f" * 64
|
|
with self.assertRaisesRegex(DAEMON.BrokerFailure, "PAYLOAD_BINDING_MISMATCH"):
|
|
self.begin(forged, construction)
|
|
|
|
|
|
class ReceiptObserverTest(BrokerFixture):
|
|
def setUp(self) -> None:
|
|
super().setUp()
|
|
observers = load_module("lease_broker_test_observer", OBSERVER_PATH)
|
|
self.observer = observers.TestReceiptObserver()
|
|
self.broker = DAEMON.Broker(self.broker.store, observer=self.observer)
|
|
|
|
def record(self, message: str) -> None:
|
|
self.observer.record_latest_assistant_message(self.session_id, 7, message)
|
|
|
|
def observe(self, challenge: str, **untrusted: object) -> dict[str, object]:
|
|
return self.broker.handle(self.peer, {
|
|
"action": "observe_receipt",
|
|
"session_id": self.session_id,
|
|
"runtime_generation": 7,
|
|
"receipt_challenge": challenge,
|
|
**untrusted,
|
|
})
|
|
|
|
def promote(self, challenge: str) -> dict[str, object]:
|
|
return self.broker.handle(self.peer, {
|
|
"action": "promote_lease",
|
|
"session_id": self.session_id,
|
|
"runtime_generation": 7,
|
|
"receipt_challenge": challenge,
|
|
})
|
|
|
|
def test_b2_echoed_request_observation_is_refused_but_observer_source_promotes(self) -> None:
|
|
construction, binding = self.construction()
|
|
cycle = self.begin(binding, construction)
|
|
challenge = cycle["receipt_challenge"]
|
|
receipt = cycle["receipt"]
|
|
|
|
with self.assertRaisesRegex(DAEMON.BrokerFailure, "INVALID_RECEIPT"):
|
|
self.observe(challenge, latest_assistant_message=receipt)
|
|
with self.assertRaisesRegex(DAEMON.BrokerFailure, "RECEIPT_OBSERVATION_UNAVAILABLE"):
|
|
self.observe(challenge)
|
|
|
|
self.record(receipt)
|
|
self.assertEqual(self.observe(challenge)["state"], DAEMON.LEASE_PENDING_PROMOTION)
|
|
self.assertEqual(self.promote(challenge)["state"], DAEMON.LEASE_VERIFIED)
|
|
self.assertEqual(self.broker.handle(self.peer, {
|
|
"action": "authorize_tool",
|
|
"session_id": self.session_id,
|
|
"runtime_generation": 7,
|
|
"runtime": "pi",
|
|
"tool_name": "bash",
|
|
})["decision"], "allow")
|
|
|
|
def test_rejected_begin_keeps_revoke_first_fence_for_all_construction_refusals(self) -> None:
|
|
construction, binding = self.construction()
|
|
refusal_cases = {
|
|
"INVALID_CONSTRUCTION": {"bad": "construction"},
|
|
"PAYLOAD_CONSTRUCTION_REFUSED": {
|
|
**construction,
|
|
"fragments": [{
|
|
**construction["fragments"][0],
|
|
"expected_sha256": "0" * 64,
|
|
}],
|
|
},
|
|
"PAYLOAD_BINDING_MISMATCH": None,
|
|
}
|
|
for expected_code, rejected_construction in refusal_cases.items():
|
|
with self.subTest(expected_code=expected_code):
|
|
verified = self.begin(binding, construction)
|
|
self.record(verified["receipt"])
|
|
self.observe(verified["receipt_challenge"])
|
|
self.assertEqual(self.promote(verified["receipt_challenge"])["state"], DAEMON.LEASE_VERIFIED)
|
|
|
|
rejected_binding = copy.deepcopy(binding)
|
|
if expected_code == "PAYLOAD_BINDING_MISMATCH":
|
|
rejected_binding["h_payload"] = "f" * 64
|
|
rejected_construction = construction
|
|
with self.assertRaisesRegex(DAEMON.BrokerFailure, expected_code):
|
|
self.begin(rejected_binding, rejected_construction)
|
|
|
|
self.assertEqual(
|
|
self.broker.leases[self.session_id]["state"], DAEMON.LEASE_UNVERIFIED
|
|
)
|
|
denied = self.broker.handle(self.peer, {
|
|
"action": "authorize_tool",
|
|
"session_id": self.session_id,
|
|
"runtime_generation": 7,
|
|
"runtime": "pi",
|
|
"tool_name": "bash",
|
|
})
|
|
self.assertEqual(denied["decision"], "deny")
|
|
self.assertEqual(denied["state"], DAEMON.LEASE_UNVERIFIED)
|
|
|
|
def test_t26_stale_epoch_receipt_cannot_promote_against_shipped_binding(self) -> None:
|
|
construction, stale_binding = self.construction()
|
|
stale = self.begin(stale_binding, construction)
|
|
current_binding = dict(stale_binding)
|
|
current_binding["compaction_epoch"] = 4
|
|
current_binding["request_epoch"] = 9
|
|
current = self.begin(current_binding, construction)
|
|
self.assertNotEqual(stale["receipt_challenge"], current["receipt_challenge"])
|
|
|
|
self.record(stale["receipt"])
|
|
with self.assertRaisesRegex(DAEMON.BrokerFailure, "RECEIPT_MISMATCH"):
|
|
self.observe(current["receipt_challenge"])
|
|
with self.assertRaisesRegex(DAEMON.BrokerFailure, "INVALID_LEASE_TRANSITION"):
|
|
self.promote(current["receipt_challenge"])
|
|
|
|
def test_t29_altered_model_hash_cannot_promote_against_shipped_binding(self) -> None:
|
|
construction, binding = self.construction()
|
|
cycle = self.begin(binding, construction)
|
|
expected = cycle["receipt"]
|
|
altered_hash = "f" * 64
|
|
self.assertNotEqual(altered_hash, cycle["binding"]["h_payload"])
|
|
altered = expected.replace(cycle["binding"]["h_payload"], altered_hash, 1)
|
|
self.assertNotEqual(altered, expected)
|
|
|
|
self.record(altered)
|
|
with self.assertRaisesRegex(DAEMON.BrokerFailure, "RECEIPT_MISMATCH"):
|
|
self.observe(cycle["receipt_challenge"])
|
|
with self.assertRaisesRegex(DAEMON.BrokerFailure, "INVALID_LEASE_TRANSITION"):
|
|
self.promote(cycle["receipt_challenge"])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|