Verify bundle-managed dataset downloads against certified SHA-256 hashes

changelog.d/495.added.md

@@ -0,0 +1 @@
+Verify bundle-managed dataset downloads against their certified SHA-256 hashes before use, failing closed on mismatch.

src/policyengine/bundle.py

@@ -136,13 +136,11 @@ def build_id(self, release: Mapping[str, Any]) -> Optional[str]:
         return str(version) if version is not None else None
 
     def verify_download(self, plan: DatasetPlan, path: Path) -> str:
-        actual_sha256 = _sha256_file(path)
-        if plan.expected_sha256 and actual_sha256 != plan.expected_sha256:
-            raise BundleError(
-                f"Downloaded {plan.country.upper()} dataset {plan.dataset} "
-                f"has sha256 {actual_sha256}, expected {plan.expected_sha256}."
-            )
-        return actual_sha256
+        return verify_file_sha256(
+            path,
+            expected_sha256=plan.expected_sha256,
+            description=f"Downloaded {plan.country.upper()} dataset {plan.dataset}",
+        )
 
     def dataset_check(
         self,
@@ -918,3 +916,19 @@ def _sha256_file(path: Path) -> str:
         for chunk in iter(lambda: file.read(1024 * 1024), b""):
             digest.update(chunk)
     return digest.hexdigest()
+
+
+def verify_file_sha256(
+    path: Path,
+    *,
+    expected_sha256: Optional[str],
+    description: str,
+) -> str:
+    """Hash a file and fail closed when its expected sha256 does not match."""
+
+    actual_sha256 = _sha256_file(path)
+    if expected_sha256 and actual_sha256 != expected_sha256:
+        raise BundleError(
+            f"{description} has sha256 {actual_sha256}, expected {expected_sha256}."
+        )
+    return actual_sha256

src/policyengine/provenance/dataset_sources.py

@@ -3,8 +3,14 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
+from pathlib import Path
 from typing import Optional
 
+from policyengine.bundle import verify_file_sha256
+from policyengine.provenance.manifest import (
+    dataset_artifact_sha256,
+    dataset_logical_name,
+)
 from policyengine.utils.google_cloud_bucket import download_file_from_gcs
 
 
@@ -78,8 +84,24 @@ def materialize_dataset_source(
     dataset_source: str,
     *,
     version: Optional[str] = None,
+    country_id: Optional[str] = None,
 ) -> str:
-    """Return a local file path for supported remote dataset URIs."""
+    """Return a local file path, verifying bundle-managed remote artifacts."""
+
+    def verify_bundle_download(local_path: str) -> None:
+        if country_id is None:
+            return
+        expected_sha256 = dataset_artifact_sha256(country_id, dataset_source)
+        if expected_sha256 is None:
+            return
+        verify_file_sha256(
+            Path(local_path),
+            expected_sha256=expected_sha256,
+            description=(
+                f"Downloaded {country_id.upper()} dataset "
+                f"{dataset_logical_name(dataset_source)}"
+            ),
+        )
 
     if dataset_source.startswith("gs://"):
         reference = parse_gs_uri(dataset_source)
@@ -88,6 +110,7 @@ def materialize_dataset_source(
             reference.path,
             version=_select_version(reference.version, version),
         )
+        verify_bundle_download(local_path)
         return local_path
 
     if dataset_source.startswith("hf://"):
@@ -97,7 +120,7 @@ def materialize_dataset_source(
 
         reference = parse_hf_uri(dataset_source)
         try:
-            return download_huggingface_dataset(
+            local_path = download_huggingface_dataset(
                 reference.repo_id,
                 reference.path,
                 version=_select_version(reference.version, version),
@@ -109,11 +132,13 @@ def materialize_dataset_source(
             # before surfacing the original failure.
             from huggingface_hub import hf_hub_download
 
-            return hf_hub_download(
+            local_path = hf_hub_download(
                 repo_id=reference.repo_id,
                 repo_type="dataset",
                 filename=reference.path,
                 revision=_select_version(reference.version, version),
             )
+        verify_bundle_download(local_path)
+        return local_path
 
     return dataset_source

src/policyengine/provenance/manifest.py

@@ -537,6 +537,33 @@ def resolve_dataset_reference(country_id: str, dataset: str) -> str:
     return artifact.uri
 
 
+def dataset_artifact_sha256(
+    country_id: str,
+    dataset_uri: str,
+) -> Optional[str]:
+    """Return the bundled sha256 for the artifact at an exact dataset URI."""
+
+    manifest = get_release_manifest(country_id)
+    for path_reference in manifest.datasets.values():
+        reference_uri = build_hf_uri(
+            repo_id=path_reference.repo_id or manifest.data_package.repo_id,
+            path_in_repo=path_reference.path,
+            revision=path_reference.revision
+            or _artifact_revision(manifest.data_package),
+        )
+        if reference_uri == dataset_uri and path_reference.sha256:
+            return path_reference.sha256
+
+    certified_artifact = manifest.certified_data_artifact
+    if (
+        certified_artifact is not None
+        and certified_artifact.uri == dataset_uri
+        and certified_artifact.sha256
+    ):
+        return certified_artifact.sha256
+    return None
+
+
 def resolve_managed_dataset_reference(
     country_id: str,
     dataset: Optional[str] = None,

src/policyengine/tax_benefit_models/uk/datasets.py

@@ -123,7 +123,10 @@ def create_datasets(
     for dataset in datasets:
         resolved_dataset = resolve_dataset_reference("uk", dataset)
         dataset_stem = dataset_logical_name(resolved_dataset)
-        runtime_dataset = materialize_dataset_source(resolved_dataset)
+        runtime_dataset = materialize_dataset_source(
+            resolved_dataset,
+            country_id="uk",
+        )
         from policyengine_uk import Microsimulation
 
         sim = Microsimulation(dataset=runtime_dataset)
@@ -251,6 +254,8 @@ def ensure_datasets(
             break
 
     if all_exist:
+        # This branch never materializes the source artifact, so runtime source
+        # byte verification only applies when create_datasets runs below.
         return load_datasets(datasets=datasets, years=years, data_folder=data_folder)
     else:
         return create_datasets(datasets=datasets, years=years, data_folder=data_folder)

src/policyengine/tax_benefit_models/uk/model.py

@@ -310,7 +310,10 @@ def managed_microsimulation(
             allow_unmanaged and dataset is not None and "://" in dataset
         ),
     )
-    runtime_dataset_source = materialize_dataset_source(dataset_source)
+    runtime_dataset_source = materialize_dataset_source(
+        dataset_source,
+        country_id="uk",
+    )
     runtime_dataset = runtime_dataset_source
     if isinstance(runtime_dataset_source, str) and "://" not in runtime_dataset_source:
         from policyengine_uk.data.dataset_schema import (

src/policyengine/tax_benefit_models/us/datasets.py

@@ -297,7 +297,10 @@ def create_datasets(
     for dataset in datasets:
         resolved_dataset = resolve_dataset_reference("us", dataset)
         dataset_stem = dataset_logical_name(resolved_dataset)
-        runtime_dataset = materialize_dataset_source(resolved_dataset)
+        runtime_dataset = materialize_dataset_source(
+            resolved_dataset,
+            country_id="us",
+        )
         sim = Microsimulation(dataset=runtime_dataset)
 
         for year in years:
@@ -1180,6 +1183,8 @@ def ensure_datasets(
             break
 
     if all_exist:
+        # This branch never materializes the source artifact, so runtime source
+        # byte verification only applies when create_datasets runs below.
         return load_datasets(datasets=datasets, years=years, data_folder=data_folder)
     else:
         return create_datasets(datasets=datasets, years=years, data_folder=data_folder)

src/policyengine/tax_benefit_models/us/model.py

@@ -453,7 +453,10 @@ def managed_microsimulation(
             allow_unmanaged and dataset is not None and "://" in dataset
         ),
     )
-    runtime_dataset_source = materialize_dataset_source(dataset_source)
+    runtime_dataset_source = materialize_dataset_source(
+        dataset_source,
+        country_id="us",
+    )
     microsim = Microsimulation(dataset=runtime_dataset_source, **kwargs)
     microsim.policyengine_bundle = _managed_release_bundle(
         dataset_uri,

tests/test_dataset_sources.py

@@ -1,12 +1,15 @@
+import hashlib
 import importlib.util
 import sys
 from pathlib import Path
-from types import SimpleNamespace
+from types import ModuleType, SimpleNamespace
 from unittest.mock import Mock
 
 import pytest
 
+from policyengine.bundle import BundleError
 from policyengine.provenance import dataset_sources
+from policyengine.provenance import manifest as manifest_module
 from policyengine.provenance.dataset_sources import (
     materialize_dataset_source,
     parse_gs_uri,
@@ -16,6 +19,67 @@
 REPO_ROOT = Path(__file__).resolve().parents[1]
 
 
+def _install_hf_downloader(monkeypatch, download):
+    core_module = ModuleType("policyengine_core")
+    core_module.__path__ = []
+    tools_module = ModuleType("policyengine_core.tools")
+    tools_module.__path__ = []
+    hf_module = ModuleType("policyengine_core.tools.hugging_face")
+    hf_module.download_huggingface_dataset = download
+    core_module.tools = tools_module
+    tools_module.hugging_face = hf_module
+    monkeypatch.setitem(sys.modules, "policyengine_core", core_module)
+    monkeypatch.setitem(sys.modules, "policyengine_core.tools", tools_module)
+    monkeypatch.setitem(
+        sys.modules,
+        "policyengine_core.tools.hugging_face",
+        hf_module,
+    )
+
+
+@pytest.fixture
+def bundle_managed_hf_download(monkeypatch, tmp_path):
+    payload = b"bundle-managed dataset bytes"
+    downloaded_path = tmp_path / "populace_uk_2023.h5"
+    country_manifest = manifest_module.get_release_manifest("uk").model_copy(deep=True)
+    dataset_name = country_manifest.default_dataset
+    path_reference = country_manifest.datasets[dataset_name]
+    dataset_uri = manifest_module.build_hf_uri(
+        repo_id=path_reference.repo_id or country_manifest.data_package.repo_id,
+        path_in_repo=path_reference.path,
+        revision=(
+            path_reference.revision
+            or country_manifest.data_package.release_manifest_revision
+            or country_manifest.data_package.version
+        ),
+    )
+    certified_artifact = country_manifest.certified_data_artifact
+    assert certified_artifact is not None
+    assert certified_artifact.uri == dataset_uri
+
+    def emit_downloaded_bytes(*args, **kwargs):
+        downloaded_path.write_bytes(payload)
+        return str(downloaded_path)
+
+    download = Mock(side_effect=emit_downloaded_bytes)
+    _install_hf_downloader(monkeypatch, download)
+    monkeypatch.setattr(
+        manifest_module,
+        "get_release_manifest",
+        Mock(return_value=country_manifest),
+    )
+    return SimpleNamespace(
+        country_manifest=country_manifest,
+        certified_artifact=certified_artifact,
+        dataset_name=dataset_name,
+        dataset_uri=dataset_uri,
+        download=download,
+        downloaded_path=downloaded_path,
+        payload=payload,
+        sha256=hashlib.sha256(payload).hexdigest(),
+    )
+
+
 def _load_module_from_path(module_name: str, path: Path):
     spec = importlib.util.spec_from_file_location(module_name, path)
     assert spec is not None
@@ -62,10 +126,7 @@ def test_materialize_dataset_source_downloads_gcs_uri(monkeypatch):
 
 def test_materialize_dataset_source_downloads_hf_uri(monkeypatch):
     download = Mock(return_value="/tmp/enhanced_cps_2024.h5")
-    monkeypatch.setattr(
-        "policyengine_core.tools.hugging_face.download_huggingface_dataset",
-        download,
-    )
+    _install_hf_downloader(monkeypatch, download)
 
     result = materialize_dataset_source(
         "hf://policyengine/policyengine-us-data/enhanced_cps_2024.h5@1.77.0"
@@ -79,6 +140,44 @@ def test_materialize_dataset_source_downloads_hf_uri(monkeypatch):
     )
 
 
+def test_bundle_managed_materialization_accepts_matching_artifact_sha256(
+    bundle_managed_hf_download,
+):
+    fixture = bundle_managed_hf_download
+    fixture.country_manifest.datasets[fixture.dataset_name].sha256 = fixture.sha256
+    fixture.certified_artifact.sha256 = "0" * 64
+
+    result = materialize_dataset_source(
+        fixture.dataset_uri,
+        country_id="uk",
+    )
+
+    assert result == str(fixture.downloaded_path)
+    assert fixture.downloaded_path.read_bytes() == fixture.payload
+    fixture.download.assert_called_once()
+
+
+def test_bundle_managed_materialization_rejects_certified_sha256_mismatch(
+    bundle_managed_hf_download,
+):
+    fixture = bundle_managed_hf_download
+    expected_sha256 = hashlib.sha256(b"expected dataset bytes").hexdigest()
+    fixture.country_manifest.datasets[fixture.dataset_name].sha256 = None
+    fixture.certified_artifact.sha256 = expected_sha256
+
+    with pytest.raises(BundleError) as exc_info:
+        materialize_dataset_source(
+            fixture.dataset_uri,
+            country_id="uk",
+        )
+
+    message = str(exc_info.value)
+    assert fixture.sha256 in message
+    assert expected_sha256 in message
+    assert "Downloaded UK dataset populace_uk_2023 has sha256" in message
+    fixture.download.assert_called_once()
+
+
 def test_materialize_dataset_source_preserves_local_path():
     assert materialize_dataset_source("/tmp/enhanced_cps_2024.h5") == (
         "/tmp/enhanced_cps_2024.h5"
@@ -116,7 +215,8 @@ def test_us_create_datasets_passes_materialized_source_to_country_package(
     )
 
     materialize.assert_called_once_with(
-        "gs://policyengine-us-data/enhanced_cps_2024.h5@1.77.0"
+        "gs://policyengine-us-data/enhanced_cps_2024.h5@1.77.0",
+        country_id="us",
     )
     microsimulation.assert_called_once_with(dataset="/tmp/enhanced_cps_2024.h5")
 
@@ -144,6 +244,7 @@ def test_uk_create_datasets_passes_materialized_source_to_country_package(
     )
 
     materialize.assert_called_once_with(
-        "gs://policyengine-uk-data-private/enhanced_frs_2023_24.h5@1.40.3"
+        "gs://policyengine-uk-data-private/enhanced_frs_2023_24.h5@1.40.3",
+        country_id="uk",
     )
     microsimulation.assert_called_once_with(dataset="/tmp/enhanced_frs_2023_24.h5")