From 09f8eeb80d32dafca43b6acda224dbf21e19facb Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Wed, 29 Jul 2026 08:09:07 +0000 Subject: [PATCH 1/5] fix(gateway): fail closed on corrupt ACME credentials --- dstack/gateway/src/distributed_certbot.rs | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/dstack/gateway/src/distributed_certbot.rs b/dstack/gateway/src/distributed_certbot.rs index 24d2613f9..bccb9d6bf 100644 --- a/dstack/gateway/src/distributed_certbot.rs +++ b/dstack/gateway/src/distributed_certbot.rs @@ -391,7 +391,9 @@ impl DistributedCertBot { // Try to load global ACME credentials from KvStore if let Some(creds) = self.kv_store.get_acme_credentials() { - if acme_url_matches(&creds.acme_credentials, acme_url) { + if acme_url_matches(&creds.acme_credentials, acme_url) + .context("invalid ACME credentials in KvStore")? + { info!("loaded global ACME account credentials from KvStore"); return AcmeClient::load( dns01_client, @@ -568,15 +570,14 @@ fn get_cert_expiry(cert_pem: &str) -> Option { Some(cert.validity().not_after.timestamp() as u64) } -fn acme_url_matches(credentials_json: &str, expected_url: &str) -> bool { +fn acme_url_matches(credentials_json: &str, expected_url: &str) -> Result { #[derive(serde::Deserialize)] struct Creds { - #[serde(default)] acme_url: String, } - serde_json::from_str::(credentials_json) - .map(|c| c.acme_url == expected_url) - .unwrap_or(false) + let credentials = serde_json::from_str::(credentials_json) + .context("failed to decode ACME credentials")?; + Ok(credentials.acme_url == expected_url) } /// Extract account_id (URI) from ACME credentials JSON From 173fa74c44ae6c781330bbab05666f9e61a51af7 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Wed, 29 Jul 2026 08:09:22 +0000 Subject: [PATCH 2/5] test(gateway): cover corrupt ACME credential handling --- dstack/gateway/src/distributed_certbot.rs | 19 +++++++++++++++++-- 1 file changed, 17 insertions(+), 2 deletions(-) diff --git a/dstack/gateway/src/distributed_certbot.rs b/dstack/gateway/src/distributed_certbot.rs index bccb9d6bf..4f0a13422 100644 --- a/dstack/gateway/src/distributed_certbot.rs +++ b/dstack/gateway/src/distributed_certbot.rs @@ -592,7 +592,6 @@ fn extract_account_uri(credentials_json: &str) -> Option { .filter(|c| !c.account_id.is_empty()) .map(|c| c.account_id) } - #[cfg(test)] mod tests { use super::*; @@ -607,7 +606,6 @@ mod tests { async fn set_caa_all_succeeds_without_configured_domains() { let data_dir = tempfile::tempdir().expect("failed to create temp dir"); let certbot = test_certbot(data_dir.path()); - // No ZT-Domain configured: nothing to reconcile and no DNS provider is contacted. certbot .set_caa_all() .await @@ -628,4 +626,21 @@ mod tests { "unexpected error: {err}" ); } + + #[test] + fn corrupt_acme_credentials_fail_closed() { + assert!(acme_url_matches("not-json", "https://acme.test/directory").is_err()); + assert!(acme_url_matches("{}", "https://acme.test/directory").is_err()); + } + + #[test] + fn valid_acme_credentials_distinguish_directory() { + let credentials = r#"{"acme_url":"https://acme.test/directory"}"#; + assert!(acme_url_matches(credentials, "https://acme.test/directory") + .expect("valid credentials rejected")); + assert!( + !acme_url_matches(credentials, "https://other.test/directory") + .expect("valid credentials rejected") + ); + } } From 888e3601a91f05d901f14f7c6037ac9d6c5b45e9 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Tue, 4 Aug 2026 19:29:16 -0700 Subject: [PATCH 3/5] feat(gateway): add safe ACME credential rotation --- dstack/gateway/rpc/proto/gateway_rpc.proto | 12 ++ dstack/gateway/src/admin_service.rs | 16 ++- dstack/gateway/src/distributed_certbot.rs | 123 +++++++++++++++++---- dstack/gateway/src/main_service.rs | 4 + 4 files changed, 130 insertions(+), 25 deletions(-) diff --git a/dstack/gateway/rpc/proto/gateway_rpc.proto b/dstack/gateway/rpc/proto/gateway_rpc.proto index b7a58f9cd..c03d9f40d 100644 --- a/dstack/gateway/rpc/proto/gateway_rpc.proto +++ b/dstack/gateway/rpc/proto/gateway_rpc.proto @@ -164,6 +164,14 @@ message AcmeInfoResponse { string account_attestation = 5; } +// Result of replacing the shared ACME account credentials. +message RotateAcmeCredentialsResponse { + // URI of the newly-created ACME account. The private credentials are never returned. + string account_uri = 1; + // Number of ZT domains whose CAA records were updated for the new account. + uint32 domains_updated = 2; +} + // Get HostInfo for associated instance id. message GetInfoRequest { string id = 1; @@ -466,6 +474,10 @@ service Admin { rpc GetCertbotConfig(google.protobuf.Empty) returns (CertbotConfigResponse) {} // Set global certbot configuration (includes ACME URL) rpc SetCertbotConfig(SetCertbotConfigRequest) returns (google.protobuf.Empty) {} + // Create a new ACME account, update every ZT-domain CAA record, and then + // replace the shared credentials. Call only one gateway at a time because + // WaveKV does not provide compare-and-swap. + rpc RotateAcmeCredentials(google.protobuf.Empty) returns (RotateAcmeCredentialsResponse) {} // ==================== Per-Instance Port Policy Override ==================== // Set an admin override for an instance's port policy. Takes precedence diff --git a/dstack/gateway/src/admin_service.rs b/dstack/gateway/src/admin_service.rs index 7b634bb67..c1e988d19 100644 --- a/dstack/gateway/src/admin_service.rs +++ b/dstack/gateway/src/admin_service.rs @@ -18,10 +18,10 @@ use dstack_gateway_rpc::{ ListCertAttestationsResponse, ListDnsCredentialsResponse, ListZtDomainsResponse, NodeStatusEntry, PeerSyncStatus as ProtoPeerSyncStatus, PortAttrs as RpcPortAttrs, PortPolicy as RpcPortPolicy, RenewCertResponse, RenewZtDomainCertRequest, - RenewZtDomainCertResponse, SetCertbotConfigRequest, SetDefaultDnsCredentialRequest, - SetInstancePortPolicyRequest, SetNodeStatusRequest, SetNodeUrlRequest, StatusResponse, - StoreSyncStatus, UpdateDnsCredentialRequest, WaveKvStatusResponse, ZtDomainCertStatus, - ZtDomainConfig as ProtoZtDomainConfig, ZtDomainInfo, + RenewZtDomainCertResponse, RotateAcmeCredentialsResponse, SetCertbotConfigRequest, + SetDefaultDnsCredentialRequest, SetInstancePortPolicyRequest, SetNodeStatusRequest, + SetNodeUrlRequest, StatusResponse, StoreSyncStatus, UpdateDnsCredentialRequest, + WaveKvStatusResponse, ZtDomainCertStatus, ZtDomainConfig as ProtoZtDomainConfig, ZtDomainInfo, }; use ra_rpc::{CallContext, RpcCall}; use tracing::info; @@ -100,6 +100,14 @@ impl AdminRpc for AdminRpcHandler { self.state.reload_all_certs_from_kvstore() } + async fn rotate_acme_credentials(self) -> Result { + let (account_uri, domains_updated) = self.state.rotate_acme_credentials().await?; + Ok(RotateAcmeCredentialsResponse { + account_uri, + domains_updated: domains_updated.try_into().unwrap_or(u32::MAX), + }) + } + async fn status(self) -> Result { self.status().await } diff --git a/dstack/gateway/src/distributed_certbot.rs b/dstack/gateway/src/distributed_certbot.rs index 4f0a13422..c7c91cccf 100644 --- a/dstack/gateway/src/distributed_certbot.rs +++ b/dstack/gateway/src/distributed_certbot.rs @@ -20,8 +20,8 @@ use tracing::{error, info, warn}; use crate::cert_store::CertResolver; use crate::kv::{ - AcmeAttestation, CertAttestation, CertCredentials, CertData, DnsProvider, KvStore, - ZtDomainConfig, + AcmeAttestation, CertAttestation, CertCredentials, CertData, DnsCredential, DnsProvider, + KvStore, ZtDomainConfig, }; /// Lock timeout for certificate renewal (10 minutes) @@ -34,11 +34,10 @@ const DEFAULT_ACME_URL: &str = "https://acme-v02.api.letsencrypt.org/directory"; pub struct DistributedCertBot { kv_store: Arc, cert_resolver: Arc, - /// Serializes CAA reconciliation within this process. + /// Serializes CAA reconciliation and credential rotation within this process. /// - /// This is deliberately not a cluster-wide lock: CAA reconciliation is a rare - /// manual operation, so a node-local guard against concurrent admin calls is - /// enough and avoids a distributed lock that could be left behind on crash. + /// This is deliberately not a cluster-wide lock because WaveKV does not + /// provide compare-and-swap. caa_lock: Mutex<()>, } @@ -51,6 +50,89 @@ impl DistributedCertBot { } } + async fn dns_client(&self, domain: &str, config: &ZtDomainConfig) -> Result { + let dns_cred = if let Some(ref cred_id) = config.dns_cred_id { + self.kv_store + .get_dns_credential(cred_id) + .context("specified DNS credential not found")? + } else { + self.kv_store + .get_default_dns_credential() + .context("no default DNS credential configured")? + }; + + match &dns_cred.provider { + DnsProvider::Cloudflare { api_token, api_url } => { + Dns01Client::new_cloudflare(domain.to_string(), api_token.clone(), api_url.clone()) + .await + } + } + } + + /// Rotate the shared ACME account without interrupting certificate serving. + /// + /// CAA records are updated before the new credentials are published. WaveKV + /// has no CAS operation, so the lock only serializes calls handled by this + /// node; operators must not rotate through multiple nodes concurrently. + pub async fn rotate_acme_credentials(&self) -> Result<(String, usize)> { + let Ok(_guard) = self.caa_lock.try_lock() else { + bail!("ACME credential rotation or CAA reconciliation is already in progress"); + }; + let configs = self.kv_store.list_zt_domain_configs(); + let first = configs + .first() + .context("no ZT-Domain configured for ACME credential rotation")?; + let certbot_config = self.config(); + let acme_url = if certbot_config.acme_url.is_empty() { + DEFAULT_ACME_URL + } else { + &certbot_config.acme_url + }; + + let first_dns_cred = dns_credential_for(&self.kv_store, first)?; + let dns_client = self.dns_client(&first.domain, first).await?; + let client = AcmeClient::new_account( + acme_url, + dns_client, + first_dns_cred.max_dns_wait, + first_dns_cred.dns_txt_ttl, + ) + .await + .context("failed to create replacement ACME account")?; + let credentials = client + .dump_credentials() + .context("failed to encode replacement ACME credentials")?; + let account_uri = client.account_id().to_string(); + + for config in &configs { + let dns_cred = dns_credential_for(&self.kv_store, config)?; + let dns_client = self.dns_client(&config.domain, config).await?; + let client = AcmeClient::load( + dns_client, + &credentials, + dns_cred.max_dns_wait, + dns_cred.dns_txt_ttl, + ) + .await + .with_context(|| format!("failed to prepare ACME client for {}", config.domain))?; + client + .set_caa_records(&[format!("*.{}", config.domain)]) + .await + .with_context(|| format!("failed to update CAA for {}", config.domain))?; + } + + // Publish only after every CAA update succeeds. Readers create an ACME + // client per operation, so all nodes recover on their next retry after + // WaveKV propagates this value. + self.kv_store.save_acme_credentials(&CertCredentials { + acme_credentials: credentials, + })?; + if let Err(err) = self.generate_and_save_acme_attestation(&account_uri).await { + warn!("failed to attest rotated ACME account: {err:?}"); + } + Ok((account_uri, configs.len())) + } + /// Get the current certbot configuration from KV store fn config(&self) -> crate::kv::GlobalCertbotConfig { self.kv_store.get_certbot_config() @@ -363,23 +445,10 @@ impl DistributedCertBot { config: &ZtDomainConfig, ) -> Result { // Get DNS credential (from config or default) - let dns_cred = if let Some(ref cred_id) = config.dns_cred_id { - self.kv_store - .get_dns_credential(cred_id) - .context("specified DNS credential not found")? - } else { - self.kv_store - .get_default_dns_credential() - .context("no default DNS credential configured")? - }; + let dns_cred = dns_credential_for(&self.kv_store, config)?; // Create DNS client based on provider - let dns01_client = match &dns_cred.provider { - DnsProvider::Cloudflare { api_token, api_url } => { - Dns01Client::new_cloudflare(domain.to_string(), api_token.clone(), api_url.clone()) - .await? - } - }; + let dns01_client = self.dns_client(domain, config).await?; // Use ACME URL from certbot config, fall back to default if not set let config = self.config(); @@ -556,6 +625,18 @@ impl DistributedCertBot { } } +fn dns_credential_for(kv_store: &KvStore, config: &ZtDomainConfig) -> Result { + if let Some(ref cred_id) = config.dns_cred_id { + kv_store + .get_dns_credential(cred_id) + .context("specified DNS credential not found") + } else { + kv_store + .get_default_dns_credential() + .context("no default DNS credential configured") + } +} + fn now_secs() -> u64 { SystemTime::now() .duration_since(UNIX_EPOCH) diff --git a/dstack/gateway/src/main_service.rs b/dstack/gateway/src/main_service.rs index 9605016bf..d38cdb7fa 100644 --- a/dstack/gateway/src/main_service.rs +++ b/dstack/gateway/src/main_service.rs @@ -430,6 +430,10 @@ impl Proxy { } } + pub(crate) async fn rotate_acme_credentials(&self) -> Result<(String, usize)> { + self.certbot.rotate_acme_credentials().await + } + /// Get ACME info for all managed domains (or a specific domain) pub(crate) fn acme_info(&self, domain: Option<&str>) -> Result { let kv_store = self.kv_store.clone(); From 68d56180c59eb5da8660e3dbf855604c1fb3d1f1 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Wed, 5 Aug 2026 01:36:11 -0700 Subject: [PATCH 4/5] fix(gateway): make ACME rotation convergent and cross-node guarded Rework RotateAcmeCredentials so a partial failure no longer strands the cluster on an unpublished account: - validate every domain's DNS credential before registering the replacement account, so a misconfigured domain aborts with no side effects and no rate-limited ACME registration consumed - publish the new credentials before re-pinning CAA; a partial re-pin now converges by rerunning SetCaa instead of registering yet another account on every retry - re-pin CAA best-effort across all domains and name the failed ones, with the recovery path, in the returned error - serialize rotation across nodes with a best-effort TTL lock in WaveKV; without CAS it narrows the race window rather than guaranteeing mutual exclusion, and the comments say so - refuse to auto-register a fresh account on ACME URL mismatch, which left every domain's CAA pinned to the old account; point at RotateAcmeCredentials, which re-pins CAA along with the switch - derive AcmeInfo.account_uri from the published credentials so the dashboard reflects a rotation even when the best-effort attestation write is skipped; attach the quote only when it matches that account - add the rotation step to the staging-to-production tutorial Verification: - cargo test -p dstack-gateway: 74 passed (7 in distributed_certbot) - cargo clippy -p dstack-gateway --all-features -- -D warnings --allow unused_variables: clean --- docs/tutorials/gateway-service-setup.md | 14 ++ dstack/gateway/rpc/proto/gateway_rpc.proto | 11 +- dstack/gateway/src/distributed_certbot.rs | 235 +++++++++++++++------ dstack/gateway/src/kv/mod.rs | 46 ++++ dstack/gateway/src/main_service.rs | 18 +- 5 files changed, 258 insertions(+), 66 deletions(-) diff --git a/docs/tutorials/gateway-service-setup.md b/docs/tutorials/gateway-service-setup.md index 9570975c3..10b592311 100644 --- a/docs/tutorials/gateway-service-setup.md +++ b/docs/tutorials/gateway-service-setup.md @@ -319,6 +319,20 @@ curl -sf -X POST "http://$ADMIN_ADDR/prpc/SetCertbotConfig" \ }' && echo "Certbot config set (PRODUCTION)" ``` +The stored ACME account still belongs to the staging directory, so rotate the +shared credentials. This registers a production account and re-pins every ZT +domain's CAA records to it in one step (renewals refuse to run while the +stored account and the configured ACME URL disagree): + +```bash +curl -sf -X POST "http://$ADMIN_ADDR/prpc/RotateAcmeCredentials" \ + -H "Content-Type: application/json" -d '{}' && echo "ACME account rotated" +``` + +> If the rotation reports that CAA re-pinning failed for some domains, the new +> account is already published — rerun `SetCaa` until it succeeds instead of +> rotating again (each rotation registers a new rate-limited ACME account). + After switching the ACME URL, the renewal loop may report "does not need renewal" because the staging cert is still valid. Force a renewal for each ZT domain to get production certificates immediately: ```bash diff --git a/dstack/gateway/rpc/proto/gateway_rpc.proto b/dstack/gateway/rpc/proto/gateway_rpc.proto index c03d9f40d..68202032d 100644 --- a/dstack/gateway/rpc/proto/gateway_rpc.proto +++ b/dstack/gateway/rpc/proto/gateway_rpc.proto @@ -474,9 +474,14 @@ service Admin { rpc GetCertbotConfig(google.protobuf.Empty) returns (CertbotConfigResponse) {} // Set global certbot configuration (includes ACME URL) rpc SetCertbotConfig(SetCertbotConfigRequest) returns (google.protobuf.Empty) {} - // Create a new ACME account, update every ZT-domain CAA record, and then - // replace the shared credentials. Call only one gateway at a time because - // WaveKV does not provide compare-and-swap. + // Create a new ACME account, publish the shared credentials, and re-pin + // every ZT-domain CAA record to the new account. If CAA re-pinning fails for + // some domains, the new credentials are already published; rerun SetCaa + // until it succeeds instead of retrying the rotation (each rotation + // registers a new rate-limited ACME account). Rotation is serialized across + // nodes by a best-effort lock (WaveKV has no compare-and-swap), so still + // avoid rotating from multiple gateways concurrently. This re-pins issuance + // to the new account; it does not deactivate the old ACME account at the CA. rpc RotateAcmeCredentials(google.protobuf.Empty) returns (RotateAcmeCredentialsResponse) {} // ==================== Per-Instance Port Policy Override ==================== diff --git a/dstack/gateway/src/distributed_certbot.rs b/dstack/gateway/src/distributed_certbot.rs index c7c91cccf..d34959c91 100644 --- a/dstack/gateway/src/distributed_certbot.rs +++ b/dstack/gateway/src/distributed_certbot.rs @@ -27,6 +27,9 @@ use crate::kv::{ /// Lock timeout for certificate renewal (10 minutes) const RENEW_LOCK_TIMEOUT_SECS: u64 = 600; +/// Lock timeout for ACME credential rotation (10 minutes) +const ROTATION_LOCK_TIMEOUT_SECS: u64 = 600; + /// Default ACME URL (Let's Encrypt production) const DEFAULT_ACME_URL: &str = "https://acme-v02.api.letsencrypt.org/directory"; @@ -36,8 +39,8 @@ pub struct DistributedCertBot { cert_resolver: Arc, /// Serializes CAA reconciliation and credential rotation within this process. /// - /// This is deliberately not a cluster-wide lock because WaveKV does not - /// provide compare-and-swap. + /// Credential rotation is additionally guarded across nodes by a + /// best-effort lock in WaveKV; see [`KvStore::try_acquire_rotation_lock`]. caa_lock: Mutex<()>, } @@ -50,17 +53,7 @@ impl DistributedCertBot { } } - async fn dns_client(&self, domain: &str, config: &ZtDomainConfig) -> Result { - let dns_cred = if let Some(ref cred_id) = config.dns_cred_id { - self.kv_store - .get_dns_credential(cred_id) - .context("specified DNS credential not found")? - } else { - self.kv_store - .get_default_dns_credential() - .context("no default DNS credential configured")? - }; - + async fn dns_client(&self, domain: &str, dns_cred: &DnsCredential) -> Result { match &dns_cred.provider { DnsProvider::Cloudflare { api_token, api_url } => { Dns01Client::new_cloudflare(domain.to_string(), api_token.clone(), api_url.clone()) @@ -71,17 +64,45 @@ impl DistributedCertBot { /// Rotate the shared ACME account without interrupting certificate serving. /// - /// CAA records are updated before the new credentials are published. WaveKV - /// has no CAS operation, so the lock only serializes calls handled by this - /// node; operators must not rotate through multiple nodes concurrently. + /// The sequence is: validate every domain's DNS credential, create the + /// replacement account, publish the new credentials, then re-pin every + /// domain's CAA record to the new account. Publishing before re-pinning + /// makes the failure mode convergent: if some domains fail to re-pin, the + /// cluster is already on the new account and rerunning `SetCaa` finishes + /// the switch without registering yet another account (Let's Encrypt caps + /// new registrations per IP). + /// + /// Between publishing and re-pinning, a renewal on another node may pick up + /// the new account while a domain's CAA still pins the old one; that + /// issuance fails and the periodic renewal task retries. Re-pinning briefly + /// installs `;` guard CAA records, so a failure can leave a domain blocked + /// from issuance until a later `SetCaa` run succeeds (the same hazard as + /// [`Self::set_caa_all`]). + /// + /// This RPC re-pins issuance to the new account; it does not deactivate the + /// old ACME account at the CA. pub async fn rotate_acme_credentials(&self) -> Result<(String, usize)> { let Ok(_guard) = self.caa_lock.try_lock() else { bail!("ACME credential rotation or CAA reconciliation is already in progress"); }; + if !self + .kv_store + .try_acquire_rotation_lock(ROTATION_LOCK_TIMEOUT_SECS) + { + bail!("another node is rotating ACME credentials; retry after it finishes"); + } + let result = self.do_rotate_acme_credentials().await; + if let Err(err) = self.kv_store.release_rotation_lock() { + error!("failed to release ACME rotation lock: {err:?}"); + } + result + } + + async fn do_rotate_acme_credentials(&self) -> Result<(String, usize)> { let configs = self.kv_store.list_zt_domain_configs(); - let first = configs - .first() - .context("no ZT-Domain configured for ACME credential rotation")?; + if configs.is_empty() { + bail!("no ZT-Domain configured for ACME credential rotation"); + } let certbot_config = self.config(); let acme_url = if certbot_config.acme_url.is_empty() { DEFAULT_ACME_URL @@ -89,11 +110,26 @@ impl DistributedCertBot { &certbot_config.acme_url }; + // Validate every domain's DNS credential up front: constructing a DNS + // client resolves the zone through an authenticated API call, so a + // misconfigured domain aborts the rotation here with no side effects + // and no ACME account consumed. + let mut prepared = Vec::with_capacity(configs.len()); + for config in &configs { + let dns_cred = dns_credential_for(&self.kv_store, config)?; + let dns_client = self + .dns_client(&config.domain, &dns_cred) + .await + .with_context(|| format!("DNS credential check failed for {}", config.domain))?; + prepared.push((&config.domain, dns_cred, dns_client)); + } + + let first = configs.first().expect("configs is non-empty"); let first_dns_cred = dns_credential_for(&self.kv_store, first)?; - let dns_client = self.dns_client(&first.domain, first).await?; + let account_dns_client = self.dns_client(&first.domain, &first_dns_cred).await?; let client = AcmeClient::new_account( acme_url, - dns_client, + account_dns_client, first_dns_cred.max_dns_wait, first_dns_cred.dns_txt_ttl, ) @@ -104,33 +140,55 @@ impl DistributedCertBot { .context("failed to encode replacement ACME credentials")?; let account_uri = client.account_id().to_string(); - for config in &configs { - let dns_cred = dns_credential_for(&self.kv_store, config)?; - let dns_client = self.dns_client(&config.domain, config).await?; - let client = AcmeClient::load( - dns_client, - &credentials, - dns_cred.max_dns_wait, - dns_cred.dns_txt_ttl, - ) - .await - .with_context(|| format!("failed to prepare ACME client for {}", config.domain))?; - client - .set_caa_records(&[format!("*.{}", config.domain)]) - .await - .with_context(|| format!("failed to update CAA for {}", config.domain))?; - } - - // Publish only after every CAA update succeeds. Readers create an ACME - // client per operation, so all nodes recover on their next retry after - // WaveKV propagates this value. + // Publish immediately. From here the cluster converges on the new + // account, and recovering from a partial re-pin below never needs to + // register another account. Readers create an ACME client per + // operation, so all nodes pick this up after WaveKV propagates it. self.kv_store.save_acme_credentials(&CertCredentials { - acme_credentials: credentials, + acme_credentials: credentials.clone(), })?; if let Err(err) = self.generate_and_save_acme_attestation(&account_uri).await { warn!("failed to attest rotated ACME account: {err:?}"); } - Ok((account_uri, configs.len())) + + // Re-pin every domain's CAA to the new account, best effort across all + // domains: one failing domain must not block re-pinning the rest. + let total = prepared.len(); + let mut failed = Vec::new(); + for (domain, dns_cred, dns_client) in prepared { + let result = async { + let client = AcmeClient::load( + dns_client, + &credentials, + dns_cred.max_dns_wait, + dns_cred.dns_txt_ttl, + ) + .await + .context("failed to prepare ACME client")?; + client + .set_caa_records(std::slice::from_ref(domain)) + .await + .context("failed to update CAA records") + } + .await; + match result { + Ok(()) => info!("cert[{domain}]: CAA re-pinned to {account_uri}"), + Err(err) => { + error!("cert[{domain}]: failed to re-pin CAA: {err:?}"); + failed.push(domain.clone()); + } + } + } + if !failed.is_empty() { + bail!( + "rotated to {account_uri} and published the new credentials, but failed to \ + re-pin CAA for {}/{total} domains: {}; rerun SetCaa until it succeeds — \ + retrying the rotation would register yet another account", + failed.len(), + failed.join(", ") + ); + } + Ok((account_uri, total)) } /// Get the current certbot configuration from KV store @@ -186,7 +244,7 @@ impl DistributedCertBot { /// the same domain; the periodic renewal task retries, so that is transient. pub async fn set_caa_all(&self) -> Result<()> { let Ok(_guard) = self.caa_lock.try_lock() else { - bail!("CAA reconciliation is already in progress"); + bail!("ACME credential rotation or CAA reconciliation is already in progress"); }; let configs = self.kv_store.list_zt_domain_configs(); if configs.is_empty() { @@ -448,7 +506,7 @@ impl DistributedCertBot { let dns_cred = dns_credential_for(&self.kv_store, config)?; // Create DNS client based on provider - let dns01_client = self.dns_client(domain, config).await?; + let dns01_client = self.dns_client(domain, &dns_cred).await?; // Use ACME URL from certbot config, fall back to default if not set let config = self.config(); @@ -460,20 +518,26 @@ impl DistributedCertBot { // Try to load global ACME credentials from KvStore if let Some(creds) = self.kv_store.get_acme_credentials() { - if acme_url_matches(&creds.acme_credentials, acme_url) - .context("invalid ACME credentials in KvStore")? - { - info!("loaded global ACME account credentials from KvStore"); - return AcmeClient::load( - dns01_client, - &creds.acme_credentials, - dns_cred.max_dns_wait, - dns_cred.dns_txt_ttl, - ) - .await - .context("failed to load ACME client from KvStore credentials"); + if !acme_url_matches(&creds.acme_credentials, acme_url).context( + "invalid ACME credentials in KvStore; call RotateAcmeCredentials to replace them", + )? { + // Registering a fresh account here would leave every domain's + // CAA pinned to the old account and block issuance; rotation + // re-pins CAA along with the switch. + bail!( + "stored ACME credentials are for a different ACME directory; \ + call RotateAcmeCredentials to switch directories" + ); } - warn!("ACME URL mismatch in KvStore credentials, will create new account"); + info!("loaded global ACME account credentials from KvStore"); + return AcmeClient::load( + dns01_client, + &creds.acme_credentials, + dns_cred.max_dns_wait, + dns_cred.dns_txt_ttl, + ) + .await + .context("failed to load ACME client from KvStore credentials"); } // Create new global ACME account @@ -662,7 +726,7 @@ fn acme_url_matches(credentials_json: &str, expected_url: &str) -> Result } /// Extract account_id (URI) from ACME credentials JSON -fn extract_account_uri(credentials_json: &str) -> Option { +pub(crate) fn extract_account_uri(credentials_json: &str) -> Option { #[derive(serde::Deserialize)] struct Creds { #[serde(default)] @@ -673,6 +737,7 @@ fn extract_account_uri(credentials_json: &str) -> Option { .filter(|c| !c.account_id.is_empty()) .map(|c| c.account_id) } + #[cfg(test)] mod tests { use super::*; @@ -687,6 +752,7 @@ mod tests { async fn set_caa_all_succeeds_without_configured_domains() { let data_dir = tempfile::tempdir().expect("failed to create temp dir"); let certbot = test_certbot(data_dir.path()); + // No ZT-Domain configured: nothing to reconcile and no DNS provider is contacted. certbot .set_caa_all() .await @@ -708,6 +774,57 @@ mod tests { ); } + #[tokio::test] + async fn rotate_acme_credentials_rejects_concurrent_runs() { + let data_dir = tempfile::tempdir().expect("failed to create temp dir"); + let certbot = test_certbot(data_dir.path()); + let _guard = certbot.caa_lock.lock().await; + let err = certbot + .rotate_acme_credentials() + .await + .expect_err("a concurrent run should be rejected"); + assert!( + err.to_string().contains("already in progress"), + "unexpected error: {err}" + ); + } + + #[tokio::test] + async fn rotate_acme_credentials_rejects_when_another_node_holds_the_lock() { + let data_dir = tempfile::tempdir().expect("failed to create temp dir"); + let certbot = test_certbot(data_dir.path()); + assert!(certbot + .kv_store + .try_acquire_rotation_lock(ROTATION_LOCK_TIMEOUT_SECS)); + let err = certbot + .rotate_acme_credentials() + .await + .expect_err("rotation should be rejected while the KV lock is held"); + assert!( + err.to_string() + .contains("another node is rotating ACME credentials"), + "unexpected error: {err}" + ); + } + + #[tokio::test] + async fn rotate_acme_credentials_requires_a_configured_domain() { + let data_dir = tempfile::tempdir().expect("failed to create temp dir"); + let certbot = test_certbot(data_dir.path()); + let err = certbot + .rotate_acme_credentials() + .await + .expect_err("rotation without domains should fail"); + assert!( + err.to_string().contains("no ZT-Domain configured"), + "unexpected error: {err}" + ); + // The failed rotation must release the KV lock so a later run can proceed. + assert!(certbot + .kv_store + .try_acquire_rotation_lock(ROTATION_LOCK_TIMEOUT_SECS)); + } + #[test] fn corrupt_acme_credentials_fail_closed() { assert!(acme_url_matches("not-json", "https://acme.test/directory").is_err()); diff --git a/dstack/gateway/src/kv/mod.rs b/dstack/gateway/src/kv/mod.rs index 7315bc1a6..bd975b628 100644 --- a/dstack/gateway/src/kv/mod.rs +++ b/dstack/gateway/src/kv/mod.rs @@ -268,6 +268,7 @@ pub mod keys { pub const GLOBAL_CERTBOT_CONFIG: &str = "global/certbot_config"; pub const GLOBAL_ACME_CREDENTIALS: &str = "global/acme_credentials"; pub const GLOBAL_ACME_ATTESTATION: &str = "global/acme_attestation"; + pub const GLOBAL_ACME_ROTATION_LOCK: &str = "global/acme_rotation_lock"; pub fn inst(instance_id: &str) -> String { format!("{INST_PREFIX}{instance_id}") @@ -979,6 +980,51 @@ impl KvStore { Ok(()) } + /// Try to acquire the global ACME credential rotation lock. + /// + /// Best-effort only: WaveKV is last-writer-wins without compare-and-swap, + /// so two nodes can both acquire during a replication gap. This narrows the + /// window for concurrent rotation from the full rotation duration to the + /// replication latency; it is not mutual exclusion. A crashed holder is + /// covered by the timeout. + pub fn try_acquire_rotation_lock(&self, lock_timeout_secs: u64) -> bool { + let now = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_secs(); + + if let Some(existing) = self.get_rotation_lock() { + // Check if lock is still valid (not expired) + if now < existing.started_at + lock_timeout_secs { + return false; + } + } + + let lock = CertRenewLock { + started_at: now, + started_by: self.my_node_id, + }; + self.persistent + .write() + .put_encoded(keys::GLOBAL_ACME_ROTATION_LOCK.to_string(), &lock) + .is_ok() + } + + /// Get the global ACME credential rotation lock + pub fn get_rotation_lock(&self) -> Option { + self.persistent + .read() + .decode(keys::GLOBAL_ACME_ROTATION_LOCK) + } + + /// Release the global ACME credential rotation lock + pub fn release_rotation_lock(&self) -> Result<()> { + self.persistent + .write() + .delete(keys::GLOBAL_ACME_ROTATION_LOCK.to_string())?; + Ok(()) + } + // ==================== Certificate Attestation ==================== /// Get the latest attestation for a domain diff --git a/dstack/gateway/src/main_service.rs b/dstack/gateway/src/main_service.rs index d38cdb7fa..b98bb6d99 100644 --- a/dstack/gateway/src/main_service.rs +++ b/dstack/gateway/src/main_service.rs @@ -450,10 +450,20 @@ impl Proxy { .collect(), }; - // Get account_uri, account_quote and account_attestation from global ACME attestation - let (account_uri, account_quote, account_attestation) = kv_store - .get_acme_attestation() - .map(|att| (att.account_uri, att.quote, att.attestation)) + // The account URI comes from the published credentials; the attestation + // record is written best-effort and may lag behind a rotation, so it + // only supplies the quote when it matches the current account. + let attestation = kv_store.get_acme_attestation(); + let account_uri = kv_store + .get_acme_credentials() + .and_then(|creds| { + crate::distributed_certbot::extract_account_uri(&creds.acme_credentials) + }) + .or_else(|| attestation.as_ref().map(|att| att.account_uri.clone())) + .unwrap_or_default(); + let (account_quote, account_attestation) = attestation + .filter(|att| att.account_uri == account_uri) + .map(|att| (att.quote, att.attestation)) .unwrap_or_default(); for domain in &domains { From c065d197bc43a54e4ae9decc940ac4836b5a8be6 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Wed, 5 Aug 2026 02:03:08 -0700 Subject: [PATCH 5/5] fix(gateway): fail closed on corrupt KV records and harden rotation locking Address review findings on the rotation rework: - get_acme_credentials now distinguishes a missing or deleted key from a record that no longer decodes; corruption at the WaveKV encoding layer previously read as absence and silently registered a fresh ACME account that the account-bound CAA records refuse. Issuance and AcmeInfo now fail with a pointer at RotateAcmeCredentials instead. - move account attestation after CAA re-pinning: its agent round trips do not gate issuance and must not widen the window where the published account and the CAA records disagree; still run it on partial re-pin failure so the new account is recorded - release the rotation lock only when the visible value is the one this rotation wrote, so a holder that outlived the timeout cannot delete the lock of the node that took over - re-pin the first domain through the registration client instead of constructing a second DNS client and re-fetching the ACME directory; this also removes an expect() that failed CI's panic lints (-D clippy::expect_used) Verification: - cargo test -p dstack-gateway: 77 passed (new: corrupt-record fail-closed at the KV layer, tombstone vs corruption, stale-holder lock release guard) - cargo clippy -- -D warnings -D clippy::expect_used -D clippy::unwrap_used --allow unused_variables: clean (workspace) - cargo fmt --all -- --check: clean --- dstack/gateway/src/distributed_certbot.rs | 110 ++++++++++++++++------ dstack/gateway/src/kv/mod.rs | 105 +++++++++++++++++++-- dstack/gateway/src/main_service.rs | 1 + 3 files changed, 179 insertions(+), 37 deletions(-) diff --git a/dstack/gateway/src/distributed_certbot.rs b/dstack/gateway/src/distributed_certbot.rs index d34959c91..a7dc07a95 100644 --- a/dstack/gateway/src/distributed_certbot.rs +++ b/dstack/gateway/src/distributed_certbot.rs @@ -85,14 +85,14 @@ impl DistributedCertBot { let Ok(_guard) = self.caa_lock.try_lock() else { bail!("ACME credential rotation or CAA reconciliation is already in progress"); }; - if !self + let Some(rotation_lock) = self .kv_store .try_acquire_rotation_lock(ROTATION_LOCK_TIMEOUT_SECS) - { + else { bail!("another node is rotating ACME credentials; retry after it finishes"); - } + }; let result = self.do_rotate_acme_credentials().await; - if let Err(err) = self.kv_store.release_rotation_lock() { + if let Err(err) = self.kv_store.release_rotation_lock(&rotation_lock) { error!("failed to release ACME rotation lock: {err:?}"); } result @@ -100,9 +100,6 @@ impl DistributedCertBot { async fn do_rotate_acme_credentials(&self) -> Result<(String, usize)> { let configs = self.kv_store.list_zt_domain_configs(); - if configs.is_empty() { - bail!("no ZT-Domain configured for ACME credential rotation"); - } let certbot_config = self.config(); let acme_url = if certbot_config.acme_url.is_empty() { DEFAULT_ACME_URL @@ -123,15 +120,17 @@ impl DistributedCertBot { .with_context(|| format!("DNS credential check failed for {}", config.domain))?; prepared.push((&config.domain, dns_cred, dns_client)); } + let total = prepared.len(); + let mut prepared = prepared.into_iter(); + let Some((first_domain, first_cred, first_client)) = prepared.next() else { + bail!("no ZT-Domain configured for ACME credential rotation"); + }; - let first = configs.first().expect("configs is non-empty"); - let first_dns_cred = dns_credential_for(&self.kv_store, first)?; - let account_dns_client = self.dns_client(&first.domain, &first_dns_cred).await?; let client = AcmeClient::new_account( acme_url, - account_dns_client, - first_dns_cred.max_dns_wait, - first_dns_cred.dns_txt_ttl, + first_client, + first_cred.max_dns_wait, + first_cred.dns_txt_ttl, ) .await .context("failed to create replacement ACME account")?; @@ -147,14 +146,26 @@ impl DistributedCertBot { self.kv_store.save_acme_credentials(&CertCredentials { acme_credentials: credentials.clone(), })?; - if let Err(err) = self.generate_and_save_acme_attestation(&account_uri).await { - warn!("failed to attest rotated ACME account: {err:?}"); - } // Re-pin every domain's CAA to the new account, best effort across all - // domains: one failing domain must not block re-pinning the rest. - let total = prepared.len(); + // domains: one failing domain must not block re-pinning the rest. The + // first domain reuses the registration client, which is already bound + // to its DNS client and the new credentials. let mut failed = Vec::new(); + let mut record = |domain: &String, result: Result<()>| match result { + Ok(()) => info!("cert[{domain}]: CAA re-pinned to {account_uri}"), + Err(err) => { + error!("cert[{domain}]: failed to re-pin CAA: {err:?}"); + failed.push(domain.clone()); + } + }; + record( + first_domain, + client + .set_caa_records(std::slice::from_ref(first_domain)) + .await + .context("failed to update CAA records"), + ); for (domain, dns_cred, dns_client) in prepared { let result = async { let client = AcmeClient::load( @@ -171,14 +182,18 @@ impl DistributedCertBot { .context("failed to update CAA records") } .await; - match result { - Ok(()) => info!("cert[{domain}]: CAA re-pinned to {account_uri}"), - Err(err) => { - error!("cert[{domain}]: failed to re-pin CAA: {err:?}"); - failed.push(domain.clone()); - } - } + record(domain, result); } + + // Attest the new account only after CAA re-pinning: attestation does + // not gate issuance, so its agent round trips must not widen the + // window where the published account and the CAA records disagree. + // Run it even when some domains failed so the new account is still + // recorded. + if let Err(err) = self.generate_and_save_acme_attestation(&account_uri).await { + warn!("failed to attest rotated ACME account: {err:?}"); + } + if !failed.is_empty() { bail!( "rotated to {account_uri} and published the new credentials, but failed to \ @@ -516,8 +531,15 @@ impl DistributedCertBot { &config.acme_url }; - // Try to load global ACME credentials from KvStore - if let Some(creds) = self.kv_store.get_acme_credentials() { + // Try to load global ACME credentials from KvStore. A corrupt record + // is an error, not absence: falling through to account registration + // would silently create an account that the account-bound CAA records + // refuse, and burn a rate-limited registration. + let stored_creds = self + .kv_store + .get_acme_credentials() + .context("call RotateAcmeCredentials to replace the stored ACME credentials")?; + if let Some(creds) = stored_creds { if !acme_url_matches(&creds.acme_credentials, acme_url).context( "invalid ACME credentials in KvStore; call RotateAcmeCredentials to replace them", )? { @@ -795,7 +817,8 @@ mod tests { let certbot = test_certbot(data_dir.path()); assert!(certbot .kv_store - .try_acquire_rotation_lock(ROTATION_LOCK_TIMEOUT_SECS)); + .try_acquire_rotation_lock(ROTATION_LOCK_TIMEOUT_SECS) + .is_some()); let err = certbot .rotate_acme_credentials() .await @@ -822,7 +845,36 @@ mod tests { // The failed rotation must release the KV lock so a later run can proceed. assert!(certbot .kv_store - .try_acquire_rotation_lock(ROTATION_LOCK_TIMEOUT_SECS)); + .try_acquire_rotation_lock(ROTATION_LOCK_TIMEOUT_SECS) + .is_some()); + } + + #[tokio::test] + async fn stale_rotation_holder_does_not_release_a_newer_lock() { + let data_dir = tempfile::tempdir().expect("failed to create temp dir"); + let certbot = test_certbot(data_dir.path()); + let current = certbot + .kv_store + .try_acquire_rotation_lock(ROTATION_LOCK_TIMEOUT_SECS) + .expect("lock should be free"); + // Simulate a holder that exceeded the timeout and was superseded. + let stale = crate::kv::CertRenewLock { + started_at: current.started_at.saturating_sub(100), + started_by: 99, + }; + certbot + .kv_store + .release_rotation_lock(&stale) + .expect("stale release should be a no-op, not an error"); + assert!( + certbot.kv_store.get_rotation_lock().is_some(), + "the newer holder's lock must remain in place" + ); + certbot + .kv_store + .release_rotation_lock(¤t) + .expect("owner release should succeed"); + assert!(certbot.kv_store.get_rotation_lock().is_none()); } #[test] diff --git a/dstack/gateway/src/kv/mod.rs b/dstack/gateway/src/kv/mod.rs index bd975b628..213be0b2e 100644 --- a/dstack/gateway/src/kv/mod.rs +++ b/dstack/gateway/src/kv/mod.rs @@ -915,9 +915,24 @@ impl KvStore { // ==================== Global ACME Credentials ==================== - /// Get global ACME credentials (shared across all domains) - pub fn get_acme_credentials(&self) -> Option { - self.persistent.read().decode(keys::GLOBAL_ACME_CREDENTIALS) + /// Get global ACME credentials (shared across all domains). + /// + /// Fails closed on a corrupt record: a missing or deleted key is + /// `Ok(None)`, but a stored value that no longer decodes is an error. + /// Treating corruption as absence would silently register a fresh ACME + /// account that the existing account-bound CAA records refuse. + pub fn get_acme_credentials(&self) -> Result> { + let state = self.persistent.read(); + let Some(entry) = state.get(keys::GLOBAL_ACME_CREDENTIALS) else { + return Ok(None); + }; + // A `None` value is a tombstone: the key was deliberately deleted. + let Some(value) = entry.value.as_ref() else { + return Ok(None); + }; + decode(value) + .map(Some) + .context("corrupt ACME credentials record in KvStore") } /// Save global ACME credentials @@ -982,12 +997,16 @@ impl KvStore { /// Try to acquire the global ACME credential rotation lock. /// + /// Returns the lock value that was written; pass it back to + /// [`Self::release_rotation_lock`] so a rotation that outlived the timeout + /// cannot delete the lock of the node that took over. + /// /// Best-effort only: WaveKV is last-writer-wins without compare-and-swap, /// so two nodes can both acquire during a replication gap. This narrows the /// window for concurrent rotation from the full rotation duration to the /// replication latency; it is not mutual exclusion. A crashed holder is /// covered by the timeout. - pub fn try_acquire_rotation_lock(&self, lock_timeout_secs: u64) -> bool { + pub fn try_acquire_rotation_lock(&self, lock_timeout_secs: u64) -> Option { let now = std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) .unwrap_or_default() @@ -996,7 +1015,7 @@ impl KvStore { if let Some(existing) = self.get_rotation_lock() { // Check if lock is still valid (not expired) if now < existing.started_at + lock_timeout_secs { - return false; + return None; } } @@ -1007,7 +1026,8 @@ impl KvStore { self.persistent .write() .put_encoded(keys::GLOBAL_ACME_ROTATION_LOCK.to_string(), &lock) - .is_ok() + .ok()?; + Some(lock) } /// Get the global ACME credential rotation lock @@ -1017,8 +1037,25 @@ impl KvStore { .decode(keys::GLOBAL_ACME_ROTATION_LOCK) } - /// Release the global ACME credential rotation lock - pub fn release_rotation_lock(&self) -> Result<()> { + /// Release the global ACME credential rotation lock. + /// + /// Only deletes the lock when the currently visible value is the one that + /// `acquired` wrote: a rotation that outlived the lock timeout must not + /// delete the lock of the node that took over (which would let a third + /// rotation start concurrently). Like acquisition, the check is + /// best-effort under WaveKV's last-writer-wins replication. + pub fn release_rotation_lock(&self, acquired: &CertRenewLock) -> Result<()> { + if let Some(current) = self.get_rotation_lock() { + if current.started_by != acquired.started_by + || current.started_at != acquired.started_at + { + warn!( + "not releasing ACME rotation lock: node {} took it over after this rotation exceeded the lock timeout", + current.started_by + ); + return Ok(()); + } + } self.persistent .write() .delete(keys::GLOBAL_ACME_ROTATION_LOCK.to_string())?; @@ -1096,6 +1133,58 @@ fn validate_peer_url(url: &str) -> Result<()> { Ok(()) } +#[cfg(test)] +mod acme_credentials_tests { + use super::*; + + fn test_kv(data_dir: &std::path::Path) -> KvStore { + KvStore::new(1, vec![], data_dir).expect("failed to create kv store") + } + + #[test] + fn missing_and_deleted_credentials_are_absent_not_errors() { + let dir = tempfile::tempdir().expect("failed to create temp dir"); + let kv = test_kv(dir.path()); + assert!(kv + .get_acme_credentials() + .expect("missing key should not error") + .is_none()); + + kv.save_acme_credentials(&CertCredentials { + acme_credentials: "{}".to_string(), + }) + .expect("save should succeed"); + kv.persistent + .write() + .delete(keys::GLOBAL_ACME_CREDENTIALS.to_string()) + .expect("delete should succeed"); + assert!(kv + .get_acme_credentials() + .expect("tombstone should not error") + .is_none()); + } + + #[test] + fn corrupt_credentials_record_fails_closed() { + let dir = tempfile::tempdir().expect("failed to create temp dir"); + let kv = test_kv(dir.path()); + kv.persistent + .write() + .put( + keys::GLOBAL_ACME_CREDENTIALS.to_string(), + b"not-messagepack".to_vec(), + ) + .expect("raw put should succeed"); + let err = kv + .get_acme_credentials() + .expect_err("corrupt record must not read as absent"); + assert!( + err.to_string().contains("corrupt ACME credentials"), + "unexpected error: {err}" + ); + } +} + #[cfg(test)] mod peer_url_tests { use super::validate_peer_url; diff --git a/dstack/gateway/src/main_service.rs b/dstack/gateway/src/main_service.rs index b98bb6d99..6167eaf1c 100644 --- a/dstack/gateway/src/main_service.rs +++ b/dstack/gateway/src/main_service.rs @@ -456,6 +456,7 @@ impl Proxy { let attestation = kv_store.get_acme_attestation(); let account_uri = kv_store .get_acme_credentials() + .context("call RotateAcmeCredentials to replace the stored ACME credentials")? .and_then(|creds| { crate::distributed_certbot::extract_account_uri(&creds.acme_credentials) })