fix(security): isolate reference execution and bound source inputs
This commit is contained in:
@@ -0,0 +1,194 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
from types import SimpleNamespace
|
||||
import unittest
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from govoplan_core.auth import ApiPrincipal
|
||||
from govoplan_core.core.access import PrincipalRef
|
||||
from govoplan_core.core.datasources import (
|
||||
CAPABILITY_DATASOURCE_CATALOGUE,
|
||||
DatasourceAccessError,
|
||||
DatasourceCatalogueProvider,
|
||||
DatasourceDescriptor,
|
||||
DatasourceReadRequest,
|
||||
DatasourceReadResult,
|
||||
)
|
||||
from govoplan_core.security.bounded_process import run_bounded_operation
|
||||
from govoplan_dataflow.backend.backends.reference import ReferenceExecutionBackend
|
||||
from govoplan_dataflow.backend.executor import EXECUTOR_VERSION, PipelineExecutionError
|
||||
from govoplan_dataflow.backend.graph import validate_graph
|
||||
from govoplan_dataflow.backend.schemas import GraphEdge, GraphNode, GraphPosition, PipelineGraph
|
||||
from govoplan_dataflow.backend.service import _execute_pipeline_preview
|
||||
|
||||
|
||||
OUTER_ROWS = ({"id": "outer", "amount": 15}, {"id": "outer-low", "amount": 5})
|
||||
INNER_ROWS = ({"id": "inner", "amount": 25}, {"id": "inner-low", "amount": 2})
|
||||
|
||||
|
||||
def node(node_id: str, node_type: str, config: dict) -> GraphNode:
|
||||
return GraphNode(
|
||||
id=node_id, type=node_type, label=node_id,
|
||||
position=GraphPosition(x=0, y=0), config=config,
|
||||
)
|
||||
|
||||
|
||||
def nested_source_graph() -> PipelineGraph:
|
||||
# Both external sources deliberately share their node ID and logical name.
|
||||
# The pinned subflow still has exactly one distinct inline input binding.
|
||||
nested = PipelineGraph(
|
||||
nodes=[
|
||||
node("input", "source.inline", {
|
||||
"source_name": "bound_input", "rows": [], "input_binding": True,
|
||||
}),
|
||||
node("shared", "source.reference", {
|
||||
"source_name": "records", "source_ref": {"$parameter": "source_ref"},
|
||||
"expected_fingerprint": "inner-pinned", "consistency": "frozen",
|
||||
}),
|
||||
node("union", "combine.union", {"mode": "all"}),
|
||||
node("minimum", "filter.expression", {"expression": "amount >= ${minimum}"}),
|
||||
node("output", "output", {}),
|
||||
],
|
||||
edges=[
|
||||
GraphEdge(id="input-union", source="input", target="union"),
|
||||
GraphEdge(id="shared-union", source="shared", target="union"),
|
||||
GraphEdge(id="union-minimum", source="union", target="minimum"),
|
||||
GraphEdge(id="minimum-output", source="minimum", target="output"),
|
||||
],
|
||||
)
|
||||
return PipelineGraph(
|
||||
nodes=[
|
||||
node("shared", "source.reference", {
|
||||
"source_name": "records", "source_ref": "datasource:outer",
|
||||
"expected_fingerprint": "outer-pinned", "consistency": "current",
|
||||
}),
|
||||
node("nested", "subflow", {
|
||||
"template_ref": "fixture-nested-source", "template_version": "1",
|
||||
"parameters": {"source_ref": "datasource:inner", "minimum": 10},
|
||||
"graph": nested.model_dump(mode="python"),
|
||||
}),
|
||||
node("output", "output", {}),
|
||||
],
|
||||
edges=[
|
||||
GraphEdge(id="shared-nested", source="shared", target="nested"),
|
||||
GraphEdge(id="nested-output", source="nested", target="output"),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
class ReferenceSubflowProcessTests(unittest.TestCase):
|
||||
def setUp(self) -> None:
|
||||
self.graph = nested_source_graph()
|
||||
self.assertEqual([], [item.model_dump() for item in validate_graph(self.graph) if item.severity == "error"])
|
||||
self.session = Mock(spec=Session)
|
||||
self.principal = ApiPrincipal(
|
||||
principal=PrincipalRef(
|
||||
account_id="fixture-account", membership_id="fixture-membership",
|
||||
tenant_id="fixture-tenant", scopes=frozenset(),
|
||||
),
|
||||
account=object(), user=object(),
|
||||
)
|
||||
self.provider = Mock(spec=DatasourceCatalogueProvider)
|
||||
self.provider.read_datasource.side_effect = self.read_source
|
||||
self.registry = SimpleNamespace(
|
||||
has_capability=lambda name: name == CAPABILITY_DATASOURCE_CATALOGUE,
|
||||
capability=lambda name: self.provider if name == CAPABILITY_DATASOURCE_CATALOGUE else None,
|
||||
)
|
||||
self.reads: list[tuple[int, DatasourceReadRequest]] = []
|
||||
self.denied = False
|
||||
|
||||
def read_source(self, session, principal, *, request: DatasourceReadRequest) -> DatasourceReadResult:
|
||||
self.assertIs(session, self.session)
|
||||
self.assertIs(principal, self.principal)
|
||||
self.reads.append((os.getpid(), request))
|
||||
if request.datasource_ref == "datasource:inner" and self.denied:
|
||||
raise DatasourceAccessError("Current principal cannot read datasource:inner.")
|
||||
rows, fingerprint = {
|
||||
"datasource:outer": (OUTER_ROWS, "outer-pinned"),
|
||||
"datasource:inner": (INNER_ROWS, "inner-pinned"),
|
||||
}[request.datasource_ref]
|
||||
self.assertEqual(fingerprint, request.expected_fingerprint)
|
||||
return DatasourceReadResult(
|
||||
datasource=DatasourceDescriptor(
|
||||
ref=request.datasource_ref, source_name="records", name="Fixture source",
|
||||
kind="custom", mode="static", shape="tabular",
|
||||
fingerprint=fingerprint, provider="fixture-catalogue",
|
||||
),
|
||||
rows=rows, total_rows=len(rows), truncated=False,
|
||||
)
|
||||
|
||||
def preview(self, *, row_limit: int = 10):
|
||||
return _execute_pipeline_preview(
|
||||
self.graph, session=self.session, principal=self.principal,
|
||||
registry=self.registry, backend="reference", row_limit=row_limit,
|
||||
preview_node_id="nested",
|
||||
)
|
||||
|
||||
def test_nested_parameterized_sources_execute_in_real_child_without_id_aliasing(self) -> None:
|
||||
with (
|
||||
patch.object(ReferenceExecutionBackend, "_execute_in_process", side_effect=AssertionError("parent evaluation")),
|
||||
patch("govoplan_dataflow.backend.backends.reference.run_bounded_operation", wraps=run_bounded_operation) as worker,
|
||||
):
|
||||
result, version = self.preview()
|
||||
worker.assert_called_once()
|
||||
self.assertEqual(EXECUTOR_VERSION, version)
|
||||
self.assertEqual([OUTER_ROWS[0], INNER_ROWS[0]], result.rows)
|
||||
self.assertEqual(2, result.total_rows)
|
||||
self.assertFalse(result.truncated)
|
||||
self.assertEqual(2, result.input_row_count) # Root input count, not the nested binding again.
|
||||
self.assertIsNotNone(result.node_preview)
|
||||
self.assertEqual("nested", result.node_preview.node_id)
|
||||
self.assertEqual(result.rows, result.node_preview.rows)
|
||||
self.assertEqual(2, result.node_preview.total_rows)
|
||||
self.assertEqual(
|
||||
[("shared", "succeeded", 0, 2), ("nested", "succeeded", 2, 2), ("output", "succeeded", 2, 2)],
|
||||
[(item.node_id, item.status, item.input_rows, item.output_rows) for item in result.node_diagnostics],
|
||||
)
|
||||
binding_hash = hashlib.sha256(json.dumps(list(OUTER_ROWS), sort_keys=True, separators=(",", ":")).encode()).hexdigest()
|
||||
self.assertEqual([
|
||||
{"node_id": "shared", "source_ref": "datasource:outer", "source_name": "records", "kind": "datasource", "provider": "fixture-catalogue", "fingerprint": "outer-pinned", "row_count": 2, "preview_rows": 2, "truncated": False},
|
||||
{"node_id": "input", "source_name": "bound_input", "kind": "inline", "fingerprint": binding_hash, "row_count": 2, "subflow_node_id": "nested"},
|
||||
{"node_id": "shared", "source_ref": "datasource:inner", "source_name": "records", "kind": "datasource", "provider": "fixture-catalogue", "fingerprint": "inner-pinned", "row_count": 2, "preview_rows": 2, "truncated": False, "subflow_node_id": "nested"},
|
||||
], result.source_fingerprints)
|
||||
self.assertEqual(
|
||||
[("datasource:outer", "current", "outer-pinned"), ("datasource:inner", "frozen", "inner-pinned")],
|
||||
[(request.datasource_ref, request.consistency, request.expected_fingerprint) for _, request in self.reads],
|
||||
)
|
||||
self.assertTrue(all(pid == os.getpid() for pid, _ in self.reads))
|
||||
self.assertTrue(all(request.limit <= 500 and request.offset == 0 for _, request in self.reads))
|
||||
self.assertEqual([], self.session.mock_calls)
|
||||
|
||||
def test_nested_result_and_node_preview_keep_full_totals_when_output_is_bounded(self) -> None:
|
||||
result, _ = self.preview(row_limit=1)
|
||||
self.assertEqual([OUTER_ROWS[0]], result.rows)
|
||||
self.assertEqual(2, result.total_rows)
|
||||
self.assertTrue(result.truncated)
|
||||
self.assertEqual([OUTER_ROWS[0]], result.node_preview.rows)
|
||||
self.assertEqual(2, result.node_preview.total_rows)
|
||||
self.assertTrue(result.node_preview.truncated)
|
||||
self.assertEqual({"outer-pinned", "inner-pinned"}, {
|
||||
item["fingerprint"] for item in result.source_fingerprints if item["kind"] == "datasource"
|
||||
})
|
||||
self.assertEqual([], self.session.mock_calls)
|
||||
|
||||
def test_denied_nested_datasource_stops_before_worker_or_persistence(self) -> None:
|
||||
self.denied = True
|
||||
with (
|
||||
patch("govoplan_dataflow.backend.service.execute_typed_graph") as execute,
|
||||
patch("govoplan_dataflow.backend.backends.reference.run_bounded_operation") as worker,
|
||||
self.assertRaises(PipelineExecutionError) as caught,
|
||||
):
|
||||
self.preview()
|
||||
execute.assert_not_called()
|
||||
worker.assert_not_called()
|
||||
self.assertEqual("shared", caught.exception.node_id)
|
||||
self.assertEqual("Current principal cannot read datasource:inner.", str(caught.exception))
|
||||
self.assertFalse(caught.exception.retryable)
|
||||
self.assertIsInstance(caught.exception.__cause__, DatasourceAccessError)
|
||||
self.assertEqual(["datasource:outer", "datasource:inner"], [request.datasource_ref for _, request in self.reads])
|
||||
self.assertEqual([], self.session.mock_calls)
|
||||
Reference in New Issue
Block a user