kms: Support for the new mr_image model

Author: kvinwang

Cargo.lock

@@ -39,6 +39,7 @@ members = [
     "lspci",
     "sdk/rust",
     "sodiumbox",
+    "serde-duration",
 ]
 resolver = "2"
 
@@ -66,6 +67,7 @@ dstack-types = { path = "dstack-types" }
 cert-client = { path = "cert-client" }
 lspci = { path = "lspci" }
 sodiumbox = { path = "sodiumbox" }
+serde-duration = { path = "serde-duration" }
 
 # Core dependencies
 anyhow = "1.0.97"

Cargo.toml

@@ -10,8 +10,13 @@ use ra_tls::{
 use tdx_attest::{eventlog::read_event_logs, get_quote};
 
 pub enum CertRequestClient {
-    Local { ca: CaCert },
-    Kms { client: KmsClient<RaClient> },
+    Local {
+        ca: CaCert,
+    },
+    Kms {
+        client: KmsClient<RaClient>,
+        vm_config: String,
+    },
 }
 
 impl CertRequestClient {
@@ -27,11 +32,13 @@ impl CertRequestClient {
                     .context("Failed to sign certificate")?;
                 Ok(vec![cert.pem(), ca.pem_cert.clone()])
             }
-            CertRequestClient::Kms { client } => {
+            CertRequestClient::Kms { client, vm_config } => {
                 let response = client
                     .sign_cert(SignCertRequest {
+                        api_version: 1,
                         csr: csr.to_vec(),
                         signature: signature.to_vec(),
+                        vm_config: vm_config.clone(),
                     })
                     .await?;
                 Ok(response.certificate_chain)
@@ -42,11 +49,15 @@ impl CertRequestClient {
     pub async fn get_root_ca(&self) -> Result<String> {
         match self {
             CertRequestClient::Local { ca } => Ok(ca.pem_cert.clone()),
-            CertRequestClient::Kms { client } => Ok(client.get_meta().await?.ca_cert),
+            CertRequestClient::Kms { client, .. } => Ok(client.get_meta().await?.ca_cert),
         }
     }
 
-    pub async fn create(keys: &AppKeys, pccs_url: Option<&str>) -> Result<CertRequestClient> {
+    pub async fn create(
+        keys: &AppKeys,
+        pccs_url: Option<&str>,
+        vm_config: String,
+    ) -> Result<CertRequestClient> {
         match &keys.key_provider {
             KeyProvider::Local { key } => {
                 let ca = CaCert::new(keys.ca_cert.clone(), key.clone())
@@ -74,7 +85,7 @@ impl CertRequestClient {
                     .into_client()
                     .context("Failed to create RA client")?;
                 let client = KmsClient::new(ra_client);
-                Ok(CertRequestClient::Kms { client })
+                Ok(CertRequestClient::Kms { client, vm_config })
             }
         }
     }

cert-client/src/lib.rs

@@ -118,8 +118,16 @@ pub struct SysConfig {
     pub pccs_url: Option<String>,
     pub docker_registry: Option<String>,
     pub host_api_url: String,
-    #[serde(with = "hex_bytes", default)]
+    // JSON serialized VmConfig
+    pub vm_config: String,
+}
+
+#[derive(Deserialize, Serialize, Debug, Clone)]
+pub struct VmConfig {
+    #[serde(with = "hex_bytes")]
     pub mr_image: Vec<u8>,
+    pub cpu_count: u32,
+    pub memory_size: u64,
 }
 
 #[derive(Serialize, Deserialize, Debug, Clone)]

dstack-types/src/lib.rs

@@ -5,7 +5,7 @@ use std::{
 };
 
 use anyhow::{anyhow, bail, Context, Result};
-use dstack_kms_rpc::GetAppKeyRequest;
+use dstack_kms_rpc as rpc;
 use dstack_types::{
     shared_filenames::{
         APP_COMPOSE, APP_KEYS, DECRYPTED_ENV, DECRYPTED_ENV_JSON, ENCRYPTED_ENV,
@@ -319,9 +319,16 @@ impl<'a> Stage0<'a> {
             .context("Failed to create client")?;
         let kms_client = dstack_kms_rpc::kms_client::KmsClient::new(ra_client);
         let response = kms_client
-            .get_app_key(GetAppKeyRequest {})
+            .get_app_key(rpc::GetAppKeyRequest {
+                api_version: 1,
+                vm_config: self.shared.sys_config.vm_config.clone(),
+            })
             .await
             .context("Failed to get app key")?;
+
+        extend_rtmr3("mr-image", &response.mr_image)
+            .context("Failed to extend mr-image to RTMR3")?;
+
         let keys = AppKeys {
             ca_cert: tmp_ca.ca_cert,
             disk_crypt_key: response.disk_crypt_key,
@@ -659,10 +666,13 @@ impl Stage1<'_> {
             usage_client_auth: true,
             ext_quote: true,
         };
-        let cert_client =
-            CertRequestClient::create(&self.keys, self.shared.sys_config.pccs_url.as_deref())
-                .await
-                .context("Failed to create cert client")?;
+        let cert_client = CertRequestClient::create(
+            &self.keys,
+            self.shared.sys_config.pccs_url.as_deref(),
+            self.shared.sys_config.vm_config.clone(),
+        )
+        .await
+        .context("Failed to create cert client")?;
         let client_key =
             KeyPair::generate_for(&PKCS_ECDSA_P256_SHA256).context("Failed to generate key")?;
         let client_certs = cert_client

dstack-util/src/system_setup.rs

@@ -41,6 +41,7 @@ dstack-guest-agent-rpc.workspace = true
 http-client = { workspace = true, features = ["prpc"] }
 sha2.workspace = true
 dstack-types.workspace = true
+serde-duration.workspace = true
 
 [target.'cfg(unix)'.dependencies]
 nix = { workspace = true, features = ["resource"] }

gateway/Cargo.toml

@@ -107,59 +107,6 @@ pub struct SyncConfig {
     pub bootnode: String,
 }
 
-mod serde_duration {
-    use serde::{Deserialize, Deserializer, Serializer};
-    use std::time::Duration;
-
-    pub fn serialize<S>(duration: &Duration, serializer: S) -> Result<S::Ok, S::Error>
-    where
-        S: Serializer,
-    {
-        if duration == &Duration::MAX {
-            return serializer.serialize_str("never");
-        }
-        let (value, unit) = if duration.as_secs() % (24 * 3600) == 0 {
-            (duration.as_secs() / (24 * 3600), "d")
-        } else if duration.as_secs() % 3600 == 0 {
-            (duration.as_secs() / 3600, "h")
-        } else if duration.as_secs() % 60 == 0 {
-            (duration.as_secs() / 60, "m")
-        } else {
-            (duration.as_secs(), "s")
-        };
-        serializer.serialize_str(&format!("{}{}", value, unit))
-    }
-
-    pub fn deserialize<'de, D>(deserializer: D) -> Result<Duration, D::Error>
-    where
-        D: Deserializer<'de>,
-    {
-        let s = String::deserialize(deserializer)?;
-        if s.is_empty() {
-            return Err(serde::de::Error::custom("Duration string cannot be empty"));
-        }
-        if s == "never" {
-            return Ok(Duration::MAX);
-        }
-        let (value, unit) = s.split_at(s.len() - 1);
-        let value = value.parse::<u64>().map_err(serde::de::Error::custom)?;
-
-        let seconds = match unit {
-            "s" => value,
-            "m" => value * 60,
-            "h" => value * 3600,
-            "d" => value * 24 * 3600,
-            _ => {
-                return Err(serde::de::Error::custom(
-                    "Invalid time unit. Use s, m, h, or d",
-                ))
-            }
-        };
-
-        Ok(Duration::from_secs(seconds))
-    }
-}
-
 #[derive(Debug, Clone, Deserialize)]
 pub struct Config {
     pub wg: WgConfig,

gateway/src/config.rs

@@ -10,6 +10,7 @@ log_level = "debug"
 app_name = ""
 keys_file = "/dstack/.host-shared/.appkeys.json"
 compose_file = "/dstack/.host-shared/app-compose.json"
+sys_config_file = "/dstack/.host-shared/.sys-config.json"
 public_logs = true
 public_sysinfo = true
 data_disks = ["/"]

guest-agent/dstack.toml

@@ -199,6 +199,8 @@ message WorkerInfo {
   string key_provider_info = 12;
   // Compose hash
   bytes compose_hash = 13;
+  // VM config
+  string vm_config = 14;
 }
 
 // The response to a WorkerInfo request

guest-agent/rpc/proto/agent_rpc.proto

@@ -22,6 +22,7 @@ pub struct Config {
     pub public_logs: bool,
     pub public_sysinfo: bool,
     pub compose_file: String,
+    pub sys_config_file: PathBuf,
     #[serde(default)]
     pub pccs_url: Option<String>,
     pub simulator: Simulator,