[STACKED on #930] fix(kms): preserve CA certificates across restart

dstack/kms/src/main.rs

@@ -63,13 +63,18 @@ async fn run_onboard_service(kms_config: KmsConfig, figment: Figment) -> Result<
 
     // Remove section tls
 
-    let _ = rocket::custom(figment)
+    let rocket = rocket::custom(figment)
         .mount("/", rocket::routes![index, finish])
         .mount(
             "/prpc",
             ra_rpc::prpc_routes!(OnboardState, OnboardHandler, trim: "Onboard."),
         )
-        .manage(state)
+        .manage(state.clone())
+        .ignite()
+        .await
+        .map_err(|err| anyhow!(err.to_string()))?;
+    state.set_shutdown(rocket.shutdown())?;
+    let _ = rocket
         .launch()
         .await
         .map_err(|err| anyhow!(err.to_string()))?;

dstack/kms/src/onboard_service.rs

@@ -2,7 +2,9 @@
 //
 // SPDX-License-Identifier: Apache-2.0
 
-use std::sync::{Arc, Mutex};
+use std::{
+    sync::{Arc, Mutex},
+};
 
 use anyhow::{bail, Context, Result};
 use dstack_kms_rpc::{
@@ -45,6 +47,7 @@ pub struct OnboardState {
     config: KmsConfig,
     attestation_verifier: Arc<AttestationVerifier>,
     bootstrap_lock: Arc<AsyncMutex<()>>,
+    shutdown: Arc<Mutex<Option<rocket::Shutdown>>>,
 }
 
 impl OnboardState {
@@ -57,8 +60,17 @@ impl OnboardState {
             config,
             attestation_verifier,
             bootstrap_lock: Arc::new(AsyncMutex::new(())),
+            shutdown: Arc::new(Mutex::new(None)),
         })
     }
+
+    pub fn set_shutdown(&self, shutdown: rocket::Shutdown) -> Result<()> {
+        *self
+            .shutdown
+            .lock()
+            .map_err(|_| anyhow::anyhow!("onboard shutdown lock poisoned"))? = Some(shutdown);
+        Ok(())
+    }
 }
 
 pub struct OnboardHandler {
@@ -126,23 +138,27 @@ impl OnboardRpc for OnboardHandler {
 
     async fn onboard(self, request: OnboardRequest) -> Result<OnboardResponse> {
         validate_onboarding_domain(&request.domain)?;
+        let _bootstrap_guard = self.state.bootstrap_lock.lock().await;
+        let cfg = &self.state.config;
+        if cfg.root_ca_key().exists() || cfg.k256_key().exists() {
+            bail!("KMS has already been onboarded");
+        }
         let source_url = request.source_url.trim_end_matches('/').to_string();
         let source_url = if source_url.ends_with("/prpc") {
             source_url
         } else {
             format!("{source_url}/prpc")
         };
         let keys = Keys::onboard(
-            &self.state.config,
+            cfg,
             &source_url,
             &request.domain,
             self.state.attestation_verifier.clone(),
         )
         .await
         .context("Failed to onboard")?;
         let k256_pubkey = keys.k256_key.verifying_key().to_sec1_bytes().to_vec();
-        keys.store(&self.state.config)
-            .context("Failed to store keys")?;
+        keys.store(cfg).context("Failed to store keys")?;
         Ok(OnboardResponse { k256_pubkey })
     }
 
@@ -199,7 +215,15 @@ impl OnboardRpc for OnboardHandler {
     }
 
     async fn finish(self) -> anyhow::Result<()> {
-        std::process::exit(0);
+        let shutdown = self
+            .state
+            .shutdown
+            .lock()
+            .map_err(|_| anyhow::anyhow!("onboard shutdown lock poisoned"))?
+            .clone()
+            .context("onboard shutdown handle is unavailable")?;
+        shutdown.notify();
+        Ok(())
     }
 }
 
@@ -616,10 +640,11 @@ pub(crate) async fn update_certs(cfg: &KmsConfig) -> Result<()> {
     .await
     .context("Failed to regenerate certificates")?;
 
-    // Write the new certificates to files. This runs on every start, so a
-    // hand-placed certificate is replaced -- say so, because the old silence
-    // made that look like the file had survived.
-    keys.store_certs(cfg)?;
+// Root and temporary CA certificates are persistent trust anchors. A normal
+    // service restart must not replace them merely because their private keys
+    // were loaded again. Only the RPC leaf depends on the refreshed domain and
+    // platform attestation.
+    safe_write(cfg.rpc_cert(), keys.rpc_cert.pem())?;
     info!("Reissued the KMS RPC certificate for {domain}");
 
     Ok(())

dstack/ra-rpc/src/rocket_helper.rs

@@ -42,6 +42,14 @@ pub struct RpcResponse {
     body: Vec<u8>,
 }
 
+fn normalize_json_response_body(is_json: bool, body: Vec<u8>) -> Vec<u8> {
+    if is_json && body.as_slice() == b"null" {
+        Vec::new()
+    } else {
+        body
+    }
+}
+
 impl<'r> Responder<'r, 'static> for RpcResponse {
     fn respond_to(self, request: &'r Request<'_>) -> rocket::response::Result<'static> {
         use rocket::http::ContentType;
@@ -50,13 +58,38 @@ impl<'r> Responder<'r, 'static> for RpcResponse {
         } else {
             ContentType::Binary
         };
-        let response = Custom(self.status, self.body).respond_to(request)?;
+        // prpc maps google.protobuf.Empty / Rust unit to JSON `null`. Case
+        // contracts and many clients expect an empty success body instead.
+        let body = normalize_json_response_body(self.is_json, self.body);
+        let response = Custom(self.status, body).respond_to(request)?;
         rocket::Response::build_from(response)
             .header(content_type)
             .ok()
     }
 }
 
+#[cfg(test)]
+mod response_tests {
+    use super::normalize_json_response_body;
+
+    #[test]
+    fn json_unit_response_has_an_empty_body() {
+        assert!(normalize_json_response_body(true, b"null".to_vec()).is_empty());
+    }
+
+    #[test]
+    fn non_unit_and_binary_responses_are_unchanged() {
+        assert_eq!(
+            normalize_json_response_body(true, br#"{"value":null}"#.to_vec()),
+            br#"{"value":null}"#
+        );
+        assert_eq!(
+            normalize_json_response_body(false, b"null".to_vec()),
+            b"null"
+        );
+    }
+}
+
 #[derive(Debug, Clone)]
 struct UnixPeerEndpoint {
     path: PathBuf,