feat: enforce container provenance in docker and gpu executors

Author: haasonsaas

src/nimbus/runners/docker.py

@@ -17,6 +17,7 @@
 
 from ..common.schemas import JobAssignment
 from ..common.settings import HostAgentSettings
+from ..common.supply_chain import ImagePolicy, ensure_provenance
 from ..common.networking import (
     MetadataEndpointDenylist,
     EgressPolicyPack,
@@ -37,6 +38,9 @@ def __init__(self, settings: Optional[HostAgentSettings] = None) -> None:
         self._job_workspaces: dict[int, Path] = {}
         self._egress_enforcer: Optional[OfflineEgressEnforcer] = None
         self._container_user: Optional[str] = None
+        self._image_policy: Optional[ImagePolicy] = None
+        self._cosign_key: Optional[Path] = None
+        self._require_provenance: bool = False
 
     def initialize(self, settings: HostAgentSettings) -> None:
         """Initialize the executor with settings."""
@@ -58,6 +62,11 @@ def initialize(self, settings: HostAgentSettings) -> None:
         # Create workspace directory
         settings.docker_workspace_path.mkdir(parents=True, exist_ok=True)
         self._container_user = settings.docker_container_user
+        self._image_policy = ImagePolicy.from_paths(
+            settings.image_allow_list_path, settings.image_deny_list_path
+        )
+        self._cosign_key = settings.cosign_certificate_authority
+        self._require_provenance = settings.provenance_required
 
         # Initialize egress enforcer (similar to HostAgent)
         metadata_denylist = MetadataEndpointDenylist(settings.metadata_endpoint_denylist)
@@ -137,6 +146,7 @@ async def run(
         try:
             # Determine container image
             image = self._get_container_image(job)
+            self._verify_image(image)
 
             # Pull image if not present locally (for better performance)
             await self._ensure_image(image)
@@ -275,6 +285,17 @@ async def _ensure_image(self, image: str) -> None:
             except Exception as exc:
                 LOGGER.error("Failed to pull image", image=image, error=str(exc))
                 raise RuntimeError(f"Failed to pull image {image}: {exc}") from exc
+    def _verify_image(self, image: str) -> None:
+        if not self._image_policy:
+            self._image_policy = ImagePolicy.from_paths(
+                self._settings.image_allow_list_path, self._settings.image_deny_list_path  # type: ignore[union-attr]
+            )
+        ensure_provenance(
+            image,
+            self._image_policy,
+            public_key_path=self._cosign_key,
+            require_provenance=self._require_provenance,
+        )
 
     async def prepare_warm_instance(self, instance_id: str) -> dict:
         """Prepare a warm Docker instance (primarily pre-pulled images)."""

src/nimbus/runners/gpu.py

@@ -19,6 +19,7 @@
 
 from ..common.schemas import JobAssignment
 from ..common.settings import HostAgentSettings
+from ..common.supply_chain import ImagePolicy, ensure_provenance
 from .base import Executor, RunResult
 
 LOGGER = structlog.get_logger("nimbus.runners.gpu")
@@ -53,6 +54,9 @@ def __init__(self, settings: Optional[HostAgentSettings] = None) -> None:
         self._available_gpus: Dict[str, GPUInfo] = {}
         self._gpu_allocations: Dict[int, List[str]] = {}  # job_id -> gpu_uuids
         self._container_user: Optional[str] = None
+        self._image_policy: Optional[ImagePolicy] = None
+        self._cosign_key: Optional[Path] = None
+        self._require_provenance: bool = False
 
     def initialize(self, settings: HostAgentSettings) -> None:
         """Initialize the GPU executor."""
@@ -78,6 +82,11 @@ def initialize(self, settings: HostAgentSettings) -> None:
         # Create workspace directory
         settings.docker_workspace_path.mkdir(parents=True, exist_ok=True)
         self._container_user = settings.docker_container_user
+        self._image_policy = ImagePolicy.from_paths(
+            settings.image_allow_list_path, settings.image_deny_list_path
+        )
+        self._cosign_key = settings.cosign_certificate_authority
+        self._require_provenance = settings.provenance_required
 
     @property
     def name(self) -> str:
@@ -226,6 +235,7 @@ async def run(
         try:
             # Get container image
             image = self._get_gpu_container_image(job)
+            self._verify_image(image)
 
             # Ensure image is available
             await self._ensure_image(image)
@@ -449,3 +459,15 @@ async def _ensure_image(self, image: str) -> None:
             except Exception as exc:
                 LOGGER.error("Failed to pull GPU image", image=image, error=str(exc))
                 raise RuntimeError(f"Failed to pull GPU image {image}: {exc}") from exc
+
+    def _verify_image(self, image: str) -> None:
+        if not self._image_policy and self._settings:
+            self._image_policy = ImagePolicy.from_paths(
+                self._settings.image_allow_list_path, self._settings.image_deny_list_path
+            )
+        ensure_provenance(
+            image,
+            self._image_policy or ImagePolicy(set(), set()),
+            public_key_path=self._cosign_key,
+            require_provenance=self._require_provenance,
+        )

tests/test_docker_executor.py

@@ -25,6 +25,10 @@ def mock_settings():
     settings.egress_policy_pack = None
     settings.offline_mode = False
     settings.artifact_registry_allow_list = []
+    settings.image_allow_list_path = None
+    settings.image_deny_list_path = None
+    settings.cosign_certificate_authority = None
+    settings.provenance_required = False
     return settings
 
 
@@ -183,6 +187,39 @@ def test_get_container_image_default(mock_settings, sample_job):
     assert image == mock_settings.docker_default_image
 
 
+@pytest.mark.asyncio
+@patch('src.nimbus.runners.docker.docker.DockerClient')
+async def test_docker_executor_enforces_provenance(mock_docker_client, tmp_path, mock_settings, sample_job, monkeypatch):
+    cosign_key = tmp_path / "cosign.pub"
+    cosign_key.write_text("public key", encoding="utf-8")
+
+    mock_settings.image_allow_list_path = None
+    mock_settings.image_deny_list_path = None
+    mock_settings.cosign_certificate_authority = cosign_key
+    mock_settings.provenance_required = True
+
+    executor = DockerExecutor()
+
+    mock_client = Mock()
+    mock_docker_client.return_value = mock_client
+    mock_client.ping.return_value = True
+    mock_client.networks.get.return_value = Mock()
+    mock_client.containers.create.return_value = Mock(wait=lambda timeout: {'StatusCode': 0}, logs=lambda **kw: b'')
+
+    with patch('src.nimbus.runners.docker.EgressPolicyPack'), patch('src.nimbus.runners.docker.MetadataEndpointDenylist'):
+        executor.initialize(mock_settings)
+
+    sample_job.labels = ["nimbus", "docker"]
+
+    with patch('src.nimbus.runners.docker.ensure_provenance') as ensure_mock:
+        ensure_mock.return_value = None
+        with patch.object(executor, '_ensure_image'):
+            await executor.prepare(sample_job)
+            await executor.run(sample_job, timeout_seconds=1)
+
+        ensure_mock.assert_called()
+
+
 def test_build_environment_variables(mock_settings, sample_job):
     """Test environment variable building."""
     executor = DockerExecutor(mock_settings)

tests/test_executor_integration.py

@@ -32,6 +32,10 @@ def mock_settings():
     settings.egress_policy_pack = None
     settings.offline_mode = False
     settings.artifact_registry_allow_list = []
+    settings.image_allow_list_path = None
+    settings.image_deny_list_path = None
+    settings.cosign_certificate_authority = None
+    settings.provenance_required = False
     return settings
 
 

tests/test_gpu_executor.py

@@ -18,6 +18,10 @@ def mock_settings():
     settings.docker_socket_path = "/var/run/docker.sock"
     settings.docker_workspace_path = Path("/tmp/test-gpu-workspaces")
     settings.docker_container_user = None
+    settings.image_allow_list_path = None
+    settings.image_deny_list_path = None
+    settings.cosign_certificate_authority = None
+    settings.provenance_required = False
     return settings
 
 
@@ -88,6 +92,47 @@ def test_gpu_executor_properties():
     assert "container" in capabilities
 
 
+@pytest.mark.asyncio
+@patch('src.nimbus.runners.gpu.docker.DockerClient')
+@patch('src.nimbus.runners.gpu.subprocess.run')
+async def test_gpu_executor_enforces_provenance(mock_run, mock_docker_client, tmp_path, mock_settings, sample_gpu_job):
+    mock_run.return_value = Mock(returncode=0, stdout='{"nvidia": {"path": "runtime"}}')
+
+    def run_side_effect(cmd, **kwargs):
+        if "nvidia-smi" in cmd:
+            return Mock(returncode=0, stdout="0, Tesla V100, GPU-12345, 32768, 30720, 7.0, 470.82")
+        return Mock(returncode=0, stdout='{"nvidia": {"path": "runtime"}}')
+
+    mock_run.side_effect = run_side_effect
+
+    mock_client = Mock()
+    mock_docker_client.return_value = mock_client
+    mock_client.ping.return_value = True
+    mock_client.containers.create.return_value = Mock(
+        wait=lambda timeout: {'StatusCode': 0},
+        logs=lambda **kw: b'',
+    )
+
+    cosign_key = tmp_path / "cosign.pub"
+    cosign_key.write_text("public key", encoding="utf-8")
+
+    mock_settings.cosign_certificate_authority = cosign_key
+    mock_settings.provenance_required = True
+
+    executor = GPUExecutor()
+    executor.initialize(mock_settings)
+
+    sample_gpu_job.labels = ["gpu", "pytorch"]
+
+    with patch('src.nimbus.runners.gpu.ensure_provenance') as ensure_mock:
+        ensure_mock.return_value = None
+        with patch.object(executor, '_ensure_image'):
+            await executor.prepare(sample_gpu_job)
+            await executor.run(sample_gpu_job, timeout_seconds=1)
+
+        ensure_mock.assert_called()
+
+
 @patch('subprocess.run')
 def test_check_nvidia_docker_available(mock_run):
     """Test nvidia-docker availability check when available."""