From ce9024eafc96a3d5d726e2dfde2e2cbc16a97cfe Mon Sep 17 00:00:00 2001 From: Leechael Yim Date: Fri, 13 Feb 2026 18:19:02 +0800 Subject: [PATCH 1/3] feat(kms): add auto-onboard support via config Add `auto_onboard_url` config to automate KMS onboarding from an existing instance, removing the need for manual Web UI interaction. When set, the new KMS automatically fetches keys from the source KMS on startup. On failure, the process exits so docker restart can retry. --- kms/kms.toml | 2 ++ kms/src/config.rs | 1 + kms/src/main.rs | 5 ++++- kms/src/onboard_service.rs | 19 +++++++++++++++++++ 4 files changed, 26 insertions(+), 1 deletion(-) diff --git a/kms/kms.toml b/kms/kms.toml index 70b2b7177..d39761505 100644 --- a/kms/kms.toml +++ b/kms/kms.toml @@ -45,5 +45,7 @@ gateway_app_id = "any" [core.onboard] enabled = true auto_bootstrap_domain = "" +auto_onboard_url = "" +quote_enabled = true address = "0.0.0.0" port = 8000 diff --git a/kms/src/config.rs b/kms/src/config.rs index 3eaa2d117..a4c27d977 100644 --- a/kms/src/config.rs +++ b/kms/src/config.rs @@ -119,4 +119,5 @@ pub(crate) struct Dev { pub(crate) struct OnboardConfig { pub enabled: bool, pub auto_bootstrap_domain: String, + pub auto_onboard_url: String, } diff --git a/kms/src/main.rs b/kms/src/main.rs index eddfbdc95..bfab1e772 100644 --- a/kms/src/main.rs +++ b/kms/src/main.rs @@ -52,7 +52,10 @@ async fn run_onboard_service(kms_config: KmsConfig, figment: Figment) -> Result< "OK" } - if !kms_config.onboard.auto_bootstrap_domain.is_empty() { + if !kms_config.onboard.auto_onboard_url.is_empty() { + onboard_service::auto_onboard_keys(&kms_config).await?; + return Ok(()); + } else if !kms_config.onboard.auto_bootstrap_domain.is_empty() { onboard_service::bootstrap_keys(&kms_config).await?; return Ok(()); } diff --git a/kms/src/onboard_service.rs b/kms/src/onboard_service.rs index 64b2390b3..5ba849105 100644 --- a/kms/src/onboard_service.rs +++ b/kms/src/onboard_service.rs @@ -336,6 +336,25 @@ pub(crate) async fn update_certs(cfg: &KmsConfig) -> Result<()> { Ok(()) } +pub(crate) async fn auto_onboard_keys(cfg: &KmsConfig) -> Result<()> { + let source_url = cfg.onboard.auto_onboard_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( + &source_url, + &cfg.onboard.auto_bootstrap_domain, + cfg.onboard.quote_enabled, + cfg.pccs_url.clone(), + ) + .await + .context("failed to auto-onboard from source KMS")?; + keys.store(cfg)?; + Ok(()) +} + pub(crate) async fn bootstrap_keys(cfg: &KmsConfig) -> Result<()> { ensure_self_kms_allowed(cfg) .await From 07faea5b3bd52089c9ac4ddd357d9c5111a9d45c Mon Sep 17 00:00:00 2001 From: Leechael Yim Date: Fri, 13 Feb 2026 18:35:48 +0800 Subject: [PATCH 2/3] style: apply cargo fmt --- kms/src/onboard_service.rs | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/kms/src/onboard_service.rs b/kms/src/onboard_service.rs index 5ba849105..8e077cadb 100644 --- a/kms/src/onboard_service.rs +++ b/kms/src/onboard_service.rs @@ -337,7 +337,11 @@ pub(crate) async fn update_certs(cfg: &KmsConfig) -> Result<()> { } pub(crate) async fn auto_onboard_keys(cfg: &KmsConfig) -> Result<()> { - let source_url = cfg.onboard.auto_onboard_url.trim_end_matches('/').to_string(); + let source_url = cfg + .onboard + .auto_onboard_url + .trim_end_matches('/') + .to_string(); let source_url = if source_url.ends_with("/prpc") { source_url } else { From b6eca71a9e7cf0f9448fef521d9d481713f8a2e8 Mon Sep 17 00:00:00 2001 From: Leechael Yim Date: Tue, 24 Feb 2026 20:40:10 +0800 Subject: [PATCH 3/3] fix(kms): validate auto bootstrap domain before key onboarding --- kms/src/onboard_service.rs | 71 ++++++++++++++++++++++++++++---------- 1 file changed, 52 insertions(+), 19 deletions(-) diff --git a/kms/src/onboard_service.rs b/kms/src/onboard_service.rs index 8e077cadb..c872ecebd 100644 --- a/kms/src/onboard_service.rs +++ b/kms/src/onboard_service.rs @@ -159,12 +159,12 @@ struct Keys { } impl Keys { - async fn generate(domain: &str) -> Result { + async fn generate(domain: &str, quote_enabled: bool) -> Result { let tmp_ca_key = KeyPair::generate_for(&PKCS_ECDSA_P256_SHA256)?; let ca_key = KeyPair::generate_for(&PKCS_ECDSA_P256_SHA256)?; let rpc_key = KeyPair::generate_for(&PKCS_ECDSA_P256_SHA256)?; let k256_key = SigningKey::random(&mut rand::rngs::OsRng); - Self::from_keys(tmp_ca_key, ca_key, rpc_key, k256_key, domain).await + Self::from_keys(tmp_ca_key, ca_key, rpc_key, k256_key, domain, quote_enabled).await } async fn from_keys( @@ -173,6 +173,7 @@ impl Keys { rpc_key: KeyPair, k256_key: SigningKey, domain: &str, + quote_enabled: bool, ) -> Result { let tmp_ca_cert = CertRequest::builder() .org_name("Dstack") @@ -190,20 +191,25 @@ impl Keys { .key(&ca_key) .build() .self_signed()?; - let pubkey = rpc_key.public_key_der(); - let report_data = QuoteContentType::RaTlsCert.to_report_data(&pubkey); - let response = app_attest(report_data.to_vec()) - .await - .context("Failed to get quote")?; - let attestation = VersionedAttestation::from_scale(&response.attestation) - .context("Invalid attestation")?; + let attestation = if quote_enabled { + let pubkey = rpc_key.public_key_der(); + let report_data = QuoteContentType::RaTlsCert.to_report_data(&pubkey); + let response = app_attest(report_data.to_vec()) + .await + .context("Failed to get quote")?; + let attestation = VersionedAttestation::from_scale(&response.attestation) + .context("Invalid attestation")?; + Some(attestation) + } else { + None + }; // Sign WWW server cert with KMS cert let rpc_cert = CertRequest::builder() .subject(domain) .alt_names(&[domain.to_string()]) .special_usage("kms:rpc") - .maybe_attestation(Some(&attestation)) + .maybe_attestation(attestation.as_ref()) .key(&rpc_key) .build() .signed_by(&ca_cert, &ca_key)?; @@ -308,6 +314,16 @@ impl Keys { } } +fn validate_domain(domain: &str, source: &str) -> Result { + let domain = domain.trim(); + if domain.is_empty() { + return Err(anyhow::anyhow!( + "invalid domain from {source}: empty or whitespace-only" + )); + } + Ok(domain.to_string()) +} + 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())?)?; @@ -318,17 +334,26 @@ pub(crate) async fn update_certs(cfg: &KmsConfig) -> Result<()> { let k256_key_bytes = fs::read(cfg.k256_key())?; let k256_key = SigningKey::from_slice(&k256_key_bytes)?; - let domain = if cfg.onboard.auto_bootstrap_domain.is_empty() { - fs::read_to_string(cfg.rpc_domain())? + let domain = if cfg.onboard.auto_bootstrap_domain.trim().is_empty() { + validate_domain(&fs::read_to_string(cfg.rpc_domain())?, "stored rpc_domain")? } else { - cfg.onboard.auto_bootstrap_domain.clone() + validate_domain( + &cfg.onboard.auto_bootstrap_domain, + "core.onboard.auto_bootstrap_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) - .await - .context("Failed to regenerate certificates")?; + 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)?; @@ -347,9 +372,13 @@ pub(crate) async fn auto_onboard_keys(cfg: &KmsConfig) -> Result<()> { } else { format!("{source_url}/prpc") }; + let domain = validate_domain( + &cfg.onboard.auto_bootstrap_domain, + "core.onboard.auto_bootstrap_domain", + )?; let keys = Keys::onboard( &source_url, - &cfg.onboard.auto_bootstrap_domain, + &domain, cfg.onboard.quote_enabled, cfg.pccs_url.clone(), ) @@ -363,7 +392,11 @@ pub(crate) async fn bootstrap_keys(cfg: &KmsConfig) -> Result<()> { ensure_self_kms_allowed(cfg) .await .context("KMS is not allowed to auto-bootstrap")?; - let keys = Keys::generate(&cfg.onboard.auto_bootstrap_domain) + let domain = validate_domain( + &cfg.onboard.auto_bootstrap_domain, + "core.onboard.auto_bootstrap_domain", + )?; + let keys = Keys::generate(&domain, cfg.onboard.quote_enabled) .await .context("Failed to generate keys")?; keys.store(cfg)?;