feat: add native SEV-SNP hardware key provider (#1526)

AMD hosts cannot use dstack's SGX local key provider. Add `key_provider: sev-snp` to derive a hardware-bound seed inside the SNP guest and reuse the existing app-key and disk-unlock flow, so a KMS CVM can recover its encrypted state and stored root keys after restart.

Signed-off-by: Hang Yin <h4x3rotab@gmail.com>

Author: h4x3rotab

docs/amd-sev-snp.md

@@ -50,20 +50,55 @@ preflight does not replace firmware, kernel, QEMU, or OVMF validation.
 
 ## Install dstack on an SNP host
 
-Pull a current unified image and select the platform explicitly:
+Use a release containing the `sev-snp` key provider in the guest OS and VMM. The
+KMS container and client attestation verifiers must recognize the new provider
+metadata. Set `VERSION` to the OS release version and `KMS_IMAGE` to the KMS
+container reference, then install:
 
 ```bash
-VERSION=0.6.0
 sudo dstackup image pull --version "$VERSION"
-sudo dstackup install --platform amd-sev-snp --image "dstack-$VERSION"
+sudo dstackup install --platform amd-sev-snp \
+  --image "dstack-$VERSION" --kms-image "$KMS_IMAGE"
 sudo dstackup status
 ```
 
-The default `--platform auto` mode also selects SEV-SNP when the host CPU flags
-advertise `sev_snp`. Explicit selection is preferable while commissioning a
-host because it fails immediately when the required SNP device is absent.
+Confirm that installation reports `booted` for the KMS CVM. `dstackup` selects
+the `sev-snp` provider for the KMS CVM on AMD hosts. The default
+`--platform auto` also detects SEV-SNP from the host CPU flags. Explicit
+selection checks for the required SNP device before installation. Review the
+[KMS key-release policy](#kms-key-release-policy) before serving applications.
 
-Unlike the TDX path, SEV-SNP does not use the local SGX key provider.
+## Run KMS with persistent local keys
+
+The `sev-snp` provider lets a KMS CVM recover its encrypted disk and stored KMS
+roots after restart on the same AMD hardware. For a manually managed KMS CVM,
+set this field in its app-compose document:
+
+```json
+{"key_provider": "sev-snp"}
+```
+
+The deployment CLI also accepts `dstack deploy --key-provider sev-snp`. Application
+CVMs continue using `key_provider: kms` and the existing KMS API.
+
+The guest derives a seed through `SNP_GET_DERIVED_KEY` using legacy VCEK at
+VMPL0. The seed binds to launch measurement, guest policy, HOST_DATA, launch
+identity, and a pinned TCB. The guest requires debugging disabled and encrypted
+storage, and verifies the configuration binding before deriving the seed.
+
+Before initializing a new disk, the guest persists public derivation parameters
+in `.instance_info` through the host API. Preserve this file, the encrypted
+disk, and the original launch configuration together. Missing parameters for
+an existing disk, unknown versions, or derivation/unlock failures stop startup.
+
+Recovery requires the original hardware and bound launch inputs. For OS or
+compose changes, or hardware replacement, preserve KMS roots through the
+existing onboarding/replication flow. Firmware updates must support the stored
+derivation TCB and satisfy attestation policy.
+
+The entire KMS CVM, including its OS, is trusted. Protect KMS root and sealing
+keys by keeping the KMS OS and workload under operator control. Application
+key release continues to follow the existing KMS authorization policy.
 
 ## Attestation and image identity
 

dstack/Cargo.lock

@@ -253,6 +253,7 @@ schnorrkel = "0.11.4"
 sha2 = { version = "0.10.8", default-features = false }
 sha3 = "0.10.8"
 subtle = "2"
+zeroize = "1.8"
 blake2 = "0.10.6"
 tokio-rustls = { version = "0.26.2", features = ["ring"] }
 x25519-dalek = { version = "2.0.1", features = ["static_secrets"] }

dstack/Cargo.toml

@@ -213,6 +213,7 @@ impl CertRequestClient {
         match &keys.key_provider {
             KeyProvider::None { key }
             | KeyProvider::Local { key, .. }
+            | KeyProvider::SevSnp { key }
             | KeyProvider::Tpm { key, .. } => {
                 let ca = CaCert::new(keys.ca_cert.clone(), key.clone())
                     .context("Failed to create CA")?;

dstack/cert-client/src/lib.rs

@@ -5,6 +5,7 @@
 //! build the app-compose manifest — the JSON document the VMM hashes (to derive
 //! the app id) and deploys. The raw docker-compose YAML is embedded as a string.
 
+use dstack_types::KeyProviderKind;
 use serde_json::json;
 
 /// build a minimal app-compose manifest from a docker-compose YAML body
@@ -34,7 +35,11 @@ pub fn build_app_compose_with_runtime(
     build_app_compose_with_runtime_and_volumes(
         name,
         docker_compose_yaml,
-        kms_enabled,
+        if kms_enabled {
+            KeyProviderKind::Kms
+        } else {
+            KeyProviderKind::None
+        },
         runner,
         snapshotter,
         &[],
@@ -45,7 +50,7 @@ pub fn build_app_compose_with_runtime(
 pub fn build_app_compose_with_runtime_and_volumes(
     name: &str,
     docker_compose_yaml: &str,
-    kms_enabled: bool,
+    key_provider: KeyProviderKind,
     runner: &str,
     snapshotter: Option<&str>,
     verity_volumes: &[dstack_types::VerityVolume],
@@ -55,7 +60,7 @@ pub fn build_app_compose_with_runtime_and_volumes(
         "name": name,
         "runner": runner,
         "docker_compose_file": docker_compose_yaml,
-        "kms_enabled": kms_enabled,
+        "kms_enabled": key_provider.is_kms(),
         "gateway_enabled": false,
         "local_key_provider_enabled": false,
         "public_logs": true,
@@ -74,7 +79,48 @@ pub fn build_app_compose_with_runtime_and_volumes(
     if !verity_volumes.is_empty() {
         manifest["verity_volumes"] = json!(verity_volumes);
     }
+    // Preserve legacy KMS/None manifest bytes, which determine app identity.
+    if !matches!(key_provider, KeyProviderKind::Kms | KeyProviderKind::None) {
+        manifest["key_provider"] = json!(key_provider);
+    }
     // pretty-print via Value's Display (`{:#}`) — infallible, and byte-identical
     // to serde_json::to_string_pretty (avoids an expect on an unfailable Result).
     format!("{manifest:#}")
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn typed_provider_matches_previous_cli_manifest_encoding() {
+        let volumes = [dstack_types::VerityVolume {
+            source: "test.img".into(),
+            verity_root: [1; 32],
+            target: "/test".into(),
+        }];
+        for volumes in [&[][..], &volumes[..]] {
+            let build = |provider| {
+                build_app_compose_with_runtime_and_volumes(
+                    "test",
+                    "",
+                    provider,
+                    "nerdctl-compose",
+                    Some("soci"),
+                    volumes,
+                )
+            };
+            for provider in [
+                KeyProviderKind::SevSnp,
+                KeyProviderKind::Local,
+                KeyProviderKind::Tpm,
+            ] {
+                // The old CLI serialized a legacy manifest, then inserted the provider.
+                let mut legacy: serde_json::Value =
+                    serde_json::from_str(&build(KeyProviderKind::None)).unwrap();
+                legacy["key_provider"] = json!(provider);
+                assert_eq!(build(provider), format!("{legacy:#}"));
+            }
+        }
+    }
+}

dstack/crates/dstack-cli-core/src/compose.rs

@@ -210,10 +210,8 @@ pub fn auth_allowlist_json(cfg: &HostConfig) -> String {
 /// compose hash.
 pub const DEFAULT_KMS_IMAGE: &str = "dstacktee/dstack-kms:0.5.11@sha256:84b793feed825a5b5e70d04386e931e0e110461492793f17ab2128e39808d989";
 
-/// build the KMS-in-CVM app-compose manifest. An init script writes the
-/// rendered `kms.toml` into the guest and the KMS container mounts it. On TDX
-/// the CVM uses the SGX local key provider to seal the KMS root key; AMD
-/// SEV-SNP has no such provider, so it's disabled there.
+/// Build the KMS app-compose manifest with an embedded `kms.toml` init script.
+/// Select SGX local provisioning on TDX and native SNP provisioning on AMD.
 pub fn kms_app_compose(kms_toml: &str, kms_image: &str, platform: Platform) -> String {
     let docker_compose = format!(
         r#"services:
@@ -234,7 +232,7 @@ volumes:
     let init_script = format!(
         "mkdir -p /dstack/kms-config\ncat > /dstack/kms-config/kms.toml <<'KMSTOML'\n{kms_toml}\nKMSTOML\ntrue\n"
     );
-    let manifest = json!({
+    let mut manifest = json!({
         "manifest_version": 2,
         "name": "dstack-kms",
         "runner": "docker-compose",
@@ -250,6 +248,9 @@ volumes:
         "secure_time": false,
         "allowed_envs": []
     });
+    if platform == Platform::AmdSevSnp {
+        manifest["key_provider"] = json!("sev-snp");
+    }
     // infallible pretty-print via Value's Display; see compose::build_app_compose.
     format!("{manifest:#}")
 }
@@ -419,7 +420,7 @@ address = "vsock:2"
 port = {host_api_port}
 
 [key_provider]
-enabled = true
+enabled = {use_sgx_key_provider}
 address = "{kp_addr}"
 port = {kp_port}
 "#,
@@ -448,6 +449,7 @@ port = {kp_port}
         host_api_port = r.host_api_port,
         kp_addr = r.key_provider_addr,
         kp_port = r.key_provider_port,
+        use_sgx_key_provider = r.platform == Platform::Tdx,
         auth_enabled = r.auth_enabled,
         auth_tokens = if r.auth_token.is_empty() {
             String::new()
@@ -525,11 +527,12 @@ mod tests {
         assert!(tdx_vmm.contains(r#"platform = "tdx""#));
         assert!(tdx_vmm.contains(r#"host_share_mode = "9p""#));
         assert!(tdx_vmm.contains("use_mrconfigid = false"));
+        assert!(tdx_vmm.contains("[key_provider]\nenabled = true"));
         toml::from_str::<toml::Value>(&tdx_vmm).expect("tdx vmm.toml valid");
         assert!(kms_app_compose("x", "img", Platform::Tdx)
             .contains(r#""local_key_provider_enabled": true"#));
 
-        // SNP: key-release gate set, vhd share, mrconfigid on, no local provider.
+        // SNP: key-release gate set, vhd share, mrconfigid on, native sealing.
         let snp = kms_toml(&HostConfig {
             platform: Platform::AmdSevSnp,
             ..Default::default()
@@ -543,9 +546,14 @@ mod tests {
         assert!(snp_vmm.contains(r#"platform = "amd-sev-snp""#));
         assert!(snp_vmm.contains(r#"host_share_mode = "vhd""#));
         assert!(snp_vmm.contains("use_mrconfigid = true"));
+        assert!(snp_vmm.contains("[key_provider]\nenabled = false"));
         toml::from_str::<toml::Value>(&snp_vmm).expect("snp vmm.toml valid");
         assert!(kms_app_compose("x", "img", Platform::AmdSevSnp)
             .contains(r#""local_key_provider_enabled": false"#));
+        let snp_compose: serde_json::Value =
+            serde_json::from_str(&kms_app_compose("x", "img", Platform::AmdSevSnp)).unwrap();
+        assert_eq!(snp_compose["key_provider"], "sev-snp");
+        assert_eq!(snp_compose["kms_enabled"], false);
     }
 
     #[test]