diff --git a/changelog.d/495.added.md b/changelog.d/495.added.md new file mode 100644 index 00000000..c4b0d50c --- /dev/null +++ b/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. diff --git a/src/policyengine/bundle.py b/src/policyengine/bundle.py index eeb45376..696ae665 100644 --- a/src/policyengine/bundle.py +++ b/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 diff --git a/src/policyengine/provenance/dataset_sources.py b/src/policyengine/provenance/dataset_sources.py index dfbdb311..094d92b3 100644 --- a/src/policyengine/provenance/dataset_sources.py +++ b/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 diff --git a/src/policyengine/provenance/manifest.py b/src/policyengine/provenance/manifest.py index 088676fd..178979a5 100644 --- a/src/policyengine/provenance/manifest.py +++ b/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, diff --git a/src/policyengine/tax_benefit_models/uk/datasets.py b/src/policyengine/tax_benefit_models/uk/datasets.py index 82522111..96c4aac3 100644 --- a/src/policyengine/tax_benefit_models/uk/datasets.py +++ b/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) diff --git a/src/policyengine/tax_benefit_models/uk/model.py b/src/policyengine/tax_benefit_models/uk/model.py index dfa54985..8b24b933 100644 --- a/src/policyengine/tax_benefit_models/uk/model.py +++ b/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 ( diff --git a/src/policyengine/tax_benefit_models/us/datasets.py b/src/policyengine/tax_benefit_models/us/datasets.py index 4b2dc2ba..b982bdf7 100644 --- a/src/policyengine/tax_benefit_models/us/datasets.py +++ b/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) diff --git a/src/policyengine/tax_benefit_models/us/model.py b/src/policyengine/tax_benefit_models/us/model.py index 090ac514..c95e252f 100644 --- a/src/policyengine/tax_benefit_models/us/model.py +++ b/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, diff --git a/tests/test_dataset_sources.py b/tests/test_dataset_sources.py index ff50865c..b0088f71 100644 --- a/tests/test_dataset_sources.py +++ b/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")