From c8395fbfb6017f2171eee792b1e39f2488baecc5 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Mon, 27 Jul 2026 18:54:30 +0000 Subject: [PATCH 1/8] fix(kms): return Finish response before exit Defect: Onboard.Finish terminated the process inside the RPC handler before Rocket could flush the documented Empty response. Observed symptom: clients received RemoteDisconnected and the graceful-response regression failed. Fix: return success immediately and schedule process exit after a short response-drain interval. --- dstack/kms/src/onboard_service.rs | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/dstack/kms/src/onboard_service.rs b/dstack/kms/src/onboard_service.rs index 334a6c6db..2160272e2 100644 --- a/dstack/kms/src/onboard_service.rs +++ b/dstack/kms/src/onboard_service.rs @@ -2,7 +2,10 @@ // // SPDX-License-Identifier: Apache-2.0 -use std::sync::{Arc, Mutex}; +use std::{ + sync::{Arc, Mutex}, + time::Duration, +}; use anyhow::{bail, Context, Result}; use dstack_kms_rpc::{ @@ -199,7 +202,11 @@ impl OnboardRpc for OnboardHandler { } async fn finish(self) -> anyhow::Result<()> { - std::process::exit(0); + tokio::spawn(async { + tokio::time::sleep(Duration::from_millis(250)).await; + std::process::exit(0); + }); + Ok(()) } } From 026a1d870ce7775b419bcd0c15ca9f7552f23f67 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Mon, 27 Jul 2026 19:03:12 +0000 Subject: [PATCH 2/8] fix(rpc): emit empty JSON unit responses --- dstack/ra-rpc/src/rocket_helper.rs | 35 +++++++++++++++++++++++++++++- 1 file changed, 34 insertions(+), 1 deletion(-) diff --git a/dstack/ra-rpc/src/rocket_helper.rs b/dstack/ra-rpc/src/rocket_helper.rs index 2cc170630..939c38239 100644 --- a/dstack/ra-rpc/src/rocket_helper.rs +++ b/dstack/ra-rpc/src/rocket_helper.rs @@ -42,6 +42,14 @@ pub struct RpcResponse { body: Vec, } +fn normalize_json_response_body(is_json: bool, body: Vec) -> Vec { + 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, From d92c2d8267734c4954ad2573d9216be6d75ba74d Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Mon, 27 Jul 2026 19:04:29 +0000 Subject: [PATCH 3/8] fix(kms): shut down after Finish response --- dstack/kms/src/main.rs | 9 +++++++-- dstack/kms/src/onboard_service.rs | 23 ++++++++++++++++++----- 2 files changed, 25 insertions(+), 7 deletions(-) diff --git a/dstack/kms/src/main.rs b/dstack/kms/src/main.rs index 478c5c52a..74921a698 100644 --- a/dstack/kms/src/main.rs +++ b/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()))?; diff --git a/dstack/kms/src/onboard_service.rs b/dstack/kms/src/onboard_service.rs index 2160272e2..cae39e1d5 100644 --- a/dstack/kms/src/onboard_service.rs +++ b/dstack/kms/src/onboard_service.rs @@ -4,7 +4,6 @@ use std::{ sync::{Arc, Mutex}, - time::Duration, }; use anyhow::{bail, Context, Result}; @@ -48,6 +47,7 @@ pub struct OnboardState { config: KmsConfig, attestation_verifier: Arc, bootstrap_lock: Arc>, + shutdown: Arc>>, } impl OnboardState { @@ -60,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 { @@ -202,10 +211,14 @@ impl OnboardRpc for OnboardHandler { } async fn finish(self) -> anyhow::Result<()> { - tokio::spawn(async { - tokio::time::sleep(Duration::from_millis(250)).await; - 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(()) } } From dd852f29e1f2fa177e73386eaf09d1583905dea4 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Wed, 29 Jul 2026 05:40:25 +0000 Subject: [PATCH 4/8] fix(kms): reject repeated onboarding --- dstack/kms/src/onboard_service.rs | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/dstack/kms/src/onboard_service.rs b/dstack/kms/src/onboard_service.rs index cae39e1d5..9a2e48d15 100644 --- a/dstack/kms/src/onboard_service.rs +++ b/dstack/kms/src/onboard_service.rs @@ -138,6 +138,11 @@ impl OnboardRpc for OnboardHandler { async fn onboard(self, request: OnboardRequest) -> Result { 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 @@ -145,7 +150,7 @@ impl OnboardRpc for OnboardHandler { format!("{source_url}/prpc") }; let keys = Keys::onboard( - &self.state.config, + cfg, &source_url, &request.domain, self.state.attestation_verifier.clone(), @@ -153,8 +158,7 @@ impl OnboardRpc for OnboardHandler { .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 }) } From 9a0b2df461eb59b5ca1d9cc0e81500f5ff46e284 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Wed, 29 Jul 2026 17:14:35 +0000 Subject: [PATCH 5/8] fix(kms): preserve CA certificates across restart --- dstack/kms/src/onboard_service.rs | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/dstack/kms/src/onboard_service.rs b/dstack/kms/src/onboard_service.rs index 9a2e48d15..b2b06edef 100644 --- a/dstack/kms/src/onboard_service.rs +++ b/dstack/kms/src/onboard_service.rs @@ -640,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(()) From cdfc719e654927289460bf98ee3b1a87d8c0f84c Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Wed, 29 Jul 2026 17:17:48 +0000 Subject: [PATCH 6/8] feat(kms): return configured historical root keys --- dstack/kms/kms.toml | 4 ++++ dstack/kms/src/config.rs | 10 ++++++++++ dstack/kms/src/main_service.rs | 35 +++++++++++++++++++++++++++++----- 3 files changed, 44 insertions(+), 5 deletions(-) diff --git a/dstack/kms/kms.toml b/dstack/kms/kms.toml index 08e954662..25355c3bc 100644 --- a/dstack/kms/kms.toml +++ b/dstack/kms/kms.toml @@ -24,6 +24,10 @@ mandatory = false [core] cert_dir = "/etc/kms/certs" +# Previous root-key pairs for authorized handover, ordered newest to oldest: +# [[core.historical_keys]] +# ca_key = "/etc/kms/history/previous/root-ca.key" +# k256_key = "/etc/kms/history/previous/root-k256.key" subject_postfix = ".dstack" site_name = "" # Whether trusted RPCs require the KMS to first attest itself to its own diff --git a/dstack/kms/src/config.rs b/dstack/kms/src/config.rs index 4d95b0e17..11da915da 100644 --- a/dstack/kms/src/config.rs +++ b/dstack/kms/src/config.rs @@ -23,6 +23,12 @@ const RPC_DOMAIN: &str = "rpc-domain"; const K256_KEY: &str = "root-k256.key"; const BOOTSTRAP_INFO: &str = "bootstrap-info.json"; +#[derive(Debug, Clone, Deserialize)] +pub(crate) struct HistoricalKeyConfig { + pub ca_key: PathBuf, + pub k256_key: PathBuf, +} + #[derive(Debug, Clone, Deserialize)] pub(crate) struct ImageConfig { pub verify: bool, @@ -35,6 +41,10 @@ pub(crate) struct ImageConfig { #[derive(Debug, Clone, Deserialize)] pub(crate) struct KmsConfig { pub cert_dir: PathBuf, + /// Previous root-key pairs returned after the current pair during an + /// authorized KMS handover. Configuration order is rotation order. + #[serde(default)] + pub historical_keys: Vec, #[serde(default)] pub attestation: AttestationVerifierConfig, pub auth_api: AuthApi, diff --git a/dstack/kms/src/main_service.rs b/dstack/kms/src/main_service.rs index 339fbd3bc..99043a3bf 100644 --- a/dstack/kms/src/main_service.rs +++ b/dstack/kms/src/main_service.rs @@ -32,7 +32,7 @@ use tracing::{info, warn}; use upgrade_authority::{build_boot_info, ensure_app_id_len, local_kms_boot_info, BootInfo}; use crate::{ - config::KmsConfig, + config::{HistoricalKeyConfig, KmsConfig}, crypto::{derive_k256_key, sign_message, sign_message_with_timestamp}, }; @@ -56,6 +56,7 @@ pub struct KmsStateInner { config: KmsConfig, root_ca: CaCert, k256_key: SigningKey, + historical_keys: Vec, temp_ca_cert: String, temp_ca_key: String, verifier: CvmVerifier, @@ -120,6 +121,24 @@ fn remove_cache(parent_dir: &Path, sub_dir: &str) -> Result<()> { Ok(()) } +fn load_historical_keys(configs: &[HistoricalKeyConfig]) -> Result> { + configs + .iter() + .enumerate() + .map(|(index, config)| { + let ca_key = fs::read_to_string(&config.ca_key) + .with_context(|| format!("Failed to read historical CA key {index}"))?; + ra_tls::rcgen::KeyPair::from_pem(&ca_key) + .with_context(|| format!("Failed to parse historical CA key {index}"))?; + let k256_key = fs::read(&config.k256_key) + .with_context(|| format!("Failed to read historical ECDSA key {index}"))?; + SigningKey::from_slice(&k256_key) + .with_context(|| format!("Failed to parse historical ECDSA key {index}"))?; + Ok(KmsKeys { ca_key, k256_key }) + }) + .collect() +} + impl KmsState { /// clear cached image and measurement material for the given hashes. Used by /// the admin `ClearImageCache` RPC; authorization is enforced by the admin @@ -139,6 +158,8 @@ impl KmsState { let key_bytes = fs::read(config.k256_key()).context("Failed to read ECDSA root key")?; let k256_key = SigningKey::from_slice(&key_bytes).context("Failed to load ECDSA root key")?; + let historical_keys = load_historical_keys(&config.historical_keys) + .context("Failed to load historical root keys")?; let temp_ca_key = fs::read_to_string(config.tmp_ca_key()).context("Faeild to read temp ca key")?; let temp_ca_cert = @@ -163,6 +184,7 @@ impl KmsState { config, root_ca, k256_key, + historical_keys, temp_ca_cert, temp_ca_key, verifier, @@ -485,12 +507,15 @@ impl KmsRpc for RpcHandler { self.state.config.sev_snp_key_release, self.state.config.aws_nitro_tpm_key_release, )?; + let mut keys = Vec::with_capacity(1 + self.state.inner.historical_keys.len()); + keys.push(KmsKeys { + ca_key: self.state.inner.root_ca.key.serialize_pem(), + k256_key: self.state.inner.k256_key.to_bytes().to_vec(), + }); + keys.extend(self.state.inner.historical_keys.iter().cloned()); Ok(KmsKeyResponse { temp_ca_key: self.state.inner.temp_ca_key.clone(), - keys: vec![KmsKeys { - ca_key: self.state.inner.root_ca.key.serialize_pem(), - k256_key: self.state.inner.k256_key.to_bytes().to_vec(), - }], + keys, }) } From eb8203fdca54dd1e2d2c7efc67aa999684819d06 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Wed, 29 Jul 2026 17:18:34 +0000 Subject: [PATCH 7/8] test(kms): cover historical root key inventory --- dstack/kms/src/main_service.rs | 62 ++++++++++++++++++++++++++++++++++ 1 file changed, 62 insertions(+) diff --git a/dstack/kms/src/main_service.rs b/dstack/kms/src/main_service.rs index 99043a3bf..6ab1c3bc3 100644 --- a/dstack/kms/src/main_service.rs +++ b/dstack/kms/src/main_service.rs @@ -610,6 +610,68 @@ mod tests { use cc_eventlog::RuntimeEvent; use sha2::{Digest, Sha256, Sha384}; + fn historical_key_fixture( + directory: &Path, + name: &str, + scalar: u8, + ) -> (HistoricalKeyConfig, String, Vec) { + let key_dir = directory.join(name); + fs::create_dir_all(&key_dir).unwrap(); + let ca_key = ra_tls::rcgen::KeyPair::generate_for(&ra_tls::rcgen::PKCS_ECDSA_P256_SHA256) + .unwrap() + .serialize_pem(); + let k256_key = SigningKey::from_slice(&[scalar; 32]) + .unwrap() + .to_bytes() + .to_vec(); + let ca_path = key_dir.join("root-ca.key"); + let k256_path = key_dir.join("root-k256.key"); + fs::write(&ca_path, &ca_key).unwrap(); + fs::write(&k256_path, &k256_key).unwrap(); + ( + HistoricalKeyConfig { + ca_key: ca_path, + k256_key: k256_path, + }, + ca_key, + k256_key, + ) + } + + #[test] + fn historical_root_keys_preserve_configured_rotation_order() { + let directory = tempfile::tempdir().unwrap(); + let (newer, newer_ca, newer_k256) = historical_key_fixture(directory.path(), "newer", 0x21); + let (older, older_ca, older_k256) = historical_key_fixture(directory.path(), "older", 0x22); + let loaded = load_historical_keys(&[newer, older]).unwrap(); + + assert_eq!(loaded.len(), 2); + assert_eq!(loaded[0].ca_key, newer_ca); + assert_eq!(loaded[0].k256_key, newer_k256); + assert_eq!(loaded[1].ca_key, older_ca); + assert_eq!(loaded[1].k256_key, older_k256); + } + + #[test] + fn historical_root_keys_fail_closed_on_missing_or_malformed_material() { + let directory = tempfile::tempdir().unwrap(); + let missing = HistoricalKeyConfig { + ca_key: directory.path().join("missing-ca.key"), + k256_key: directory.path().join("missing-k256.key"), + }; + assert!(load_historical_keys(&[missing]).is_err()); + + let ca_path = directory.path().join("bad-ca.key"); + let k256_path = directory.path().join("bad-k256.key"); + fs::write(&ca_path, "not a PEM key").unwrap(); + fs::write(&k256_path, [0u8; 31]).unwrap(); + let malformed = HistoricalKeyConfig { + ca_key: ca_path, + k256_key: k256_path, + }; + assert!(load_historical_keys(&[malformed]).is_err()); + } + #[test] fn remove_cache_only_deletes_the_named_hex_entry() { let dir = tempfile::tempdir().unwrap(); From af6c3bcb6ad636cee23461c1ce149ede2fd748d8 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Wed, 29 Jul 2026 17:30:17 +0000 Subject: [PATCH 8/8] docs(kms): define cold backup recovery procedure --- dstack/kms/README.md | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/dstack/kms/README.md b/dstack/kms/README.md index 4ea96fbda..4d9241eb0 100644 --- a/dstack/kms/README.md +++ b/dstack/kms/README.md @@ -232,3 +232,21 @@ The `SignCert` RPC is used by the dstack app to sign a TLS certificate. In this - Verify the CSR signature - Query the smart contract to check if the app is authorized - If authorized, sign the CSR with the CA root key and return the certificate chain to the app + +## Cold backup and recovery + +KMS root material is backed up as a complete, offline copy of `core.cert_dir`. +There is no online backup RPC. Stop the KMS before taking or restoring a copy, +preserve file modes and ownership, and protect the backup with the same controls +as the live private keys. Do not copy only one key: the root CA, root K256 key, +temporary CA, RPC identity, domain, and public certificates form one identity +set. + +A supported recovery restores the complete directory while KMS is stopped, +verifies that every private-key file remains owner-only (`0600`), and then +starts KMS with the unchanged configuration. Compare `KMS.GetMeta` before the +backup and after recovery: `ca_cert` and `k256_pubkey` must match exactly. KMS +must fail closed when a required key is missing or malformed; never generate a +new identity to repair a partial restore. Orphaned `*.private-tmp` files from an +interrupted atomic write are not trust anchors and may be removed only after the +final key file has been verified.