guest-agent: Request demo cert lazily

Author: kvinwang

Cargo.toml

@@ -218,3 +218,6 @@ serde_yaml2 = "0.1.2"
 
 luks2 = "0.5.0"
 scopeguard = "1.2.0"
+
+[profile.release]
+panic = "abort"

basefiles/docker.service.d/dstack-guest-agent.conf

@@ -0,0 +1,7 @@
+[Unit]
+Wants=dstack-guest-agent.service
+After=dstack-guest-agent.service
+
+[Service]
+TimeoutStartSec=0
+ExecStartPre=/bin/sh -c 'while ! systemctl is-active --quiet dstack-guest-agent.service || [ ! -S /var/run/dstack.sock ] || [ ! -S /var/run/tappd.sock ]; do sleep 1; done'

cert-client/src/lib.rs

@@ -68,15 +68,13 @@ impl CertRequestClient {
                     .context("Failed to create CA")?;
                 Ok(CertRequestClient::Local { ca })
             }
-            KeyProvider::Kms { url, .. } => {
-                let tmp_client =
-                    RaClient::new(url.into(), true).context("Failed to create RA client")?;
-                let tmp_client = KmsClient::new(tmp_client);
-                let tmp_ca = tmp_client
-                    .get_temp_ca_cert()
-                    .await
-                    .context("Failed to get temp CA cert")?;
-                let client_cert = generate_ra_cert(tmp_ca.temp_ca_cert, tmp_ca.temp_ca_key)
+            KeyProvider::Kms {
+                url,
+                tmp_ca_key,
+                tmp_ca_cert,
+                ..
+            } => {
+                let client_cert = generate_ra_cert(tmp_ca_cert.clone(), tmp_ca_key.clone())
                     .context("Failed to generate RA cert")?;
                 let ra_client = RaClientConfig::builder()
                     .remote_uri(url.clone())

dstack-types/src/lib.rs

@@ -192,6 +192,8 @@ pub enum KeyProvider {
         url: String,
         #[serde(with = "hex_bytes")]
         pubkey: Vec<u8>,
+        tmp_ca_key: String,
+        tmp_ca_cert: String,
     },
 }
 

dstack-util/src/system_setup.rs

@@ -561,7 +561,7 @@ impl<'a> Stage0<'a> {
                 .await
                 .context("Failed to get temp ca cert")?
         };
-        let cert_pair = generate_ra_cert(tmp_ca.temp_ca_cert, tmp_ca.temp_ca_key)?;
+        let cert_pair = generate_ra_cert(tmp_ca.temp_ca_cert.clone(), tmp_ca.temp_ca_key.clone())?;
         let ra_client = RaClientConfig::builder()
             .tls_no_check(false)
             .tls_built_in_root_certs(false)
@@ -619,6 +619,8 @@ impl<'a> Stage0<'a> {
             key_provider: KeyProvider::Kms {
                 url: kms_url,
                 pubkey: root_pubkey,
+                tmp_ca_key: tmp_ca.temp_ca_key,
+                tmp_ca_cert: tmp_ca.temp_ca_cert,
             },
         };
         Ok(keys)

guest-agent/src/rpc_service.rs

@@ -2,7 +2,7 @@
 //
 // SPDX-License-Identifier: Apache-2.0
 
-use std::sync::Arc;
+use std::sync::{Arc, RwLock};
 
 use anyhow::{Context, Result};
 use cert_client::CertRequestClient;
@@ -28,6 +28,7 @@ use ring::rand::{SecureRandom, SystemRandom};
 use serde_json::json;
 use sha3::{Digest, Keccak256};
 use tdx_attest::eventlog::read_event_logs;
+use tracing::error;
 
 use crate::config::Config;
 
@@ -41,23 +42,14 @@ struct AppStateInner {
     keys: AppKeys,
     vm_config: String,
     cert_client: CertRequestClient,
-    demo_cert: String,
+    demo_cert: RwLock<String>,
 }
 
-impl AppState {
-    pub async fn new(config: Config) -> Result<Self> {
-        let keys: AppKeys = serde_json::from_str(&fs::read_to_string(&config.keys_file)?)
-            .context("Failed to parse app keys")?;
-        let sys_config: SysConfig =
-            serde_json::from_str(&fs::read_to_string(&config.sys_config_file)?)
-                .context("Failed to parse VM config")?;
-        let vm_config = sys_config.vm_config;
-        let cert_client =
-            CertRequestClient::create(&keys, config.pccs_url.as_deref(), vm_config.clone())
-                .await
-                .context("Failed to create cert signer")?;
+impl AppStateInner {
+    async fn request_demo_cert(&self) -> Result<String> {
         let key = KeyPair::generate().context("Failed to generate demo key")?;
-        let demo_cert = cert_client
+        let demo_cert = self
+            .cert_client
             .request_cert(
                 &key,
                 CertConfig {
@@ -68,20 +60,55 @@ impl AppState {
                     usage_client_auth: true,
                     ext_quote: true,
                 },
-                config.simulator.enabled,
+                self.config.simulator.enabled,
             )
             .await
             .context("Failed to get app cert")?
             .join("\n");
-        Ok(Self {
+        Ok(demo_cert)
+    }
+}
+
+impl AppState {
+    fn maybe_request_demo_cert(&self) {
+        let state = self.inner.clone();
+        if !state.demo_cert.read().unwrap().is_empty() {
+            return;
+        }
+        tokio::spawn(async move {
+            match state.request_demo_cert().await {
+                Ok(demo_cert) => {
+                    *state.demo_cert.write().unwrap() = demo_cert;
+                }
+                Err(e) => {
+                    error!("Failed to request demo cert: {e}");
+                }
+            }
+        });
+    }
+
+    pub async fn new(config: Config) -> Result<Self> {
+        let keys: AppKeys = serde_json::from_str(&fs::read_to_string(&config.keys_file)?)
+            .context("Failed to parse app keys")?;
+        let sys_config: SysConfig =
+            serde_json::from_str(&fs::read_to_string(&config.sys_config_file)?)
+                .context("Failed to parse VM config")?;
+        let vm_config = sys_config.vm_config;
+        let cert_client =
+            CertRequestClient::create(&keys, config.pccs_url.as_deref(), vm_config.clone())
+                .await
+                .context("Failed to create cert signer")?;
+        let me = Self {
             inner: Arc::new(AppStateInner {
                 config,
                 keys,
                 cert_client,
-                demo_cert,
+                demo_cert: RwLock::new(String::new()),
                 vm_config,
             }),
-        })
+        };
+        me.maybe_request_demo_cert();
+        Ok(me)
     }
 
     pub fn config(&self) -> &Config {
@@ -136,6 +163,7 @@ pub async fn get_info(state: &AppState, external: bool) -> Result<AppInfo> {
     } else {
         state.inner.vm_config.clone()
     };
+    state.maybe_request_demo_cert();
     Ok(AppInfo {
         app_name: state.config().app_compose.name.clone(),
         app_id: app_info.app_id,
@@ -145,7 +173,7 @@ pub async fn get_info(state: &AppState, external: bool) -> Result<AppInfo> {
         os_image_hash: app_info.os_image_hash.clone(),
         key_provider_info: String::from_utf8(app_info.key_provider_info).unwrap_or_default(),
         compose_hash: app_info.compose_hash.clone(),
-        app_cert: state.inner.demo_cert.clone(),
+        app_cert: state.inner.demo_cert.read().unwrap().clone(),
         tcb_info,
         vm_config,
     })