feat(vmm): select TEE simulator per instance

Signed-off-by: Kevin Wang <wy721@qq.com>

Author: kvinwang

docs/development-without-tee.md

@@ -89,7 +89,6 @@ qemu_path = "/usr/bin/qemu-system-x86_64"
 mode = "user"
 
 [cvm.tee_simulator]
-platform = "dstack-tdx"
 # Development credential only. Use a different random 32-byte hex seed for
 # each isolated test environment; never use it for production secrets.
 mock_attestation_seed = "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef"
@@ -99,9 +98,11 @@ collateral_base_url = "http://10.0.2.2:18088"
 exe = "/home/USER/.local/bin/supervisor"
 ```
 
-The VMM writes these settings to the guest's development-only
-`.tee-simulator.json`. The guest starts the simulator only when that file is
-present in the host share.
+These node-local settings provide development credentials and collateral URLs;
+they do not enable simulation globally. Each deployment selects its simulated
+platform with `--simulated-tee`. The VMM then writes an instance-specific
+`.tee-simulator.json`, and the guest starts the simulator only when that file
+is present in the host share.
 
 Start the VMM from a stable working directory because its default API socket is
 relative to that directory:
@@ -153,15 +154,16 @@ dstack compose \
   --output app-compose.json
 ```
 
-Deploy with the development image, no KMS, and no TEE:
+Deploy with the development image, no KMS, and an instance-specific simulated
+TEE platform:
 
 ```bash
 dstack deploy \
   --name swtpm-persistence \
   --image dstack-dev-<version> \
   --compose app-compose.json \
   --vcpu 2 --memory 3G --disk 10G \
-  --no-tee
+  --simulated-tee dstack-tdx
 ```
 
 Successful output contains a VM ID. `dstack info VM_ID` should eventually show

dstack/vmm/rpc/proto/vmm_rpc.proto

@@ -119,6 +119,9 @@ message VmConfiguration {
   // Per-VM networking overrides. If set, wins over singular networking.
   repeated NetworkingConfig networks = 19;
   reserved 20;
+  // Development-only TEE ABI to simulate for this VM. When set, the VMM
+  // launches a non-TEE VM and provides an instance-specific simulator config.
+  optional string simulated_tee = 21;
 }
 
 // Per-VM networking configuration.

dstack/vmm/src/app.rs

@@ -110,6 +110,8 @@ pub struct Manifest {
     pub gateway_urls: Vec<String>,
     #[serde(default)]
     pub no_tee: bool,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub simulated_tee: Option<dstack_types::TeeVariant>,
     #[serde(default, skip_serializing_if = "Vec::is_empty")]
     pub networks: Vec<Networking>,
     #[serde(default)]
@@ -1108,7 +1110,8 @@ impl App {
         )?;
         fs::write(shared_dir.join(SYS_CONFIG), &sys_config_str)
             .context("Failed to write vm config")?;
-        sync_tee_simulator_config(&shared_dir, cfg.cvm.tee_simulator.as_ref(), &sys_config_str)?;
+        let simulator_config = simulator_config_for_manifest(&self.config.cvm, &manifest)?;
+        sync_tee_simulator_config(&shared_dir, simulator_config.as_ref(), &sys_config_str)?;
         Ok(())
     }
 
@@ -1244,7 +1247,22 @@ fn rotate_serial_log(work_dir: &VmWorkDir, max_bytes: u64) {
     }
 }
 
