kms: Auto update certs on start

Author: kvinwang

dstack-util/src/system_setup.rs

@@ -163,7 +163,7 @@ impl HostShared {
             if src_size > max_size {
                 bail!("Source file {src} is too large, max size is {max_size} bytes");
             }
-            std::fs::copy(src_path, dst_path)?;
+            fs_err::copy(src_path, dst_path)?;
             Ok(())
         };
         cmd! {

kms/src/config.rs

@@ -14,6 +14,7 @@ const ROOT_CA_CERT: &str = "root-ca.crt";
 const ROOT_CA_KEY: &str = "root-ca.key";
 const RPC_CERT: &str = "rpc.crt";
 const RPC_KEY: &str = "rpc.key";
+const RPC_DOMAIN: &str = "rpc-domain";
 const K256_KEY: &str = "root-k256.key";
 const BOOTSTRAP_INFO: &str = "bootstrap-info.json";
 
@@ -60,6 +61,10 @@ impl KmsConfig {
         self.cert_dir.join(RPC_KEY)
     }
 
+    pub fn rpc_domain(&self) -> PathBuf {
+        self.cert_dir.join(RPC_DOMAIN)
+    }
+
     pub fn k256_key(&self) -> PathBuf {
         self.cert_dir.join(K256_KEY)
     }

kms/src/main.rs

@@ -9,7 +9,7 @@ use rocket::{
     response::content::RawHtml,
     Shutdown,
 };
-use tracing::info;
+use tracing::{info, warn};
 
 mod config;
 // mod ct_log;
@@ -93,6 +93,11 @@ async fn main() -> Result<()> {
         }
     }
 
+    info!("Updating certs");
+    if let Err(err) = onboard_service::update_certs(&config).await {
+        warn!("Failed to update certs: {err}");
+    };
+
     info!("Starting KMS");
     info!("Supported methods:");
     for method in main_service::rpc_methods() {

kms/src/onboard_service.rs

@@ -7,6 +7,7 @@ use dstack_kms_rpc::{
     onboard_server::{OnboardRpc, OnboardServer},
     BootstrapRequest, BootstrapResponse, OnboardRequest, OnboardResponse,
 };
+use fs_err as fs;
 use http_client::prpc::PrpcClient;
 use k256::ecdsa::SigningKey;
 use ra_rpc::{client::RaClient, CallContext, RpcCall};
@@ -102,6 +103,7 @@ struct Keys {
     ca_cert: Certificate,
     rpc_key: KeyPair,
     rpc_cert: Certificate,
+    rpc_domain: String,
 }
 
 impl Keys {
@@ -169,6 +171,7 @@ impl Keys {
             ca_cert,
             rpc_key,
             rpc_cert,
+            rpc_domain: domain.to_string(),
         })
     }
 
@@ -216,25 +219,59 @@ impl Keys {
     }
 
     fn store(&self, cfg: &KmsConfig) -> Result<()> {
-        // Store the temporary CA cert and key
-        safe_write(cfg.tmp_ca_cert(), self.tmp_ca_cert.pem())?;
-        safe_write(cfg.tmp_ca_key(), self.tmp_ca_key.serialize_pem())?;
+        self.store_keys(cfg)?;
+        self.store_certs(cfg)?;
+        safe_write(cfg.rpc_domain(), self.rpc_domain.as_bytes())?;
+        Ok(())
+    }
 
-        // Store the root CA cert and key
-        safe_write(cfg.root_ca_cert(), self.ca_cert.pem())?;
+    fn store_keys(&self, cfg: &KmsConfig) -> Result<()> {
+        safe_write(cfg.tmp_ca_key(), self.tmp_ca_key.serialize_pem())?;
         safe_write(cfg.root_ca_key(), self.ca_key.serialize_pem())?;
-
-        // Store the RPC cert and key
-        safe_write(cfg.rpc_cert(), self.rpc_cert.pem())?;
         safe_write(cfg.rpc_key(), self.rpc_key.serialize_pem())?;
-
-        // Store the ECDSA root key
         safe_write(cfg.k256_key(), self.k256_key.to_bytes())?;
+        Ok(())
+    }
 
+    fn store_certs(&self, cfg: &KmsConfig) -> Result<()> {
+        safe_write(cfg.tmp_ca_cert(), self.tmp_ca_cert.pem())?;
+        safe_write(cfg.root_ca_cert(), self.ca_cert.pem())?;
+        safe_write(cfg.rpc_cert(), self.rpc_cert.pem())?;
         Ok(())
     }
 }
 
+pub(crate) async fn update_certs(cfg: &KmsConfig) -> Result<()> {
+    // Read existing keys
+    let tmp_ca_key = KeyPair::from_pem(&fs::read_to_string(cfg.tmp_ca_key())?)?;
+    let ca_key = KeyPair::from_pem(&fs::read_to_string(cfg.root_ca_key())?)?;
+    let rpc_key = KeyPair::from_pem(&fs::read_to_string(cfg.rpc_key())?)?;
+
+    // Read k256 key
+    let k256_key_bytes = fs::read(cfg.k256_key())?;
+    let k256_key = SigningKey::from_slice(&k256_key_bytes)?;
+
+    let domain = fs::read_to_string(cfg.rpc_domain())?;
+    let domain = domain.trim();
+
+    // Regenerate certificates using existing keys
+    let keys = Keys::from_keys(
+        tmp_ca_key,
+        ca_key,
+        rpc_key,
+        k256_key,
+        domain,
+        cfg.onboard.quote_enabled,
+    )
+    .await
+    .context("Failed to regenerate certificates")?;
+
+    // Write the new certificates to files
+    keys.store_certs(cfg)?;
+
+    Ok(())
+}
+
 pub(crate) async fn bootstrap_keys(cfg: &KmsConfig) -> Result<()> {
     let keys = Keys::generate(
         &cfg.onboard.auto_bootstrap_domain,

supervisor/src/main.rs

@@ -98,7 +98,7 @@ async fn async_main(args: Args) -> Result<()> {
     if let Some(uds) = args.uds {
         mk_parents(&uds)?;
         if args.remove_existing_uds {
-            std::fs::remove_file(&uds).ok();
+            fs_err::remove_file(&uds).ok();
         }
         figment = figment.join(("address", format!("unix:{uds}")));
     } else if let Some(address) = args.address {