Implement mr_config_id v2

cert-client/src/lib.rs

@@ -59,12 +59,12 @@ impl CertRequestClient {
         vm_config: String,
     ) -> Result<CertRequestClient> {
         match &keys.key_provider {
-            KeyProvider::Local { key } => {
+            KeyProvider::None { key } | KeyProvider::Local { key, .. } => {
                 let ca = CaCert::new(keys.ca_cert.clone(), key.clone())
                     .context("Failed to create CA")?;
                 Ok(CertRequestClient::Local { ca })
             }
-            KeyProvider::Kms { url } => {
+            KeyProvider::Kms { url, .. } => {
                 let tmp_client =
                     RaClient::new(url.into(), true).context("Failed to create RA client")?;
                 let tmp_client = KmsClient::new(tmp_client);

dstack-types/Cargo.toml

@@ -8,3 +8,4 @@ license.workspace = true
 [dependencies]
 serde = { workspace = true, features = ["derive"] }
 serde-human-bytes.workspace = true
+sha3.workspace = true

dstack-types/src/lib.rs

@@ -151,8 +151,37 @@ pub struct AppKeys {
 
 #[derive(Serialize, Deserialize, Debug, Clone)]
 pub enum KeyProvider {
-    Local { key: String },
-    Kms { url: String },
+    None {
+        key: String,
+    },
+    Local {
+        key: String,
+        #[serde(with = "hex_bytes")]
+        mr: Vec<u8>,
+    },
+    Kms {
+        url: String,
+        #[serde(with = "hex_bytes")]
+        pubkey: Vec<u8>,
+    },
+}
+
+impl KeyProvider {
+    pub fn kind(&self) -> KeyProviderKind {
+        match self {
+            KeyProvider::None { .. } => KeyProviderKind::None,
+            KeyProvider::Local { .. } => KeyProviderKind::Local,
+            KeyProvider::Kms { .. } => KeyProviderKind::Kms,
+        }
+    }
+
+    pub fn id(&self) -> &[u8] {
+        match self {
+            KeyProvider::None { .. } => &[],
+            KeyProvider::Local { mr, .. } => mr,
+            KeyProvider::Kms { pubkey, .. } => pubkey,
+        }
+    }
 }
 
 #[derive(Serialize, Deserialize, Debug, Clone)]
@@ -167,6 +196,7 @@ impl KeyProviderInfo {
     }
 }
 
+pub mod mr_config;
 pub mod shared_filenames;
 
 /// Get the address of the dstack agent

dstack-types/src/mr_config.rs

@@ -0,0 +1,50 @@
+use sha3::{Digest, Keccak256};
+
+use crate::KeyProviderKind;
+
+pub enum MrConfig<'a> {
+    V1 {
+        compose_hash: &'a [u8; 32],
+    },
+    V2 {
+        compose_hash: &'a [u8; 32],
+        app_id: &'a [u8; 20],
+        key_provider: KeyProviderKind,
+        key_provider_id: &'a [u8],
+    },
+}
+
+impl MrConfig<'_> {
+    pub fn to_mr_config_id(&self) -> [u8; 48] {
+        match self {
+            MrConfig::V1 { compose_hash } => {
+                let mut config_id = [0u8; 48];
+                config_id[0] = 1;
+                config_id[1..33].copy_from_slice(*compose_hash);
+                config_id
+            }
+            MrConfig::V2 {
+                compose_hash,
+                app_id,
+                key_provider,
+                key_provider_id,
+            } => {
+                let kp_kind = match key_provider {
+                    KeyProviderKind::None => 0_u8,
+                    KeyProviderKind::Local => 1,
+                    KeyProviderKind::Kms => 2,
+                };
+                let mut hasher = Keccak256::new();
+                hasher.update(compose_hash);
+                hasher.update(app_id);
+                hasher.update([kp_kind]);
+                hasher.update(key_provider_id);
+                let digest = hasher.finalize();
+                let mut config_id = [0u8; 48];
+                config_id[0] = 2;
+                config_id[1..33].copy_from_slice(digest.as_slice());
+                config_id
+            }
+        }
+    }
+}

dstack-util/src/main.rs

@@ -339,25 +339,26 @@ fn cmd_gen_app_keys(args: GenAppKeysArgs) -> Result<()> {
     let key = KeyPair::generate_for(&PKCS_ECDSA_P256_SHA256)?;
     let disk_key = KeyPair::generate_for(&PKCS_ECDSA_P256_SHA256)?;
     let k256_key = SigningKey::random(&mut rand::thread_rng());
-    let app_keys = make_app_keys(key, disk_key, k256_key, args.ca_level)?;
+    let app_keys = make_app_keys(key, disk_key, k256_key, args.ca_level, None)?;
     let app_keys = serde_json::to_string(&app_keys).context("Failed to serialize app keys")?;
     fs::write(&args.output, app_keys).context("Failed to write app keys")?;
     Ok(())
 }
 
-fn gen_app_keys_from_seed(seed: &[u8]) -> Result<AppKeys> {
+fn gen_app_keys_from_seed(seed: &[u8], mr: Option<Vec<u8>>) -> Result<AppKeys> {
     let key = derive_ecdsa_key_pair_from_bytes(seed, &["app-key".as_bytes()])?;
     let disk_key = derive_ecdsa_key_pair_from_bytes(seed, &["app-disk-key".as_bytes()])?;
     let k256_key = derive_ecdsa_key(seed, &["app-k256-key".as_bytes()], 32)?;
     let k256_key = SigningKey::from_bytes(&k256_key).context("Failed to parse k256 key")?;
-    make_app_keys(key, disk_key, k256_key, 1)
+    make_app_keys(key, disk_key, k256_key, 1, mr)
 }
 
 fn make_app_keys(
     app_key: KeyPair,
     disk_key: KeyPair,
     k256_key: SigningKey,
     ca_level: u8,
+    mr: Option<Vec<u8>>,
 ) -> Result<AppKeys> {
     use ra_tls::cert::CertRequest;
     let pubkey = app_key.public_key_der();
@@ -383,8 +384,14 @@ fn make_app_keys(
         k256_signature: vec![],
         gateway_app_id: "".to_string(),
         ca_cert: cert.pem(),
-        key_provider: KeyProvider::Local {
-            key: app_key.serialize_pem(),
+        key_provider: match mr {
+            Some(mr) => KeyProvider::Local {
+                mr,
+                key: app_key.serialize_pem(),
+            },
+            None => KeyProvider::None {
+                key: app_key.serialize_pem(),
+            },
         },
     })
 }