guest: gate boot on GPU TEE attestation via requirements.verify_gpu

Author: kvinwang

dstack-types/src/lib.rs

@@ -143,6 +143,16 @@ pub struct Requirements {
     /// (e.g. 32 random alphanumeric characters).
     #[serde(skip_serializing_if = "Option::is_none")]
     pub launch_token_hash: Option<String>,
+    /// GPU TEE attestation requirement, defaults to `true` when omitted.
+    ///
+    /// On guests with an NVIDIA GPU attached, `true` means the guest runs
+    /// local GPU attestation (nvattest) during system setup and refuses to
+    /// boot — before key provisioning — if the GPU fails to attest (e.g. a
+    /// non-CC GPU or CC mode disabled by the host). `false` skips attestation
+    /// and sets the GPU ready state directly. Guests without a GPU attached
+    /// are unaffected either way.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub verify_gpu: Option<bool>,
 }
 
 impl Requirements {
@@ -151,6 +161,7 @@ impl Requirements {
             && self.platforms.is_none()
             && self.tdx_measure_acpi_tables.is_none()
             && self.launch_token_hash.is_none()
+            && self.verify_gpu.is_none()
     }
 }
 
@@ -345,6 +356,16 @@ impl AppCompose {
             }
         }
     }
+
+    /// Whether an attached GPU must pass local TEE attestation before the
+    /// guest continues booting. Defaults to `true` when
+    /// `requirements.verify_gpu` is omitted.
+    pub fn verify_gpu(&self) -> bool {
+        self.requirements
+            .as_ref()
+            .and_then(|r| r.verify_gpu)
+            .unwrap_or(true)
+    }
 }
 
 #[cfg(test)]
@@ -501,6 +522,41 @@ mod app_compose_tests {
         assert!(!requirements.is_empty());
     }
 