-fn sync_tee_simulator_config(
+pub(crate) fn simulator_config_for_manifest(
+    cvm: &crate::config::CvmConfig,
+    manifest: &Manifest,
+) -> Result<Option<dstack_types::TeeSimulatorConfig>> {
+    let Some(platform) = manifest.simulated_tee else {
+        return Ok(None);
+    };
+    let mut config = cvm
+        .tee_simulator
+        .clone()
+        .context("tee simulator credentials are not configured on this VMM")?;
+    config.platform = platform;
+    Ok(Some(config))
+}
+
+pub(crate) fn sync_tee_simulator_config(
     shared_dir: &Path,
     simulator_config: Option<&dstack_types::TeeSimulatorConfig>,
     sys_config: &str,
@@ -1520,6 +1538,28 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn instance_platform_overrides_node_simulator_template() -> Result<()> {
+        let mut config = test_tdx_config()?;
+        config.cvm.tee_simulator = Some(dstack_types::TeeSimulatorConfig {
+            platform: dstack_types::TeeVariant::DstackTdx,
+            mock_attestation_seed: Some("11".repeat(32)),
+            ..Default::default()
+        });
+        let mut manifest = test_manifest(2048);
+        assert!(simulator_config_for_manifest(&config.cvm, &manifest)?.is_none());
+        manifest.simulated_tee = Some(dstack_types::TeeVariant::DstackNitroEnclave);
+
+        let resolved = simulator_config_for_manifest(&config.cvm, &manifest)?
+            .context("simulator config should be enabled")?;
+
+        assert_eq!(
+            resolved.platform,
+            dstack_types::TeeVariant::DstackNitroEnclave
+        );
+        Ok(())
+    }
+
     #[test]
     fn gpu_config_has_gpus_only_when_resolved_gpu_list_is_non_empty() {
         assert!(!GpuConfig::default().has_gpus());
@@ -1682,6 +1722,7 @@ mod tests {
             kms_urls: vec![],
             gateway_urls: vec![],
             no_tee: false,
+            simulated_tee: None,
             networks: vec![],
             volumes: vec![],
         }
@@ -1969,6 +2010,7 @@ mod tests {
             kms_urls: vec![],
             gateway_urls: vec![],
             no_tee: false,
+            simulated_tee: None,
             networks: vec![],
             volumes: vec![],
         };

dstack/vmm/src/app/qemu.rs

@@ -19,9 +19,7 @@ use crate::{
 };
 use anyhow::{bail, Context, Result};
 use bon::Builder;
-use dstack_types::{
-    shared_filenames::HOST_SHARED_DISK_LABEL, KeyProviderKind, TeeSimulatorConfig, TeeVariant,
-};
+use dstack_types::{shared_filenames::HOST_SHARED_DISK_LABEL, KeyProviderKind, TeeVariant};
 use fs_err as fs;
 use nix::unistd::User;
 use serde::Serialize;
@@ -179,15 +177,16 @@ struct PreparedVolume {
     source: String,
 }
 
-fn needs_qemu_swtpm(key_provider: KeyProviderKind, simulator: Option<&TeeSimulatorConfig>) -> bool {
-    if !matches!(key_provider, KeyProviderKind::Tpm) {
-        return false;
-    }
-    !matches!(
-        simulator.map(|config| config.platform),
+fn simulated_platform_provides_tpm(platform: Option<TeeVariant>) -> bool {
+    matches!(
+        platform,
         Some(TeeVariant::DstackGcpTdx | TeeVariant::DstackAwsNitroTpm)
     )
 }
+
+fn needs_qemu_swtpm(key_provider: KeyProviderKind, simulated_tee: Option<TeeVariant>) -> bool {
+    matches!(key_provider, KeyProviderKind::Tpm) && !simulated_platform_provides_tpm(simulated_tee)
+}
 struct PreparedQemuLaunch {
     workdir: VmWorkDir,
     platform: CvmPlatform,
@@ -246,7 +245,7 @@ impl PreparedQemuLaunch {
             None
         };
         let (swtpm_socket, swtpm_path) =
-            if needs_qemu_swtpm(app_compose.key_provider(), cfg.tee_simulator.as_ref()) {
+            if needs_qemu_swtpm(app_compose.key_provider(), vm.manifest.simulated_tee) {
                 let swtpm_path = which::which("swtpm")
                     .context("tpm key provider requested but swtpm is not installed")?;
                 let state_dir = workdir.swtpm_state_dir();
@@ -989,20 +988,19 @@ mod tests {
     use crate::app::image::{Image, ImageInfo};
     use crate::app::{GpuConfig, Manifest, PortMapping, VmVolume, VmWorkDir};
     use crate::config::{Config, CvmPlatform, Protocol, DEFAULT_CONFIG};
-    use dstack_types::{KeyProviderKind, TeeSimulatorConfig, TeeVariant};
+    use dstack_types::{KeyProviderKind, TeeVariant};
 
     #[test]
     fn qemu_swtpm_is_omitted_when_simulator_provides_the_tpm() {
         for platform in [TeeVariant::DstackGcpTdx, TeeVariant::DstackAwsNitroTpm] {
-            let simulator = TeeSimulatorConfig {
-                platform,
-                ..Default::default()
-            };
-            assert!(!needs_qemu_swtpm(KeyProviderKind::Tpm, Some(&simulator)));
+            assert!(!needs_qemu_swtpm(KeyProviderKind::Tpm, Some(platform)));
+            assert!(!needs_qemu_swtpm(KeyProviderKind::Kms, Some(platform)));
         }
 
-        let tdx = TeeSimulatorConfig::default();
-        assert!(needs_qemu_swtpm(KeyProviderKind::Tpm, Some(&tdx)));
+        assert!(needs_qemu_swtpm(
+            KeyProviderKind::Tpm,
+            Some(TeeVariant::DstackTdx)
+        ));
         assert!(needs_qemu_swtpm(KeyProviderKind::Tpm, None));
         assert!(!needs_qemu_swtpm(KeyProviderKind::Kms, None));
     }
@@ -1070,6 +1068,7 @@ mod tests {
                 kms_urls: vec![],
                 gateway_urls: vec![],
                 no_tee: true,
+                simulated_tee: None,
                 networks: vec![],
                 volumes: vec![VmVolume {
                     source: "/does-not-exist/volume.img".into(),

dstack/vmm/src/app/vm_info.rs

@@ -167,6 +167,10 @@ impl VmInfo {
                     gateway_urls: custom_gateway_urls.clone(),
                     stopped,
                     no_tee,
+                    simulated_tee: self
+                        .manifest
+                        .simulated_tee
+                        .map(|platform| platform.as_str().to_string()),
                     networking: configured_networking,
                     networks: configured_networks,
                 })

dstack/vmm/src/config.rs

@@ -260,7 +260,8 @@ pub struct CvmConfig {
     /// The URL of the PCCS server
     #[serde(default)]
     pub pccs_url: String,
-    /// Development-image simulator configuration copied into guest sys-config.
+    /// Node-local credentials and collateral settings used when an individual
+    /// VM requests a simulated TEE platform.
     #[serde(default)]
     pub tee_simulator: Option<dstack_types::TeeSimulatorConfig>,
     /// Optional NVIDIA OCSP/RIM cache passed to guests in sys-config.

dstack/vmm/src/main_service.rs

@@ -189,6 +189,16 @@ pub fn create_manifest_from_vm_config(
     dstack_types::validate_verity_volumes(&verity_volumes).map_err(anyhow::Error::msg)?;
     let volumes = resolve_volumes(&verity_volumes, cvm_config)?;
 
+    let simulated_tee = request
+        .simulated_tee
+        .as_deref()
+        .map(str::parse)
+        .transpose()
+        .map_err(anyhow::Error::msg)?;
+    if simulated_tee.is_some() && cvm_config.tee_simulator.is_none() {
+        bail!("tee simulator credentials are not configured on this VMM");
+    }
+
     Ok(Manifest {
         id,
         name: request.name.clone(),
@@ -204,7 +214,8 @@ pub fn create_manifest_from_vm_config(
         gpus: Some(gpus),
         kms_urls: request.kms_urls.clone(),
         gateway_urls: request.gateway_urls.clone(),
-        no_tee: request.no_tee,
+        no_tee: request.no_tee || simulated_tee.is_some(),
+        simulated_tee,
         networks: networks_from_vm_config(&request, cvm_config)?,
         volumes,
     })
@@ -941,6 +952,7 @@ mod tests {
             gateway_urls: vec![],
             stopped: false,
             no_tee: false,
+            simulated_tee: None,
             networking: None,
             networks: vec![],
         }
@@ -954,6 +966,47 @@ mod tests {
         assert!(manifest.networks.is_empty());
     }
 
+    #[test]
+    fn simulated_tee_is_selected_per_instance_and_implies_no_tee() {
+        let mut request = test_vm_configuration();
+        request.simulated_tee = Some("dstack-amd-sev-snp".into());
+        let mut config = test_cvm_config();
+        config.tee_simulator = Some(dstack_types::TeeSimulatorConfig {
+            mock_attestation_seed: Some("11".repeat(32)),
+            ..Default::default()
+        });
+
+        let manifest = create_manifest_from_vm_config(request, &config).unwrap();
+
+        assert_eq!(
+            manifest.simulated_tee,
+            Some(dstack_types::TeeVariant::DstackAmdSevSnp)
+        );
+        assert!(manifest.no_tee);
+    }
+
+    #[test]
+    fn simulated_tee_requires_node_credentials() {
+        let mut request = test_vm_configuration();
+        request.simulated_tee = Some("dstack-tdx".into());
+
+        let err = create_manifest_from_vm_config(request, &test_cvm_config()).unwrap_err();
+
+        assert!(err
+            .to_string()
+            .contains("tee simulator credentials are not configured"));
+    }
+
+    #[test]
+    fn invalid_simulated_tee_is_rejected() {
+        let mut request = test_vm_configuration();
+        request.simulated_tee = Some("not-a-platform".into());
+
+        let err = create_manifest_from_vm_config(request, &test_cvm_config()).unwrap_err();
+
+        assert!(err.to_string().contains("unsupported TEE variant"));
+    }
+
     #[test]
     fn volume_extraction_keeps_other_compose_fields_opaque() -> Result<()> {
         assert!(extract_verity_volumes("not json")?.is_empty());

dstack/vmm/src/one_shot.rs

@@ -2,7 +2,10 @@
 //
 // SPDX-License-Identifier: Apache-2.0
 
-use crate::app::{make_sys_config, Image, VmConfig, VmWorkDir};
+use crate::app::{
+    make_sys_config, simulator_config_for_manifest, sync_tee_simulator_config, Image, VmConfig,
+    VmWorkDir,
+};
 use crate::config::Config;
 use crate::main_service;
 use anyhow::{Context, Result};
@@ -265,7 +268,13 @@ Compose file content (first 200 chars):
         app_compose.requirements.as_ref(),
     )?;
     let sys_config_path = vm_work_dir.shared_dir().join(".sys-config.json");
-    fs_err::write(&sys_config_path, sys_config_str).context("Failed to write sys config")?;
+    fs_err::write(&sys_config_path, &sys_config_str).context("Failed to write sys config")?;
+    let simulator_config = simulator_config_for_manifest(&config.cvm, &manifest)?;
+    sync_tee_simulator_config(
+        &vm_work_dir.shared_dir(),
+        simulator_config.as_ref(),
+        &sys_config_str,
+    )?;
 
     // Create vm-state.json with initial state
     vm_work_dir

dstack/vmm/src/vmm-cli.py

@@ -901,6 +901,8 @@ def create_vm(self, args) -> None:
             "stopped": args.stopped,
             "no_tee": args.no_tee,
         }
+        if args.simulated_tee:
+            params["simulated_tee"] = args.simulated_tee
         if args.swap is not None:
             swap_bytes = max(0, int(round(args.swap)) * 1024 * 1024)
             if swap_bytes > 0:
@@ -1810,6 +1812,17 @@ def _patched_format_help():
         help="Force-enable Intel TDX (default)",
     )
     deploy_parser.set_defaults(no_tee=False)
+    deploy_parser.add_argument(
+        "--simulated-tee",
+        choices=[
+            "dstack-tdx",
+            "dstack-gcp-tdx",
+            "dstack-amd-sev-snp",
+            "dstack-nitro-enclave",
+            "dstack-aws-nitro-tpm",
+        ],
+        help="Simulate the selected TEE ABI for this VM (development images only)",
+    )
     deploy_parser.add_argument(
         "--net",
         choices=["bridge", "user"],