From 2035e0effc240fe7b3044f31c7ec41ac16140115 Mon Sep 17 00:00:00 2001 From: Imran Siddique Date: Sat, 6 Jun 2026 12:15:52 -0700 Subject: [PATCH] feat(inspection): stage 2 response schema validation (#74) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Implements output_schema validation in the inspection pipeline. Supports three per-entry modes — redact (default), strict, and log — driven by schema_validation_mode on CatalogEntry. Redact mode strips surplus fields and propagates modified_response to the caller; strict denies; log passes through with audit metadata. Non-JSON responses and entries without output_schema skip the stage. Co-Authored-By: Claude Sonnet 4.6 --- src/cmcp_gateway/catalog/loader.py | 14 +- src/cmcp_gateway/inspection/pipeline.py | 101 +++++++++++- tests/unit/test_stage2_schema.py | 195 ++++++++++++++++++++++++ 3 files changed, 305 insertions(+), 5 deletions(-) create mode 100644 tests/unit/test_stage2_schema.py diff --git a/src/cmcp_gateway/catalog/loader.py b/src/cmcp_gateway/catalog/loader.py index 8f641037..c11391cd 100644 --- a/src/cmcp_gateway/catalog/loader.py +++ b/src/cmcp_gateway/catalog/loader.py @@ -4,9 +4,9 @@ import hashlib import json -from dataclasses import dataclass +from dataclasses import dataclass, field from pathlib import Path -from typing import Any +from typing import Any, Literal import jsonschema @@ -49,6 +49,7 @@ class CatalogEntry: added_at: str approved_by: str catalog_exception: bool = False + schema_validation_mode: Literal["redact", "strict", "log"] = field(default="redact") @dataclass @@ -157,6 +158,14 @@ def load_catalog(catalog_path: str, expected_hash: str | None = None) -> ToolCat f"Stored: {raw['definition_hash']}, computed: {computed_def_hash}" ) + raw_mode = raw.get("schema_validation_mode", "redact") + if raw_mode not in ("redact", "strict", "log"): + raise ConfigError( + f"Catalog entry '{tool_name}': invalid schema_validation_mode '{raw_mode}'; " + "must be 'redact', 'strict', or 'log'" + ) + schema_validation_mode: Literal["redact", "strict", "log"] = raw_mode + entries[tool_name] = CatalogEntry( tool_name=tool_name, server=server, @@ -168,6 +177,7 @@ def load_catalog(catalog_path: str, expected_hash: str | None = None) -> ToolCat added_at=raw.get("added_at", ""), approved_by=raw.get("approved_by", ""), catalog_exception=raw.get("catalog_exception", False), + schema_validation_mode=schema_validation_mode, ) computed_hash = _catalog_hash(raw_list) diff --git a/src/cmcp_gateway/inspection/pipeline.py b/src/cmcp_gateway/inspection/pipeline.py index aece7bfe..9cb25ccb 100644 --- a/src/cmcp_gateway/inspection/pipeline.py +++ b/src/cmcp_gateway/inspection/pipeline.py @@ -13,10 +13,13 @@ from __future__ import annotations import hashlib +import json import re from dataclasses import dataclass, field from typing import Any +import jsonschema + from cmcp_gateway.catalog.loader import CatalogEntry # ── AGT components (optional — fall back gracefully) ───────────────────────── @@ -84,6 +87,89 @@ def _stage1_size_check(response_bytes: bytes, max_bytes: int) -> StageResult: return StageResult(stage="size", decision="allow") +def _stage2_schema_validation( + response_bytes: bytes, + catalog_entry: CatalogEntry, +) -> tuple[StageResult, bytes]: + """ + Stage 2: validate response against the catalog entry's approved output_schema. + + Mode is per catalog entry (catalog_entry.schema_validation_mode, default "redact"). + + Returns a (StageResult, response_bytes) tuple where response_bytes may be + modified (surplus fields stripped) in redact mode. + """ + output_schema = catalog_entry.approved_definition.output_schema + if output_schema is None: + return StageResult(stage="schema", decision="skip"), response_bytes + + # Only validate JSON responses; non-JSON passes through + try: + payload: Any = json.loads(response_bytes) + except (json.JSONDecodeError, ValueError): + return StageResult(stage="schema", decision="allow", reason="non-JSON response; schema check skipped"), response_bytes + + # Identify surplus fields at the top level only + surplus: list[str] = [] + if isinstance(payload, dict) and isinstance(output_schema.get("properties"), dict): + allowed_props: set[str] = set(output_schema["properties"].keys()) + surplus = [k for k in payload if k not in allowed_props] + + mode = catalog_entry.schema_validation_mode + + if not surplus: + # No surplus fields — still run jsonschema for type/required violations + try: + jsonschema.validate(payload, output_schema) + except jsonschema.ValidationError as exc: + return ( + StageResult( + stage="schema", + decision="deny", + reason=f"RESPONSE_SCHEMA_VIOLATION: {exc.message}", + ), + response_bytes, + ) + return StageResult(stage="schema", decision="allow"), response_bytes + + # Surplus fields present — mode determines action + if mode == "strict": + return ( + StageResult( + stage="schema", + decision="deny", + reason="RESPONSE_SCHEMA_VIOLATION_STRICT", + stripped_fields=surplus, + ), + response_bytes, + ) + + if mode == "log": + return ( + StageResult( + stage="schema", + decision="allow", + reason="surplus fields logged", + stripped_fields=surplus, + ), + response_bytes, + ) + + # mode == "redact" (default): strip surplus fields and return modified bytes + assert isinstance(payload, dict) # guaranteed by surplus check above + redacted = {k: v for k, v in payload.items() if k not in surplus} + modified_bytes = json.dumps(redacted, separators=(",", ":"), ensure_ascii=False).encode() + return ( + StageResult( + stage="schema", + decision="allow", + reason="surplus fields redacted", + stripped_fields=surplus, + ), + modified_bytes, + ) + + def _stage4_injection_detection( response_text: str, custom_patterns: list[tuple[re.Pattern[str], str]] | None = None, @@ -214,6 +300,7 @@ def run( stripped_fields: list[str] | None = None injection_pattern: str | None = None sensitivity_tags: list[str] = [] + modified_response: bytes | None = None # Stage 1: size check s1 = _stage1_size_check(response_bytes, self._max_bytes) @@ -221,8 +308,16 @@ def run( if s1.decision == "deny": deny_reasons.append(s1.reason or "size exceeded") - # Stage 2: schema validation (Phase 1 GA — issue #74; skipped here) - stage_results["schema"] = "skip" + # Stage 2: schema validation (issue #74) + s2, response_bytes = _stage2_schema_validation(response_bytes, catalog_entry) + stage_results["schema"] = s2.decision + if s2.decision == "deny": + deny_reasons.append(s2.reason or "schema violation") + if s2.stripped_fields: + stripped_fields = s2.stripped_fields + if s2.decision == "allow" and s2.stripped_fields and s2.reason == "surplus fields redacted": + # Redact mode modified the bytes — expose to caller + modified_response = response_bytes # Stage 3: sensitivity classification (AGT CredentialRedactor + catalog) try: @@ -289,5 +384,5 @@ def run( injection_pattern_matched=injection_pattern, stage_results=stage_results, response_payload_hash=response_payload_hash, - modified_response=None, + modified_response=modified_response, ) diff --git a/tests/unit/test_stage2_schema.py b/tests/unit/test_stage2_schema.py new file mode 100644 index 00000000..78b12de0 --- /dev/null +++ b/tests/unit/test_stage2_schema.py @@ -0,0 +1,195 @@ +"""Tests for Stage 2 response schema validation (issue #74).""" + +from __future__ import annotations + +import json + +from cmcp_gateway.catalog.loader import ApprovedDefinition, CatalogEntry, ServerIdentity +from cmcp_gateway.inspection.pipeline import InspectionPipeline, _stage2_schema_validation + + +def _make_entry( + output_schema: dict | None = None, + schema_validation_mode: str = "redact", +) -> CatalogEntry: + return CatalogEntry( + tool_name="test.tool", + server=ServerIdentity( + display_name="Test", + url="https://test.example.com", + tls_fingerprint="SHA256:AAAA/BBBB==", + spiffe_id=None, + transport="http-sse", + rotation_mode="key-pinned", + ), + approved_definition=ApprovedDefinition( + description="test tool", + input_schema={}, + output_schema=output_schema, + ), + definition_hash="sha256:" + "0" * 64, + compliance_domain="external", + requires_baa=False, + sensitivity_level="public", + added_at="2026-06-01T00:00:00Z", + approved_by="test", + schema_validation_mode=schema_validation_mode, # type: ignore[arg-type] + ) + + +_SCHEMA_WITH_PROPERTIES = { + "type": "object", + "properties": { + "result": {"type": "string"}, + "count": {"type": "integer"}, + }, +} + +_VALID_PAYLOAD = {"result": "ok", "count": 3} +_SURPLUS_PAYLOAD = {"result": "ok", "count": 3, "internal_id": "secret", "debug": True} + + +# ── Skip when output_schema is None ────────────────────────────────────────── + +def test_skip_when_no_output_schema(): + entry = _make_entry(output_schema=None) + result, out_bytes = _stage2_schema_validation(b'{"result": "ok"}', entry) + assert result.decision == "skip" + assert result.stage == "schema" + + +def test_pipeline_records_skip_when_no_schema(): + pipeline = InspectionPipeline() + entry = _make_entry(output_schema=None) + ir = pipeline.run("c1", entry, b'{"result": "ok"}') + assert ir.stage_results["schema"] == "skip" + assert ir.modified_response is None + + +# ── Non-JSON response passes through ───────────────────────────────────────── + +def test_non_json_response_allows(): + entry = _make_entry(output_schema=_SCHEMA_WITH_PROPERTIES) + result, out_bytes = _stage2_schema_validation(b"plain text response", entry) + assert result.decision == "allow" + assert out_bytes == b"plain text response" + + +def test_pipeline_non_json_passes_through(): + pipeline = InspectionPipeline() + entry = _make_entry(output_schema=_SCHEMA_WITH_PROPERTIES) + ir = pipeline.run("c1", entry, b"plain text response") + assert ir.final_decision == "allow" + assert ir.stage_results["schema"] == "allow" + assert ir.modified_response is None + + +# ── Valid response matching schema: allow ───────────────────────────────────── + +def test_valid_response_allows(): + entry = _make_entry(output_schema=_SCHEMA_WITH_PROPERTIES) + payload_bytes = json.dumps(_VALID_PAYLOAD).encode() + result, out_bytes = _stage2_schema_validation(payload_bytes, entry) + assert result.decision == "allow" + assert result.stripped_fields is None or result.stripped_fields == [] + assert out_bytes == payload_bytes + + +def test_pipeline_valid_response_no_modification(): + pipeline = InspectionPipeline() + entry = _make_entry(output_schema=_SCHEMA_WITH_PROPERTIES) + ir = pipeline.run("c1", entry, json.dumps(_VALID_PAYLOAD).encode()) + assert ir.final_decision == "allow" + assert ir.modified_response is None + + +# ── Redact mode: strips surplus fields ─────────────────────────────────────── + +def test_redact_mode_strips_surplus_fields(): + entry = _make_entry(output_schema=_SCHEMA_WITH_PROPERTIES, schema_validation_mode="redact") + payload_bytes = json.dumps(_SURPLUS_PAYLOAD).encode() + result, out_bytes = _stage2_schema_validation(payload_bytes, entry) + assert result.decision == "allow" + assert result.stripped_fields is not None + assert set(result.stripped_fields) == {"internal_id", "debug"} + redacted = json.loads(out_bytes) + assert "internal_id" not in redacted + assert "debug" not in redacted + assert redacted["result"] == "ok" + assert redacted["count"] == 3 + + +def test_pipeline_redact_mode_sets_modified_response(): + pipeline = InspectionPipeline() + entry = _make_entry(output_schema=_SCHEMA_WITH_PROPERTIES, schema_validation_mode="redact") + ir = pipeline.run("c1", entry, json.dumps(_SURPLUS_PAYLOAD).encode()) + assert ir.final_decision == "allow" + assert ir.modified_response is not None + redacted = json.loads(ir.modified_response) + assert "internal_id" not in redacted + assert "debug" not in redacted + assert ir.stripped_fields is not None + assert set(ir.stripped_fields) == {"internal_id", "debug"} + + +def test_pipeline_redact_mode_no_modified_response_when_no_surplus(): + pipeline = InspectionPipeline() + entry = _make_entry(output_schema=_SCHEMA_WITH_PROPERTIES, schema_validation_mode="redact") + ir = pipeline.run("c1", entry, json.dumps(_VALID_PAYLOAD).encode()) + assert ir.modified_response is None + + +# ── Strict mode: deny on surplus fields ────────────────────────────────────── + +def test_strict_mode_denies_on_surplus(): + entry = _make_entry(output_schema=_SCHEMA_WITH_PROPERTIES, schema_validation_mode="strict") + payload_bytes = json.dumps(_SURPLUS_PAYLOAD).encode() + result, out_bytes = _stage2_schema_validation(payload_bytes, entry) + assert result.decision == "deny" + assert result.reason == "RESPONSE_SCHEMA_VIOLATION_STRICT" + assert result.stripped_fields is not None + assert set(result.stripped_fields) == {"internal_id", "debug"} + # Response bytes unchanged in strict mode + assert out_bytes == payload_bytes + + +def test_pipeline_strict_mode_denies(): + pipeline = InspectionPipeline() + entry = _make_entry(output_schema=_SCHEMA_WITH_PROPERTIES, schema_validation_mode="strict") + ir = pipeline.run("c1", entry, json.dumps(_SURPLUS_PAYLOAD).encode()) + assert ir.final_decision == "deny" + assert ir.stage_results["schema"] == "deny" + assert "RESPONSE_SCHEMA_VIOLATION_STRICT" in (ir.deny_reason or "") + + +def test_strict_mode_allows_valid_response(): + entry = _make_entry(output_schema=_SCHEMA_WITH_PROPERTIES, schema_validation_mode="strict") + payload_bytes = json.dumps(_VALID_PAYLOAD).encode() + result, out_bytes = _stage2_schema_validation(payload_bytes, entry) + assert result.decision == "allow" + + +# ── Log mode: passes through but records surplus_fields ────────────────────── + +def test_log_mode_allows_with_surplus_logged(): + entry = _make_entry(output_schema=_SCHEMA_WITH_PROPERTIES, schema_validation_mode="log") + payload_bytes = json.dumps(_SURPLUS_PAYLOAD).encode() + result, out_bytes = _stage2_schema_validation(payload_bytes, entry) + assert result.decision == "allow" + assert result.stripped_fields is not None + assert set(result.stripped_fields) == {"internal_id", "debug"} + # Response bytes unmodified in log mode + assert out_bytes == payload_bytes + + +def test_pipeline_log_mode_passes_through_unmodified(): + pipeline = InspectionPipeline() + entry = _make_entry(output_schema=_SCHEMA_WITH_PROPERTIES, schema_validation_mode="log") + payload_bytes = json.dumps(_SURPLUS_PAYLOAD).encode() + ir = pipeline.run("c1", entry, payload_bytes) + assert ir.final_decision == "allow" + assert ir.stage_results["schema"] == "allow" + # Log mode does not set modified_response (bytes unchanged) + assert ir.modified_response is None + assert ir.stripped_fields is not None + assert set(ir.stripped_fields) == {"internal_id", "debug"}