From a9844654aff55070c0070a816d22a4f127937aaf Mon Sep 17 00:00:00 2001 From: Gary <59334078+garrettallen14@users.noreply.github.com> Date: Tue, 24 Mar 2026 17:43:22 -0700 Subject: [PATCH 1/4] Add layerlens.attestation: cryptographic hash chains for tamper-evident traces --- pyproject.toml | 1 + src/layerlens/attestation/__init__.py | 24 ++++ src/layerlens/attestation/_chain.py | 84 ++++++++++++++ src/layerlens/attestation/_envelope.py | 31 ++++++ src/layerlens/attestation/_hash.py | 39 +++++++ src/layerlens/attestation/_verify.py | 115 +++++++++++++++++++ src/layerlens/instrument/_recorder.py | 55 ++++++++- src/layerlens/instrument/_upload.py | 24 +++- tests/attestation/__init__.py | 0 tests/attestation/test_chain.py | 132 ++++++++++++++++++++++ tests/attestation/test_hash.py | 62 +++++++++++ tests/attestation/test_integration.py | 148 +++++++++++++++++++++++++ tests/attestation/test_verify.py | 127 +++++++++++++++++++++ 13 files changed, 834 insertions(+), 8 deletions(-) create mode 100644 src/layerlens/attestation/__init__.py create mode 100644 src/layerlens/attestation/_chain.py create mode 100644 src/layerlens/attestation/_envelope.py create mode 100644 src/layerlens/attestation/_hash.py create mode 100644 src/layerlens/attestation/_verify.py create mode 100644 tests/attestation/__init__.py create mode 100644 tests/attestation/test_chain.py create mode 100644 tests/attestation/test_hash.py create mode 100644 tests/attestation/test_integration.py create mode 100644 tests/attestation/test_verify.py diff --git a/pyproject.toml b/pyproject.toml index 16cb9a1..d0fabba 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -141,6 +141,7 @@ known-first-party = ["openai", "tests"] "scripts/**.py" = ["T201", "T203"] "tests/**.py" = ["T201", "T203"] "tests/instrument/**.py" = ["T201", "T203", "ARG"] +"tests/attestation/**.py" = ["T201", "T203", "ARG"] "examples/**.py" = ["T201", "T203"] "src/layerlens/cli/**" = ["T201", "T203"] "src/layerlens/instrument/adapters/frameworks/langchain.py" = ["ARG002"] diff --git a/src/layerlens/attestation/__init__.py b/src/layerlens/attestation/__init__.py new file mode 100644 index 0000000..e8d8481 --- /dev/null +++ b/src/layerlens/attestation/__init__.py @@ -0,0 +1,24 @@ +from __future__ import annotations + +from ._hash import compute_hash +from ._chain import HashChain +from ._verify import ( + TamperingResult, + ChainVerification, + verify_chain, + verify_trial, + detect_tampering, +) +from ._envelope import HashScope, AttestationEnvelope + +__all__ = [ + "AttestationEnvelope", + "ChainVerification", + "HashChain", + "HashScope", + "TamperingResult", + "compute_hash", + "detect_tampering", + "verify_chain", + "verify_trial", +] diff --git a/src/layerlens/attestation/_chain.py b/src/layerlens/attestation/_chain.py new file mode 100644 index 0000000..6e3fc3f --- /dev/null +++ b/src/layerlens/attestation/_chain.py @@ -0,0 +1,84 @@ +from __future__ import annotations + +from typing import Any, Dict, List, Optional + +from ._hash import compute_hash +from ._envelope import HashScope, AttestationEnvelope + + +class HashChain: + """Builds a linear hash chain over a sequence of events. + + Each event is hashed and linked to the previous hash, forming + a tamper-evident chain. If any event is modified after the fact, + the chain breaks at that point. + """ + + def __init__(self) -> None: + self._chain: List[AttestationEnvelope] = [] + self._last_hash: Optional[str] = None + self._terminated: bool = False + self._terminate_reason: Optional[str] = None + + @property + def envelopes(self) -> List[AttestationEnvelope]: + return list(self._chain) + + @property + def is_terminated(self) -> bool: + return self._terminated + + def _check_active(self) -> None: + if self._terminated: + raise RuntimeError(f"Hash chain terminated: {self._terminate_reason}. No further events can be added.") + + def add_event(self, data: Dict[str, Any]) -> AttestationEnvelope: + """Hash an event and append it to the chain.""" + self._check_active() + # Include previous_hash in the hashed payload for chaining + payload = {**data, "_previous_hash": self._last_hash} + event_hash = compute_hash(payload) + envelope = AttestationEnvelope( + hash=event_hash, + scope=HashScope.EVENT, + previous_hash=self._last_hash, + ) + self._chain.append(envelope) + self._last_hash = event_hash + return envelope + + def terminate(self, reason: str) -> None: + """Permanently stop the chain. No further events or finalization allowed.""" + self._terminated = True + self._terminate_reason = reason + + def finalize(self) -> AttestationEnvelope: + """Compute a trial-level root hash over all event hashes and seal the chain.""" + if self._terminated: + raise RuntimeError( + f"Cannot finalize terminated hash chain. Trial is non-attestable due to: {self._terminate_reason}" + ) + if not self._chain: + raise RuntimeError("Cannot finalize empty hash chain.") + event_hashes = [e.hash for e in self._chain] + root_hash = compute_hash({"event_hashes": event_hashes}) + trial_envelope = AttestationEnvelope( + hash=root_hash, + scope=HashScope.TRIAL, + previous_hash=self._last_hash, + ) + # Seal — no more events after finalization + self._terminated = True + self._terminate_reason = "chain finalized" + return trial_envelope + + def to_dict(self) -> Dict[str, Any]: + """Serialize the chain for inclusion in trace uploads.""" + result: Dict[str, Any] = { + "events": [e.to_dict() for e in self._chain], + } + # Only include termination details when the chain was stopped + # due to a policy violation (not normal finalization). + if self._terminated and self._terminate_reason != "chain finalized": + result["terminated_reason"] = self._terminate_reason + return result diff --git a/src/layerlens/attestation/_envelope.py b/src/layerlens/attestation/_envelope.py new file mode 100644 index 0000000..bb054c1 --- /dev/null +++ b/src/layerlens/attestation/_envelope.py @@ -0,0 +1,31 @@ +from __future__ import annotations + +from enum import Enum +from typing import Any, Dict, Optional +from datetime import datetime, timezone +from dataclasses import field, dataclass + + +class HashScope(Enum): + """Level at which a hash was computed.""" + + EVENT = "event" + TRIAL = "trial" + + +@dataclass +class AttestationEnvelope: + """Single entry in a hash chain.""" + + hash: str + scope: HashScope + previous_hash: Optional[str] = None + timestamp: datetime = field(default_factory=lambda: datetime.now(timezone.utc)) + + def to_dict(self) -> Dict[str, Any]: + return { + "hash": self.hash, + "scope": self.scope.value, + "previous_hash": self.previous_hash, + "timestamp": self.timestamp.isoformat(), + } diff --git a/src/layerlens/attestation/_hash.py b/src/layerlens/attestation/_hash.py new file mode 100644 index 0000000..f6284cf --- /dev/null +++ b/src/layerlens/attestation/_hash.py @@ -0,0 +1,39 @@ +from __future__ import annotations + +import json +import hashlib +from enum import Enum +from typing import Any +from datetime import datetime +from dataclasses import asdict + + +def _json_default(obj: Any) -> Any: + """Handle non-standard types for canonical JSON serialization.""" + if isinstance(obj, datetime): + return obj.isoformat() + if isinstance(obj, Enum): + return obj.value + if hasattr(obj, "to_dict"): + return obj.to_dict() + if hasattr(obj, "__dataclass_fields__"): + return asdict(obj) + raise TypeError(f"Object of type {type(obj).__name__} is not JSON serializable") + + +def canonical_json(data: Any) -> str: + """Serialize data to canonical JSON: sorted keys, compact, deterministic.""" + return json.dumps( + data, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=True, + default=_json_default, + ) + + +def compute_hash(data: Any) -> str: + """Compute SHA-256 hash of canonicalized data. Returns 'sha256:<64 hex chars>'.""" + raw = canonical_json(data) + digest = hashlib.sha256(raw.encode("utf-8")).hexdigest() + return f"sha256:{digest}" diff --git a/src/layerlens/attestation/_verify.py b/src/layerlens/attestation/_verify.py new file mode 100644 index 0000000..35579e3 --- /dev/null +++ b/src/layerlens/attestation/_verify.py @@ -0,0 +1,115 @@ +from __future__ import annotations + +from typing import Any, Dict, List, Optional +from dataclasses import field, dataclass + +from ._hash import compute_hash +from ._envelope import HashScope, AttestationEnvelope + + +@dataclass +class ChainVerification: + """Result of verifying a hash chain's integrity.""" + + valid: bool + break_index: Optional[int] = None + error: Optional[str] = None + + +@dataclass +class TamperingResult: + """Result of checking whether trace data was modified after hashing.""" + + tampered: bool + modified_indices: List[int] = field(default_factory=list) + chain_broken: bool = False + + +def verify_chain(envelopes: List[AttestationEnvelope]) -> ChainVerification: + """Verify that a hash chain is continuous and unbroken. + + Checks: + - First envelope has previous_hash=None + - Each subsequent envelope's previous_hash matches the prior envelope's hash + """ + if not envelopes: + return ChainVerification(valid=True) + + if envelopes[0].previous_hash is not None: + return ChainVerification( + valid=False, + break_index=0, + error="First envelope must have previous_hash=None", + ) + + for i in range(1, len(envelopes)): + if envelopes[i].previous_hash != envelopes[i - 1].hash: + return ChainVerification( + valid=False, + break_index=i, + error=f"Chain broken at index {i}: " + f"expected previous_hash={envelopes[i - 1].hash!r}, " + f"got {envelopes[i].previous_hash!r}", + ) + + return ChainVerification(valid=True) + + +def verify_trial( + envelopes: List[AttestationEnvelope], + trial_envelope: AttestationEnvelope, +) -> ChainVerification: + """Verify a trial envelope against its event chain. + + Checks chain integrity, then verifies the trial hash is correctly + computed over all event hashes. + """ + chain_result = verify_chain(envelopes) + if not chain_result.valid: + return chain_result + + if trial_envelope.scope != HashScope.TRIAL: + return ChainVerification( + valid=False, + error=f"Trial envelope has wrong scope: {trial_envelope.scope}", + ) + + event_hashes = [e.hash for e in envelopes] + expected_hash = compute_hash({"event_hashes": event_hashes}) + if trial_envelope.hash != expected_hash: + return ChainVerification( + valid=False, + error="Trial hash does not match event hashes", + ) + + return ChainVerification(valid=True) + + +def detect_tampering( + envelopes: List[AttestationEnvelope], + original_data: List[Dict[str, Any]], +) -> TamperingResult: + """Detect which events were modified after being hashed. + + Recomputes the hash for each event (using its stored previous_hash + for chain linkage) and compares against the stored hash. + """ + if len(envelopes) != len(original_data): + return TamperingResult( + tampered=True, + chain_broken=True, + ) + + modified: List[int] = [] + for i, (envelope, data) in enumerate(zip(envelopes, original_data)): + payload = {**data, "_previous_hash": envelope.previous_hash} + recomputed = compute_hash(payload) + if recomputed != envelope.hash: + modified.append(i) + + chain_result = verify_chain(envelopes) + return TamperingResult( + tampered=len(modified) > 0 or not chain_result.valid, + modified_indices=modified, + chain_broken=not chain_result.valid, + ) diff --git a/src/layerlens/instrument/_recorder.py b/src/layerlens/instrument/_recorder.py index dba6a45..dce18b3 100644 --- a/src/layerlens/instrument/_recorder.py +++ b/src/layerlens/instrument/_recorder.py @@ -1,22 +1,71 @@ from __future__ import annotations -from typing import Any, Optional +import logging +from typing import Any, Dict, List, Optional + +from layerlens.attestation import HashChain from ._types import SpanData from ._upload import upload_trace, async_upload_trace +log: logging.Logger = logging.getLogger(__name__) + + +def _collect_spans(span: SpanData) -> List[Dict[str, Any]]: + """Walk the span tree depth-first and return a flat list of span dicts. + + Uses SpanData.to_dict() to capture every field — structure, inputs, + outputs, metadata, and errors. Children are excluded because we + flatten the tree ourselves; any future SpanData fields are automatically + included in the hash. + """ + result: List[Dict[str, Any]] = [] + span_dict = span.to_dict() + span_dict.pop("children") + result.append(span_dict) + for child in span.children: + result.extend(_collect_spans(child)) + return result + class TraceRecorder: def __init__(self, client: Any) -> None: self._client = client self.root: Optional[SpanData] = None + def _build_attestation(self) -> Dict[str, Any]: + """Build a hash chain from the span tree and return attestation data.""" + if self.root is None: + return {} + chain = HashChain() + spans = _collect_spans(self.root) + for span_dict in spans: + chain.add_event(span_dict) + trial = chain.finalize() + return { + "chain": chain.to_dict(), + "root_hash": trial.hash, + "schema_version": "1.0", + } + def flush(self) -> None: if self.root is None: return - upload_trace(self._client, self.root.to_dict()) + trace_data = self.root.to_dict() + try: + attestation = self._build_attestation() + except Exception: + log.debug("Failed to build attestation chain", exc_info=True) + attestation = {} + upload_trace(self._client, trace_data, attestation) async def async_flush(self) -> None: if self.root is None: return - await async_upload_trace(self._client, self.root.to_dict()) + trace_data = self.root.to_dict() + try: + attestation = self._build_attestation() + except Exception: + log.debug("Failed to build attestation chain", exc_info=True) + attestation = {} + await async_upload_trace(self._client, trace_data, attestation) diff --git a/src/layerlens/instrument/_upload.py b/src/layerlens/instrument/_upload.py index 6597970..e6cd3a1 100644 --- a/src/layerlens/instrument/_upload.py +++ b/src/layerlens/instrument/_upload.py @@ -4,16 +4,23 @@ import json import logging import tempfile -from typing import Any, Dict +from typing import Any, Dict, Optional log: logging.Logger = logging.getLogger(__name__) -def upload_trace(client: Any, trace_data: Dict[str, Any]) -> None: +def upload_trace( + client: Any, + trace_data: Dict[str, Any], + attestation: Optional[Dict[str, Any]] = None, +) -> None: + payload = trace_data + if attestation: + payload = {**trace_data, "attestation": attestation} fd, path = tempfile.mkstemp(suffix=".json", prefix="layerlens_trace_") try: with os.fdopen(fd, "w") as f: - json.dump([trace_data], f, default=str) + json.dump([payload], f, default=str) client.traces.upload(path) finally: try: @@ -22,11 +29,18 @@ def upload_trace(client: Any, trace_data: Dict[str, Any]) -> None: log.debug("Failed to remove temp trace file: %s", path) -async def async_upload_trace(client: Any, trace_data: Dict[str, Any]) -> None: +async def async_upload_trace( + client: Any, + trace_data: Dict[str, Any], + attestation: Optional[Dict[str, Any]] = None, +) -> None: + payload = trace_data + if attestation: + payload = {**trace_data, "attestation": attestation} fd, path = tempfile.mkstemp(suffix=".json", prefix="layerlens_trace_") try: with os.fdopen(fd, "w") as f: - json.dump([trace_data], f, default=str) + json.dump([payload], f, default=str) await client.traces.upload(path) finally: try: diff --git a/tests/attestation/__init__.py b/tests/attestation/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/attestation/test_chain.py b/tests/attestation/test_chain.py new file mode 100644 index 0000000..b204578 --- /dev/null +++ b/tests/attestation/test_chain.py @@ -0,0 +1,132 @@ +from __future__ import annotations + +import pytest + +from layerlens.attestation._chain import HashChain +from layerlens.attestation._envelope import HashScope + + +class TestHashChainBuilding: + def test_single_event(self): + chain = HashChain() + env = chain.add_event({"name": "span-1"}) + assert env.previous_hash is None + assert env.scope == HashScope.EVENT + assert env.hash.startswith("sha256:") + + def test_chain_linking(self): + """Each event links to the previous hash.""" + chain = HashChain() + e1 = chain.add_event({"name": "span-1"}) + e2 = chain.add_event({"name": "span-2"}) + e3 = chain.add_event({"name": "span-3"}) + + assert e1.previous_hash is None + assert e2.previous_hash == e1.hash + assert e3.previous_hash == e2.hash + + def test_different_data_different_hashes(self): + chain = HashChain() + e1 = chain.add_event({"name": "a"}) + e2 = chain.add_event({"name": "b"}) + assert e1.hash != e2.hash + + def test_envelopes_property(self): + chain = HashChain() + chain.add_event({"name": "span-1"}) + chain.add_event({"name": "span-2"}) + assert len(chain.envelopes) == 2 + + +class TestHashChainFinalization: + def test_finalize_produces_trial_scope(self): + chain = HashChain() + chain.add_event({"name": "span-1"}) + trial = chain.finalize() + assert trial.scope == HashScope.TRIAL + + def test_finalize_root_hash_deterministic(self): + """Same events in same order produce the same root hash.""" + + def build(): + c = HashChain() + c.add_event({"name": "a"}) + c.add_event({"name": "b"}) + return c.finalize() + + assert build().hash == build().hash + + def test_finalize_seals_chain(self): + """No events can be added after finalization.""" + chain = HashChain() + chain.add_event({"name": "span-1"}) + chain.finalize() + with pytest.raises(RuntimeError, match="terminated"): + chain.add_event({"name": "span-2"}) + + def test_finalize_empty_chain_raises(self): + chain = HashChain() + with pytest.raises(RuntimeError, match="empty"): + chain.finalize() + + def test_finalize_links_to_last_event(self): + chain = HashChain() + chain.add_event({"name": "a"}) + last = chain.add_event({"name": "b"}) + trial = chain.finalize() + assert trial.previous_hash == last.hash + + +class TestHashChainTermination: + def test_terminate_blocks_add(self): + chain = HashChain() + chain.add_event({"name": "span-1"}) + chain.terminate("policy_violation") + with pytest.raises(RuntimeError, match="terminated"): + chain.add_event({"name": "span-2"}) + + def test_terminate_blocks_finalize(self): + chain = HashChain() + chain.add_event({"name": "span-1"}) + chain.terminate("policy_violation") + with pytest.raises(RuntimeError, match="non-attestable"): + chain.finalize() + + def test_is_terminated_flag(self): + chain = HashChain() + assert not chain.is_terminated + chain.terminate("test") + assert chain.is_terminated + + def test_terminate_reason_in_error(self): + chain = HashChain() + chain.terminate("safety_check_failed") + with pytest.raises(RuntimeError, match="safety_check_failed"): + chain.add_event({"name": "span-1"}) + + +class TestHashChainSerialization: + def test_to_dict(self): + chain = HashChain() + chain.add_event({"name": "span-1"}) + d = chain.to_dict() + assert "events" in d + assert len(d["events"]) == 1 + assert d["events"][0]["scope"] == "event" + assert d["events"][0]["hash"].startswith("sha256:") + + def test_to_dict_finalized_is_clean(self): + """Normal finalization should not include termination details.""" + chain = HashChain() + chain.add_event({"name": "span-1"}) + chain.finalize() + d = chain.to_dict() + assert "terminated_reason" not in d + + def test_to_dict_terminated_includes_reason(self): + """Policy violation termination should include the reason.""" + chain = HashChain() + chain.add_event({"name": "span-1"}) + chain.terminate("policy_violation") + d = chain.to_dict() + assert d["terminated_reason"] == "policy_violation" diff --git a/tests/attestation/test_hash.py b/tests/attestation/test_hash.py new file mode 100644 index 0000000..1f5e9ad --- /dev/null +++ b/tests/attestation/test_hash.py @@ -0,0 +1,62 @@ +from __future__ import annotations + +import re +from enum import Enum +from datetime import datetime, timezone + +from layerlens.attestation._hash import compute_hash, canonical_json + + +class TestCanonicalJson: + def test_sorted_keys(self): + """Key order must not affect output.""" + a = canonical_json({"b": 2, "a": 1}) + b = canonical_json({"a": 1, "b": 2}) + assert a == b + + def test_compact_format(self): + """No whitespace in output.""" + result = canonical_json({"a": 1, "b": [2, 3]}) + assert " " not in result + assert result == '{"a":1,"b":[2,3]}' + + def test_nested_structures(self): + """Nested dicts and lists are handled deterministically.""" + data = {"z": {"y": 1, "x": 2}, "a": [3, 2, 1]} + result = canonical_json(data) + assert result == '{"a":[3,2,1],"z":{"x":2,"y":1}}' + + def test_datetime_serialization(self): + dt = datetime(2026, 3, 23, 12, 0, 0, tzinfo=timezone.utc) + result = canonical_json({"ts": dt}) + assert "2026-03-23" in result + + def test_enum_serialization(self): + class Color(Enum): + RED = "red" + + result = canonical_json({"color": Color.RED}) + assert '"red"' in result + + +class TestComputeHash: + def test_format(self): + """Hash must be 'sha256:' followed by 64 hex chars.""" + h = compute_hash({"test": "data"}) + assert re.match(r"^sha256:[0-9a-f]{64}$", h) + + def test_deterministic(self): + """Same data always produces the same hash.""" + data = {"key": "value", "num": 42} + assert compute_hash(data) == compute_hash(data) + + def test_key_order_irrelevant(self): + """Different key orders produce the same hash.""" + assert compute_hash({"b": 2, "a": 1}) == compute_hash({"a": 1, "b": 2}) + + def test_different_data_different_hash(self): + assert compute_hash({"a": 1}) != compute_hash({"a": 2}) + + def test_empty_dict(self): + h = compute_hash({}) + assert re.match(r"^sha256:[0-9a-f]{64}$", h) diff --git a/tests/attestation/test_integration.py b/tests/attestation/test_integration.py new file mode 100644 index 0000000..5f8c436 --- /dev/null +++ b/tests/attestation/test_integration.py @@ -0,0 +1,148 @@ +from __future__ import annotations + +import json +from unittest.mock import Mock + +from layerlens.instrument import span, trace +from layerlens.attestation import verify_chain, detect_tampering +from layerlens.attestation._envelope import HashScope, AttestationEnvelope + + +def _make_client(): + """Create a mock client that captures the uploaded trace JSON.""" + client = Mock() + client.traces = Mock() + uploaded = {} + + def capture(path): + with open(path) as f: + uploaded["data"] = json.load(f) + + client.traces.upload = Mock(side_effect=capture) + return client, uploaded + + +class TestTraceAttestation: + def test_trace_includes_attestation(self): + """@trace should include attestation data in the upload.""" + client, uploaded = _make_client() + + @trace(client) + def my_agent(query: str): + return f"answer to {query}" + + my_agent("hello") + + payload = uploaded["data"][0] + assert "attestation" in payload + att = payload["attestation"] + assert "chain" in att + assert "root_hash" in att + assert att["root_hash"].startswith("sha256:") + assert att["schema_version"] == "1.0" + + def test_trace_with_child_spans(self): + """Attestation chain should include all spans in the tree.""" + client, uploaded = _make_client() + + @trace(client) + def my_agent(query: str): + with span("step-1", kind="tool") as s: + s.output = "result-1" + with span("step-2", kind="llm") as s: + s.output = "result-2" + return "done" + + my_agent("test") + + att = uploaded["data"][0]["attestation"] + chain_events = att["chain"]["events"] + # Root span + 2 child spans = 3 events in the chain + assert len(chain_events) == 3 + + def test_chain_events_are_linked(self): + """Verify the chain in the uploaded payload is valid.""" + client, uploaded = _make_client() + + @trace(client) + def my_agent(query: str): + with span("step-1") as s: + s.output = "r1" + with span("step-2") as s: + s.output = "r2" + return "done" + + my_agent("test") + + chain_events = uploaded["data"][0]["attestation"]["chain"]["events"] + # Reconstruct envelopes and verify chain integrity + envelopes = [ + AttestationEnvelope( + hash=e["hash"], + scope=HashScope(e["scope"]), + previous_hash=e["previous_hash"], + ) + for e in chain_events + ] + result = verify_chain(envelopes) + assert result.valid + + def test_trace_error_still_has_attestation(self): + """Even when the traced function raises, attestation should be present.""" + client, uploaded = _make_client() + + @trace(client) + def failing_agent(): + with span("step-1") as s: + s.output = "ok" + raise ValueError("boom") + + try: + failing_agent() + except ValueError: + pass + + payload = uploaded["data"][0] + assert "attestation" in payload + assert payload["attestation"]["root_hash"].startswith("sha256:") + + def test_modifying_output_breaks_chain(self): + """Changing what the agent said must invalidate the attestation.""" + client, uploaded = _make_client() + + @trace(client) + def my_agent(query: str): + with span("llm-call", kind="llm") as s: + s.output = "the real answer" + return "done" + + my_agent("test") + + att = uploaded["data"][0]["attestation"] + envelopes = [ + AttestationEnvelope( + hash=e["hash"], + scope=HashScope(e["scope"]), + previous_hash=e["previous_hash"], + ) + for e in att["chain"]["events"] + ] + + # Build the original span dicts that were hashed (root + child) + payload = uploaded["data"][0] + original_spans = [] + for s in [payload] + payload.get("children", []): + d = {k: v for k, v in s.items() if k not in ("children", "attestation")} + original_spans.append(d) + + # Verify clean data passes + clean = detect_tampering(envelopes, original_spans) + assert not clean.tampered + + # Tamper: change the LLM output + tampered_spans = [dict(d) for d in original_spans] + tampered_spans[1] = {**tampered_spans[1], "output": "a forged answer"} + + tampered = detect_tampering(envelopes, tampered_spans) + assert tampered.tampered + assert 1 in tampered.modified_indices diff --git a/tests/attestation/test_verify.py b/tests/attestation/test_verify.py new file mode 100644 index 0000000..f7567d8 --- /dev/null +++ b/tests/attestation/test_verify.py @@ -0,0 +1,127 @@ +from __future__ import annotations + +from layerlens.attestation._chain import HashChain +from layerlens.attestation._verify import ( + verify_chain, + verify_trial, + detect_tampering, +) +from layerlens.attestation._envelope import HashScope + + +class TestVerifyChain: + def test_valid_chain(self): + chain = HashChain() + chain.add_event({"name": "a"}) + chain.add_event({"name": "b"}) + chain.add_event({"name": "c"}) + result = verify_chain(chain.envelopes) + assert result.valid + assert result.break_index is None + + def test_empty_chain_valid(self): + result = verify_chain([]) + assert result.valid + + def test_single_event_valid(self): + chain = HashChain() + chain.add_event({"name": "a"}) + result = verify_chain(chain.envelopes) + assert result.valid + + def test_broken_first_link(self): + """First envelope must have previous_hash=None.""" + chain = HashChain() + chain.add_event({"name": "a"}) + envelopes = chain.envelopes + # Tamper: set previous_hash on first event + envelopes[0].previous_hash = "sha256:fake" + result = verify_chain(envelopes) + assert not result.valid + assert result.break_index == 0 + + def test_broken_middle_link(self): + chain = HashChain() + chain.add_event({"name": "a"}) + chain.add_event({"name": "b"}) + chain.add_event({"name": "c"}) + envelopes = chain.envelopes + # Tamper: break the link between event 1 and 2 + envelopes[2].previous_hash = "sha256:fake" + result = verify_chain(envelopes) + assert not result.valid + assert result.break_index == 2 + + +class TestVerifyTrial: + def test_valid_trial(self): + chain = HashChain() + chain.add_event({"name": "a"}) + chain.add_event({"name": "b"}) + envelopes = chain.envelopes + trial = chain.finalize() + result = verify_trial(envelopes, trial) + assert result.valid + + def test_wrong_scope_rejected(self): + chain = HashChain() + chain.add_event({"name": "a"}) + envelopes = chain.envelopes + trial = chain.finalize() + trial.scope = HashScope.EVENT # Wrong scope + result = verify_trial(envelopes, trial) + assert not result.valid + assert "scope" in (result.error or "") + + def test_tampered_trial_hash(self): + chain = HashChain() + chain.add_event({"name": "a"}) + envelopes = chain.envelopes + trial = chain.finalize() + trial.hash = "sha256:" + "0" * 64 # Wrong hash + result = verify_trial(envelopes, trial) + assert not result.valid + assert "does not match" in (result.error or "") + + +class TestDetectTampering: + def test_no_tampering(self): + data = [{"name": "a"}, {"name": "b"}, {"name": "c"}] + chain = HashChain() + for d in data: + chain.add_event(d) + result = detect_tampering(chain.envelopes, data) + assert not result.tampered + assert result.modified_indices == [] + assert not result.chain_broken + + def test_detect_modified_event(self): + data = [{"name": "a"}, {"name": "b"}, {"name": "c"}] + chain = HashChain() + for d in data: + chain.add_event(d) + # Tamper with the second event's data + tampered_data = [{"name": "a"}, {"name": "CHANGED"}, {"name": "c"}] + result = detect_tampering(chain.envelopes, tampered_data) + assert result.tampered + assert 1 in result.modified_indices + + def test_detect_multiple_modifications(self): + data = [{"name": "a"}, {"name": "b"}, {"name": "c"}] + chain = HashChain() + for d in data: + chain.add_event(d) + tampered = [{"name": "X"}, {"name": "b"}, {"name": "Z"}] + result = detect_tampering(chain.envelopes, tampered) + assert result.tampered + assert 0 in result.modified_indices + assert 2 in result.modified_indices + + def test_detect_count_mismatch(self): + data = [{"name": "a"}, {"name": "b"}] + chain = HashChain() + for d in data: + chain.add_event(d) + result = detect_tampering(chain.envelopes, [{"name": "a"}]) + assert result.tampered + assert result.chain_broken From abf41511189d0ebb380708d09f1629b2405b51ce Mon Sep 17 00:00:00 2001 From: Gary <59334078+garrettallen14@users.noreply.github.com> Date: Thu, 26 Mar 2026 14:35:36 -0700 Subject: [PATCH 2/4] feat: signing keys --- src/layerlens/_client.py | 13 + src/layerlens/attestation/__init__.py | 5 + src/layerlens/attestation/_chain.py | 22 +- src/layerlens/attestation/_envelope.py | 9 +- src/layerlens/attestation/_signing.py | 19 + src/layerlens/attestation/_verify.py | 80 +++- src/layerlens/instrument/__init__.py | 3 +- src/layerlens/instrument/_decorator.py | 7 +- src/layerlens/instrument/_recorder.py | 125 +++++- .../resources/signing_keys/__init__.py | 3 + .../resources/signing_keys/signing_keys.py | 153 +++++++ tests/attestation/test_hash.py | 9 + tests/attestation/test_signing.py | 190 ++++++++ tests/attestation/test_verify.py | 30 +- tests/instrument/test_signing_autofetch.py | 412 ++++++++++++++++++ 15 files changed, 1045 insertions(+), 35 deletions(-) create mode 100644 src/layerlens/attestation/_signing.py create mode 100644 src/layerlens/resources/signing_keys/__init__.py create mode 100644 src/layerlens/resources/signing_keys/signing_keys.py create mode 100644 tests/attestation/test_signing.py create mode 100644 tests/instrument/test_signing_autofetch.py diff --git a/src/layerlens/_client.py b/src/layerlens/_client.py index 032e15b..c679990 100644 --- a/src/layerlens/_client.py +++ b/src/layerlens/_client.py @@ -25,6 +25,7 @@ from .resources.benchmarks import Benchmarks, AsyncBenchmarks from .resources.evaluations import Evaluations, AsyncEvaluations from .resources.integrations import Integrations, AsyncIntegrations + from .resources.signing_keys import SigningKeys, AsyncSigningKeys from .resources.evaluation_spaces import EvaluationSpaces, AsyncEvaluationSpaces from .resources.trace_evaluations import TraceEvaluations, AsyncTraceEvaluations from .resources.judge_optimizations import JudgeOptimizations, AsyncJudgeOptimizations @@ -139,6 +140,12 @@ def scorers(self) -> Scorers: return Scorers(self) + @cached_property + def signing_keys(self) -> SigningKeys: + from .resources.signing_keys import SigningKeys + + return SigningKeys(self) + @cached_property def evaluation_spaces(self) -> EvaluationSpaces: from .resources.evaluation_spaces import EvaluationSpaces @@ -326,6 +333,12 @@ def scorers(self) -> AsyncScorers: return AsyncScorers(self) + @cached_property + def signing_keys(self) -> AsyncSigningKeys: + from .resources.signing_keys import AsyncSigningKeys + + return AsyncSigningKeys(self) + @cached_property def evaluation_spaces(self) -> AsyncEvaluationSpaces: from .resources.evaluation_spaces import AsyncEvaluationSpaces diff --git a/src/layerlens/attestation/__init__.py b/src/layerlens/attestation/__init__.py index e8d8481..eda36cd 100644 --- a/src/layerlens/attestation/__init__.py +++ b/src/layerlens/attestation/__init__.py @@ -5,10 +5,12 @@ from ._verify import ( TamperingResult, ChainVerification, + TrialVerification, verify_chain, verify_trial, detect_tampering, ) +from ._signing import hmac_sign, hmac_verify from ._envelope import HashScope, AttestationEnvelope __all__ = [ @@ -17,8 +19,11 @@ "HashChain", "HashScope", "TamperingResult", + "TrialVerification", "compute_hash", "detect_tampering", + "hmac_sign", + "hmac_verify", "verify_chain", "verify_trial", ] diff --git a/src/layerlens/attestation/_chain.py b/src/layerlens/attestation/_chain.py index 6e3fc3f..93a3ace 100644 --- a/src/layerlens/attestation/_chain.py +++ b/src/layerlens/attestation/_chain.py @@ -3,6 +3,7 @@ from typing import Any, Dict, List, Optional from ._hash import compute_hash +from ._signing import hmac_sign from ._envelope import HashScope, AttestationEnvelope @@ -12,13 +13,24 @@ class HashChain: Each event is hashed and linked to the previous hash, forming a tamper-evident chain. If any event is modified after the fact, the chain breaks at that point. + + If ``signing_secret`` is provided, each envelope's hash is + HMAC-SHA256 signed for authenticity on top of integrity. """ - def __init__(self) -> None: + def __init__( + self, + signing_key_id: Optional[str] = None, + signing_secret: Optional[bytes] = None, + ) -> None: + if signing_secret is not None and not signing_key_id: + raise ValueError("signing_key_id is required when signing_secret is provided") self._chain: List[AttestationEnvelope] = [] self._last_hash: Optional[str] = None self._terminated: bool = False self._terminate_reason: Optional[str] = None + self._signing_key_id = signing_key_id + self._signing_secret = signing_secret @property def envelopes(self) -> List[AttestationEnvelope]: @@ -32,6 +44,12 @@ def _check_active(self) -> None: if self._terminated: raise RuntimeError(f"Hash chain terminated: {self._terminate_reason}. No further events can be added.") + def _sign_envelope(self, envelope: AttestationEnvelope) -> None: + """Sign an envelope's hash if a signing secret is configured.""" + if self._signing_secret is not None: + envelope.signature = hmac_sign(self._signing_secret, envelope.hash.encode("utf-8")) + envelope.signing_key_id = self._signing_key_id + def add_event(self, data: Dict[str, Any]) -> AttestationEnvelope: """Hash an event and append it to the chain.""" self._check_active() @@ -43,6 +61,7 @@ def add_event(self, data: Dict[str, Any]) -> AttestationEnvelope: scope=HashScope.EVENT, previous_hash=self._last_hash, ) + self._sign_envelope(envelope) self._chain.append(envelope) self._last_hash = event_hash return envelope @@ -67,6 +86,7 @@ def finalize(self) -> AttestationEnvelope: scope=HashScope.TRIAL, previous_hash=self._last_hash, ) + self._sign_envelope(trial_envelope) # Seal — no more events after finalization self._terminated = True self._terminate_reason = "chain finalized" diff --git a/src/layerlens/attestation/_envelope.py b/src/layerlens/attestation/_envelope.py index bb054c1..7fbc978 100644 --- a/src/layerlens/attestation/_envelope.py +++ b/src/layerlens/attestation/_envelope.py @@ -21,11 +21,18 @@ class AttestationEnvelope: scope: HashScope previous_hash: Optional[str] = None timestamp: datetime = field(default_factory=lambda: datetime.now(timezone.utc)) + signature: Optional[str] = None + signing_key_id: Optional[str] = None def to_dict(self) -> Dict[str, Any]: - return { + d: Dict[str, Any] = { "hash": self.hash, "scope": self.scope.value, "previous_hash": self.previous_hash, "timestamp": self.timestamp.isoformat(), } + if self.signature is not None: + d["signature"] = self.signature + if self.signing_key_id is not None: + d["signing_key_id"] = self.signing_key_id + return d diff --git a/src/layerlens/attestation/_signing.py b/src/layerlens/attestation/_signing.py new file mode 100644 index 0000000..6164c14 --- /dev/null +++ b/src/layerlens/attestation/_signing.py @@ -0,0 +1,19 @@ +"""HMAC-SHA256 signing for attestation envelopes.""" + +from __future__ import annotations + +import hmac as hmac_mod +import base64 +import hashlib + + +def hmac_sign(secret: bytes, data: bytes) -> str: + """Sign data with HMAC-SHA256, returning a base64-encoded signature.""" + sig = hmac_mod.new(secret, data, hashlib.sha256).digest() + return base64.b64encode(sig).decode("ascii") + + +def hmac_verify(secret: bytes, data: bytes, signature: str) -> bool: + """Verify a base64-encoded HMAC-SHA256 signature. Timing-safe.""" + expected = hmac_sign(secret, data) + return hmac_mod.compare_digest(signature, expected) diff --git a/src/layerlens/attestation/_verify.py b/src/layerlens/attestation/_verify.py index 35579e3..33b595d 100644 --- a/src/layerlens/attestation/_verify.py +++ b/src/layerlens/attestation/_verify.py @@ -4,6 +4,7 @@ from dataclasses import field, dataclass from ._hash import compute_hash +from ._signing import hmac_verify from ._envelope import HashScope, AttestationEnvelope @@ -16,6 +17,17 @@ class ChainVerification: error: Optional[str] = None +@dataclass +class TrialVerification: + """Result of verifying a full trial: chain + root hash + signatures.""" + + valid: bool + chain_valid: bool = True + trial_hash_valid: bool = True + signatures_valid: bool = True + errors: List[str] = field(default_factory=list) + + @dataclass class TamperingResult: """Result of checking whether trace data was modified after hashing.""" @@ -58,31 +70,61 @@ def verify_chain(envelopes: List[AttestationEnvelope]) -> ChainVerification: def verify_trial( envelopes: List[AttestationEnvelope], trial_envelope: AttestationEnvelope, -) -> ChainVerification: + signing_secret: Optional[bytes] = None, +) -> TrialVerification: """Verify a trial envelope against its event chain. - Checks chain integrity, then verifies the trial hash is correctly - computed over all event hashes. + Checks chain integrity, trial hash correctness, and (optionally) signatures. + Pass ``signing_secret`` to verify HMAC-SHA256 signatures. """ + errors: List[str] = [] + + # 1. Chain continuity chain_result = verify_chain(envelopes) - if not chain_result.valid: - return chain_result + chain_valid = chain_result.valid + if not chain_valid: + errors.append(f"Chain integrity failed: {chain_result.error}") + # 2. Trial scope + hash + trial_hash_valid = True if trial_envelope.scope != HashScope.TRIAL: - return ChainVerification( - valid=False, - error=f"Trial envelope has wrong scope: {trial_envelope.scope}", - ) - - event_hashes = [e.hash for e in envelopes] - expected_hash = compute_hash({"event_hashes": event_hashes}) - if trial_envelope.hash != expected_hash: - return ChainVerification( - valid=False, - error="Trial hash does not match event hashes", - ) - - return ChainVerification(valid=True) + trial_hash_valid = False + errors.append(f"Trial envelope has wrong scope: {trial_envelope.scope}") + else: + event_hashes = [e.hash for e in envelopes] + expected_hash = compute_hash({"event_hashes": event_hashes}) + if trial_envelope.hash != expected_hash: + trial_hash_valid = False + errors.append("Trial hash does not match event hashes") + + # 3. Signatures (only if a signing secret is provided) + signatures_valid = True + if signing_secret is not None: + for i, envelope in enumerate(envelopes): + if not envelope.signature: + signatures_valid = False + errors.append(f"Missing signature on event {i}") + else: + if not hmac_verify(signing_secret, envelope.hash.encode("utf-8"), envelope.signature): + signatures_valid = False + errors.append(f"Invalid signature on event {i}") + + if not trial_envelope.signature: + signatures_valid = False + errors.append("Missing signature on trial envelope") + else: + if not hmac_verify(signing_secret, trial_envelope.hash.encode("utf-8"), trial_envelope.signature): + signatures_valid = False + errors.append("Invalid signature on trial envelope") + + valid = chain_valid and trial_hash_valid and signatures_valid + return TrialVerification( + valid=valid, + chain_valid=chain_valid, + trial_hash_valid=trial_hash_valid, + signatures_valid=signatures_valid, + errors=errors, + ) def detect_tampering( diff --git a/src/layerlens/instrument/__init__.py b/src/layerlens/instrument/__init__.py index 2e11b51..8dde6d0 100644 --- a/src/layerlens/instrument/__init__.py +++ b/src/layerlens/instrument/__init__.py @@ -2,12 +2,13 @@ from ._span import span from ._types import SpanData -from ._recorder import TraceRecorder +from ._recorder import TraceRecorder, clear_signing_key_cache from ._decorator import trace __all__ = [ "SpanData", "TraceRecorder", + "clear_signing_key_cache", "span", "trace", ] diff --git a/src/layerlens/instrument/_decorator.py b/src/layerlens/instrument/_decorator.py index 4f4644f..85ca673 100644 --- a/src/layerlens/instrument/_decorator.py +++ b/src/layerlens/instrument/_decorator.py @@ -6,7 +6,7 @@ from ._types import SpanData from ._context import _current_span, _current_recorder -from ._recorder import TraceRecorder +from ._recorder import _SENTINEL, TraceRecorder def trace( @@ -14,6 +14,7 @@ def trace( *, name: Optional[str] = None, metadata: Optional[Dict[str, Any]] = None, + signing_service: Any = _SENTINEL, ) -> Callable[..., Any]: def decorator(fn: Callable[..., Any]) -> Callable[..., Any]: span_name = name or fn.__name__ @@ -22,7 +23,7 @@ def decorator(fn: Callable[..., Any]) -> Callable[..., Any]: @functools.wraps(fn) async def async_wrapper(*args: Any, **kwargs: Any) -> Any: - recorder = TraceRecorder(client) + recorder = TraceRecorder(client, signing_service=signing_service) root = SpanData( name=span_name, kind="chain", @@ -52,7 +53,7 @@ async def async_wrapper(*args: Any, **kwargs: Any) -> Any: @functools.wraps(fn) def sync_wrapper(*args: Any, **kwargs: Any) -> Any: - recorder = TraceRecorder(client) + recorder = TraceRecorder(client, signing_service=signing_service) root = SpanData( name=span_name, kind="chain", diff --git a/src/layerlens/instrument/_recorder.py b/src/layerlens/instrument/_recorder.py index dce18b3..10561de 100644 --- a/src/layerlens/instrument/_recorder.py +++ b/src/layerlens/instrument/_recorder.py @@ -1,7 +1,10 @@ from __future__ import annotations +import base64 import logging -from typing import Any, Dict, List, Optional +import weakref +import threading +from typing import Any, Dict, List, Tuple, Optional from layerlens.attestation import HashChain @@ -10,6 +13,97 @@ log: logging.Logger = logging.getLogger(__name__) +# Per-client cache for auto-resolved signing keys. +# Uses weakref to the client so entries are evicted when the client is GC'd, +# preventing stale keys from being served to a new client at the same address. +_signing_key_cache: Dict[int, Tuple[Any, Optional[Tuple[str, bytes]]]] = {} # (weakref.ref | callable, value) +_cache_lock = threading.Lock() + +_SENTINEL = object() # distinguishes "not passed" from "passed as None" +_NOT_RESOLVED = object() # cache miss marker + + +def _cache_get(client: Any) -> Any: + """Look up cached signing key for a client. Returns _NOT_RESOLVED on miss.""" + entry = _signing_key_cache.get(id(client), None) + if entry is None: + return _NOT_RESOLVED + ref, value = entry + # If the weakref is dead, the original client was GC'd and a new object + # now occupies the same id(). Evict the stale entry. + if ref() is None: + del _signing_key_cache[id(client)] + return _NOT_RESOLVED + return value + + +def _cache_put(client: Any, value: Optional[Tuple[str, bytes]]) -> None: + """Store signing key in cache, keyed by client identity.""" + try: + ref = weakref.ref(client) + except TypeError: + # Client doesn't support weakrefs (e.g. some Mock objects). + # Fall back to caching without liveness check. + ref = lambda: client # type: ignore[assignment] + _signing_key_cache[id(client)] = (ref, value) + + +def _resolve_signing_key(client: Any) -> Optional[Tuple[str, bytes]]: + """Fetch the org's active signing key, or auto-create one if none exists. + + Returns (key_id, secret_bytes) or None. Result is cached per client + instance so we only hit the API once. If the org has no signing key, + the SDK will attempt to create one automatically. + """ + with _cache_lock: + cached = _cache_get(client) + if cached is not _NOT_RESOLVED: + return cached # type: ignore[no-any-return] + + # Fetch outside the lock to avoid holding it during I/O. + result: Optional[Tuple[str, bytes]] = None + try: + if hasattr(client, "signing_keys"): + key_data = client.signing_keys.get_active() + if not _is_valid_key_data(key_data): + # No active key — auto-create one for the org. + log.info("No active signing key found, auto-creating one for attestation") + key_data = client.signing_keys.create() + if _is_valid_key_data(key_data): + secret_bytes = base64.b64decode(key_data["secret"]) + result = (key_data["key_id"], secret_bytes) + log.info("Attestation signing key resolved: %s", key_data["key_id"]) + else: + log.info("Could not resolve or create signing key — traces will be unsigned") + except Exception: + log.warning("Failed to resolve signing key, traces will be unsigned", exc_info=True) + + with _cache_lock: + # Another thread may have populated while we were fetching — first writer wins. + existing = _cache_get(client) + if existing is not _NOT_RESOLVED: + return existing # type: ignore[no-any-return] + _cache_put(client, result) + + return result + + +def _is_valid_key_data(data: Any) -> bool: + """Check that key data is a dict with both 'key_id' and 'secret'.""" + return isinstance(data, dict) and "secret" in data and "key_id" in data + + +def clear_signing_key_cache(client: Any = None) -> None: + """Clear cached signing keys. Call after key rotation. + + Pass a specific client to clear only its cache, or None to clear all. + """ + with _cache_lock: + if client is None: + _signing_key_cache.clear() + else: + _signing_key_cache.pop(id(client), None) + def _collect_spans(span: SpanData) -> List[Dict[str, Any]]: """Walk the span tree depth-first and return a flat list of span dicts. @@ -29,15 +123,36 @@ def _collect_spans(span: SpanData) -> List[Dict[str, Any]]: class TraceRecorder: - def __init__(self, client: Any) -> None: + def __init__( + self, + client: Any, + signing_service: Any = _SENTINEL, + ) -> None: self._client = client + + if signing_service is _SENTINEL: + # Auto-resolve: fetch the org's active signing key + self._signing_key = _resolve_signing_key(client) + elif signing_service is None: + # Explicit None: no signing + self._signing_key = None + else: + # Explicit (key_id, secret) tuple + self._signing_key = signing_service + self.root: Optional[SpanData] = None def _build_attestation(self) -> Dict[str, Any]: """Build a hash chain from the span tree and return attestation data.""" if self.root is None: return {} - chain = HashChain() + + if self._signing_key is not None: + key_id, secret = self._signing_key + chain = HashChain(signing_key_id=key_id, signing_secret=secret) + else: + chain = HashChain() + spans = _collect_spans(self.root) for span_dict in spans: chain.add_event(span_dict) @@ -55,7 +170,7 @@ def flush(self) -> None: try: attestation = self._build_attestation() except Exception: - log.debug("Failed to build attestation chain", exc_info=True) + log.warning("Failed to build attestation chain", exc_info=True) attestation = {} upload_trace(self._client, trace_data, attestation) @@ -66,6 +181,6 @@ async def async_flush(self) -> None: try: attestation = self._build_attestation() except Exception: - log.debug("Failed to build attestation chain", exc_info=True) + log.warning("Failed to build attestation chain", exc_info=True) attestation = {} await async_upload_trace(self._client, trace_data, attestation) diff --git a/src/layerlens/resources/signing_keys/__init__.py b/src/layerlens/resources/signing_keys/__init__.py new file mode 100644 index 0000000..7a96b5a --- /dev/null +++ b/src/layerlens/resources/signing_keys/__init__.py @@ -0,0 +1,3 @@ +from .signing_keys import SigningKeys, AsyncSigningKeys + +__all__ = ["SigningKeys", "AsyncSigningKeys"] diff --git a/src/layerlens/resources/signing_keys/signing_keys.py b/src/layerlens/resources/signing_keys/signing_keys.py new file mode 100644 index 0000000..a4ff1c7 --- /dev/null +++ b/src/layerlens/resources/signing_keys/signing_keys.py @@ -0,0 +1,153 @@ +from __future__ import annotations + +import logging +from typing import Any, Dict, List, Union, Optional + +import httpx + +from ..._resource import SyncAPIResource, AsyncAPIResource +from ..._constants import DEFAULT_TIMEOUT + +log: logging.Logger = logging.getLogger(__name__) + + +def _unwrap(resp: Any) -> Any: + if isinstance(resp, dict) and "data" in resp and "status" in resp: + return resp["data"] + return resp + + +class SigningKeys(SyncAPIResource): + def _base_url(self) -> str: + org_id = self._client.organization_id + if not org_id: + raise ValueError("Client has no organization_id configured") + return f"/organizations/{org_id}/signing-keys" + + def get_active( + self, + *, + timeout: float | httpx.Timeout | None = DEFAULT_TIMEOUT, + ) -> Optional[Dict[str, Any]]: + """Fetch the active signing key (key_id, name, secret). + + Returns None if no active signing key exists (404). + """ + try: + resp = self._get( + f"{self._base_url()}/active", + timeout=timeout, + cast_to=dict, + ) + data = _unwrap(resp) + return data if isinstance(data, dict) else None + except Exception: + log.debug("No active signing key found", exc_info=True) + return None + + def create( + self, + *, + name: str = "default", + timeout: float | httpx.Timeout | None = DEFAULT_TIMEOUT, + ) -> Optional[Dict[str, Any]]: + """Create a new signing key for the organization. + + Returns the key data (key_id, name, secret) or None on failure. + """ + try: + resp = self._post( + self._base_url(), + body={"name": name}, + timeout=timeout, + cast_to=dict, + ) + data = _unwrap(resp) + return data if isinstance(data, dict) else None + except Exception: + log.debug("Failed to create signing key", exc_info=True) + return None + + def list( + self, + *, + timeout: float | httpx.Timeout | None = DEFAULT_TIMEOUT, + ) -> Optional[Union[List[Dict[str, Any]], Dict[str, Any]]]: + """List signing key metadata (no secrets).""" + try: + resp = self._get(self._base_url(), timeout=timeout, cast_to=dict) + data = _unwrap(resp) + if isinstance(data, (dict, list)): + return data + return None + except Exception: + log.debug("Failed to list signing keys", exc_info=True) + return None + + +class AsyncSigningKeys(AsyncAPIResource): + def _base_url(self) -> str: + org_id = self._client.organization_id + if not org_id: + raise ValueError("Client has no organization_id configured") + return f"/organizations/{org_id}/signing-keys" + + async def get_active( + self, + *, + timeout: float | httpx.Timeout | None = DEFAULT_TIMEOUT, + ) -> Optional[Dict[str, Any]]: + """Fetch the active signing key (key_id, name, secret). + + Returns None if no active signing key exists (404). + """ + try: + resp = await self._get( + f"{self._base_url()}/active", + timeout=timeout, + cast_to=dict, + ) + data = _unwrap(resp) + return data if isinstance(data, dict) else None + except Exception: + log.debug("No active signing key found", exc_info=True) + return None + + async def create( + self, + *, + name: str = "default", + timeout: float | httpx.Timeout | None = DEFAULT_TIMEOUT, + ) -> Optional[Dict[str, Any]]: + """Create a new signing key for the organization. + + Returns the key data (key_id, name, secret) or None on failure. + """ + try: + resp = await self._post( + self._base_url(), + body={"name": name}, + timeout=timeout, + cast_to=dict, + ) + data = _unwrap(resp) + return data if isinstance(data, dict) else None + except Exception: + log.debug("Failed to create signing key", exc_info=True) + return None + + async def list( + self, + *, + timeout: float | httpx.Timeout | None = DEFAULT_TIMEOUT, + ) -> Optional[Union[List[Dict[str, Any]], Dict[str, Any]]]: + """List signing key metadata (no secrets).""" + try: + resp = await self._get(self._base_url(), timeout=timeout, cast_to=dict) + data = _unwrap(resp) + if isinstance(data, (dict, list)): + return data + return None + except Exception: + log.debug("Failed to list signing keys", exc_info=True) + return None diff --git a/tests/attestation/test_hash.py b/tests/attestation/test_hash.py index 1f5e9ad..203c6e6 100644 --- a/tests/attestation/test_hash.py +++ b/tests/attestation/test_hash.py @@ -60,3 +60,12 @@ def test_different_data_different_hash(self): def test_empty_dict(self): h = compute_hash({}) assert re.match(r"^sha256:[0-9a-f]{64}$", h) + + def test_cross_language_vector(self): + """Pinned vector shared with Go backend (TestComputeCanonicalHash_CrossLanguageVector). + + If this test fails, Python and Go will produce different root hashes + for the same trace, breaking attestation verification. + """ + h = compute_hash({"event_hashes": ["sha256:aaa", "sha256:bbb"]}) + assert h == "sha256:b930d0a2cbda5171b8a12d17445c38b8c0842344f2d691a00d24b3359a854db5" diff --git a/tests/attestation/test_signing.py b/tests/attestation/test_signing.py new file mode 100644 index 0000000..13ca357 --- /dev/null +++ b/tests/attestation/test_signing.py @@ -0,0 +1,190 @@ +from __future__ import annotations + +from layerlens.attestation import ( + HashChain, + hmac_sign, + hmac_verify, + verify_trial, +) + + +class TestHMACSigning: + def test_sign_produces_base64(self): + sig = hmac_sign(b"test-key", b"sha256:" + b"a" * 64) + assert sig # non-empty string + assert isinstance(sig, str) + + def test_sign_deterministic(self): + data = b"sha256:" + b"a" * 64 + assert hmac_sign(b"test-key", data) == hmac_sign(b"test-key", data) + + def test_different_data_different_signatures(self): + s1 = hmac_sign(b"test-key", b"sha256:" + b"a" * 64) + s2 = hmac_sign(b"test-key", b"sha256:" + b"b" * 64) + assert s1 != s2 + + def test_different_keys_different_signatures(self): + data = b"sha256:" + b"a" * 64 + assert hmac_sign(b"key-1", data) != hmac_sign(b"key-2", data) + + def test_verify_valid(self): + data = b"sha256:" + b"a" * 64 + sig = hmac_sign(b"test-key", data) + assert hmac_verify(b"test-key", data, sig) + + def test_verify_invalid(self): + assert not hmac_verify(b"test-key", b"sha256:" + b"a" * 64, "bogus") + + def test_verify_wrong_data(self): + sig = hmac_sign(b"test-key", b"sha256:" + b"a" * 64) + assert not hmac_verify(b"test-key", b"sha256:" + b"b" * 64, sig) + + def test_verify_wrong_key(self): + sig = hmac_sign(b"key-1", b"data") + assert not hmac_verify(b"key-2", b"data", sig) + + +class TestChainWithSigning: + def test_signed_envelopes(self): + chain = HashChain(signing_key_id="org-123", signing_secret=b"test-key") + e1 = chain.add_event({"name": "span-1"}) + e2 = chain.add_event({"name": "span-2"}) + + assert e1.signature is not None + assert e1.signing_key_id == "org-123" + assert e2.signature is not None + assert e2.signing_key_id == "org-123" + # Different events get different signatures + assert e1.signature != e2.signature + + def test_trial_signed(self): + chain = HashChain(signing_key_id="org-123", signing_secret=b"test-key") + chain.add_event({"name": "span-1"}) + trial = chain.finalize() + + assert trial.signature is not None + assert trial.signing_key_id == "org-123" + + def test_unsigned_chain_has_no_signatures(self): + chain = HashChain() + e1 = chain.add_event({"name": "span-1"}) + trial = chain.finalize() + + assert e1.signature is None + assert e1.signing_key_id is None + assert trial.signature is None + + def test_to_dict_includes_signature_fields(self): + chain = HashChain(signing_key_id="org-123", signing_secret=b"test-key") + chain.add_event({"name": "span-1"}) + d = chain.to_dict() + + event = d["events"][0] + assert "signature" in event + assert "signing_key_id" in event + + def test_to_dict_omits_signature_when_unsigned(self): + chain = HashChain() + chain.add_event({"name": "span-1"}) + d = chain.to_dict() + + event = d["events"][0] + assert "signature" not in event + assert "signing_key_id" not in event + + +class TestVerifyTrialWithSigning: + def test_valid_signed_trial(self): + secret = b"test-key" + chain = HashChain(signing_key_id="org-123", signing_secret=secret) + chain.add_event({"name": "a"}) + chain.add_event({"name": "b"}) + envelopes = chain.envelopes + trial = chain.finalize() + + result = verify_trial(envelopes, trial, signing_secret=secret) + assert result.valid + assert result.chain_valid + assert result.trial_hash_valid + assert result.signatures_valid + assert result.errors == [] + + def test_tampered_signature_detected(self): + secret = b"test-key" + chain = HashChain(signing_key_id="org-123", signing_secret=secret) + chain.add_event({"name": "a"}) + envelopes = chain.envelopes + trial = chain.finalize() + + # Tamper with the event signature + envelopes[0].signature = "dGFtcGVyZWQ=" # base64("tampered") + + result = verify_trial(envelopes, trial, signing_secret=secret) + assert not result.valid + assert not result.signatures_valid + assert result.chain_valid # chain structure is still fine + assert result.trial_hash_valid # trial hash is still fine + + def test_wrong_key_rejects(self): + chain = HashChain(signing_key_id="org-123", signing_secret=b"key-1") + chain.add_event({"name": "a"}) + envelopes = chain.envelopes + trial = chain.finalize() + + result = verify_trial(envelopes, trial, signing_secret=b"key-2") + assert not result.valid + assert not result.signatures_valid + + def test_unsigned_chain_passes_without_secret(self): + """verify_trial without signing_secret ignores missing signatures.""" + chain = HashChain() + chain.add_event({"name": "a"}) + envelopes = chain.envelopes + trial = chain.finalize() + + result = verify_trial(envelopes, trial) + assert result.valid + assert result.signatures_valid # vacuously true + + def test_stripped_signatures_detected(self): + """When signing_secret is provided, missing signatures should fail.""" + secret = b"test-key" + chain = HashChain(signing_key_id="org-123", signing_secret=secret) + chain.add_event({"name": "a"}) + envelopes = chain.envelopes + trial = chain.finalize() + + # Strip signatures + envelopes[0].signature = None + trial.signature = None + + result = verify_trial(envelopes, trial, signing_secret=secret) + assert not result.valid + assert not result.signatures_valid + assert any("Missing signature" in e for e in result.errors) + + def test_backward_compat_old_verify_trial(self): + """Old-style verify_trial (no signing_secret) still returns valid for valid chains.""" + chain = HashChain() + chain.add_event({"name": "a"}) + chain.add_event({"name": "b"}) + envelopes = chain.envelopes + trial = chain.finalize() + + result = verify_trial(envelopes, trial) + assert result.valid + + def test_single_event_signed_chain(self): + """Signed chain with exactly one event works correctly.""" + secret = b"test-key" + chain = HashChain(signing_key_id="org-1", signing_secret=secret) + chain.add_event({"name": "only"}) + envelopes = chain.envelopes + trial = chain.finalize() + + assert len(envelopes) == 1 + assert envelopes[0].signature is not None + + result = verify_trial(envelopes, trial, signing_secret=secret) + assert result.valid + assert result.signatures_valid diff --git a/tests/attestation/test_verify.py b/tests/attestation/test_verify.py index f7567d8..60f14c8 100644 --- a/tests/attestation/test_verify.py +++ b/tests/attestation/test_verify.py @@ -1,12 +1,12 @@ from __future__ import annotations -from layerlens.attestation._chain import HashChain -from layerlens.attestation._verify import ( +from layerlens.attestation import ( + HashChain, + HashScope, verify_chain, verify_trial, detect_tampering, ) -from layerlens.attestation._envelope import HashScope class TestVerifyChain: @@ -71,7 +71,8 @@ def test_wrong_scope_rejected(self): trial.scope = HashScope.EVENT # Wrong scope result = verify_trial(envelopes, trial) assert not result.valid - assert "scope" in (result.error or "") + assert not result.trial_hash_valid + assert any("scope" in e for e in result.errors) def test_tampered_trial_hash(self): chain = HashChain() @@ -81,7 +82,8 @@ def test_tampered_trial_hash(self): trial.hash = "sha256:" + "0" * 64 # Wrong hash result = verify_trial(envelopes, trial) assert not result.valid - assert "does not match" in (result.error or "") + assert not result.trial_hash_valid + assert any("does not match" in e for e in result.errors) class TestDetectTampering: @@ -125,3 +127,21 @@ def test_detect_count_mismatch(self): result = detect_tampering(chain.envelopes, [{"name": "a"}]) assert result.tampered assert result.chain_broken + + def test_detect_tampering_with_signed_chain(self): + """detect_tampering works correctly on chains that were signed.""" + data = [{"name": "a"}, {"name": "b"}, {"name": "c"}] + chain = HashChain(signing_key_id="org-1", signing_secret=b"test-key") + for d in data: + chain.add_event(d) + + # No tampering — should pass + result = detect_tampering(chain.envelopes, data) + assert not result.tampered + assert result.modified_indices == [] + + # Tamper with one event + tampered = [{"name": "a"}, {"name": "CHANGED"}, {"name": "c"}] + result = detect_tampering(chain.envelopes, tampered) + assert result.tampered + assert 1 in result.modified_indices diff --git a/tests/instrument/test_signing_autofetch.py b/tests/instrument/test_signing_autofetch.py new file mode 100644 index 0000000..928372f --- /dev/null +++ b/tests/instrument/test_signing_autofetch.py @@ -0,0 +1,412 @@ +"""Tests for automatic signing key fetch in @trace decorator.""" + +from __future__ import annotations + +import json +import base64 +import asyncio +import threading +from unittest.mock import Mock, AsyncMock + +import pytest + +from layerlens.instrument import trace, clear_signing_key_cache +from layerlens.instrument._recorder import _signing_key_cache, _resolve_signing_key + + +@pytest.fixture(autouse=True) +def _clear_cache(): + """Clear signing key cache before and after each test.""" + _signing_key_cache.clear() + yield + _signing_key_cache.clear() + + +def _make_client(*, signing_key_response=None, create_key_response=None): + """Create a mock client with optional signing_keys.get_active() response. + + If create_key_response is provided, signing_keys.create() returns it. + Otherwise create() returns None (simulating backend failure or no-op). + """ + client = Mock() + client.traces = Mock() + client.traces.upload = Mock() + client.signing_keys = Mock() + if signing_key_response is not None: + client.signing_keys.get_active = Mock(return_value=signing_key_response) + else: + client.signing_keys.get_active = Mock(return_value=None) + if create_key_response is not None: + client.signing_keys.create = Mock(return_value=create_key_response) + else: + client.signing_keys.create = Mock(return_value=None) + return client + + +def _capture_upload(client): + """Set up trace capture on the mock client. Returns dict that gets populated.""" + uploaded = {} + + def _capture(path): + with open(path) as f: + uploaded["trace"] = json.load(f) + + client.traces.upload.side_effect = _capture + return uploaded + + +class TestAutoFetchSigningKey: + def test_auto_fetches_and_signs(self): + """When no signing_service passed, auto-fetch from client and sign.""" + secret = b"test-auto-key-32-bytes-long!!!!!" + client = _make_client( + signing_key_response={ + "key_id": "sk_auto_123", + "name": "auto-key", + "secret": base64.b64encode(secret).decode(), + } + ) + uploaded = _capture_upload(client) + + @trace(client) + def my_agent(): + return "hello" + + my_agent() + + client.signing_keys.get_active.assert_called_once() + payload = uploaded["trace"][0] + assert "attestation" in payload + att = payload["attestation"] + # Chain events should have signatures + events = att["chain"]["events"] + assert len(events) > 0 + for event in events: + assert "signature" in event, "Event should be signed" + assert event["signing_key_id"] == "sk_auto_123" + + def test_auto_creates_key_when_none_exists(self): + """When org has no active key, SDK auto-creates one and signs.""" + secret = b"auto-created-key-32-bytes!!!!!!!" + client = _make_client( + signing_key_response=None, + create_key_response={ + "key_id": "sk_auto_created", + "name": "default", + "secret": base64.b64encode(secret).decode(), + }, + ) + uploaded = _capture_upload(client) + + @trace(client) + def my_agent(): + return "hello" + + my_agent() + + client.signing_keys.get_active.assert_called_once() + client.signing_keys.create.assert_called_once() + payload = uploaded["trace"][0] + assert "attestation" in payload + events = payload["attestation"]["chain"]["events"] + assert len(events) > 0 + for event in events: + assert "signature" in event, "Event should be signed with auto-created key" + assert event["signing_key_id"] == "sk_auto_created" + + def test_no_signing_key_and_create_fails_produces_unsigned(self): + """When org has no active key AND create fails, traces are unsigned.""" + client = _make_client(signing_key_response=None, create_key_response=None) + uploaded = _capture_upload(client) + + @trace(client) + def my_agent(): + return "hello" + + my_agent() + + client.signing_keys.get_active.assert_called_once() + client.signing_keys.create.assert_called_once() + payload = uploaded["trace"][0] + assert "attestation" in payload + events = payload["attestation"]["chain"]["events"] + assert len(events) > 0 + for event in events: + assert "signature" not in event, "Event should NOT be signed" + + def test_caches_across_traces(self): + """Signing key is fetched once and reused across multiple @trace calls.""" + secret = b"cached-key-32-bytes-long!!!!!!!!" + client = _make_client( + signing_key_response={ + "key_id": "sk_cached", + "name": "cached", + "secret": base64.b64encode(secret).decode(), + } + ) + all_uploads: list = [] + + def _capture_all(path): + with open(path) as f: + all_uploads.append(json.load(f)) + + client.traces.upload.side_effect = _capture_all + + @trace(client) + def agent_a(): + return "a" + + @trace(client) + def agent_b(): + return "b" + + agent_a() + agent_b() + + # Only one API call despite two traces + client.signing_keys.get_active.assert_called_once() + # Both traces should be signed + assert len(all_uploads) == 2 + for upload in all_uploads: + events = upload[0]["attestation"]["chain"]["events"] + assert "signature" in events[0] + assert events[0]["signing_key_id"] == "sk_cached" + + def test_clear_cache_forces_refetch(self): + """clear_signing_key_cache() causes next trace to refetch.""" + secret = b"key-before-rotation!!!!!!!!!!!!!!" + client = _make_client( + signing_key_response={ + "key_id": "sk_old", + "name": "old", + "secret": base64.b64encode(secret).decode(), + } + ) + _capture_upload(client) + + @trace(client) + def my_agent(): + return "hello" + + my_agent() + assert client.signing_keys.get_active.call_count == 1 + + # Simulate key rotation + clear_signing_key_cache(client) + new_secret = b"key-after-rotation!!!!!!!!!!!!!!!" + client.signing_keys.get_active.return_value = { + "key_id": "sk_new", + "name": "new", + "secret": base64.b64encode(new_secret).decode(), + } + + my_agent() + assert client.signing_keys.get_active.call_count == 2 + + def test_explicit_signing_key_skips_autofetch(self): + """Passing signing_service= explicitly bypasses auto-fetch entirely.""" + client = _make_client( + signing_key_response={ + "key_id": "sk_should_not_fetch", + "name": "nope", + "secret": base64.b64encode(b"nope").decode(), + } + ) + uploaded = _capture_upload(client) + + @trace(client, signing_service=("explicit-key", b"explicit-secret")) + def my_agent(): + return "hello" + + my_agent() + + # Auto-fetch should NOT be called + client.signing_keys.get_active.assert_not_called() + # But traces should still be signed with the explicit key + payload = uploaded["trace"][0] + events = payload["attestation"]["chain"]["events"] + assert events[0]["signing_key_id"] == "explicit-key" + + def test_explicit_none_disables_signing(self): + """Passing signing_service=None explicitly disables signing (no auto-fetch).""" + client = _make_client( + signing_key_response={ + "key_id": "sk_should_not_fetch", + "name": "nope", + "secret": base64.b64encode(b"nope").decode(), + } + ) + uploaded = _capture_upload(client) + + @trace(client, signing_service=None) + def my_agent(): + return "hello" + + my_agent() + + # Auto-fetch should NOT be called + client.signing_keys.get_active.assert_not_called() + # Traces should be unsigned + payload = uploaded["trace"][0] + events = payload["attestation"]["chain"]["events"] + for event in events: + assert "signature" not in event + + def test_fetch_failure_degrades_to_unsigned(self): + """If get_active() throws, traces are uploaded unsigned (not broken).""" + client = _make_client() + client.signing_keys.get_active = Mock(side_effect=RuntimeError("network error")) + uploaded = _capture_upload(client) + + @trace(client) + def my_agent(): + return "hello" + + my_agent() + + payload = uploaded["trace"][0] + assert "attestation" in payload + events = payload["attestation"]["chain"]["events"] + for event in events: + assert "signature" not in event + + def test_client_without_signing_keys_attr(self): + """Clients that don't have signing_keys (e.g. old SDK) degrade gracefully.""" + client = Mock(spec=["traces"]) + client.traces = Mock() + client.traces.upload = Mock() + uploaded = {} + + def _capture(path): + with open(path) as f: + uploaded["trace"] = json.load(f) + + client.traces.upload.side_effect = _capture + + @trace(client) + def my_agent(): + return "hello" + + my_agent() + + payload = uploaded["trace"][0] + assert "attestation" in payload + # Should be unsigned, no crash + events = payload["attestation"]["chain"]["events"] + for event in events: + assert "signature" not in event + + def test_malformed_response_missing_key_id(self): + """get_active() returns dict with secret but no key_id — falls back to create().""" + client = _make_client( + signing_key_response={ + "name": "broken-key", + "secret": base64.b64encode(b"some-secret").decode(), + # Missing "key_id" + }, + create_key_response=None, # create also fails + ) + uploaded = _capture_upload(client) + + @trace(client) + def my_agent(): + return "hello" + + my_agent() + + client.signing_keys.create.assert_called_once() + payload = uploaded["trace"][0] + events = payload["attestation"]["chain"]["events"] + for event in events: + assert "signature" not in event, "Should be unsigned when key_id is missing and create fails" + + def test_malformed_response_missing_secret(self): + """get_active() returns dict with key_id but no secret — falls back to create().""" + client = _make_client( + signing_key_response={ + "key_id": "sk_123", + "name": "broken-key", + # Missing "secret" + }, + create_key_response=None, # create also fails + ) + uploaded = _capture_upload(client) + + @trace(client) + def my_agent(): + return "hello" + + my_agent() + + payload = uploaded["trace"][0] + events = payload["attestation"]["chain"]["events"] + for event in events: + assert "signature" not in event, "Should be unsigned when secret is missing" + + def test_concurrent_cache_access(self): + """Multiple threads resolving the same client only fetch once.""" + secret = b"concurrent-key-32-bytes-long!!!!!" + client = _make_client( + signing_key_response={ + "key_id": "sk_concurrent", + "name": "concurrent", + "secret": base64.b64encode(secret).decode(), + } + ) + + results = [] + errors = [] + + def resolve(): + try: + result = _resolve_signing_key(client) + results.append(result) + except Exception as e: + errors.append(e) + + threads = [threading.Thread(target=resolve) for _ in range(10)] + for t in threads: + t.start() + for t in threads: + t.join() + + assert not errors + assert len(results) == 10 + # All threads got the same result + for r in results: + assert r is not None + assert r[0] == "sk_concurrent" + + +class TestAsyncSigningFlush: + def test_async_trace_signs_correctly(self): + """Async @trace path produces signed attestation.""" + secret = b"async-test-key-32-bytes-long!!!!!" + client = _make_client( + signing_key_response={ + "key_id": "sk_async", + "name": "async-key", + "secret": base64.b64encode(secret).decode(), + } + ) + uploaded = {} + + async def _async_capture(path): + with open(path) as f: + uploaded["trace"] = json.load(f) + + client.traces.upload = AsyncMock(side_effect=_async_capture) + + @trace(client) + async def my_async_agent(): + return "hello async" + + asyncio.run(my_async_agent()) + + payload = uploaded["trace"][0] + assert "attestation" in payload + events = payload["attestation"]["chain"]["events"] + assert len(events) > 0 + for event in events: + assert "signature" in event + assert event["signing_key_id"] == "sk_async" From 4c447317cc1ee1abf14a881caf7bd656603ade41 Mon Sep 17 00:00:00 2001 From: Gary <59334078+garrettallen14@users.noreply.github.com> Date: Fri, 27 Mar 2026 13:44:46 -0700 Subject: [PATCH 3/4] refactor: remove client-side signing, delegate to server-side attestation --- src/layerlens/_client.py | 13 - src/layerlens/attestation/_chain.py | 23 +- src/layerlens/instrument/__init__.py | 3 +- src/layerlens/instrument/_decorator.py | 7 +- src/layerlens/instrument/_recorder.py | 126 +----- .../resources/signing_keys/__init__.py | 3 - .../resources/signing_keys/signing_keys.py | 153 ------- tests/attestation/test_signing.py | 96 ++-- tests/attestation/test_verify.py | 6 +- tests/instrument/test_signing_autofetch.py | 412 ------------------ 10 files changed, 55 insertions(+), 787 deletions(-) delete mode 100644 src/layerlens/resources/signing_keys/__init__.py delete mode 100644 src/layerlens/resources/signing_keys/signing_keys.py delete mode 100644 tests/instrument/test_signing_autofetch.py diff --git a/src/layerlens/_client.py b/src/layerlens/_client.py index c679990..032e15b 100644 --- a/src/layerlens/_client.py +++ b/src/layerlens/_client.py @@ -25,7 +25,6 @@ from .resources.benchmarks import Benchmarks, AsyncBenchmarks from .resources.evaluations import Evaluations, AsyncEvaluations from .resources.integrations import Integrations, AsyncIntegrations - from .resources.signing_keys import SigningKeys, AsyncSigningKeys from .resources.evaluation_spaces import EvaluationSpaces, AsyncEvaluationSpaces from .resources.trace_evaluations import TraceEvaluations, AsyncTraceEvaluations from .resources.judge_optimizations import JudgeOptimizations, AsyncJudgeOptimizations @@ -140,12 +139,6 @@ def scorers(self) -> Scorers: return Scorers(self) - @cached_property - def signing_keys(self) -> SigningKeys: - from .resources.signing_keys import SigningKeys - - return SigningKeys(self) - @cached_property def evaluation_spaces(self) -> EvaluationSpaces: from .resources.evaluation_spaces import EvaluationSpaces @@ -333,12 +326,6 @@ def scorers(self) -> AsyncScorers: return AsyncScorers(self) - @cached_property - def signing_keys(self) -> AsyncSigningKeys: - from .resources.signing_keys import AsyncSigningKeys - - return AsyncSigningKeys(self) - @cached_property def evaluation_spaces(self) -> AsyncEvaluationSpaces: from .resources.evaluation_spaces import AsyncEvaluationSpaces diff --git a/src/layerlens/attestation/_chain.py b/src/layerlens/attestation/_chain.py index 93a3ace..b6b45d2 100644 --- a/src/layerlens/attestation/_chain.py +++ b/src/layerlens/attestation/_chain.py @@ -3,7 +3,6 @@ from typing import Any, Dict, List, Optional from ._hash import compute_hash -from ._signing import hmac_sign from ._envelope import HashScope, AttestationEnvelope @@ -14,23 +13,15 @@ class HashChain: a tamper-evident chain. If any event is modified after the fact, the chain breaks at that point. - If ``signing_secret`` is provided, each envelope's hash is - HMAC-SHA256 signed for authenticity on top of integrity. + Signing is handled server-side at trace ingestion. The SDK builds + the hash chain for integrity; the backend signs for authenticity. """ - def __init__( - self, - signing_key_id: Optional[str] = None, - signing_secret: Optional[bytes] = None, - ) -> None: - if signing_secret is not None and not signing_key_id: - raise ValueError("signing_key_id is required when signing_secret is provided") + def __init__(self) -> None: self._chain: List[AttestationEnvelope] = [] self._last_hash: Optional[str] = None self._terminated: bool = False self._terminate_reason: Optional[str] = None - self._signing_key_id = signing_key_id - self._signing_secret = signing_secret @property def envelopes(self) -> List[AttestationEnvelope]: @@ -44,12 +35,6 @@ def _check_active(self) -> None: if self._terminated: raise RuntimeError(f"Hash chain terminated: {self._terminate_reason}. No further events can be added.") - def _sign_envelope(self, envelope: AttestationEnvelope) -> None: - """Sign an envelope's hash if a signing secret is configured.""" - if self._signing_secret is not None: - envelope.signature = hmac_sign(self._signing_secret, envelope.hash.encode("utf-8")) - envelope.signing_key_id = self._signing_key_id - def add_event(self, data: Dict[str, Any]) -> AttestationEnvelope: """Hash an event and append it to the chain.""" self._check_active() @@ -61,7 +46,6 @@ def add_event(self, data: Dict[str, Any]) -> AttestationEnvelope: scope=HashScope.EVENT, previous_hash=self._last_hash, ) - self._sign_envelope(envelope) self._chain.append(envelope) self._last_hash = event_hash return envelope @@ -86,7 +70,6 @@ def finalize(self) -> AttestationEnvelope: scope=HashScope.TRIAL, previous_hash=self._last_hash, ) - self._sign_envelope(trial_envelope) # Seal — no more events after finalization self._terminated = True self._terminate_reason = "chain finalized" diff --git a/src/layerlens/instrument/__init__.py b/src/layerlens/instrument/__init__.py index 8dde6d0..2e11b51 100644 --- a/src/layerlens/instrument/__init__.py +++ b/src/layerlens/instrument/__init__.py @@ -2,13 +2,12 @@ from ._span import span from ._types import SpanData -from ._recorder import TraceRecorder, clear_signing_key_cache +from ._recorder import TraceRecorder from ._decorator import trace __all__ = [ "SpanData", "TraceRecorder", - "clear_signing_key_cache", "span", "trace", ] diff --git a/src/layerlens/instrument/_decorator.py b/src/layerlens/instrument/_decorator.py index 85ca673..4f4644f 100644 --- a/src/layerlens/instrument/_decorator.py +++ b/src/layerlens/instrument/_decorator.py @@ -6,7 +6,7 @@ from ._types import SpanData from ._context import _current_span, _current_recorder -from ._recorder import _SENTINEL, TraceRecorder +from ._recorder import TraceRecorder def trace( @@ -14,7 +14,6 @@ def trace( *, name: Optional[str] = None, metadata: Optional[Dict[str, Any]] = None, - signing_service: Any = _SENTINEL, ) -> Callable[..., Any]: def decorator(fn: Callable[..., Any]) -> Callable[..., Any]: span_name = name or fn.__name__ @@ -23,7 +22,7 @@ def decorator(fn: Callable[..., Any]) -> Callable[..., Any]: @functools.wraps(fn) async def async_wrapper(*args: Any, **kwargs: Any) -> Any: - recorder = TraceRecorder(client, signing_service=signing_service) + recorder = TraceRecorder(client) root = SpanData( name=span_name, kind="chain", @@ -53,7 +52,7 @@ async def async_wrapper(*args: Any, **kwargs: Any) -> Any: @functools.wraps(fn) def sync_wrapper(*args: Any, **kwargs: Any) -> Any: - recorder = TraceRecorder(client, signing_service=signing_service) + recorder = TraceRecorder(client) root = SpanData( name=span_name, kind="chain", diff --git a/src/layerlens/instrument/_recorder.py b/src/layerlens/instrument/_recorder.py index 10561de..758c3d8 100644 --- a/src/layerlens/instrument/_recorder.py +++ b/src/layerlens/instrument/_recorder.py @@ -1,10 +1,7 @@ from __future__ import annotations -import base64 import logging -import weakref -import threading -from typing import Any, Dict, List, Tuple, Optional +from typing import Any, Dict, List, Optional from layerlens.attestation import HashChain @@ -13,97 +10,6 @@ log: logging.Logger = logging.getLogger(__name__) -# Per-client cache for auto-resolved signing keys. -# Uses weakref to the client so entries are evicted when the client is GC'd, -# preventing stale keys from being served to a new client at the same address. -_signing_key_cache: Dict[int, Tuple[Any, Optional[Tuple[str, bytes]]]] = {} # (weakref.ref | callable, value) -_cache_lock = threading.Lock() - -_SENTINEL = object() # distinguishes "not passed" from "passed as None" -_NOT_RESOLVED = object() # cache miss marker - - -def _cache_get(client: Any) -> Any: - """Look up cached signing key for a client. Returns _NOT_RESOLVED on miss.""" - entry = _signing_key_cache.get(id(client), None) - if entry is None: - return _NOT_RESOLVED - ref, value = entry - # If the weakref is dead, the original client was GC'd and a new object - # now occupies the same id(). Evict the stale entry. - if ref() is None: - del _signing_key_cache[id(client)] - return _NOT_RESOLVED - return value - - -def _cache_put(client: Any, value: Optional[Tuple[str, bytes]]) -> None: - """Store signing key in cache, keyed by client identity.""" - try: - ref = weakref.ref(client) - except TypeError: - # Client doesn't support weakrefs (e.g. some Mock objects). - # Fall back to caching without liveness check. - ref = lambda: client # type: ignore[assignment] - _signing_key_cache[id(client)] = (ref, value) - - -def _resolve_signing_key(client: Any) -> Optional[Tuple[str, bytes]]: - """Fetch the org's active signing key, or auto-create one if none exists. - - Returns (key_id, secret_bytes) or None. Result is cached per client - instance so we only hit the API once. If the org has no signing key, - the SDK will attempt to create one automatically. - """ - with _cache_lock: - cached = _cache_get(client) - if cached is not _NOT_RESOLVED: - return cached # type: ignore[no-any-return] - - # Fetch outside the lock to avoid holding it during I/O. - result: Optional[Tuple[str, bytes]] = None - try: - if hasattr(client, "signing_keys"): - key_data = client.signing_keys.get_active() - if not _is_valid_key_data(key_data): - # No active key — auto-create one for the org. - log.info("No active signing key found, auto-creating one for attestation") - key_data = client.signing_keys.create() - if _is_valid_key_data(key_data): - secret_bytes = base64.b64decode(key_data["secret"]) - result = (key_data["key_id"], secret_bytes) - log.info("Attestation signing key resolved: %s", key_data["key_id"]) - else: - log.info("Could not resolve or create signing key — traces will be unsigned") - except Exception: - log.warning("Failed to resolve signing key, traces will be unsigned", exc_info=True) - - with _cache_lock: - # Another thread may have populated while we were fetching — first writer wins. - existing = _cache_get(client) - if existing is not _NOT_RESOLVED: - return existing # type: ignore[no-any-return] - _cache_put(client, result) - - return result - - -def _is_valid_key_data(data: Any) -> bool: - """Check that key data is a dict with both 'key_id' and 'secret'.""" - return isinstance(data, dict) and "secret" in data and "key_id" in data - - -def clear_signing_key_cache(client: Any = None) -> None: - """Clear cached signing keys. Call after key rotation. - - Pass a specific client to clear only its cache, or None to clear all. - """ - with _cache_lock: - if client is None: - _signing_key_cache.clear() - else: - _signing_key_cache.pop(id(client), None) - def _collect_spans(span: SpanData) -> List[Dict[str, Any]]: """Walk the span tree depth-first and return a flat list of span dicts. @@ -123,36 +29,20 @@ def _collect_spans(span: SpanData) -> List[Dict[str, Any]]: class TraceRecorder: - def __init__( - self, - client: Any, - signing_service: Any = _SENTINEL, - ) -> None: + def __init__(self, client: Any) -> None: self._client = client - - if signing_service is _SENTINEL: - # Auto-resolve: fetch the org's active signing key - self._signing_key = _resolve_signing_key(client) - elif signing_service is None: - # Explicit None: no signing - self._signing_key = None - else: - # Explicit (key_id, secret) tuple - self._signing_key = signing_service - self.root: Optional[SpanData] = None def _build_attestation(self) -> Dict[str, Any]: - """Build a hash chain from the span tree and return attestation data.""" + """Build an unsigned hash chain from the span tree. + + The chain provides integrity (tamper-evidence). Signing is + handled server-side at trace ingestion for authenticity. + """ if self.root is None: return {} - if self._signing_key is not None: - key_id, secret = self._signing_key - chain = HashChain(signing_key_id=key_id, signing_secret=secret) - else: - chain = HashChain() - + chain = HashChain() spans = _collect_spans(self.root) for span_dict in spans: chain.add_event(span_dict) diff --git a/src/layerlens/resources/signing_keys/__init__.py b/src/layerlens/resources/signing_keys/__init__.py deleted file mode 100644 index 7a96b5a..0000000 --- a/src/layerlens/resources/signing_keys/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from .signing_keys import SigningKeys, AsyncSigningKeys - -__all__ = ["SigningKeys", "AsyncSigningKeys"] diff --git a/src/layerlens/resources/signing_keys/signing_keys.py b/src/layerlens/resources/signing_keys/signing_keys.py deleted file mode 100644 index a4ff1c7..0000000 --- a/src/layerlens/resources/signing_keys/signing_keys.py +++ /dev/null @@ -1,153 +0,0 @@ -from __future__ import annotations - -import logging -from typing import Any, Dict, List, Union, Optional - -import httpx - -from ..._resource import SyncAPIResource, AsyncAPIResource -from ..._constants import DEFAULT_TIMEOUT - -log: logging.Logger = logging.getLogger(__name__) - - -def _unwrap(resp: Any) -> Any: - if isinstance(resp, dict) and "data" in resp and "status" in resp: - return resp["data"] - return resp - - -class SigningKeys(SyncAPIResource): - def _base_url(self) -> str: - org_id = self._client.organization_id - if not org_id: - raise ValueError("Client has no organization_id configured") - return f"/organizations/{org_id}/signing-keys" - - def get_active( - self, - *, - timeout: float | httpx.Timeout | None = DEFAULT_TIMEOUT, - ) -> Optional[Dict[str, Any]]: - """Fetch the active signing key (key_id, name, secret). - - Returns None if no active signing key exists (404). - """ - try: - resp = self._get( - f"{self._base_url()}/active", - timeout=timeout, - cast_to=dict, - ) - data = _unwrap(resp) - return data if isinstance(data, dict) else None - except Exception: - log.debug("No active signing key found", exc_info=True) - return None - - def create( - self, - *, - name: str = "default", - timeout: float | httpx.Timeout | None = DEFAULT_TIMEOUT, - ) -> Optional[Dict[str, Any]]: - """Create a new signing key for the organization. - - Returns the key data (key_id, name, secret) or None on failure. - """ - try: - resp = self._post( - self._base_url(), - body={"name": name}, - timeout=timeout, - cast_to=dict, - ) - data = _unwrap(resp) - return data if isinstance(data, dict) else None - except Exception: - log.debug("Failed to create signing key", exc_info=True) - return None - - def list( - self, - *, - timeout: float | httpx.Timeout | None = DEFAULT_TIMEOUT, - ) -> Optional[Union[List[Dict[str, Any]], Dict[str, Any]]]: - """List signing key metadata (no secrets).""" - try: - resp = self._get(self._base_url(), timeout=timeout, cast_to=dict) - data = _unwrap(resp) - if isinstance(data, (dict, list)): - return data - return None - except Exception: - log.debug("Failed to list signing keys", exc_info=True) - return None - - -class AsyncSigningKeys(AsyncAPIResource): - def _base_url(self) -> str: - org_id = self._client.organization_id - if not org_id: - raise ValueError("Client has no organization_id configured") - return f"/organizations/{org_id}/signing-keys" - - async def get_active( - self, - *, - timeout: float | httpx.Timeout | None = DEFAULT_TIMEOUT, - ) -> Optional[Dict[str, Any]]: - """Fetch the active signing key (key_id, name, secret). - - Returns None if no active signing key exists (404). - """ - try: - resp = await self._get( - f"{self._base_url()}/active", - timeout=timeout, - cast_to=dict, - ) - data = _unwrap(resp) - return data if isinstance(data, dict) else None - except Exception: - log.debug("No active signing key found", exc_info=True) - return None - - async def create( - self, - *, - name: str = "default", - timeout: float | httpx.Timeout | None = DEFAULT_TIMEOUT, - ) -> Optional[Dict[str, Any]]: - """Create a new signing key for the organization. - - Returns the key data (key_id, name, secret) or None on failure. - """ - try: - resp = await self._post( - self._base_url(), - body={"name": name}, - timeout=timeout, - cast_to=dict, - ) - data = _unwrap(resp) - return data if isinstance(data, dict) else None - except Exception: - log.debug("Failed to create signing key", exc_info=True) - return None - - async def list( - self, - *, - timeout: float | httpx.Timeout | None = DEFAULT_TIMEOUT, - ) -> Optional[Union[List[Dict[str, Any]], Dict[str, Any]]]: - """List signing key metadata (no secrets).""" - try: - resp = await self._get(self._base_url(), timeout=timeout, cast_to=dict) - data = _unwrap(resp) - if isinstance(data, (dict, list)): - return data - return None - except Exception: - log.debug("Failed to list signing keys", exc_info=True) - return None diff --git a/tests/attestation/test_signing.py b/tests/attestation/test_signing.py index 13ca357..5a2b8f7 100644 --- a/tests/attestation/test_signing.py +++ b/tests/attestation/test_signing.py @@ -44,27 +44,7 @@ def test_verify_wrong_key(self): assert not hmac_verify(b"key-2", b"data", sig) -class TestChainWithSigning: - def test_signed_envelopes(self): - chain = HashChain(signing_key_id="org-123", signing_secret=b"test-key") - e1 = chain.add_event({"name": "span-1"}) - e2 = chain.add_event({"name": "span-2"}) - - assert e1.signature is not None - assert e1.signing_key_id == "org-123" - assert e2.signature is not None - assert e2.signing_key_id == "org-123" - # Different events get different signatures - assert e1.signature != e2.signature - - def test_trial_signed(self): - chain = HashChain(signing_key_id="org-123", signing_secret=b"test-key") - chain.add_event({"name": "span-1"}) - trial = chain.finalize() - - assert trial.signature is not None - assert trial.signing_key_id == "org-123" - +class TestUnsignedChainHasNoSignatures: def test_unsigned_chain_has_no_signatures(self): chain = HashChain() e1 = chain.add_event({"name": "span-1"}) @@ -74,15 +54,6 @@ def test_unsigned_chain_has_no_signatures(self): assert e1.signing_key_id is None assert trial.signature is None - def test_to_dict_includes_signature_fields(self): - chain = HashChain(signing_key_id="org-123", signing_secret=b"test-key") - chain.add_event({"name": "span-1"}) - d = chain.to_dict() - - event = d["events"][0] - assert "signature" in event - assert "signing_key_id" in event - def test_to_dict_omits_signature_when_unsigned(self): chain = HashChain() chain.add_event({"name": "span-1"}) @@ -94,14 +65,34 @@ def test_to_dict_omits_signature_when_unsigned(self): class TestVerifyTrialWithSigning: - def test_valid_signed_trial(self): - secret = b"test-key" - chain = HashChain(signing_key_id="org-123", signing_secret=secret) + """Verify that verify_trial still works with externally-signed envelopes. + + In the server-side signing model, the backend signs the chain after + ingestion. These tests simulate that by manually signing envelopes + and verifying them with verify_trial(). + """ + + def _build_and_sign(self, secret: bytes, key_id: str = "org-123"): + """Build an unsigned chain, then manually sign each envelope.""" + chain = HashChain() chain.add_event({"name": "a"}) chain.add_event({"name": "b"}) envelopes = chain.envelopes trial = chain.finalize() + # Simulate server-side signing + for env in envelopes: + env.signature = hmac_sign(secret, env.hash.encode("utf-8")) + env.signing_key_id = key_id + trial.signature = hmac_sign(secret, trial.hash.encode("utf-8")) + trial.signing_key_id = key_id + + return envelopes, trial + + def test_valid_signed_trial(self): + secret = b"test-key" + envelopes, trial = self._build_and_sign(secret) + result = verify_trial(envelopes, trial, signing_secret=secret) assert result.valid assert result.chain_valid @@ -111,10 +102,7 @@ def test_valid_signed_trial(self): def test_tampered_signature_detected(self): secret = b"test-key" - chain = HashChain(signing_key_id="org-123", signing_secret=secret) - chain.add_event({"name": "a"}) - envelopes = chain.envelopes - trial = chain.finalize() + envelopes, trial = self._build_and_sign(secret) # Tamper with the event signature envelopes[0].signature = "dGFtcGVyZWQ=" # base64("tampered") @@ -122,14 +110,11 @@ def test_tampered_signature_detected(self): result = verify_trial(envelopes, trial, signing_secret=secret) assert not result.valid assert not result.signatures_valid - assert result.chain_valid # chain structure is still fine - assert result.trial_hash_valid # trial hash is still fine + assert result.chain_valid + assert result.trial_hash_valid def test_wrong_key_rejects(self): - chain = HashChain(signing_key_id="org-123", signing_secret=b"key-1") - chain.add_event({"name": "a"}) - envelopes = chain.envelopes - trial = chain.finalize() + envelopes, trial = self._build_and_sign(b"key-1") result = verify_trial(envelopes, trial, signing_secret=b"key-2") assert not result.valid @@ -149,10 +134,7 @@ def test_unsigned_chain_passes_without_secret(self): def test_stripped_signatures_detected(self): """When signing_secret is provided, missing signatures should fail.""" secret = b"test-key" - chain = HashChain(signing_key_id="org-123", signing_secret=secret) - chain.add_event({"name": "a"}) - envelopes = chain.envelopes - trial = chain.finalize() + envelopes, trial = self._build_and_sign(secret) # Strip signatures envelopes[0].signature = None @@ -163,25 +145,21 @@ def test_stripped_signatures_detected(self): assert not result.signatures_valid assert any("Missing signature" in e for e in result.errors) - def test_backward_compat_old_verify_trial(self): - """Old-style verify_trial (no signing_secret) still returns valid for valid chains.""" - chain = HashChain() - chain.add_event({"name": "a"}) - chain.add_event({"name": "b"}) - envelopes = chain.envelopes - trial = chain.finalize() - - result = verify_trial(envelopes, trial) - assert result.valid - def test_single_event_signed_chain(self): """Signed chain with exactly one event works correctly.""" secret = b"test-key" - chain = HashChain(signing_key_id="org-1", signing_secret=secret) + chain = HashChain() chain.add_event({"name": "only"}) envelopes = chain.envelopes trial = chain.finalize() + # Manually sign + for env in envelopes: + env.signature = hmac_sign(secret, env.hash.encode("utf-8")) + env.signing_key_id = "org-1" + trial.signature = hmac_sign(secret, trial.hash.encode("utf-8")) + trial.signing_key_id = "org-1" + assert len(envelopes) == 1 assert envelopes[0].signature is not None diff --git a/tests/attestation/test_verify.py b/tests/attestation/test_verify.py index 60f14c8..b2f34c8 100644 --- a/tests/attestation/test_verify.py +++ b/tests/attestation/test_verify.py @@ -128,10 +128,10 @@ def test_detect_count_mismatch(self): assert result.tampered assert result.chain_broken - def test_detect_tampering_with_signed_chain(self): - """detect_tampering works correctly on chains that were signed.""" + def test_detect_tampering_with_multi_event_chain(self): + """detect_tampering works correctly on multi-event chains.""" data = [{"name": "a"}, {"name": "b"}, {"name": "c"}] - chain = HashChain(signing_key_id="org-1", signing_secret=b"test-key") + chain = HashChain() for d in data: chain.add_event(d) diff --git a/tests/instrument/test_signing_autofetch.py b/tests/instrument/test_signing_autofetch.py deleted file mode 100644 index 928372f..0000000 --- a/tests/instrument/test_signing_autofetch.py +++ /dev/null @@ -1,412 +0,0 @@ -"""Tests for automatic signing key fetch in @trace decorator.""" - -from __future__ import annotations - -import json -import base64 -import asyncio -import threading -from unittest.mock import Mock, AsyncMock - -import pytest - -from layerlens.instrument import trace, clear_signing_key_cache -from layerlens.instrument._recorder import _signing_key_cache, _resolve_signing_key - - -@pytest.fixture(autouse=True) -def _clear_cache(): - """Clear signing key cache before and after each test.""" - _signing_key_cache.clear() - yield - _signing_key_cache.clear() - - -def _make_client(*, signing_key_response=None, create_key_response=None): - """Create a mock client with optional signing_keys.get_active() response. - - If create_key_response is provided, signing_keys.create() returns it. - Otherwise create() returns None (simulating backend failure or no-op). - """ - client = Mock() - client.traces = Mock() - client.traces.upload = Mock() - client.signing_keys = Mock() - if signing_key_response is not None: - client.signing_keys.get_active = Mock(return_value=signing_key_response) - else: - client.signing_keys.get_active = Mock(return_value=None) - if create_key_response is not None: - client.signing_keys.create = Mock(return_value=create_key_response) - else: - client.signing_keys.create = Mock(return_value=None) - return client - - -def _capture_upload(client): - """Set up trace capture on the mock client. Returns dict that gets populated.""" - uploaded = {} - - def _capture(path): - with open(path) as f: - uploaded["trace"] = json.load(f) - - client.traces.upload.side_effect = _capture - return uploaded - - -class TestAutoFetchSigningKey: - def test_auto_fetches_and_signs(self): - """When no signing_service passed, auto-fetch from client and sign.""" - secret = b"test-auto-key-32-bytes-long!!!!!" - client = _make_client( - signing_key_response={ - "key_id": "sk_auto_123", - "name": "auto-key", - "secret": base64.b64encode(secret).decode(), - } - ) - uploaded = _capture_upload(client) - - @trace(client) - def my_agent(): - return "hello" - - my_agent() - - client.signing_keys.get_active.assert_called_once() - payload = uploaded["trace"][0] - assert "attestation" in payload - att = payload["attestation"] - # Chain events should have signatures - events = att["chain"]["events"] - assert len(events) > 0 - for event in events: - assert "signature" in event, "Event should be signed" - assert event["signing_key_id"] == "sk_auto_123" - - def test_auto_creates_key_when_none_exists(self): - """When org has no active key, SDK auto-creates one and signs.""" - secret = b"auto-created-key-32-bytes!!!!!!!" - client = _make_client( - signing_key_response=None, - create_key_response={ - "key_id": "sk_auto_created", - "name": "default", - "secret": base64.b64encode(secret).decode(), - }, - ) - uploaded = _capture_upload(client) - - @trace(client) - def my_agent(): - return "hello" - - my_agent() - - client.signing_keys.get_active.assert_called_once() - client.signing_keys.create.assert_called_once() - payload = uploaded["trace"][0] - assert "attestation" in payload - events = payload["attestation"]["chain"]["events"] - assert len(events) > 0 - for event in events: - assert "signature" in event, "Event should be signed with auto-created key" - assert event["signing_key_id"] == "sk_auto_created" - - def test_no_signing_key_and_create_fails_produces_unsigned(self): - """When org has no active key AND create fails, traces are unsigned.""" - client = _make_client(signing_key_response=None, create_key_response=None) - uploaded = _capture_upload(client) - - @trace(client) - def my_agent(): - return "hello" - - my_agent() - - client.signing_keys.get_active.assert_called_once() - client.signing_keys.create.assert_called_once() - payload = uploaded["trace"][0] - assert "attestation" in payload - events = payload["attestation"]["chain"]["events"] - assert len(events) > 0 - for event in events: - assert "signature" not in event, "Event should NOT be signed" - - def test_caches_across_traces(self): - """Signing key is fetched once and reused across multiple @trace calls.""" - secret = b"cached-key-32-bytes-long!!!!!!!!" - client = _make_client( - signing_key_response={ - "key_id": "sk_cached", - "name": "cached", - "secret": base64.b64encode(secret).decode(), - } - ) - all_uploads: list = [] - - def _capture_all(path): - with open(path) as f: - all_uploads.append(json.load(f)) - - client.traces.upload.side_effect = _capture_all - - @trace(client) - def agent_a(): - return "a" - - @trace(client) - def agent_b(): - return "b" - - agent_a() - agent_b() - - # Only one API call despite two traces - client.signing_keys.get_active.assert_called_once() - # Both traces should be signed - assert len(all_uploads) == 2 - for upload in all_uploads: - events = upload[0]["attestation"]["chain"]["events"] - assert "signature" in events[0] - assert events[0]["signing_key_id"] == "sk_cached" - - def test_clear_cache_forces_refetch(self): - """clear_signing_key_cache() causes next trace to refetch.""" - secret = b"key-before-rotation!!!!!!!!!!!!!!" - client = _make_client( - signing_key_response={ - "key_id": "sk_old", - "name": "old", - "secret": base64.b64encode(secret).decode(), - } - ) - _capture_upload(client) - - @trace(client) - def my_agent(): - return "hello" - - my_agent() - assert client.signing_keys.get_active.call_count == 1 - - # Simulate key rotation - clear_signing_key_cache(client) - new_secret = b"key-after-rotation!!!!!!!!!!!!!!!" - client.signing_keys.get_active.return_value = { - "key_id": "sk_new", - "name": "new", - "secret": base64.b64encode(new_secret).decode(), - } - - my_agent() - assert client.signing_keys.get_active.call_count == 2 - - def test_explicit_signing_key_skips_autofetch(self): - """Passing signing_service= explicitly bypasses auto-fetch entirely.""" - client = _make_client( - signing_key_response={ - "key_id": "sk_should_not_fetch", - "name": "nope", - "secret": base64.b64encode(b"nope").decode(), - } - ) - uploaded = _capture_upload(client) - - @trace(client, signing_service=("explicit-key", b"explicit-secret")) - def my_agent(): - return "hello" - - my_agent() - - # Auto-fetch should NOT be called - client.signing_keys.get_active.assert_not_called() - # But traces should still be signed with the explicit key - payload = uploaded["trace"][0] - events = payload["attestation"]["chain"]["events"] - assert events[0]["signing_key_id"] == "explicit-key" - - def test_explicit_none_disables_signing(self): - """Passing signing_service=None explicitly disables signing (no auto-fetch).""" - client = _make_client( - signing_key_response={ - "key_id": "sk_should_not_fetch", - "name": "nope", - "secret": base64.b64encode(b"nope").decode(), - } - ) - uploaded = _capture_upload(client) - - @trace(client, signing_service=None) - def my_agent(): - return "hello" - - my_agent() - - # Auto-fetch should NOT be called - client.signing_keys.get_active.assert_not_called() - # Traces should be unsigned - payload = uploaded["trace"][0] - events = payload["attestation"]["chain"]["events"] - for event in events: - assert "signature" not in event - - def test_fetch_failure_degrades_to_unsigned(self): - """If get_active() throws, traces are uploaded unsigned (not broken).""" - client = _make_client() - client.signing_keys.get_active = Mock(side_effect=RuntimeError("network error")) - uploaded = _capture_upload(client) - - @trace(client) - def my_agent(): - return "hello" - - my_agent() - - payload = uploaded["trace"][0] - assert "attestation" in payload - events = payload["attestation"]["chain"]["events"] - for event in events: - assert "signature" not in event - - def test_client_without_signing_keys_attr(self): - """Clients that don't have signing_keys (e.g. old SDK) degrade gracefully.""" - client = Mock(spec=["traces"]) - client.traces = Mock() - client.traces.upload = Mock() - uploaded = {} - - def _capture(path): - with open(path) as f: - uploaded["trace"] = json.load(f) - - client.traces.upload.side_effect = _capture - - @trace(client) - def my_agent(): - return "hello" - - my_agent() - - payload = uploaded["trace"][0] - assert "attestation" in payload - # Should be unsigned, no crash - events = payload["attestation"]["chain"]["events"] - for event in events: - assert "signature" not in event - - def test_malformed_response_missing_key_id(self): - """get_active() returns dict with secret but no key_id — falls back to create().""" - client = _make_client( - signing_key_response={ - "name": "broken-key", - "secret": base64.b64encode(b"some-secret").decode(), - # Missing "key_id" - }, - create_key_response=None, # create also fails - ) - uploaded = _capture_upload(client) - - @trace(client) - def my_agent(): - return "hello" - - my_agent() - - client.signing_keys.create.assert_called_once() - payload = uploaded["trace"][0] - events = payload["attestation"]["chain"]["events"] - for event in events: - assert "signature" not in event, "Should be unsigned when key_id is missing and create fails" - - def test_malformed_response_missing_secret(self): - """get_active() returns dict with key_id but no secret — falls back to create().""" - client = _make_client( - signing_key_response={ - "key_id": "sk_123", - "name": "broken-key", - # Missing "secret" - }, - create_key_response=None, # create also fails - ) - uploaded = _capture_upload(client) - - @trace(client) - def my_agent(): - return "hello" - - my_agent() - - payload = uploaded["trace"][0] - events = payload["attestation"]["chain"]["events"] - for event in events: - assert "signature" not in event, "Should be unsigned when secret is missing" - - def test_concurrent_cache_access(self): - """Multiple threads resolving the same client only fetch once.""" - secret = b"concurrent-key-32-bytes-long!!!!!" - client = _make_client( - signing_key_response={ - "key_id": "sk_concurrent", - "name": "concurrent", - "secret": base64.b64encode(secret).decode(), - } - ) - - results = [] - errors = [] - - def resolve(): - try: - result = _resolve_signing_key(client) - results.append(result) - except Exception as e: - errors.append(e) - - threads = [threading.Thread(target=resolve) for _ in range(10)] - for t in threads: - t.start() - for t in threads: - t.join() - - assert not errors - assert len(results) == 10 - # All threads got the same result - for r in results: - assert r is not None - assert r[0] == "sk_concurrent" - - -class TestAsyncSigningFlush: - def test_async_trace_signs_correctly(self): - """Async @trace path produces signed attestation.""" - secret = b"async-test-key-32-bytes-long!!!!!" - client = _make_client( - signing_key_response={ - "key_id": "sk_async", - "name": "async-key", - "secret": base64.b64encode(secret).decode(), - } - ) - uploaded = {} - - async def _async_capture(path): - with open(path) as f: - uploaded["trace"] = json.load(f) - - client.traces.upload = AsyncMock(side_effect=_async_capture) - - @trace(client) - async def my_async_agent(): - return "hello async" - - asyncio.run(my_async_agent()) - - payload = uploaded["trace"][0] - assert "attestation" in payload - events = payload["attestation"]["chain"]["events"] - assert len(events) > 0 - for event in events: - assert "signature" in event - assert event["signing_key_id"] == "sk_async" From a6d9bbf3e22d580ffdb3d99ba6eefdee0883b328 Mon Sep 17 00:00:00 2001 From: Gary <59334078+garrettallen14@users.noreply.github.com> Date: Fri, 27 Mar 2026 14:14:40 -0700 Subject: [PATCH 4/4] fix: attestation chain integrity: error propagation, async I/O, envelope immutability --- src/layerlens/attestation/_chain.py | 3 ++- src/layerlens/instrument/_recorder.py | 8 ++++---- src/layerlens/instrument/_upload.py | 17 +++++++++++------ 3 files changed, 17 insertions(+), 11 deletions(-) diff --git a/src/layerlens/attestation/_chain.py b/src/layerlens/attestation/_chain.py index b6b45d2..ae6a3d1 100644 --- a/src/layerlens/attestation/_chain.py +++ b/src/layerlens/attestation/_chain.py @@ -1,5 +1,6 @@ from __future__ import annotations +from copy import copy from typing import Any, Dict, List, Optional from ._hash import compute_hash @@ -25,7 +26,7 @@ def __init__(self) -> None: @property def envelopes(self) -> List[AttestationEnvelope]: - return list(self._chain) + return [copy(e) for e in self._chain] @property def is_terminated(self) -> bool: diff --git a/src/layerlens/instrument/_recorder.py b/src/layerlens/instrument/_recorder.py index 758c3d8..9960577 100644 --- a/src/layerlens/instrument/_recorder.py +++ b/src/layerlens/instrument/_recorder.py @@ -59,9 +59,9 @@ def flush(self) -> None: trace_data = self.root.to_dict() try: attestation = self._build_attestation() - except Exception: + except Exception as exc: log.warning("Failed to build attestation chain", exc_info=True) - attestation = {} + attestation = {"attestation_error": str(exc)} upload_trace(self._client, trace_data, attestation) async def async_flush(self) -> None: @@ -70,7 +70,7 @@ async def async_flush(self) -> None: trace_data = self.root.to_dict() try: attestation = self._build_attestation() - except Exception: + except Exception as exc: log.warning("Failed to build attestation chain", exc_info=True) - attestation = {} + attestation = {"attestation_error": str(exc)} await async_upload_trace(self._client, trace_data, attestation) diff --git a/src/layerlens/instrument/_upload.py b/src/layerlens/instrument/_upload.py index e6cd3a1..020d990 100644 --- a/src/layerlens/instrument/_upload.py +++ b/src/layerlens/instrument/_upload.py @@ -2,6 +2,7 @@ import os import json +import asyncio import logging import tempfile from typing import Any, Dict, Optional @@ -9,6 +10,14 @@ log: logging.Logger = logging.getLogger(__name__) +def _write_trace_file(payload: Dict[str, Any]) -> str: + """Write trace payload to a temp file and return its path.""" + fd, path = tempfile.mkstemp(suffix=".json", prefix="layerlens_trace_") + with os.fdopen(fd, "w") as f: + json.dump([payload], f, default=str) + return path + + def upload_trace( client: Any, trace_data: Dict[str, Any], @@ -17,10 +26,8 @@ def upload_trace( payload = trace_data if attestation: payload = {**trace_data, "attestation": attestation} - fd, path = tempfile.mkstemp(suffix=".json", prefix="layerlens_trace_") + path = _write_trace_file(payload) try: - with os.fdopen(fd, "w") as f: - json.dump([payload], f, default=str) client.traces.upload(path) finally: try: @@ -37,10 +44,8 @@ async def async_upload_trace( payload = trace_data if attestation: payload = {**trace_data, "attestation": attestation} - fd, path = tempfile.mkstemp(suffix=".json", prefix="layerlens_trace_") + path = await asyncio.to_thread(_write_trace_file, payload) try: - with os.fdopen(fd, "w") as f: - json.dump([payload], f, default=str) await client.traces.upload(path) finally: try: