ソースを参照

fix: harden governed rule execution contracts

马小龙 4 週間 前
コミット
d823d8c7b5

+ 155 - 30
app/core/data_rules/execution_contracts.py

@@ -5,6 +5,7 @@ from __future__ import annotations
 import copy
 import hashlib
 import json
+import math
 import re
 from typing import Any
 
@@ -16,6 +17,22 @@ BACKENDS = {"sql_pushdown", "polars_batch", "quality_check"}
 OBJECT_KINDS = {"table", "view", "query", "parquet_artifact"}
 ACCESS_MODES = {"read", "write", "read_write"}
 PLAN_STATUSES = {"compiled", "tested", "published", "revoked"}
+SQL_DIALECTS = {"postgresql", "mysql"}
+POLARS_OPERATIONS = {
+    "aggregate",
+    "assert",
+    "cast",
+    "deduplicate",
+    "derive",
+    "fill_null",
+    "filter",
+    "lookup_join",
+    "map_values",
+    "mask",
+    "normalize_text",
+    "regex_replace",
+}
+QUALITY_CHECKS = {"assert"}
 
 _SECRET_KEY_TOKENS = {
     "apikey",
@@ -63,7 +80,28 @@ _EXECUTION_PLAN_KEYS = {
     "operations",
     "status",
 }
-_OPERATION_KEYS = {"kind", "statement", "parameters"}
+_SQL_OPERATION_KEYS = {"kind", "dialect", "statement", "parameters"}
+_POLARS_OPERATION_KEYS = {"kind", "op", "arguments"}
+_QUALITY_OPERATION_KEYS = {"kind", "check", "arguments"}
+_COMPILER_VERSION_PATTERNS = {
+    "sql_pushdown": re.compile(r"^dataops-sqlglot-[1-9][0-9]*(?:\.[0-9]+)+$"),
+    "polars_batch": re.compile(r"^dataops-polars-[1-9][0-9]*(?:\.[0-9]+)+$"),
+    "quality_check": re.compile(r"^dataops-quality-[1-9][0-9]*(?:\.[0-9]+)+$"),
+}
+_QUERY_REFERENCE = re.compile(r"^query://([0-9a-fA-F-]{36})$")
+_PARAMETER_NAME = re.compile(r"^[A-Za-z_][A-Za-z0-9_]{0,99}$")
+_FORBIDDEN_RUNTIME_CONTENT = re.compile(
+    r"\b(?:generated_python|python|pickle(?:d)?|callable|module|file)\b"
+    r"|[A-Za-z][A-Za-z0-9+.-]*://",
+    re.IGNORECASE,
+)
+_SQL_COMMENT = re.compile(r"--|/\*|\*/")
+_SQL_DISALLOWED = re.compile(
+    r"\b(?:alter|analyze|begin|call|commit|copy|create|deallocate|do|drop|"
+    r"execute|grant|listen|notify|prepare|reset|revoke|rollback|set|show|"
+    r"truncate|use|vacuum)\b",
+    re.IGNORECASE,
+)
 
 
 def _normalized_key(value: Any) -> str:
@@ -137,6 +175,23 @@ def _canonical_json(value: Any) -> str:
         raise ValueError("execution contract must be JSON serializable") from exc
 
 
+def _normalized_schema_fields(value: Any) -> list[dict[str, Any]]:
+    if not isinstance(value, list) or not value or len(value) > 1_000:
+        raise ValueError("schema fields must be a non-empty bounded array")
+    fields = [_validate_schema_field(item) for item in value]
+    names = [item["name"] for item in fields]
+    if len(names) != len(set(names)):
+        raise ValueError("schema field names must be unique")
+    return sorted(fields, key=lambda item: item["name"])
+
+
+def canonical_schema_hash(fields: Any) -> str:
+    """Hash normalized schema fields with the runtime's canonical JSON form."""
+
+    normalized = _normalized_schema_fields(fields)
+    return hashlib.sha256(_canonical_json(normalized).encode("utf-8")).hexdigest()
+
+
 def _validate_schema_field(value: Any) -> dict[str, Any]:
     field = _closed_object(value, _SCHEMA_FIELD_KEYS, "schema field")
     _required_fields(field, {"name", "type", "nullable"}, "schema field")
@@ -168,14 +223,9 @@ def validate_schema_snapshot(value: Any) -> dict:
     snapshot["source_revision"] = _required_string(
         snapshot["source_revision"], "source_revision", 200
     )
-    raw_fields = snapshot["fields"]
-    if not isinstance(raw_fields, list) or not raw_fields or len(raw_fields) > 1_000:
-        raise ValueError("schema fields must be a non-empty bounded array")
-    fields = [_validate_schema_field(item) for item in raw_fields]
-    names = [item["name"] for item in fields]
-    if len(names) != len(set(names)):
-        raise ValueError("schema field names must be unique")
-    snapshot["fields"] = fields
+    snapshot["fields"] = _normalized_schema_fields(snapshot["fields"])
+    if snapshot["schema_hash"] != canonical_schema_hash(snapshot["fields"]):
+        raise ValueError("schema_hash does not match normalized schema fields")
     return snapshot
 
 
@@ -192,7 +242,14 @@ def validate_dataset_binding(value: Any) -> dict:
     binding["object_ref"] = _required_string(
         binding["object_ref"], "object_ref", 1_000
     )
-    if "://" in binding["object_ref"]:
+    if binding["object_kind"] == "query":
+        match = _QUERY_REFERENCE.fullmatch(binding["object_ref"])
+        if match is None:
+            raise ValueError(
+                "query object_ref must be an opaque query reference"
+            )
+        binding["object_ref"] = f"query://{_uid(match.group(1), 'query reference')}"
+    elif "://" in binding["object_ref"]:
         raise ValueError("dataset object_ref cannot contain a connection string")
     binding["schema_snapshot_id"] = _uid(
         binding["schema_snapshot_id"], "schema_snapshot_id"
@@ -206,32 +263,91 @@ def validate_dataset_binding(value: Any) -> dict:
     return binding
 
 
-def _validate_operation(value: Any) -> dict[str, Any]:
-    operation = _closed_object(value, _OPERATION_KEYS, "execution operation")
-    _required_fields(operation, {"kind"}, "execution operation")
-    operation["kind"] = _required_string(
-        operation["kind"], "execution operation kind", 30
-    )
-    if operation["kind"] == "sql":
-        operation["statement"] = _required_string(
-            operation.get("statement"), "sql execution statement", 20_000
-        )
-    elif "statement" in operation:
-        operation["statement"] = _required_string(
-            operation["statement"], "execution operation statement", 20_000
-        )
-    if "parameters" in operation and not isinstance(operation["parameters"], dict):
-        raise ValueError("execution operation parameters must be an object")
+def _reject_forbidden_runtime_content(value: Any) -> None:
+    if isinstance(value, dict):
+        for key, item in value.items():
+            _reject_forbidden_runtime_content(key)
+            _reject_forbidden_runtime_content(item)
+    elif isinstance(value, list):
+        for item in value:
+            _reject_forbidden_runtime_content(item)
+    elif isinstance(value, str) and _FORBIDDEN_RUNTIME_CONTENT.search(value):
+        raise ValueError("forbidden runtime content is not allowed in an execution plan")
+
+
+def _validate_parameters(value: Any) -> dict[str, Any]:
+    if not isinstance(value, dict) or len(value) > 100:
+        raise ValueError("SQL operation parameters must be a bounded object")
+    normalized: dict[str, Any] = {}
+    for key, item in value.items():
+        if not isinstance(key, str) or not _PARAMETER_NAME.fullmatch(key):
+            raise ValueError("SQL parameter names must be identifiers")
+        if isinstance(item, float) and not math.isfinite(item):
+            raise ValueError("SQL parameter value must be finite")
+        if not isinstance(item, (type(None), bool, int, float, str)):
+            raise ValueError("SQL parameter value must be a JSON scalar")
+        if isinstance(item, str) and len(item) > 4_000:
+            raise ValueError("SQL parameter value exceeds 4000 characters")
+        normalized[key] = item
+    _reject_forbidden_runtime_content(normalized)
+    return normalized
+
+
+def _validate_sql_statement(value: Any) -> str:
+    statement = _required_string(value, "sql execution statement", 20_000)
+    _reject_forbidden_runtime_content(statement)
+    if ";" in statement:
+        raise ValueError("SQL operation must contain one single statement")
+    if _SQL_COMMENT.search(statement):
+        raise ValueError("SQL comments are not allowed in execution plans")
+    if _SQL_DISALLOWED.search(statement):
+        raise ValueError("SQL DDL or control statements are not allowed")
+    if not re.match(r"^(?:SELECT|INSERT|UPDATE|DELETE|WITH)\b", statement, re.IGNORECASE):
+        raise ValueError("SQL execution statement must start with a supported data statement")
+    return statement
+
+
+def _validate_sql_operation(value: Any) -> dict[str, Any]:
+    operation = _closed_object(value, _SQL_OPERATION_KEYS, "SQL execution operation")
+    _required_fields(operation, _SQL_OPERATION_KEYS, "SQL execution operation")
+    if operation["kind"] != "sql":
+        raise ValueError("unsupported sql_pushdown operation kind")
+    if operation["dialect"] not in SQL_DIALECTS:
+        raise ValueError("unsupported SQL dialect")
+    operation["statement"] = _validate_sql_statement(operation["statement"])
+    operation["parameters"] = _validate_parameters(operation["parameters"])
+    return operation
+
+
+def _validate_polars_operation(value: Any) -> dict[str, Any]:
+    operation = _closed_object(value, _POLARS_OPERATION_KEYS, "Polars execution operation")
+    _required_fields(operation, _POLARS_OPERATION_KEYS, "Polars execution operation")
+    if operation["kind"] != "polars" or operation["op"] not in POLARS_OPERATIONS:
+        raise ValueError("unsupported polars_batch operation kind")
+    operation["arguments"] = _validate_parameters(operation["arguments"])
+    return operation
+
+
+def _validate_quality_operation(value: Any) -> dict[str, Any]:
+    operation = _closed_object(value, _QUALITY_OPERATION_KEYS, "quality execution operation")
+    _required_fields(operation, _QUALITY_OPERATION_KEYS, "quality execution operation")
+    if operation["kind"] != "quality" or operation["check"] not in QUALITY_CHECKS:
+        raise ValueError("unsupported quality_check operation kind")
+    operation["arguments"] = _validate_parameters(operation["arguments"])
     return operation
 
 
 def validate_execution_plan_v2(value: Any) -> dict:
-    """Validate an execution plan that contains only immutable references."""
+    """Validate a structurally safe plan; Task 4 supplies SQLGlot compilation.
+
+    This contract intentionally does not prove SQL semantics. Task 4 must
+    compile SQL through SQLGlot ASTs before any plan can reach published state.
+    """
 
     plan = _closed_object(value, _EXECUTION_PLAN_KEYS, "execution plan")
     _required_fields(
         plan,
-        _EXECUTION_PLAN_KEYS - {"status"},
+        _EXECUTION_PLAN_KEYS,
         "execution plan",
     )
     if plan["schema_version"] != SCHEMA_VERSION:
@@ -241,6 +357,10 @@ def validate_execution_plan_v2(value: Any) -> dict:
     plan["compiler_version"] = _required_string(
         plan["compiler_version"], "compiler_version", 80
     )
+    if not _COMPILER_VERSION_PATTERNS[plan["backend"]].fullmatch(
+        plan["compiler_version"]
+    ):
+        raise ValueError("compiler provenance does not match execution backend")
     for key in (
         "rule_version_id",
         "input_schema_snapshot_id",
@@ -249,7 +369,7 @@ def validate_execution_plan_v2(value: Any) -> dict:
         "output_binding_id",
     ):
         plan[key] = _uid(plan[key], key)
-    if "status" in plan and plan["status"] not in PLAN_STATUSES:
+    if plan["status"] not in PLAN_STATUSES:
         raise ValueError("unsupported execution plan status")
     raw_operations = plan["operations"]
     if (
@@ -258,7 +378,12 @@ def validate_execution_plan_v2(value: Any) -> dict:
         or len(raw_operations) > 500
     ):
         raise ValueError("execution plan operations must be a non-empty bounded array")
-    plan["operations"] = [_validate_operation(item) for item in raw_operations]
+    validator = {
+        "sql_pushdown": _validate_sql_operation,
+        "polars_batch": _validate_polars_operation,
+        "quality_check": _validate_quality_operation,
+    }[plan["backend"]]
+    plan["operations"] = [validator(item) for item in raw_operations]
     _canonical_json(plan)
     return plan
 

+ 131 - 6
tests/core/data_rules/test_execution_contracts.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
 import copy
+import hashlib
+import json
 
 import pytest
 
@@ -14,17 +16,29 @@ from app.core.data_rules.execution_contracts import (
 
 
 def schema_snapshot():
+    fields = [
+        {"name": "customer_id", "type": "string", "nullable": False},
+        {"name": "mobile", "type": "string", "nullable": True},
+    ]
     return {
         "schema_ref": "bd:customer:v7",
-        "schema_hash": "a" * 64,
-        "fields": [
-            {"name": "customer_id", "type": "string", "nullable": False},
-            {"name": "mobile", "type": "string", "nullable": True},
-        ],
+        "schema_hash": canonical_fields_hash(fields),
+        "fields": fields,
         "source_revision": "neo4j:42",
     }
 
 
+def canonical_fields_hash(fields):
+    normalized = sorted(fields, key=lambda field: field["name"])
+    encoded = json.dumps(
+        normalized,
+        sort_keys=True,
+        separators=(",", ":"),
+        ensure_ascii=False,
+    )
+    return hashlib.sha256(encoded.encode("utf-8")).hexdigest()
+
+
 def dataset_binding():
     return {
         "data_source_uid": new_governance_uid(),
@@ -47,7 +61,15 @@ def execution_plan():
         "output_schema_snapshot_id": new_governance_uid(),
         "input_binding_id": new_governance_uid(),
         "output_binding_id": new_governance_uid(),
-        "operations": [{"kind": "sql", "statement": "SELECT 1"}],
+        "status": "compiled",
+        "operations": [
+            {
+                "kind": "sql",
+                "dialect": "postgresql",
+                "statement": "SELECT 1",
+                "parameters": {},
+            }
+        ],
     }
 
 
@@ -80,6 +102,26 @@ def test_schema_snapshot_is_closed_and_preserves_typed_fields():
         validate_schema_snapshot(unsafe)
 
 
+def test_schema_snapshot_rejects_a_digest_that_does_not_match_fields():
+    snapshot = schema_snapshot()
+    snapshot["fields"][0]["nullable"] = True
+
+    with pytest.raises(ValueError, match="schema_hash does not match"):
+        validate_schema_snapshot(snapshot)
+
+
+def test_query_dataset_binding_requires_an_opaque_server_published_reference():
+    binding = dataset_binding()
+    binding["object_kind"] = "query"
+    binding["object_ref"] = "SELECT * FROM public.customer"
+
+    with pytest.raises(ValueError, match="opaque query reference"):
+        validate_dataset_binding(binding)
+
+    binding["object_ref"] = f"query://{new_governance_uid()}"
+    assert validate_dataset_binding(binding)["object_ref"] == binding["object_ref"]
+
+
 @pytest.mark.parametrize(
     ("mutation", "message"),
     [
@@ -107,3 +149,86 @@ def test_execution_plan_rejects_unpinned_or_secret_runtime_material(
 
     with pytest.raises(ValueError, match=message):
         validate_execution_plan_v2(plan)
+
+
+def test_execution_plan_requires_backend_specific_compiler_provenance_and_status():
+    plan = execution_plan()
+    plan.pop("status")
+    with pytest.raises(ValueError, match="status"):
+        validate_execution_plan_v2(plan)
+
+    plan = execution_plan()
+    plan["compiler_version"] = "untrusted-compiler-1.0"
+    with pytest.raises(ValueError, match="compiler provenance"):
+        validate_execution_plan_v2(plan)
+
+
+@pytest.mark.parametrize(
+    ("mutation", "message"),
+    [
+        (
+            lambda operation: operation.update({"kind": "shell"}),
+            "unsupported sql_pushdown operation kind",
+        ),
+        (
+            lambda operation: operation.update({"dialect": "sqlite"}),
+            "unsupported SQL dialect",
+        ),
+        (
+            lambda operation: operation.update({"unsafe": True}),
+            "unsupported fields",
+        ),
+        (
+            lambda operation: operation.update({"parameters": {"id": {"x": 1}}}),
+            "parameter value",
+        ),
+    ],
+)
+def test_sql_plan_rejects_unknown_operation_dialect_keys_and_parameter_shapes(
+    mutation, message
+):
+    plan = execution_plan()
+    mutation(plan["operations"][0])
+
+    with pytest.raises(ValueError, match=message):
+        validate_execution_plan_v2(plan)
+
+
+@pytest.mark.parametrize(
+    "statement",
+    [
+        "SELECT 1; SELECT 2",
+        "SELECT 1 -- comment",
+        "SELECT 1 /* comment */",
+        "CREATE TABLE unsafe(id integer)",
+        "BEGIN",
+    ],
+)
+def test_sql_plan_rejects_multiple_commented_ddl_and_control_statements(statement):
+    plan = execution_plan()
+    plan["operations"][0]["statement"] = statement
+
+    with pytest.raises(ValueError, match="single statement|comment|DDL|control"):
+        validate_execution_plan_v2(plan)
+
+
+@pytest.mark.parametrize(
+    "unsafe_value",
+    [
+        "generated_python",
+        "python",
+        "pickled payload",
+        "callable payload",
+        "module payload",
+        "file payload",
+        "s3://unsafe-plan",
+    ],
+)
+def test_execution_plan_rejects_code_serialization_and_url_like_content(
+    unsafe_value,
+):
+    plan = execution_plan()
+    plan["operations"][0]["parameters"] = {"value": unsafe_value}
+
+    with pytest.raises(ValueError, match="forbidden runtime content"):
+        validate_execution_plan_v2(plan)