Merge branch 'configid'

Author: kvinwang

Cargo.lock

@@ -27,6 +27,8 @@ pub struct AppCompose {
     pub local_key_provider_enabled: bool,
     #[serde(default)]
     pub key_provider: Option<KeyProviderKind>,
+    #[serde(default, with = "hex_bytes")]
+    pub key_provider_id: Vec<u8>,
     #[serde(default)]
     pub allowed_envs: Vec<String>,
     #[serde(default)]

dstack-types/src/lib.rs

@@ -5,6 +5,7 @@ use std::{
 };
 
 use anyhow::{anyhow, bail, Context, Result};
+use dcap_qvl::quote::{Quote, Report};
 use dstack_kms_rpc as rpc;
 use dstack_types::{
     shared_filenames::{
@@ -17,7 +18,7 @@ use fs_err as fs;
 use ra_rpc::client::{CertInfo, RaClient, RaClientConfig};
 use ra_tls::cert::generate_ra_cert;
 use serde::{Deserialize, Serialize};
-use tdx_attest::extend_rtmr3;
+use tdx_attest::{extend_rtmr3, get_quote};
 use tracing::{info, warn};
 
 use crate::{
@@ -364,6 +365,7 @@ impl<'a> Stage0<'a> {
             let (_, ca_pem) = x509_parser::pem::parse_x509_pem(keys.ca_cert.as_bytes())
                 .context("Failed to parse ca cert")?;
             let x509 = ca_pem.parse_x509().context("Failed to parse ca cert")?;
+            self.ensure_provider_id_matches(x509.public_key().raw)?;
             let id = hex::encode(x509.public_key().raw);
             let provider_info = KeyProviderInfo::new("kms".into(), id);
             emit_key_provider_info(&provider_info)?;
@@ -373,6 +375,20 @@ impl<'a> Stage0<'a> {
         Ok(())
     }
 
+    fn ensure_provider_id_matches(&self, provider_id: &[u8]) -> Result<()> {
+        let expected_key_provider_id = &self.shared.app_compose.key_provider_id;
+        if expected_key_provider_id.is_empty() {
+            return Ok(());
+        };
+        if expected_key_provider_id != provider_id {
+            bail!(
+                "Unexpected key provider id: {:?}, expected: {:?}",
+                hex_fmt::HexFmt(provider_id),
+                hex_fmt::HexFmt(expected_key_provider_id)
+            );
+        }
+        Ok(())
+    }
     async fn get_keys_from_local_key_provider(&self) -> Result<()> {
         info!("Getting keys from local key provider");
         let provision = self
@@ -387,6 +403,7 @@ impl<'a> Stage0<'a> {
         fs::write(self.app_keys_file(), keys_json).context("Failed to write app keys")?;
 
         // write to RTMR
+        self.ensure_provider_id_matches(&provision.mr)?;
         let provider_info = KeyProviderInfo::new("local-sgx".into(), hex::encode(provision.mr));
         emit_key_provider_info(&provider_info)?;
         Ok(())
@@ -470,6 +487,8 @@ impl<'a> Stage0<'a> {
         let key_provider = self.shared.app_compose.key_provider();
         let mut instance_info = self.shared.instance_info.clone();
 
+        validate_compose_hash(&compose_hash).context("Failed to validate compose hash")?;
+
         if instance_info.app_id.is_empty() {
             instance_info.app_id = truncated_compose_hash.to_vec();
         }
@@ -536,6 +555,25 @@ impl<'a> Stage0<'a> {
     }
 }
 
+fn validate_compose_hash(compose_hash: &[u8]) -> Result<()> {
+    // If configid is not all zero, use it as compose_hash
+    let (_, quote) = get_quote(&[0u8; 64], None).context("Failed to get quote")?;
+    let quote = Quote::parse(&quote).context("Failed to parse quote")?;
+    let configid = match quote.report {
+        Report::SgxEnclave(_report) => bail!("SGX quote is not supported"),
+        Report::TD10(report) => report.mr_config_id,
+        Report::TD15(report) => report.base.mr_config_id,
+    };
+    info!("mr_config_id: {}", hex_fmt::HexFmt(&configid));
+    if configid == [0u8; 48] {
+        return Ok(());
+    }
+    if &configid[..32] != compose_hash {
+        bail!("mr_config_id does not match compose hash");
+    }
+    Ok(())
+}
+
 impl Stage1<'_> {
     fn resolve(&self, path: &str) -> String {
         path.to_string()

dstack-util/src/system_setup.rs

@@ -42,6 +42,7 @@ key-provider-client.workspace = true
 dstack-types.workspace = true
 hex_fmt.workspace = true
 lspci.workspace = true
+base64.workspace = true
 
 [dev-dependencies]
 insta.workspace = true

vmm/Cargo.toml

@@ -14,6 +14,7 @@ use std::{
 
 use super::{image::Image, GpuConfig, VmState};
 use anyhow::{bail, Context, Result};
+use base64::prelude::*;
 use bon::Builder;
 use dstack_types::{
     shared_filenames::{APP_COMPOSE, ENCRYPTED_ENV, INSTANCE_INFO, USER_CONFIG},
@@ -301,7 +302,19 @@ impl VmConfig {
         command
             .arg("-machine")
             .arg("q35,kernel-irqchip=split,confidential-guest-support=tdx,hpet=off");
-        command.arg("-object").arg("tdx-guest,id=tdx");
+
+        let tdx_object = if cfg.use_mrconfigid {
+            let mut compose_hash = workdir
+                .app_compose_hash()
+                .context("Failed to get compose hash")?;
+            compose_hash.resize(48, 0);
+            let mrconfigid = BASE64_STANDARD.encode(&compose_hash);
+            format!("tdx-guest,id=tdx,mrconfigid={mrconfigid}")
+        } else {
+            "tdx-guest,id=tdx".to_string()
+        };
+        command.arg("-object").arg(tdx_object);
+
         command
             .arg("-device")
             .arg(format!("vhost-vsock-pci,guest-cid={}", self.cid));
@@ -611,6 +624,13 @@ impl VmWorkDir {
         self.shared_dir().join(APP_COMPOSE)
     }
 
+    pub fn app_compose_hash(&self) -> Result<Vec<u8>> {
+        use sha2::Digest;
+        let compose_path = self.app_compose_path();
+        let compose = fs::read(compose_path).context("Failed to read compose")?;
+        Ok(sha2::Sha256::new_with_prefix(&compose).finalize().to_vec())
+    }
+
     pub fn user_config_path(&self) -> PathBuf {
         self.shared_dir().join(USER_CONFIG)
     }

vmm/src/app/qemu.rs

@@ -120,6 +120,9 @@ pub struct CvmConfig {
 
     /// Auto restart configuration
     pub auto_restart: AutoRestartConfig,
+
+    /// Use mrconfigid instead of compose hash
+    pub use_mrconfigid: bool,
 }
 
 #[derive(Debug, Clone, Deserialize)]

vmm/src/config.rs

@@ -745,6 +745,11 @@ <h2>Deploy a new instance</h2>
                             </label>
                         </div>
                     </div>
+                    <div class="form-group">
+                        <label for="keyProviderId">Key Provider ID (optional)</label>
+                        <input id="keyProviderId" v-model="vmForm.key_provider_id" type="text"
+                            placeholder="Key provider ID if you want to bind to a specific key provider">
+                    </div>
 
                     <div class="form-group full-width" v-if="vmForm.kms_enabled">
                         <encrypted-env-editor :env-vars="vmForm.encryptedEnvs"></encrypted-env-editor>
@@ -1437,6 +1442,7 @@ <h3>Derive VM</h3>
                     app_id: '',
                     kms_enabled: true,
                     local_key_provider_enabled: false,
+                    key_provider_id: '',
                     gateway_enabled: true,
                     public_logs: true,
                     public_sysinfo: true,
@@ -1604,6 +1610,7 @@ <h3>Derive VM</h3>
                         "public_sysinfo": vmForm.value.public_sysinfo,
                         "public_tcbinfo": vmForm.value.public_tcbinfo,
                         "local_key_provider_enabled": vmForm.value.local_key_provider_enabled,
+                        "key_provider_id": vmForm.value.key_provider_id,
                         "allowed_envs": vmForm.value.encryptedEnvs.map(env => env.key),
                         "no_instance_id": !vmForm.value.gateway_enabled,
                         "secure_time": false,
@@ -2289,6 +2296,7 @@ <h3>Derive VM</h3>
                     () => vmForm.value.public_sysinfo,
                     () => vmForm.value.public_tcbinfo,
                     () => vmForm.value.local_key_provider_enabled,
+                    () => vmForm.value.key_provider_id,
                     () => vmForm.value.docker_config.enabled,
                     () => vmForm.value.docker_config.username,
                     () => vmForm.value.docker_config.token_key

vmm/src/console.html

@@ -404,6 +404,7 @@ def create_app_compose(self,
                            kms_enabled: bool,
                            gateway_enabled: bool,
                            local_key_provider_enabled: bool,
+                           key_provider_id: str,
                            public_logs: bool,
                            public_sysinfo: bool,
                            envs: Optional[Dict],
@@ -419,6 +420,7 @@ def create_app_compose(self,
             "kms_enabled": kms_enabled,
             "gateway_enabled": gateway_enabled,
             "local_key_provider_enabled": local_key_provider_enabled,
+            "key_provider_id": key_provider_id,
             "public_logs": public_logs,
             "public_sysinfo": public_sysinfo,
             "allowed_envs": [k for k in envs.keys()],
@@ -781,6 +783,8 @@ def main():
         '--gateway', action='store_true', help='Enable dstack-gateway')
     compose_parser.add_argument(
         '--local-key-provider', action='store_true', help='Enable local key provider')
+    compose_parser.add_argument(
+        '--key-provider-id', default=None, help='Key provider ID if you want to bind to a specific key provider')
     compose_parser.add_argument(
         '--public-logs', action='store_true', help='Enable public logs')
     compose_parser.add_argument(
@@ -872,6 +876,7 @@ def main():
             kms_enabled=args.kms,
             gateway_enabled=args.gateway,
             local_key_provider_enabled=args.local_key_provider,
+            key_provider_id=args.key_provider_id,
             public_logs=args.public_logs,
             public_sysinfo=args.public_sysinfo,
             envs=parse_env_file(args.env_file),

vmm/src/vmm-cli.py

@@ -31,6 +31,7 @@ max_allocable_memory_in_mb = 100_000 # MB
 qmp_socket = false
 # The user to run the VM as. If empty, the VM will be run as the current user.
 user = ""
+use_mrconfigid = true
 
 [cvm.port_mapping]
 enabled = false