diff --git a/src/agentrust_trace/provenance.py b/src/agentrust_trace/provenance.py index bfaa8235..a548deb2 100644 --- a/src/agentrust_trace/provenance.py +++ b/src/agentrust_trace/provenance.py @@ -66,6 +66,14 @@ def _as_object(value: Any, field: str) -> dict[str, Any]: return value +def _tool_count(catalog: dict[str, Any]) -> int: + """Return the required catalog count as a JSON integer.""" + value = catalog.get("tool_count") + if not isinstance(value, int) or isinstance(value, bool) or value < 0: + raise ProvenanceError("tool_catalog.tool_count must be a non-negative integer") + return value + + class ToolCatalogMismatch(ProvenanceError): """The server offered a tool set the record does not describe. @@ -344,6 +352,7 @@ def verify_record( catalog = _as_object(record.get("tool_catalog"), "tool_catalog") if not _DIGEST_RE.match(str(catalog.get("hash", ""))): raise ProvenanceError("tool_catalog.hash is not a sha256: digest") + _tool_count(catalog) # Freshness. `issued_at` has been required and type-checked since the format # existed, with an error message explaining that a record with no issue time @@ -443,3 +452,9 @@ def check_tool_catalog(record: dict[str, Any], tools: list[dict[str, Any]]) -> N "The signature may be perfectly valid; this is about the server, not the " "document." ) + declared_count = _tool_count(catalog) + if declared_count != len(tools): + raise ProvenanceError( + f"tool_catalog.tool_count declares {declared_count} tools, but the matching " + f"catalog contains {len(tools)}" + ) diff --git a/tests/test_provenance_tool_count.py b/tests/test_provenance_tool_count.py new file mode 100644 index 00000000..a5ce5f14 --- /dev/null +++ b/tests/test_provenance_tool_count.py @@ -0,0 +1,90 @@ +from __future__ import annotations + +import pytest + +from agentrust_trace.provenance import ( + ProvenanceError, + ToolCatalogMismatch, + build_record, + check_tool_catalog, + sign_record, + verify_record, +) +from agentrust_trace.sign import generate_key, key_to_jwk + +TOOLS = [ + { + "name": "search", + "description": "search the docs", + "input_schema": {"type": "object"}, + }, + { + "name": "fetch", + "description": "fetch a page", + "input_schema": {"type": "object"}, + }, +] +ARTIFACT = { + "package": "pkg:npm/%40acme/mcp-search@2.1.0", + "digest": "sha256:" + "a" * 64, +} +_MISSING = object() + + +def _signed_with_count(count: object = _MISSING): + key = generate_key() + record = build_record( + kind="publisher-asserted", + publisher="did:web:acme.example", + tools=TOOLS, + artifact=ARTIFACT, + ) + if count is _MISSING: + del record["tool_catalog"]["tool_count"] + else: + record["tool_catalog"]["tool_count"] = count + return sign_record(record, key), key_to_jwk(key) + + +@pytest.mark.parametrize("bad_count", [_MISSING, -1, True, "2", None]) +def test_verify_record_refuses_malformed_or_missing_tool_count(bad_count) -> None: + record, trusted = _signed_with_count(bad_count) + with pytest.raises(ProvenanceError, match="tool_catalog.tool_count must be"): + verify_record(record, trusted) + + +@pytest.mark.parametrize("bad_count", [2.0, "2", True, -1, None]) +def test_check_tool_catalog_alone_refuses_malformed_count(bad_count) -> None: + record, _ = _signed_with_count(bad_count) + with pytest.raises(ProvenanceError, match="tool_catalog.tool_count must be"): + check_tool_catalog(record, TOOLS) + + +def test_wrong_positive_count_is_detected_when_live_catalog_is_checked() -> None: + record, trusted = _signed_with_count(999) + + verify_record(record, trusted) + with pytest.raises(ProvenanceError, match="declares 999 tools") as excinfo: + check_tool_catalog(record, TOOLS) + assert not isinstance(excinfo.value, ToolCatalogMismatch) + + +def test_correct_count_and_hash_still_verify() -> None: + record, trusted = _signed_with_count(len(TOOLS)) + verify_record(record, trusted) + check_tool_catalog(record, TOOLS) + + +def test_hash_mismatch_remains_a_tool_catalog_mismatch() -> None: + record, trusted = _signed_with_count(len(TOOLS)) + verify_record(record, trusted) + + offered = [TOOLS[0]] + with pytest.raises(ToolCatalogMismatch, match="about the server, not the document"): + check_tool_catalog(record, offered) + + +def test_hash_mismatch_outranks_a_malformed_count() -> None: + record, _ = _signed_with_count("2") + with pytest.raises(ToolCatalogMismatch): + check_tool_catalog(record, [TOOLS[0]])