+    #[test]
+    fn verify_gpu_defaults_to_true() {
+        let no_requirements: AppCompose = serde_json::from_value(serde_json::json!({
+            "manifest_version": 2,
+            "name": "test",
+            "runner": "docker-compose"
+        }))
+        .unwrap();
+        assert!(no_requirements.verify_gpu());
+
+        let omitted: AppCompose = serde_json::from_value(serde_json::json!({
+            "manifest_version": "3",
+            "name": "test",
+            "runner": "docker-compose",
+            "requirements": {}
+        }))
+        .unwrap();
+        assert!(omitted.verify_gpu());
+        assert!(omitted.requirements.as_ref().unwrap().is_empty());
+
+        let disabled: AppCompose = serde_json::from_value(serde_json::json!({
+            "manifest_version": "3",
+            "name": "test",
+            "runner": "docker-compose",
+            "requirements": {
+                "verify_gpu": false
+            }
+        }))
+        .unwrap();
+        assert!(!disabled.verify_gpu());
+        let requirements = disabled.requirements.as_ref().unwrap();
+        assert_eq!(requirements.verify_gpu, Some(false));
+        assert!(!requirements.is_empty());
+    }
+
     #[test]
     fn launch_token_hash_is_domain_separated() {
         assert_eq!(

dstack-util/src/system_setup.rs

@@ -1047,10 +1047,153 @@ async fn do_sys_setup(stage0: Stage0<'_>) -> Result<()> {
     } else {
         info!("System time will be synchronized by chronyd in background");
     }
+    stage0
+        .setup_gpu()
+        .await
+        .context("Failed to verify GPU TEE attestation")?;
     let stage1 = stage0.setup_fs().await?;
     stage1.setup().await
 }
 
+/// GPU TEE attestation gate (`requirements.verify_gpu`, defaults to true).
+///
+/// Runs before key provisioning so a CVM whose GPU cannot prove it is a
+/// genuine, CC-enabled NVIDIA TEE never gets its app keys. The GPU
+/// "ready" state (`nvidia-smi conf-compute -srs 1`) is only set from here —
+/// nvidia-persistenced deliberately does not set it — so CUDA work cannot be
+/// submitted to an unverified GPU either.
+mod gpu {
+    use super::*;
+
+    const NVATTEST: &str = "/usr/bin/nvattest";
+    const ATTESTATION_OUTPUT: &str = "/run/nvidia-gpu-attestation/attestation.out";
+    const ATTESTATION_TIMEOUT: Duration = Duration::from_secs(300);
+
+    /// True if a passed-through NVIDIA GPU is present, detected via sysfs PCI
+    /// (vendor 0x10de, class VGA 0x0300xx or 3D controller 0x0302xx) so it
+    /// works even before the nvidia driver is loaded.
+    pub(super) fn nvidia_gpu_present() -> bool {
+        let Ok(entries) = fs::read_dir("/sys/bus/pci/devices") else {
+            return false;
+        };
+        entries.filter_map(|e| e.ok()).any(|dev| {
+            let read = |name: &str| {
+                fs::read_to_string(dev.path().join(name))
+                    .unwrap_or_default()
+                    .trim()
+                    .to_string()
+            };
+            read("vendor") == "0x10de"
+                && matches!(read("class").get(..6), Some("0x0300") | Some("0x0302"))
+        })
+    }
+
+    /// Mark the GPU as ready to accept work. Only meaningful (and only
+    /// succeeds) when the GPU runs in CC mode.
+    pub(super) fn set_gpu_ready_state() -> Result<()> {
+        let output = Command::new("nvidia-smi")
+            .args(["conf-compute", "-srs", "1"])
+            .output()
+            .context("failed to run nvidia-smi conf-compute")?;
+        if !output.status.success() {
+            bail!(
+                "nvidia-smi conf-compute -srs 1 failed ({}): {}",
+                output.status,
+                truncated_lossy(&output.stderr, 512),
+            );
+        }
+        info!("GPU ready state set");
+        Ok(())
+    }
+
+    /// Run local GPU attestation via nvattest with a fresh nonce, keeping the
+    /// verifier output in /run for debugging. Fails on any non-zero exit —
+    /// including a GPU that cannot produce an attestation report at all (a
+    /// non-CC GPU, or CC mode left off by the host).
+    pub(super) async fn attest_gpu() -> Result<()> {
+        if !Path::new(NVATTEST).exists() {
+            bail!("nvattest is not available in this image");
+        }
+        // Certificate/OCSP validation needs a sane clock even when
+        // secure_time is off; best-effort step chrony before attesting.
+        if let Err(err) = cmd!(chronyc makestep) {
+            warn!("failed to step system clock: {err:?}");
+        }
+        let nonce = hex::encode(rand::thread_rng().gen::<[u8; 32]>());
+        let output = tokio::time::timeout(
+            ATTESTATION_TIMEOUT,
+            tokio::process::Command::new(NVATTEST)
+                .args(["attest", "--device", "gpu", "--verifier", "local"])
+                .args(["--nonce", &nonce])
+                .output(),
+        )
+        .await
+        .context("GPU attestation timed out")?
+        .context("failed to run nvattest")?;
+        if !output.stderr.is_empty() {
+            info!("nvattest: {}", truncated_lossy(&output.stderr, 2048));
+        }
+        if !output.status.success() {
+            bail!(
+                "nvattest exited with {}: {}",
+                output.status,
+                truncated_lossy(&output.stderr, 512),
+            );
+        }
+        if let Err(err) = save_attestation_output(&output.stdout) {
+            warn!("failed to save GPU attestation output: {err:?}");
+        }
+        Ok(())
+    }
+
+    fn save_attestation_output(stdout: &[u8]) -> Result<()> {
+        let output_path = Path::new(ATTESTATION_OUTPUT);
+        if let Some(parent) = output_path.parent() {
+            fs::create_dir_all(parent)?;
+        }
+        safe_write(output_path, stdout)?;
+        fs::set_permissions(
+            output_path,
+            std::os::unix::fs::PermissionsExt::from_mode(0o600),
+        )?;
+        Ok(())
+    }
+
+    fn truncated_lossy(bytes: &[u8], limit: usize) -> String {
+        let text = String::from_utf8_lossy(bytes);
+        let text = text.trim();
+        match text.char_indices().nth(limit) {
+            Some((idx, _)) => format!("{}...", &text[..idx]),
+            None => text.to_string(),
+        }
+    }
+}
+
+impl Stage0<'_> {
+    /// Enforce `requirements.verify_gpu` (default true): attest an attached
+    /// NVIDIA GPU before continuing to key provisioning, or — when explicitly
+    /// disabled — set the GPU ready state without verification.
+    async fn setup_gpu(&self) -> Result<()> {
+        if !gpu::nvidia_gpu_present() {
+            return Ok(());
+        }
+        if !self.shared.app_compose.verify_gpu() {
+            warn!("requirements.verify_gpu is false; setting GPU ready state without attestation");
+            // Best-effort: a GPU with CC mode off has no ready state to set.
+            if let Err(err) = gpu::set_gpu_ready_state() {
+                warn!("failed to set GPU ready state: {err:?}");
+            }
+            return Ok(());
+        }
+        self.vmm.notify_q("boot.progress", "attesting GPU").await;
+        info!("verifying GPU TEE attestation");
+        gpu::attest_gpu().await?;
+        gpu::set_gpu_ready_state()?;
+        info!("GPU TEE attestation succeeded");
+        Ok(())
+    }
+}
+
 pub async fn cmd_gateway_refresh(args: GatewayRefreshArgs) -> Result<()> {
     let host_shared_dir = args.work_dir.join(HOST_SHARED_DIR_NAME);
     let shared = HostShared::load(host_shared_dir.as_path()).with_context(|| {