From c3b9189eb1836a074ea2fcf7cbbe4ed9128bc34c Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 5 Nov 2026 22:05:00 -0800 Subject: [PATCH 01/17] kms: keep a multi-epoch root keyring and select epochs in GetAppKey Every KMS now holds root keys for epochs 0..current_epoch. The current epoch's roots stay in the existing top-level files, so all paths that only use the current roots are unchanged; older epochs live under epochs// and a root-epoch marker records the stored epoch. Startup refuses a configured epoch that does not match the stored keys. GetAppKey gains a sint32 key_epoch selector (absent = epoch 0, -1 = current) and reports the resolved epoch, so old OS images keep their keys while new ones can ask for the latest epoch or an older one to migrate. Older epochs are not time-limited: their roots are already exposed, so withholding them only strands late instances. Key handover sends the full keyring tagged by epoch. Onboarding requires exactly epochs 0..N from the source, or with rotate imports 0..N-1 and generates epoch N locally, attesting the new public keys bound to the epoch. Bootstrap only creates epoch 0. Signed-off-by: Kevin Wang --- dstack/dstack-util/src/main.rs | 1 + dstack/dstack-util/src/system_setup.rs | 1 + dstack/kms/kms.toml | 6 + dstack/kms/rpc/proto/kms_rpc.proto | 16 +- dstack/kms/src/config.rs | 65 ++++++ dstack/kms/src/main_service.rs | 110 +++++++--- dstack/kms/src/onboard_service.rs | 278 +++++++++++++++++++++---- dstack/kms/src/www/onboard.html | 12 +- 8 files changed, 418 insertions(+), 71 deletions(-) diff --git a/dstack/dstack-util/src/main.rs b/dstack/dstack-util/src/main.rs index 87867085b..9b28165b9 100644 --- a/dstack/dstack-util/src/main.rs +++ b/dstack/dstack-util/src/main.rs @@ -701,6 +701,7 @@ async fn cmd_get_keys(args: GetKeysArgs) -> Result<()> { .get_app_key(dstack_kms_rpc::GetAppKeyRequest { api_version: 1, vm_config: "".to_string(), + ..Default::default() }) .await .context("Failed to get app key")?; diff --git a/dstack/dstack-util/src/system_setup.rs b/dstack/dstack-util/src/system_setup.rs index 358a8a003..35aa841ab 100644 --- a/dstack/dstack-util/src/system_setup.rs +++ b/dstack/dstack-util/src/system_setup.rs @@ -2445,6 +2445,7 @@ impl<'a> Stage0<'a> { .get_app_key(rpc::GetAppKeyRequest { api_version: 1, vm_config: self.shared.sys_config.vm_config.clone(), + ..Default::default() }) .await .context("Failed to get app key")?; diff --git a/dstack/kms/kms.toml b/dstack/kms/kms.toml index 22ee6d7b5..233fda094 100644 --- a/dstack/kms/kms.toml +++ b/dstack/kms/kms.toml @@ -124,6 +124,12 @@ timeout = "60s" [core.auth_api.dev] gateway_app_id = "any" +[core.key_rotation] +# Root key epoch served as current. Raising it requires a KMS onboarded with +# `rotate` (or from a source already at this epoch); editing it alone fails. +# Older epochs stay available so CVMs can migrate their data at any time. +current_epoch = 0 + [core.onboard] enabled = true auto_bootstrap_domain = "" diff --git a/dstack/kms/rpc/proto/kms_rpc.proto b/dstack/kms/rpc/proto/kms_rpc.proto index 6853df2be..0405d2606 100644 --- a/dstack/kms/rpc/proto/kms_rpc.proto +++ b/dstack/kms/rpc/proto/kms_rpc.proto @@ -11,6 +11,9 @@ package kms; message GetAppKeyRequest { uint32 api_version = 1; string vm_config = 2; + // Root key epoch to derive from: 0 (default) = epoch 0, N > 0 = epoch N, + // -1 = the current epoch. Older epochs remain available for migration. + sint32 key_epoch = 3; } message AppId { @@ -46,6 +49,8 @@ message AppKeyResponse { string gateway_app_id = 7; // OS Image hash bytes os_image_hash = 8; + // Root key epoch all epoch-dependent fields above were derived from. + uint32 key_epoch = 9; } message GetMetaResponse { @@ -70,6 +75,8 @@ message GetKmsKeyRequest { message KmsKeys { string ca_key = 1; bytes k256_key = 2; + // Root key epoch of this pair. + uint32 epoch = 3; } message KmsKeyResponse { @@ -158,11 +165,18 @@ message OnboardRequest { string domain = 2; // Optional bearer token for the source KMS admin listener. string source_token = 3; + // Import epochs 0..N-1 from the source and generate epoch N locally, where N + // is this KMS's configured current epoch. + bool rotate = 4; } message OnboardResponse { - // k256 public key (secp256k1) inherited from source KMS + // k256 public key (secp256k1) of the current epoch bytes k256_pubkey = 1; + // Set only by a rotation: the generated CA public key and the attestation + // binding the new epoch's public keys. + bytes ca_pubkey = 2; + bytes attestation = 3; } // Attestation info needed for on-chain KMS authorization. diff --git a/dstack/kms/src/config.rs b/dstack/kms/src/config.rs index fd2200fb8..f3869ca23 100644 --- a/dstack/kms/src/config.rs +++ b/dstack/kms/src/config.rs @@ -2,6 +2,7 @@ // // SPDX-License-Identifier: Apache-2.0 +use anyhow::{anyhow, bail, Result}; use dstack_attest::attestation::AttestationVerifierConfig; use load_config::load_config; use rocket::figment::Figment; @@ -21,6 +22,8 @@ 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 ROOT_EPOCH: &str = "root-epoch"; +const EPOCHS_DIR: &str = "epochs"; const BOOTSTRAP_INFO: &str = "bootstrap-info.json"; #[derive(Debug, Clone, Deserialize)] @@ -38,6 +41,8 @@ pub(crate) struct KmsConfig { pub attestation: AttestationVerifierConfig, pub auth_api: AuthApi, pub onboard: OnboardConfig, + #[serde(default)] + pub key_rotation: KeyRotationConfig, pub image: ImageConfig, /// Whether to enable the additional local release gate for AMD SEV-SNP /// key/cert material. This is separate from the auth API so production @@ -171,6 +176,27 @@ impl KmsConfig { self.cert_dir.join(K256_KEY) } + /// Records which epoch the top-level root keys belong to. Absent means 0. + pub fn root_epoch(&self) -> PathBuf { + self.cert_dir.join(ROOT_EPOCH) + } + + /// The current epoch's roots are the top-level root files, so every path + /// that only needs the current roots is unchanged. Older epochs live in + /// `epochs//`. + pub fn root_key_paths(&self, epoch: u32) -> RootKeyPaths { + let dir = if epoch == self.key_rotation.current_epoch { + self.cert_dir.clone() + } else { + self.cert_dir.join(EPOCHS_DIR).join(epoch.to_string()) + }; + RootKeyPaths { + ca_cert: dir.join(ROOT_CA_CERT), + ca_key: dir.join(ROOT_CA_KEY), + k256_key: dir.join(K256_KEY), + } + } + pub fn bootstrap_info(&self) -> PathBuf { self.cert_dir.join(BOOTSTRAP_INFO) } @@ -209,6 +235,34 @@ pub(crate) struct Dev { pub gateway_app_id: String, } +pub(crate) struct RootKeyPaths { + pub ca_cert: PathBuf, + pub ca_key: PathBuf, + pub k256_key: PathBuf, +} + +#[derive(Debug, Clone, Default, Deserialize)] +pub(crate) struct KeyRotationConfig { + /// Root key epoch served as current. Every KMS holds all epochs up to it. + #[serde(default)] + pub current_epoch: u32, +} + +impl KeyRotationConfig { + /// Resolve a GetAppKey epoch selector: `-1` is the current epoch. + pub fn resolve_epoch(&self, requested: i32) -> Result { + let current = self.current_epoch; + let epoch = match requested { + -1 => current, + epoch => u32::try_from(epoch).map_err(|_| anyhow!("invalid key epoch {requested}"))?, + }; + if epoch > current { + bail!("unknown key epoch {epoch}; the current epoch is {current}"); + } + Ok(epoch) + } +} + #[derive(Debug, Clone, Deserialize)] pub(crate) struct OnboardConfig { pub enabled: bool, @@ -238,6 +292,17 @@ mod tests { assert!(!config.admin.insecure_no_auth); } + #[test] + fn key_epoch_selector_resolution() { + let cfg = KeyRotationConfig { current_epoch: 2 }; + assert_eq!(cfg.resolve_epoch(-1).unwrap(), 2); + assert_eq!(cfg.resolve_epoch(0).unwrap(), 0); + assert_eq!(cfg.resolve_epoch(2).unwrap(), 2); + assert!(cfg.resolve_epoch(3).is_err()); + assert!(cfg.resolve_epoch(-2).is_err()); + assert_eq!(KeyRotationConfig::default().resolve_epoch(-1).unwrap(), 0); + } + #[test] fn rpc_cert_is_attested_by_default() { let figment = load_config_figment(None); diff --git a/dstack/kms/src/main_service.rs b/dstack/kms/src/main_service.rs index baac0147d..12274fe2d 100644 --- a/dstack/kms/src/main_service.rs +++ b/dstack/kms/src/main_service.rs @@ -34,7 +34,7 @@ use upgrade_authority::{ }; use crate::{ - config::KmsConfig, + config::{KmsConfig, RootKeyPaths}, crypto::{derive_k256_key, sign_message, sign_message_with_timestamp}, }; @@ -67,10 +67,26 @@ impl std::ops::Deref for KmsState { } } +/// The root key pair of one epoch. +struct RootKeys { + ca: CaCert, + k256: SigningKey, +} + +impl RootKeys { + fn load(paths: &RootKeyPaths) -> Result { + let ca = CaCert::load(&paths.ca_cert, &paths.ca_key) + .context("failed to load root CA certificate")?; + let key_bytes = fs::read(&paths.k256_key).context("failed to read ECDSA root key")?; + let k256 = SigningKey::from_slice(&key_bytes).context("failed to load ECDSA root key")?; + Ok(Self { ca, k256 }) + } +} + pub struct KmsStateInner { config: KmsConfig, - root_ca: CaCert, - k256_key: SigningKey, + /// Root keys of every epoch, indexed by epoch; the last is current. + roots: Vec, temp_ca_cert: String, temp_ca_key: String, verifier: CvmVerifier, @@ -128,11 +144,20 @@ impl KmsState { } pub fn new(config: KmsConfig) -> Result { - let root_ca = CaCert::load(config.root_ca_cert(), config.root_ca_key()) - .context("Failed to load root CA certificate")?; - 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 current_epoch = config.key_rotation.current_epoch; + let stored_epoch = read_root_epoch(&config)?; + if stored_epoch != current_epoch { + bail!( + "stored root keys belong to epoch {stored_epoch}, but key_rotation.current_epoch \ + is {current_epoch}; a new epoch must be imported or generated by onboarding" + ); + } + let roots = (0..=current_epoch) + .map(|epoch| { + RootKeys::load(&config.root_key_paths(epoch)) + .with_context(|| format!("failed to load root keys of epoch {epoch}")) + }) + .collect::>>()?; let temp_ca_key = fs::read_to_string(config.tmp_ca_key()).context("Faeild to read temp ca key")?; let temp_ca_cert = @@ -155,8 +180,7 @@ impl KmsState { Ok(Self { inner: Arc::new(KmsStateInner { config, - root_ca, - k256_key, + roots, temp_ca_cert, temp_ca_key, verifier, @@ -179,6 +203,19 @@ impl KmsState { pub(crate) fn attestation_verifier(&self) -> Arc { self.inner.attestation_verifier.clone() } + + /// The current epoch's root keys. + fn root(&self) -> &RootKeys { + self.roots.last().expect("the keyring always holds epoch 0") + } +} + +pub(crate) fn read_root_epoch(config: &KmsConfig) -> Result { + match fs::read_to_string(config.root_epoch()) { + Ok(epoch) => epoch.trim().parse().context("invalid root epoch file"), + Err(err) if err.kind() == std::io::ErrorKind::NotFound => Ok(0), + Err(err) => Err(err).context("failed to read root epoch file"), + } } pub struct RpcHandler { @@ -202,10 +239,14 @@ impl RpcHandler { ensure_key_release_allowed(&info, (&self.state.config).into())?; 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: (0..) + .zip(&self.state.roots) + .map(|(epoch, root)| KmsKeys { + ca_key: root.ca.key.serialize_pem(), + k256_key: root.k256.to_bytes().to_vec(), + epoch, + }) + .collect(), }) } } @@ -369,7 +410,7 @@ impl RpcHandler { fn derive_app_ca(&self, app_id: &[u8]) -> Result { let context_data = vec![app_id, b"app-ca"]; - let app_key = kdf::derive_p256_key_pair(&self.state.root_ca.key, &context_data) + let app_key = kdf::derive_p256_key_pair(&self.state.root().ca.key, &context_data) .context("Failed to derive app disk key")?; let req = CertRequest::builder() .key(&app_key) @@ -381,7 +422,8 @@ impl RpcHandler { .build(); let app_ca = self .state - .root_ca + .root() + .ca .sign(req) .context("Failed to sign App CA")?; Ok(CaCert::from_parts(app_key, app_ca)) @@ -404,29 +446,34 @@ impl KmsRpc for RpcHandler { .await .context("App not allowed")?; ensure_key_release_allowed(&boot_info, (&self.state.config).into())?; + let key_epoch = self + .state + .config + .key_rotation + .resolve_epoch(request.key_epoch)?; + let root = &self.state.roots[key_epoch as usize]; let app_id = boot_info.app_id; let instance_id = boot_info.instance_id; let os_image_hash = boot_info.os_image_hash; let context_data = vec![&app_id[..], &instance_id[..], b"app-disk-crypt-key"]; - let app_disk_key = kdf::derive_dh_secret(&self.state.root_ca.key, &context_data) + let app_disk_key = kdf::derive_dh_secret(&root.ca.key, &context_data) .context("Failed to derive app disk key")?; let env_crypt_key = { - let secret = - kdf::derive_dh_secret(&self.state.root_ca.key, &[&app_id[..], b"env-encrypt-key"]) - .context("Failed to derive env encrypt key")?; + let secret = kdf::derive_dh_secret(&root.ca.key, &[&app_id[..], b"env-encrypt-key"]) + .context("Failed to derive env encrypt key")?; let secret = x25519_dalek::StaticSecret::from(secret); secret.to_bytes() }; let (k256_key, k256_signature) = { - let (k256_app_key, signature) = derive_k256_key(&self.state.k256_key, &app_id) - .context("Failed to derive app ecdsa key")?; + let (k256_app_key, signature) = + derive_k256_key(&root.k256, &app_id).context("Failed to derive app ecdsa key")?; (k256_app_key.to_bytes().to_vec(), signature) }; Ok(AppKeyResponse { - ca_cert: self.state.root_ca.pem_cert.clone(), + ca_cert: root.ca.pem_cert.clone(), disk_crypt_key: app_disk_key.to_vec(), env_crypt_key: env_crypt_key.to_vec(), k256_key, @@ -434,6 +481,7 @@ impl KmsRpc for RpcHandler { tproxy_app_id: gateway_app_id.clone(), gateway_app_id, os_image_hash, + key_epoch, }) } @@ -443,7 +491,7 @@ impl KmsRpc for RpcHandler { .context("KMS self authorization failed")?; ensure_app_id_len(&request.app_id)?; let secret = kdf::derive_dh_secret( - &self.state.root_ca.key, + &self.state.root().ca.key, &[&request.app_id[..], "env-encrypt-key".as_bytes()], ) .context("Failed to derive env encrypt key")?; @@ -458,7 +506,7 @@ impl KmsRpc for RpcHandler { // Legacy signature (without timestamp) for backward compatibility let signature = sign_message( - &self.state.k256_key, + &self.state.root().k256, b"dstack-env-encrypt-pubkey", &request.app_id, &public_key, @@ -467,7 +515,7 @@ impl KmsRpc for RpcHandler { // New signature with timestamp to prevent replay attacks let signature_v1 = sign_message_with_timestamp( - &self.state.k256_key, + &self.state.root().k256, b"dstack-env-encrypt-pubkey", &request.app_id, timestamp, @@ -494,12 +542,12 @@ impl KmsRpc for RpcHandler { .await .map_err(|err| anyhow!("{err:#}"))?; Ok(GetMetaResponse { - ca_cert: self.state.inner.root_ca.pem_cert.clone(), + ca_cert: self.state.root().ca.pem_cert.clone(), allow_any_upgrade: self.state.inner.config.auth_api.is_dev(), k256_pubkey: self .state - .inner - .k256_key + .root() + .k256 .verifying_key() .to_sec1_bytes() .to_vec(), @@ -543,7 +591,7 @@ impl KmsRpc for RpcHandler { Ok(GetTempCaCertResponse { temp_ca_cert: self.state.inner.temp_ca_cert.clone(), temp_ca_key: self.state.inner.temp_ca_key.clone(), - ca_cert: self.state.inner.root_ca.pem_cert.clone(), + ca_cert: self.state.root().ca.pem_cert.clone(), }) } @@ -592,7 +640,7 @@ impl KmsRpc for RpcHandler { certificate_chain: vec![ cert.pem(), app_ca.pem_cert.clone(), - self.state.root_ca.pem_cert.clone(), + self.state.root().ca.pem_cert.clone(), ], }) } diff --git a/dstack/kms/src/onboard_service.rs b/dstack/kms/src/onboard_service.rs index c16b67d4d..20775bd3c 100644 --- a/dstack/kms/src/onboard_service.rs +++ b/dstack/kms/src/onboard_service.rs @@ -140,22 +140,14 @@ impl OnboardRpc for OnboardHandler { if cfg.root_keys_exist() { bail!("KMS has already been bootstrapped"); } + ensure_bootstrap_epoch(cfg)?; ensure_self_kms_allowed(cfg, &self.state.attestation_verifier) .await .context("KMS is not allowed to bootstrap")?; let keys = Keys::generate(&request.domain, self.state.config.attest_rpc_cert) .await .context("Failed to generate keys")?; - - let k256_pubkey = keys.k256_key.verifying_key().to_sec1_bytes().to_vec(); - let ca_pubkey = keys.ca_key.public_key_der(); - let attestation = attest_keys(&ca_pubkey, &k256_pubkey).await?; - - let response = BootstrapResponse { - ca_pubkey, - k256_pubkey, - attestation, - }; + let response = keys.attest().await?; // Store the bootstrap info safe_write(cfg.bootstrap_info(), serde_json::to_vec(&response)?)?; keys.store(cfg)?; @@ -180,13 +172,29 @@ impl OnboardRpc for OnboardHandler { &source_url, &request.source_token, &request.domain, + request.rotate, self.state.attestation_verifier.clone(), ) .await .context("Failed to onboard")?; - let k256_pubkey = keys.k256_key.verifying_key().to_sec1_bytes().to_vec(); + if !request.rotate { + let k256_pubkey = keys.k256_key.verifying_key().to_sec1_bytes().to_vec(); + keys.store(cfg).context("Failed to store keys")?; + return Ok(OnboardResponse { + k256_pubkey, + ..Default::default() + }); + } + // The attestation over the new epoch is what governance registers, so + // it is persisted with the keys and stays readable via GetSetupStatus. + let info = keys.attest().await?; + safe_write(cfg.bootstrap_info(), serde_json::to_vec(&info)?)?; keys.store(cfg).context("Failed to store keys")?; - Ok(OnboardResponse { k256_pubkey }) + Ok(OnboardResponse { + k256_pubkey: info.k256_pubkey, + ca_pubkey: info.ca_pubkey, + attestation: info.attestation, + }) } async fn get_attestation_info(self) -> Result { @@ -287,6 +295,7 @@ mod tests { use crate::main_service::amd_attest::{ compute_expected_measurement, MeasurementInput, OvmfSectionParam, }; + use dstack_kms_rpc::KmsKeys; use sha2::Digest; fn hex_of(byte: u8, len: usize) -> String { @@ -536,6 +545,104 @@ mod tests { assert!(status.attestation.is_empty()); } + fn handover(epochs: u32) -> KmsKeyResponse { + KmsKeyResponse { + temp_ca_key: KeyPair::generate_for(&PKCS_ECDSA_P256_SHA256) + .unwrap() + .serialize_pem(), + keys: (0..epochs) + .map(|epoch| KmsKeys { + ca_key: KeyPair::generate_for(&PKCS_ECDSA_P256_SHA256) + .unwrap() + .serialize_pem(), + k256_key: SigningKey::random(&mut rand::rngs::OsRng) + .to_bytes() + .to_vec(), + epoch, + }) + .collect(), + } + } + + fn epoch_config(cert_dir: &std::path::Path, current_epoch: u32) -> KmsConfig { + let mut cfg = auto_bootstrap_config(cert_dir); + cfg.key_rotation.current_epoch = current_epoch; + cfg + } + + #[rocket::async_test] + async fn join_imports_and_stores_every_epoch() { + let cert_dir = tempfile::tempdir().unwrap(); + let cfg = epoch_config(cert_dir.path(), 2); + let res = handover(3); + let sent: Vec<_> = res.keys.iter().map(|k| k.k256_key.clone()).collect(); + + let keys = Keys::from_key_response(&cfg, res, "kms.example.com", false) + .await + .unwrap(); + keys.store(&cfg).unwrap(); + + assert_eq!(crate::main_service::read_root_epoch(&cfg).unwrap(), 2); + for (epoch, k256) in (0..).zip(&sent) { + assert_eq!(&fs::read(cfg.root_key_paths(epoch).k256_key).unwrap(), k256); + } + assert_eq!(&fs::read(cfg.k256_key()).unwrap(), &sent[2]); + + crate::main_service::KmsState::new(cfg.clone()).unwrap(); + // Editing the configured epoch cannot relabel the stored roots. + let err = crate::main_service::KmsState::new(epoch_config(cert_dir.path(), 3)) + .err() + .unwrap(); + assert!(format!("{err:#}").contains("belong to epoch 2"), "{err:#}"); + } + + #[rocket::async_test] + async fn rotation_imports_history_and_generates_the_new_epoch() { + let cert_dir = tempfile::tempdir().unwrap(); + let cfg = epoch_config(cert_dir.path(), 2); + let res = handover(2); + let sent: Vec<_> = res.keys.iter().map(|k| k.k256_key.clone()).collect(); + + let keys = Keys::from_key_response(&cfg, res, "kms.example.com", true) + .await + .unwrap(); + keys.store(&cfg).unwrap(); + + assert_eq!(fs::read(cfg.root_key_paths(1).k256_key).unwrap(), sent[1]); + let current = fs::read(cfg.k256_key()).unwrap(); + assert!(!sent.contains(¤t)); + } + + #[rocket::async_test] + async fn onboarding_rejects_incomplete_or_extra_history() { + let cert_dir = tempfile::tempdir().unwrap(); + let cfg = epoch_config(cert_dir.path(), 1); + for (epochs, rotate) in [(1, false), (3, false), (2, true), (0, true)] { + let res = handover(epochs); + assert!( + Keys::from_key_response(&cfg, res, "kms.example.com", rotate) + .await + .is_err(), + "epochs={epochs} rotate={rotate}" + ); + } + let cfg = epoch_config(cert_dir.path(), 0); + assert!( + Keys::from_key_response(&cfg, handover(0), "kms.example.com", true) + .await + .is_err() + ); + } + + #[rocket::async_test] + async fn bootstrap_refuses_a_later_epoch() { + let cert_dir = tempfile::tempdir().unwrap(); + let cfg = epoch_config(cert_dir.path(), 1); + let verifier = AttestationVerifier::load(&cfg.attestation).unwrap(); + assert!(bootstrap_keys(&cfg, &verifier).await.is_err()); + assert!(!cfg.root_keys_exist()); + } + #[rocket::async_test] async fn setup_status_includes_bootstrap_info_when_present() { let cert_dir = tempfile::tempdir().unwrap(); @@ -607,7 +714,28 @@ mod tests { } } +/// Root keys of an epoch older than the current one. +struct HistoricalRoot { + ca_key: KeyPair, + ca_cert: Certificate, + k256_key: SigningKey, +} + +impl HistoricalRoot { + fn new(ca_key: KeyPair, k256_key: SigningKey) -> Result { + let ca_cert = self_signed_root_ca(&ca_key)?; + Ok(Self { + ca_key, + ca_cert, + k256_key, + }) + } +} + struct Keys { + /// Epoch of the root keys below; `history` holds epochs `0..epoch`. + epoch: u32, + history: Vec, k256_key: SigningKey, tmp_ca_key: KeyPair, tmp_ca_cert: Certificate, @@ -635,6 +763,19 @@ impl Keys { .await } + /// Attest the current root public keys. Epoch 0 keeps the original + /// bootstrap encoding; later epochs also bind the epoch number. + async fn attest(&self) -> Result { + let ca_pubkey = self.ca_key.public_key_der(); + let k256_pubkey = self.k256_key.verifying_key().to_sec1_bytes().to_vec(); + let attestation = attest_keys(self.epoch, &ca_pubkey, &k256_pubkey).await?; + Ok(BootstrapResponse { + ca_pubkey, + k256_pubkey, + attestation, + }) + } + async fn from_keys( tmp_ca_key: KeyPair, ca_key: KeyPair, @@ -651,14 +792,7 @@ impl Keys { .build() .self_signed()?; - // Create self-signed KMS cert - let ca_cert = CertRequest::builder() - .org_name("Dstack") - .subject("Dstack KMS CA") - .ca_level(1) - .key(&ca_key) - .build() - .self_signed()?; + let ca_cert = self_signed_root_ca(&ca_key)?; // The only place the KMS embeds its own attestation. Skipping it lets // the KMS run outside a TEE for development; it does not affect the // verification of quotes presented *to* the KMS, which is a separate @@ -689,6 +823,8 @@ impl Keys { .build() .signed_by(&ca_cert, &ca_key)?; Ok(Keys { + epoch: 0, + history: Vec::new(), k256_key, tmp_ca_key, tmp_ca_cert, @@ -705,6 +841,7 @@ impl Keys { other_kms_url: &str, source_token: &str, domain: &str, + rotate: bool, attestation_verifier: Arc, ) -> Result { let attestation_slot = Arc::new(Mutex::new(None::)); @@ -756,7 +893,7 @@ impl Keys { ensure_kms_allowed(cfg, &source_attestation()?, &attestation_verifier) .await .context("Source KMS is not allowed for onboarding")?; - return Self::from_key_response(cfg, keys_res, domain).await; + return Self::from_key_response(cfg, keys_res, domain, rotate).await; } let client = RaClientConfig::builder() @@ -789,43 +926,87 @@ impl Keys { vm_config: info.vm_config, }) .await?; - Self::from_key_response(cfg, keys_res, domain).await + Self::from_key_response(cfg, keys_res, domain, rotate).await } async fn from_key_response( cfg: &KmsConfig, keys_res: KmsKeyResponse, domain: &str, + rotate: bool, ) -> Result { - if keys_res.keys.len() != 1 { - return Err(anyhow::anyhow!("Invalid keys")); + let epoch = cfg.key_rotation.current_epoch; + // A rotation imports every epoch before the new one; a join imports all. + let expected = if rotate { + if epoch == 0 { + bail!("rotation requires key_rotation.current_epoch > 0"); + } + epoch + } else { + epoch + 1 + }; + if !keys_res.keys.iter().map(|k| k.epoch).eq(0..expected) { + bail!( + "source KMS must hand over root key epochs 0..{expected}, got {:?}", + keys_res.keys.iter().map(|k| k.epoch).collect::>() + ); } - let keys = keys_res.keys[0].clone(); - let tmp_ca_key_pem = keys_res.temp_ca_key; - let root_ca_key_pem = keys.ca_key; - let root_k256_key = keys.k256_key; + let mut roots = keys_res + .keys + .into_iter() + .map(|keys| { + let ca_key = KeyPair::from_pem(&keys.ca_key).context("Failed to parse CA key")?; + let k256_key = + SigningKey::from_slice(&keys.k256_key).context("Failed to parse ECDSA key")?; + Ok((ca_key, k256_key)) + }) + .collect::>>()?; + let (ca_key, k256_key) = if rotate { + ( + KeyPair::generate_for(&PKCS_ECDSA_P256_SHA256)?, + SigningKey::random(&mut rand::rngs::OsRng), + ) + } else { + roots.pop().context("source KMS handed over no keys")? + }; + let history = roots + .into_iter() + .map(|(ca_key, k256_key)| HistoricalRoot::new(ca_key, k256_key)) + .collect::>>()?; let rpc_key = KeyPair::generate_for(&PKCS_ECDSA_P256_SHA256)?; - let ca_key = KeyPair::from_pem(&root_ca_key_pem).context("Failed to parse CA key")?; let tmp_ca_key = - KeyPair::from_pem(&tmp_ca_key_pem).context("Failed to parse tmp CA key")?; - let ecdsa_key = - SigningKey::from_slice(&root_k256_key).context("Failed to parse ECDSA key")?; - Self::from_keys( + KeyPair::from_pem(&keys_res.temp_ca_key).context("Failed to parse tmp CA key")?; + let keys = Self::from_keys( tmp_ca_key, ca_key, rpc_key, - ecdsa_key, + k256_key, domain, cfg.attest_rpc_cert, ) - .await + .await?; + Ok(Self { + epoch, + history, + ..keys + }) } /// The root keys go last: see [`KmsConfig::root_keys_exist`]. fn store(&self, cfg: &KmsConfig) -> Result<()> { + for (epoch, root) in (0..).zip(&self.history) { + let paths = cfg.root_key_paths(epoch); + if let Some(dir) = paths.ca_key.parent() { + fs::create_dir_all(dir)?; + } + safe_write(&paths.ca_cert, root.ca_cert.pem())?; + safe_write_with_mode(&paths.ca_key, root.ca_key.serialize_pem(), 0o600)?; + safe_write_with_mode(&paths.k256_key, root.k256_key.to_bytes(), 0o600)?; + } self.store_certs(cfg)?; safe_write(cfg.rpc_domain(), self.rpc_domain.as_bytes())?; + safe_write(cfg.root_epoch(), self.epoch.to_string())?; self.store_keys(cfg)?; Ok(()) } @@ -935,6 +1116,7 @@ pub(crate) async fn bootstrap_keys(cfg: &KmsConfig, verifier: &AttestationVerifi if cfg.root_keys_exist() { bail!("KMS has already been bootstrapped"); } + ensure_bootstrap_epoch(cfg)?; ensure_self_kms_allowed(cfg, verifier) .await .context("KMS is not allowed to auto-bootstrap")?; @@ -945,10 +1127,32 @@ pub(crate) async fn bootstrap_keys(cfg: &KmsConfig, verifier: &AttestationVerifi Ok(()) } -async fn attest_keys(p256_pubkey: &[u8], k256_pubkey: &[u8]) -> Result> { +/// Bootstrapping creates epoch 0 only; later epochs come from rotation onboarding. +fn ensure_bootstrap_epoch(cfg: &KmsConfig) -> Result<()> { + if cfg.key_rotation.current_epoch != 0 { + bail!("bootstrap only creates epoch 0; onboard with rotate to create a later epoch"); + } + Ok(()) +} + +fn self_signed_root_ca(ca_key: &KeyPair) -> Result { + CertRequest::builder() + .org_name("Dstack") + .subject("Dstack KMS CA") + .ca_level(1) + .key(ca_key) + .build() + .self_signed() +} + +async fn attest_keys(epoch: u32, p256_pubkey: &[u8], k256_pubkey: &[u8]) -> Result> { let p256_hex = hex::encode(p256_pubkey); let k256_hex = hex::encode(k256_pubkey); - let content_to_quote = format!("dstack-kms-genereted-keys-v1:{p256_hex};{k256_hex};"); + let content_to_quote = if epoch == 0 { + format!("dstack-kms-genereted-keys-v1:{p256_hex};{k256_hex};") + } else { + format!("dstack-kms-generated-keys-v2:{epoch};{p256_hex};{k256_hex};") + }; let hash = keccak256(content_to_quote.as_bytes()); let report_data = pad64(hash); let res = app_attest(report_data).await?; diff --git a/dstack/kms/src/www/onboard.html b/dstack/kms/src/www/onboard.html index 4794cf724..d476e64e8 100644 --- a/dstack/kms/src/www/onboard.html +++ b/dstack/kms/src/www/onboard.html @@ -261,6 +261,12 @@

Onboard from an Existing KMS Instance

+
+ +
Onboard from an Existing KMS Instance sourceUrl: '', sourceToken: '', onboardDomain: '', + rotate: false, error: '', success: '', result: '', @@ -350,7 +357,7 @@

Onboard from an Existing KMS Instance

return; } if (data.attestation) { - this.success = 'Bootstrap successful!'; + this.success = 'New root keys generated!'; this.result = JSON.stringify({ caPubkey: '0x' + data.ca_pubkey, k256Pubkey: '0x' + data.k256_pubkey, @@ -391,7 +398,8 @@

Onboard from an Existing KMS Instance

const data = await rpcCall('Onboard', { source_url: this.sourceUrl + "/prpc", source_token: this.sourceToken, - domain: this.onboardDomain + domain: this.onboardDomain, + rotate: this.rotate }); if (data.error) throw new Error(data.error); From 352c62817aa791d7d7c4decdfec9a581b2ba9033 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 5 Nov 2026 22:22:56 -0800 Subject: [PATCH 02/17] kms: register key epochs on chain and refuse unregistered or stale keyrings DstackKms gains an append-only registry of each epoch's root public keys (registerKeyEpoch, keyEpochCount, keyEpoch). auth-eth and auth-simple list the registered epochs in their info response and report the latest one in every boot decision. The KMS checks at startup that its keyring matches the registered public keys epoch by epoch, and on every authorization that the latest registered epoch is still its own. A node left behind after a newer epoch is registered stops serving keys, certificates, and handovers. Deployments that never register epochs keep working at epoch 0. Signed-off-by: Kevin Wang --- dstack/kms/auth-eth/contracts/DstackKms.sol | 37 ++++++++- dstack/kms/auth-eth/src/ethereum.ts | 33 +++++++- dstack/kms/auth-eth/src/main.test.ts | 34 ++++++++ dstack/kms/auth-eth/src/server.ts | 3 + dstack/kms/auth-eth/src/types.ts | 7 ++ dstack/kms/auth-eth/test/DstackKms.t.sol | 35 ++++++++ dstack/kms/auth-simple/index.test.ts | 32 +++++++ dstack/kms/auth-simple/index.ts | 21 ++++- dstack/kms/src/main.rs | 4 + dstack/kms/src/main_service.rs | 59 ++++++++++++- .../kms/src/main_service/upgrade_authority.rs | 70 ++++++++++++++-- dstack/kms/src/onboard_service.rs | 83 +++++++++++++++++++ 12 files changed, 405 insertions(+), 13 deletions(-) diff --git a/dstack/kms/auth-eth/contracts/DstackKms.sol b/dstack/kms/auth-eth/contracts/DstackKms.sol index f79325b89..ae44664b1 100644 --- a/dstack/kms/auth-eth/contracts/DstackKms.sol +++ b/dstack/kms/auth-eth/contracts/DstackKms.sol @@ -44,6 +44,15 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E // DstackApp implementation contract address for factory deployment address public appImplementation; + struct KeyEpoch { + bytes caPubkey; + bytes k256Pubkey; + } + + // Root public keys of each KMS key epoch, indexed by epoch. Append-only: + // a registered epoch is never changed, so the last entry is current. + KeyEpoch[] private _keyEpochs; + // Events // appId is unindexed for backward compatibility with deployed log indexers. // slither-disable-next-line unindexed-event-address @@ -58,6 +67,7 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E event GatewayAppIdSet(string gatewayAppId); // slither-disable-next-line unindexed-event-address event AppImplementationSet(address implementation); + event KeyEpochRegistered(uint256 indexed epoch, bytes caPubkey, bytes k256Pubkey); event AppDeployedViaFactory(address indexed appId, address indexed deployer); /// @notice Additive audit event for reconstructing authorization policy. /// @dev `value` is the affected bytes32 value, address, or hash of dynamic public data. @@ -115,6 +125,31 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E _emitPolicy("kms-info", keccak256(info.k256Pubkey), true); } + // Register the root public keys of the next KMS key epoch + function registerKeyEpoch( + bytes calldata caPubkey, + bytes calldata k256Pubkey + ) + external + onlyOwner + returns (uint256 epoch) + { + require(caPubkey.length > 0 && k256Pubkey.length > 0, "Empty epoch public key"); + epoch = _keyEpochs.length; + _keyEpochs.push(KeyEpoch(caPubkey, k256Pubkey)); + emit KeyEpochRegistered(epoch, caPubkey, k256Pubkey); + _emitPolicy("key-epoch", keccak256(abi.encode(caPubkey, k256Pubkey)), true); + } + + function keyEpochCount() external view returns (uint256) { + return _keyEpochs.length; + } + + function keyEpoch(uint256 epoch) external view returns (bytes memory caPubkey, bytes memory k256Pubkey) { + KeyEpoch storage keys = _keyEpochs[epoch]; + return (keys.caPubkey, keys.k256Pubkey); + } + // Function to set KMS quote function setKmsQuote(bytes memory quote) external onlyOwner { kmsInfo.quote = quote; @@ -306,7 +341,7 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E } // Add storage gap for upgradeable contracts - uint256[50] private __gap; + uint256[49] private __gap; } function isContract(address addr) view returns (bool) { diff --git a/dstack/kms/auth-eth/src/ethereum.ts b/dstack/kms/auth-eth/src/ethereum.ts index ebd655140..3c4a696ad 100644 --- a/dstack/kms/auth-eth/src/ethereum.ts +++ b/dstack/kms/auth-eth/src/ethereum.ts @@ -3,16 +3,23 @@ // SPDX-License-Identifier: Apache-2.0 import { ethers } from 'ethers'; -import { BootInfo, BootResponse } from './types'; +import { BootInfo, BootResponse, KeyEpoch } from './types'; // Minimal ABI for DstackKms contract const DSTACK_KMS_ABI = [ "function isAppAllowed((address appId,bytes32 composeHash,address instanceId,bytes32 deviceId,bytes32 mrAggregated,bytes32 mrSystem,bytes32 osImageHash,string tcbStatus,string[] advisoryIds) bootInfo) view returns (bool, string)", "function isKmsAllowed((address appId,bytes32 composeHash,address instanceId,bytes32 deviceId,bytes32 mrAggregated,bytes32 mrSystem,bytes32 osImageHash,string tcbStatus,string[] advisoryIds) bootInfo) view returns (bool, string)", "function gatewayAppId() view returns (string)", - "function appImplementation() view returns (address)" + "function appImplementation() view returns (address)", + "function keyEpochCount() view returns (uint256)", + "function keyEpoch(uint256 epoch) view returns (bytes caPubkey, bytes k256Pubkey)" ]; +// Contracts deployed before key epochs existed lack the registry getters. +function isMissingFunction(error: unknown): boolean { + return ethers.isError(error, 'CALL_EXCEPTION') || ethers.isError(error, 'BAD_DATA'); +} + export class EthereumBackend { private provider: ethers.JsonRpcProvider; private kmsContract: ethers.Contract; @@ -60,13 +67,35 @@ export class EthereumBackend { } const [isAllowed, reason] = response; const gatewayAppId = await this.kmsContract.gatewayAppId({ blockTag }); + const epochCount = await this.getKeyEpochCount(blockTag); return { isAllowed, reason, gatewayAppId, + ...(epochCount > 0 && { keyEpoch: epochCount - 1 }), + } + } + + private async getKeyEpochCount(blockTag?: number): Promise { + try { + return Number(await this.kmsContract.keyEpochCount({ blockTag })); + } catch (error) { + if (isMissingFunction(error)) return 0; + throw error; } } + async getKeyEpochs(): Promise { + const blockTag = await this.provider.getBlockNumber(); + const count = await this.getKeyEpochCount(blockTag); + return Promise.all( + Array.from({ length: count }, async (_, epoch) => { + const [caPubkey, k256Pubkey] = await this.kmsContract.keyEpoch(epoch, { blockTag }); + return { caPubkey, k256Pubkey }; + }) + ); + } + async getGatewayAppId(): Promise { return await this.kmsContract.gatewayAppId(); } diff --git a/dstack/kms/auth-eth/src/main.test.ts b/dstack/kms/auth-eth/src/main.test.ts index 842f1e36b..43605eb7b 100644 --- a/dstack/kms/auth-eth/src/main.test.ts +++ b/dstack/kms/auth-eth/src/main.test.ts @@ -88,6 +88,40 @@ describe('Server', () => { }); }); + describe('key epochs', () => { + it('should report the latest registered key epoch', async () => { + app.ethereum.checkBoot = jest.fn().mockResolvedValue({ + isAllowed: true, + reason: '', + gatewayAppId: '', + keyEpoch: 2, + }); + + const response = await app.inject({ + method: 'POST', + url: '/bootAuth/app', + payload: mockBootInfo + }); + + expect(JSON.parse(response.payload).keyEpoch).toBe(2); + }); + + it('should list registered key epochs in the info response', async () => { + const keyEpochs = [{ caPubkey: '0xa0', k256Pubkey: '0xb0' }]; + Object.assign(app.ethereum, { + getGatewayAppId: jest.fn().mockResolvedValue('gw'), + getChainId: jest.fn().mockResolvedValue(1), + getAppImplementation: jest.fn().mockResolvedValue('0x00'), + getKeyEpochs: jest.fn().mockResolvedValue(keyEpochs), + }); + + const response = await app.inject({ method: 'GET', url: '/' }); + + expect(response.statusCode).toBe(200); + expect(JSON.parse(response.payload).keyEpochs).toEqual(keyEpochs); + }); + }); + describe('POST /bootAuth/kms', () => { it('should return 200 when backend check passes', async () => { // Mock successful response diff --git a/dstack/kms/auth-eth/src/server.ts b/dstack/kms/auth-eth/src/server.ts index 54dd14346..08a7e913c 100644 --- a/dstack/kms/auth-eth/src/server.ts +++ b/dstack/kms/auth-eth/src/server.ts @@ -50,6 +50,7 @@ export async function build(): Promise { isAllowed: { type: 'boolean' }, reason: { type: 'string' }, gatewayAppId: { type: 'string' }, + keyEpoch: { type: 'integer' }, } }); @@ -84,6 +85,7 @@ export async function build(): Promise { server.ethereum.getGatewayAppId(), server.ethereum.getChainId(), server.ethereum.getAppImplementation(), + server.ethereum.getKeyEpochs(), ]); return { status: 'ok', @@ -93,6 +95,7 @@ export async function build(): Promise { chainId: batch[1], appAuthImplementation: batch[2], // NOTE: for backward compatibility appImplementation: batch[2], + keyEpochs: batch[3], }; } catch { _request.log.error('authorization backend health check failed'); diff --git a/dstack/kms/auth-eth/src/types.ts b/dstack/kms/auth-eth/src/types.ts index fd1a89143..5442a7e69 100644 --- a/dstack/kms/auth-eth/src/types.ts +++ b/dstack/kms/auth-eth/src/types.ts @@ -18,6 +18,13 @@ export interface BootResponse { isAllowed: boolean; gatewayAppId: string; reason: string; + // Latest registered KMS key epoch; absent when none is registered. + keyEpoch?: number; +} + +export interface KeyEpoch { + caPubkey: string; + k256Pubkey: string; } // Removed KMS_CONTRACT_ABI and APP_CONTRACT_ABI since we're using typechain types now diff --git a/dstack/kms/auth-eth/test/DstackKms.t.sol b/dstack/kms/auth-eth/test/DstackKms.t.sol index 899b2b7b9..ab99f12ff 100644 --- a/dstack/kms/auth-eth/test/DstackKms.t.sol +++ b/dstack/kms/auth-eth/test/DstackKms.t.sol @@ -19,6 +19,7 @@ contract DstackKmsTest is Test { address public user; event KmsInfoSet(bytes k256Pubkey); + event KeyEpochRegistered(uint256 indexed epoch, bytes caPubkey, bytes k256Pubkey); event AppRegistered(address appId); @@ -66,6 +67,40 @@ contract DstackKmsTest is Test { assertEq(storedEventlog, eventlog); } + function test_RegisterKeyEpochAppends() public { + assertEq(kms.keyEpochCount(), 0); + + vm.startPrank(owner); + vm.expectEmit(true, true, true, true); + emit KeyEpochRegistered(0, hex"a0", hex"b0"); + assertEq(kms.registerKeyEpoch(hex"a0", hex"b0"), 0); + assertEq(kms.registerKeyEpoch(hex"a1", hex"b1"), 1); + vm.stopPrank(); + + assertEq(kms.keyEpochCount(), 2); + (bytes memory ca, bytes memory k256) = kms.keyEpoch(0); + assertEq(ca, hex"a0"); + assertEq(k256, hex"b0"); + (ca, k256) = kms.keyEpoch(1); + assertEq(ca, hex"a1"); + assertEq(k256, hex"b1"); + } + + function test_RegisterKeyEpochOnlyOwner() public { + vm.prank(user); + vm.expectRevert(abi.encodeWithSelector(OwnableUpgradeable.OwnableUnauthorizedAccount.selector, user)); + kms.registerKeyEpoch(hex"a0", hex"b0"); + } + + function test_RegisterKeyEpochRejectsEmptyKeys() public { + vm.startPrank(owner); + vm.expectRevert("Empty epoch public key"); + kms.registerKeyEpoch("", hex"b0"); + vm.expectRevert("Empty epoch public key"); + kms.registerKeyEpoch(hex"a0", ""); + vm.stopPrank(); + } + function test_SetKmsInfoOnlyOwner() public { bytes memory k256Pubkey = hex"123456"; bytes memory caPubkey = hex"789abc"; diff --git a/dstack/kms/auth-simple/index.test.ts b/dstack/kms/auth-simple/index.test.ts index fee441712..a75120fca 100644 --- a/dstack/kms/auth-simple/index.test.ts +++ b/dstack/kms/auth-simple/index.test.ts @@ -446,6 +446,38 @@ describe('auth-simple', () => { }); }); + describe('key epochs', () => { + const keyEpochs = [ + { caPubkey: '0xa0', k256Pubkey: '0xb0' }, + { caPubkey: '0xa1', k256Pubkey: '0xb1' } + ]; + + it('reports the latest key epoch in boot responses and lists all epochs', async () => { + writeTestConfig({ gatewayAppId: '0xgateway', keyEpochs }); + + const boot = await app.fetch(new Request('http://localhost/bootAuth/app', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify(baseBootInfo) + })); + expect((await boot.json()).keyEpoch).toBe(1); + + const info = await app.fetch(new Request('http://localhost/', { method: 'GET' })); + expect((await info.json()).keyEpochs).toEqual(keyEpochs); + }); + + it('omits the key epoch when none is configured', async () => { + writeTestConfig({ gatewayAppId: '0xgateway' }); + + const boot = await app.fetch(new Request('http://localhost/bootAuth/app', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify(baseBootInfo) + })); + expect((await boot.json()).keyEpoch).toBeUndefined(); + }); + }); + describe('hex normalization', () => { it('handles uppercase hex', async () => { writeTestConfig({ diff --git a/dstack/kms/auth-simple/index.ts b/dstack/kms/auth-simple/index.ts index 9da149014..d2be75515 100644 --- a/dstack/kms/auth-simple/index.ts +++ b/dstack/kms/auth-simple/index.ts @@ -24,7 +24,13 @@ const BootInfoSchema = z.object({ const BootResponseSchema = z.object({ isAllowed: z.boolean(), reason: z.string(), - gatewayAppId: z.string() + gatewayAppId: z.string(), + keyEpoch: z.number().optional() +}); + +const KeyEpochSchema = z.object({ + caPubkey: z.string(), + k256Pubkey: z.string() }); // config schema @@ -52,7 +58,9 @@ const AuthConfigSchema = z.object({ allowedTcbStatuses: z.array(z.string()).default(['UpToDate']), allowedAdvisoryIds: z.array(z.string()).default([]), kms: KmsConfigSchema.default({}), - apps: z.record(z.string(), AppConfigSchema).default({}) + apps: z.record(z.string(), AppConfigSchema).default({}), + // Root public keys of each KMS key epoch, indexed by epoch; append only. + keyEpochs: z.array(KeyEpochSchema).default([]) }); type BootInfo = z.infer; @@ -93,6 +101,12 @@ class ConfigBackend { async checkBoot(bootInfo: BootInfo, isKms: boolean): Promise { const config = this.loadConfig(); + const response = this.evaluateBoot(bootInfo, isKms, config); + const epochCount = config.keyEpochs.length; + return epochCount > 0 ? { ...response, keyEpoch: epochCount - 1 } : response; + } + + private evaluateBoot(bootInfo: BootInfo, isKms: boolean, config: AuthConfig): BootResponse { const osImageHash = normalizeHex(bootInfo.osImageHash); const deviceId = normalizeHex(bootInfo.deviceId); @@ -220,6 +234,7 @@ class ConfigBackend { kmsContractAddr: string; chainId: number; appImplementation: string; + keyEpochs: AuthConfig['keyEpochs']; }> { const config = this.loadConfig(); return { @@ -227,6 +242,7 @@ class ConfigBackend { kmsContractAddr: config.kmsContractAddr, chainId: config.chainId, appImplementation: config.appImplementation, + keyEpochs: config.keyEpochs, }; } } @@ -248,6 +264,7 @@ app.get('/', async (c) => { chainId: info.chainId, appAuthImplementation: info.appImplementation, // backward compat appImplementation: info.appImplementation, + keyEpochs: info.keyEpochs, }); } catch (error) { console.error('error in health check:', error); diff --git a/dstack/kms/src/main.rs b/dstack/kms/src/main.rs index f4bf428a8..526fd7f9b 100644 --- a/dstack/kms/src/main.rs +++ b/dstack/kms/src/main.rs @@ -177,6 +177,10 @@ async fn main() -> Result<()> { None }; let state = main_service::KmsState::new(config).context("Failed to initialize KMS state")?; + state + .ensure_keyring_registered() + .await + .context("KMS keyring does not match the registered key epochs")?; let quote_verifier = QuoteVerifier::new(state.attestation_verifier()); let admin_quote_verifier = QuoteVerifier::new(state.attestation_verifier()); let figment = figment diff --git a/dstack/kms/src/main_service.rs b/dstack/kms/src/main_service.rs index 12274fe2d..f7187e288 100644 --- a/dstack/kms/src/main_service.rs +++ b/dstack/kms/src/main_service.rs @@ -30,7 +30,8 @@ use ra_tls::{ use scale::Decode; use tracing::warn; use upgrade_authority::{ - build_boot_info, ensure_app_id_len, local_kms_boot_info, BootInfo, GetInfoResponse, + build_boot_info, ensure_app_id_len, local_kms_boot_info, BootInfo, BootResponse, + GetInfoResponse, KeyEpochPubkeys, }; use crate::{ @@ -81,6 +82,13 @@ impl RootKeys { let k256 = SigningKey::from_slice(&key_bytes).context("failed to load ECDSA root key")?; Ok(Self { ca, k256 }) } + + fn pubkeys(&self) -> KeyEpochPubkeys { + KeyEpochPubkeys { + ca_pubkey: self.ca.key.public_key_der(), + k256_pubkey: self.k256.verifying_key().to_sec1_bytes().to_vec(), + } + } } pub struct KmsStateInner { @@ -208,6 +216,52 @@ impl KmsState { fn root(&self) -> &RootKeys { self.roots.last().expect("the keyring always holds epoch 0") } + + /// Refuse to serve a keyring governance has not registered. Deployments + /// that never registered epochs keep working while they stay at epoch 0. + pub async fn ensure_keyring_registered(&self) -> Result<()> { + let current = self.config.key_rotation.current_epoch; + let info = match self.config.auth_api.get_info().await { + Ok(info) => info, + Err(err) if current == 0 => { + warn!("skipping key epoch registration check: {err:#}"); + return Ok(()); + } + Err(err) => return Err(err).context("failed to read registered key epochs"), + }; + let Some(registered) = info.key_epochs else { + return Ok(()); + }; + if registered.is_empty() && current == 0 { + return Ok(()); + } + if registered.len() != current as usize + 1 { + bail!( + "this KMS serves key epoch {current}, but {} epochs are registered", + registered.len() + ); + } + for (epoch, (root, registered)) in self.roots.iter().zip(®istered).enumerate() { + if root.pubkeys() != *registered { + bail!("root keys of epoch {epoch} do not match the registered public keys"); + } + } + Ok(()) + } + + /// Stop serving once governance registers a newer epoch than this KMS has. + pub(crate) fn ensure_epoch_current(&self, response: &BootResponse) -> Result<()> { + let current = self.config.key_rotation.current_epoch; + match response.key_epoch { + Some(epoch) if epoch != current => { + bail!("the registered key epoch is {epoch}, but this KMS serves epoch {current}") + } + None if current > 0 && !self.config.auth_api.is_dev() => { + bail!("the auth API reports no key epoch, but this KMS serves epoch {current}") + } + _ => Ok(()), + } + } } pub(crate) fn read_root_epoch(config: &KmsConfig) -> Result { @@ -344,7 +398,7 @@ impl RpcHandler { if !response.is_allowed { bail!("KMS is not allowed: {}", response.reason); } - Ok(()) + self.state.ensure_epoch_current(&response) }) .await .map_err(|err| anyhow!("{err:#}"))?; @@ -388,6 +442,7 @@ impl RpcHandler { if !response.is_allowed { bail!("Boot denied: {}", response.reason); } + self.state.ensure_epoch_current(&response)?; // SNP rootfs/app/config binding is handled by the SNP launch-measurement // helper above. The legacy OS-image verifier is TDX-oriented and still // rejects SNP quotes; keep SNP on the explicit fail-closed helper path. diff --git a/dstack/kms/src/main_service/upgrade_authority.rs b/dstack/kms/src/main_service/upgrade_authority.rs index 79f4a7c23..e9d7e97e0 100644 --- a/dstack/kms/src/main_service/upgrade_authority.rs +++ b/dstack/kms/src/main_service/upgrade_authority.rs @@ -84,6 +84,24 @@ pub(crate) struct BootResponse { pub is_allowed: bool, pub gateway_app_id: String, pub reason: String, + /// Latest KMS key epoch registered by governance, if any. + #[serde(default)] + pub key_epoch: Option, +} + +/// Root public keys governance registered for one KMS key epoch. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct KeyEpochPubkeys { + #[serde(deserialize_with = "deserialize_hex")] + pub ca_pubkey: Vec, + #[serde(deserialize_with = "deserialize_hex")] + pub k256_pubkey: Vec, +} + +fn deserialize_hex<'de, D: serde::Deserializer<'de>>(deserializer: D) -> Result, D::Error> { + let value = String::deserialize(deserializer)?; + hex::decode(value.trim_start_matches("0x")).map_err(serde::de::Error::custom) } #[derive(Debug, Serialize, Deserialize)] @@ -96,6 +114,8 @@ pub(crate) struct AuthApiInfoResponse { pub gateway_app_id: String, pub chain_id: u64, pub app_implementation: String, + #[serde(default)] + pub key_epochs: Vec, } #[derive(Debug, Clone, Serialize, Deserialize)] @@ -107,6 +127,8 @@ pub(crate) struct GetInfoResponse { pub eth_rpc_url: Option, pub chain_id: Option, pub app_implementation: Option, + /// Registered key epochs, indexed by epoch. `None` means no registry. + pub key_epochs: Option>, } const AUTH_API_CONNECT_TIMEOUT: Duration = Duration::from_secs(5); @@ -174,6 +196,7 @@ impl AuthApi { is_allowed: true, reason: "".to_string(), gateway_app_id: dev.gateway_app_id.clone(), + key_epoch: None, }), AuthApi::Webhook { webhook } => { let path = if is_kms { @@ -196,6 +219,7 @@ impl AuthApi { gateway_app_id: Some(dev.gateway_app_id.clone()), chain_id: None, app_implementation: None, + key_epochs: None, }), AuthApi::Webhook { webhook } => { let info: AuthApiInfoResponse = http_get(&webhook.url, webhook.timeout).await?; @@ -211,6 +235,7 @@ impl AuthApi { chain_id: Some(info.chain_id), gateway_app_id: Some(info.gateway_app_id.clone()), app_implementation: Some(info.app_implementation.clone()), + key_epochs: Some(info.key_epochs), }) } } @@ -295,7 +320,7 @@ pub(crate) async fn ensure_kms_allowed( } #[cfg(test)] -mod tests { +pub(crate) mod tests { use super::*; use crate::config::Webhook; use ra_tls::attestation::TeeVariant; @@ -320,7 +345,7 @@ mod tests { } } - fn serve>( + pub(crate) fn serve>( responses: Vec, ) -> (String, thread::JoinHandle>) { let responses: Vec = responses.into_iter().map(Into::into).collect(); @@ -364,9 +389,12 @@ mod tests { .nth(1) .unwrap() .to_string(); - let body = - serde_json::from_slice(&request[header_end..header_end + content_length]) - .unwrap(); + let body = &request[header_end..header_end + content_length]; + let body = if body.is_empty() { + Value::Null + } else { + serde_json::from_slice(body).unwrap() + }; requests.push((path, body)); write!( stream, @@ -381,7 +409,7 @@ mod tests { (format!("http://{address}"), handle) } - fn webhook(url: String) -> AuthApi { + pub(crate) fn webhook(url: String) -> AuthApi { AuthApi::Webhook { webhook: Webhook { url, @@ -509,6 +537,36 @@ mod tests { assert_ne!(requests[0].1["mrAggregated"], requests[1].1["mrAggregated"]); } + #[rocket::async_test] + async fn key_epochs_are_parsed_from_info_and_boot_responses() { + let (url, server) = serve(vec![ + r#"{"status":"ok","kmsContractAddr":"0x00","gatewayAppId":"gw","chainId":1, + "appImplementation":"0x00","keyEpochs":[{"caPubkey":"0xa0","k256Pubkey":"b0"}]}"#, + r#"{"isAllowed":true,"gatewayAppId":"gw","reason":"","keyEpoch":3}"#, + r#"{"isAllowed":true,"gatewayAppId":"gw","reason":""}"#, + ]); + let auth = webhook(url); + + let info = auth.get_info().await.unwrap(); + assert_eq!( + info.key_epochs, + Some(vec![KeyEpochPubkeys { + ca_pubkey: vec![0xa0], + k256_pubkey: vec![0xb0], + }]) + ); + let info = boot_info(50); + assert_eq!( + auth.is_app_allowed(&info, false).await.unwrap().key_epoch, + Some(3) + ); + assert_eq!( + auth.is_app_allowed(&info, false).await.unwrap().key_epoch, + None + ); + server.join().unwrap(); + } + #[rocket::async_test] async fn malformed_backend_response_fails_closed_and_next_request_recovers() { let (url, server) = serve(vec![ diff --git a/dstack/kms/src/onboard_service.rs b/dstack/kms/src/onboard_service.rs index 20775bd3c..764023562 100644 --- a/dstack/kms/src/onboard_service.rs +++ b/dstack/kms/src/onboard_service.rs @@ -295,6 +295,7 @@ mod tests { use crate::main_service::amd_attest::{ compute_expected_measurement, MeasurementInput, OvmfSectionParam, }; + use crate::main_service::upgrade_authority::{BootResponse, KeyEpochPubkeys}; use dstack_kms_rpc::KmsKeys; use sha2::Digest; @@ -596,6 +597,88 @@ mod tests { assert!(format!("{err:#}").contains("belong to epoch 2"), "{err:#}"); } + fn pubkeys_of(res: &KmsKeyResponse) -> Vec { + res.keys + .iter() + .map(|k| KeyEpochPubkeys { + ca_pubkey: KeyPair::from_pem(&k.ca_key).unwrap().public_key_der(), + k256_pubkey: SigningKey::from_slice(&k.k256_key) + .unwrap() + .verifying_key() + .to_sec1_bytes() + .to_vec(), + }) + .collect() + } + + async fn joined_state(res: KmsKeyResponse, url: String) -> crate::main_service::KmsState { + let cert_dir = tempfile::tempdir().unwrap(); + let mut cfg = epoch_config(cert_dir.path(), res.keys.len() as u32 - 1); + Keys::from_key_response(&cfg, res, "kms.example.com", false) + .await + .unwrap() + .store(&cfg) + .unwrap(); + cfg.auth_api = crate::main_service::upgrade_authority::tests::webhook(url); + crate::main_service::KmsState::new(cfg).unwrap() + } + + fn info_with_epochs(epochs: &[KeyEpochPubkeys]) -> String { + let epochs: Vec<_> = epochs + .iter() + .map(|e| { + serde_json::json!({ + "caPubkey": hex::encode(&e.ca_pubkey), + "k256Pubkey": hex::encode(&e.k256_pubkey), + }) + }) + .collect(); + serde_json::json!({ + "status": "ok", "kmsContractAddr": "0x00", "gatewayAppId": "", "chainId": 1, + "appImplementation": "0x00", "keyEpochs": epochs, + }) + .to_string() + } + + #[rocket::async_test] + async fn keyring_must_match_the_registered_epochs() { + use crate::main_service::upgrade_authority::tests::serve; + + type Registered = fn(Vec) -> Vec; + let cases: [(Registered, bool); 4] = [ + (|keys| keys, true), + (|keys| keys[..1].to_vec(), false), + (|keys| vec![keys[0].clone(), keys[0].clone()], false), + (|_| Vec::new(), false), + ]; + for (registered, ok) in cases { + let res = handover(2); + let epochs = registered(pubkeys_of(&res)); + let (url, server) = serve(vec![info_with_epochs(&epochs)]); + let state = joined_state(res, url).await; + assert_eq!( + state.ensure_keyring_registered().await.is_ok(), + ok, + "{epochs:?}" + ); + server.join().unwrap(); + } + } + + #[rocket::async_test] + async fn a_newer_registered_epoch_stops_a_stale_kms() { + let state = joined_state(handover(2), "http://127.0.0.1:1".into()).await; + let response = |key_epoch| BootResponse { + is_allowed: true, + gateway_app_id: String::new(), + reason: String::new(), + key_epoch, + }; + assert!(state.ensure_epoch_current(&response(Some(1))).is_ok()); + assert!(state.ensure_epoch_current(&response(Some(2))).is_err()); + assert!(state.ensure_epoch_current(&response(None)).is_err()); + } + #[rocket::async_test] async fn rotation_imports_history_and_generates_the_new_epoch() { let cert_dir = tempfile::tempdir().unwrap(); From 4db641ff9566126ae40cca2ecb2b0a896f360413 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 5 Nov 2026 22:31:18 -0800 Subject: [PATCH 03/17] dstack-util: request the latest key epoch and re-key older data disks The guest now asks the KMS for the current key epoch and records the served epoch in the KMS key provider. A KMS that predates epochs ignores the selector and serves epoch 0, which is accepted. When the data disk does not unlock with the current key, the epoch recorded in its LUKS2 header (none means 0) selects which historical disk key to request from the same KMS. luksChangeKey replaces the keyslot 0 passphrase in place, keeping the volume key, and a dstack-kms token records the new epoch. The header validator accepts exactly that token. An unrotated epoch-0 disk is left untouched so older OS images can still open it; disks formatted at a later epoch are tagged immediately. Signed-off-by: Kevin Wang --- dstack/dstack-types/src/lib.rs | 4 + dstack/dstack-util/src/main.rs | 3 +- dstack/dstack-util/src/system_setup.rs | 462 ++++++++++++++++---- dstack/dstack-util/src/system_setup/luks.rs | 14 +- 4 files changed, 403 insertions(+), 80 deletions(-) diff --git a/dstack/dstack-types/src/lib.rs b/dstack/dstack-types/src/lib.rs index 68369be18..1d7070313 100644 --- a/dstack/dstack-types/src/lib.rs +++ b/dstack/dstack-types/src/lib.rs @@ -2431,6 +2431,9 @@ pub enum KeyProvider { pubkey: Vec, tmp_ca_key: String, tmp_ca_cert: String, + /// KMS root key epoch the keys were derived from. + #[serde(default)] + key_epoch: u32, }, } @@ -2494,6 +2497,7 @@ mod key_provider_tests { pubkey: vec![0xab; 32], tmp_ca_key: String::new(), tmp_ca_cert: String::new(), + key_epoch: 0, }; assert_eq!(kms.id(), &[0xab; 32]); } diff --git a/dstack/dstack-util/src/main.rs b/dstack/dstack-util/src/main.rs index 9b28165b9..f78a77216 100644 --- a/dstack/dstack-util/src/main.rs +++ b/dstack/dstack-util/src/main.rs @@ -701,7 +701,7 @@ async fn cmd_get_keys(args: GetKeysArgs) -> Result<()> { .get_app_key(dstack_kms_rpc::GetAppKeyRequest { api_version: 1, vm_config: "".to_string(), - ..Default::default() + key_epoch: -1, }) .await .context("Failed to get app key")?; @@ -724,6 +724,7 @@ async fn cmd_get_keys(args: GetKeysArgs) -> Result<()> { pubkey: root_pubkey, tmp_ca_key: tmp_ca.temp_ca_key, tmp_ca_cert: tmp_ca.temp_ca_cert, + key_epoch: response.key_epoch, }, }; diff --git a/dstack/dstack-util/src/system_setup.rs b/dstack/dstack-util/src/system_setup.rs index 35aa841ab..b5374b3eb 100644 --- a/dstack/dstack-util/src/system_setup.rs +++ b/dstack/dstack-util/src/system_setup.rs @@ -29,7 +29,7 @@ use dstack_types::{ use fs_err as fs; use luks::{ LuksAf, LuksArea, LuksConfig, LuksDigest, LuksHeader, LuksJson, LuksKdf, LuksKeyslot, - LuksSegment, BINARY_HEADER_SIZE, + LuksSegment, BINARY_HEADER_SIZE, KEY_EPOCH_TOKEN_TYPE, }; use ra_rpc::{ client::{CertInfo, RaClient, RaClientConfig}, @@ -2353,6 +2353,87 @@ struct AppInfo { gpu_attestation_event: Option>, } +impl HostShared { + /// Request app keys of `key_epoch` (`-1` = current) from one KMS. + async fn request_kms_keys( + &self, + kms_url: String, + key_epoch: i32, + ) -> Result<(AppKeys, Vec)> { + info!("Requesting app keys of epoch {key_epoch} from KMS: {kms_url}"); + let tmp_ca = { + info!("Getting temp ca cert"); + let client = RaClient::new(kms_url.clone(), true)?; + let kms_client = dstack_kms_rpc::kms_client::KmsClient::new(client); + kms_client + .get_temp_ca_cert() + .await + .context("Failed to get temp ca cert")? + }; + let cert_pair = generate_ra_cert(tmp_ca.temp_ca_cert.clone(), tmp_ca.temp_ca_key.clone())?; + let attestation_verifier = attestation_verifier(&self.sys_config)?; + let ra_client = RaClientConfig::builder() + .tls_no_check(false) + .tls_built_in_root_certs(false) + .remote_uri(kms_url.clone()) + .tls_client_cert(cert_pair.cert_pem) + .tls_client_key(cert_pair.key_pem) + .tls_ca_cert(tmp_ca.ca_cert.clone()) + .attestation_verifier(attestation_verifier) + .cert_validator(Box::new(validate_kms_rpc_cert)) + .build() + .into_client() + .context("Failed to create client")?; + let kms_client = dstack_kms_rpc::kms_client::KmsClient::new(ra_client); + let response = kms_client + .get_app_key(rpc::GetAppKeyRequest { + api_version: 1, + vm_config: self.sys_config.vm_config.clone(), + key_epoch, + }) + .await + .context("Failed to get app key")?; + // A KMS predating key epochs ignores the selector and serves epoch 0. + if key_epoch >= 0 && response.key_epoch as i32 != key_epoch { + bail!( + "KMS returned keys of epoch {} instead of {key_epoch}", + response.key_epoch + ); + } + + let (_, ca_pem) = x509_parser::pem::parse_x509_pem(tmp_ca.ca_cert.as_bytes()) + .context("Failed to parse ca cert")?; + let x509 = ca_pem.parse_x509().context("Failed to parse ca cert")?; + let root_pubkey = x509.public_key().raw.to_vec(); + + let keys = AppKeys { + ca_cert: tmp_ca.ca_cert, + disk_crypt_key: response.disk_crypt_key, + env_crypt_key: response.env_crypt_key, + k256_key: response.k256_key, + k256_signature: response.k256_signature, + gateway_app_id: response.gateway_app_id, + key_provider: KeyProvider::Kms { + url: kms_url, + pubkey: root_pubkey, + tmp_ca_key: tmp_ca.temp_ca_key, + tmp_ca_cert: tmp_ca.temp_ca_cert, + key_epoch: response.key_epoch, + }, + }; + Ok((keys, response.os_image_hash)) + } + + /// Request keys of an older epoch from the KMS that served `keys`. + async fn request_historical_keys(&self, keys: &AppKeys, epoch: u32) -> Result { + let KeyProvider::Kms { url, .. } = &keys.key_provider else { + bail!("only KMS-provided keys have historical epochs"); + }; + let epoch = i32::try_from(epoch).context("key epoch out of range")?; + Ok(self.request_kms_keys(url.clone(), epoch).await?.0) + } +} + struct Stage0<'a> { args: &'a SetupArgs, shared: HostShared, @@ -2415,63 +2496,6 @@ impl<'a> Stage0<'a> { self.shared.dir.join(APP_KEYS) } - async fn request_app_keys_from_kms_url(&self, kms_url: String) -> Result<(AppKeys, Vec)> { - info!("Requesting app keys from KMS: {kms_url}"); - let tmp_ca = { - info!("Getting temp ca cert"); - let client = RaClient::new(kms_url.clone(), true)?; - let kms_client = dstack_kms_rpc::kms_client::KmsClient::new(client); - kms_client - .get_temp_ca_cert() - .await - .context("Failed to get temp ca cert")? - }; - let cert_pair = generate_ra_cert(tmp_ca.temp_ca_cert.clone(), tmp_ca.temp_ca_key.clone())?; - let attestation_verifier = attestation_verifier(&self.shared.sys_config)?; - let ra_client = RaClientConfig::builder() - .tls_no_check(false) - .tls_built_in_root_certs(false) - .remote_uri(kms_url.clone()) - .tls_client_cert(cert_pair.cert_pem) - .tls_client_key(cert_pair.key_pem) - .tls_ca_cert(tmp_ca.ca_cert.clone()) - .attestation_verifier(attestation_verifier) - .cert_validator(Box::new(validate_kms_rpc_cert)) - .build() - .into_client() - .context("Failed to create client")?; - let kms_client = dstack_kms_rpc::kms_client::KmsClient::new(ra_client); - let response = kms_client - .get_app_key(rpc::GetAppKeyRequest { - api_version: 1, - vm_config: self.shared.sys_config.vm_config.clone(), - ..Default::default() - }) - .await - .context("Failed to get app key")?; - - let (_, ca_pem) = x509_parser::pem::parse_x509_pem(tmp_ca.ca_cert.as_bytes()) - .context("Failed to parse ca cert")?; - let x509 = ca_pem.parse_x509().context("Failed to parse ca cert")?; - let root_pubkey = x509.public_key().raw.to_vec(); - - let keys = AppKeys { - ca_cert: tmp_ca.ca_cert, - disk_crypt_key: response.disk_crypt_key, - env_crypt_key: response.env_crypt_key, - k256_key: response.k256_key, - k256_signature: response.k256_signature, - gateway_app_id: response.gateway_app_id, - key_provider: KeyProvider::Kms { - url: kms_url, - pubkey: root_pubkey, - tmp_ca_key: tmp_ca.temp_ca_key, - tmp_ca_cert: tmp_ca.temp_ca_cert, - }, - }; - Ok((keys, response.os_image_hash)) - } - async fn request_app_keys_from_kms(&self) -> Result { if self.shared.sys_config.kms_urls.is_empty() { bail!("No KMS URLs are set"); @@ -2479,7 +2503,7 @@ impl<'a> Stage0<'a> { let mut errors = vec![]; for kms_url in &self.shared.sys_config.kms_urls { let kms_url = kms_rpc_url(kms_url); - match self.request_app_keys_from_kms_url(kms_url.clone()).await { + match self.shared.request_kms_keys(kms_url.clone(), -1).await { Ok((keys, os_image_hash)) => { // Measured here, outside the per-URL request, so that a // failover can never extend os-image-hash twice. @@ -2722,7 +2746,9 @@ impl<'a> Stage0<'a> { initialized } - async fn mount_data_disk(&self, disk_crypt_key: &str, opts: &DstackOptions) -> Result<()> { + async fn mount_data_disk(&self, keys: &AppKeys, opts: &DstackOptions) -> Result<()> { + let disk_crypt_key = &hex::encode(&keys.disk_crypt_key); + let key_epoch = kms_key_epoch(keys); let name = "dstack_data_disk"; let mount_point = &self.args.mount_point; @@ -2743,6 +2769,11 @@ impl<'a> Stage0<'a> { let has_filesystem = if opts.storage_encrypted && disk_initialized { info!("Mounting encrypted data disk"); + if let Some(epoch) = key_epoch { + self.migrate_disk_key(keys, disk_crypt_key, epoch) + .await + .context("failed to migrate the data disk key")?; + } let initializing = self.open_encrypted_volume(disk_crypt_key, name, opts.storage_discard)?; if initializing { @@ -2762,7 +2793,7 @@ impl<'a> Stage0<'a> { info!("Skipping disk encryption as requested by kernel cmdline"); } else if !disk_initialized { info!("Setting up disk encryption"); - self.luks_setup(disk_crypt_key, name, opts.storage_discard)?; + self.luks_setup(disk_crypt_key, key_epoch, name, opts.storage_discard)?; } match opts.storage_fs { @@ -2891,7 +2922,13 @@ impl<'a> Stage0<'a> { Ok(()) } - fn luks_setup(&self, disk_crypt_key: &str, name: &str, discard: bool) -> Result<()> { + fn luks_setup( + &self, + disk_crypt_key: &str, + key_epoch: Option, + name: &str, + discard: bool, + ) -> Result<()> { let root_hd = &self.args.device; let sector_offset = PAYLOAD_OFFSET / 512; info!("Formatting encrypted disk"); @@ -2929,10 +2966,46 @@ impl<'a> Stage0<'a> { { bail!("Failed to setup luks volume"); } + if let Some(epoch) = key_epoch.filter(|epoch| *epoch > 0) { + write_luks_key_epoch(root_hd, epoch, false)?; + } self.open_encrypted_volume(disk_crypt_key, name, discard)?; Ok(()) } + /// Re-key a data disk still protected by an older KMS epoch. + /// + /// Only the keyslot passphrase changes; the volume key stays. The epoch + /// is recorded in the header so the old key can be requested next time + /// the KMS rotates. An unrotated epoch-0 disk is left untouched, so it + /// stays readable by older OS images. + async fn migrate_disk_key( + &self, + keys: &AppKeys, + current_key: &str, + current: u32, + ) -> Result<()> { + let device = &self.args.device; + let stored = read_luks_key_epoch(device)?; + if !luks_unlocks(device, current_key)? { + let epoch = stored.unwrap_or(0); + if epoch >= current { + bail!("the data disk does not unlock with the key of epoch {current}"); + } + info!("re-keying the data disk from key epoch {epoch} to {current}"); + let old = self.shared.request_historical_keys(keys, epoch).await?; + let old_key = hex::encode(&old.disk_crypt_key); + luks_change_key(device, &old_key, current_key, Path::new(REKEY_DIR))?; + if !luks_unlocks(device, current_key)? { + bail!("the data disk does not unlock after re-keying"); + } + } + if stored.unwrap_or(0) != current { + write_luks_key_epoch(device, current, stored.is_some())?; + } + Ok(()) + } + /// Returns whether the header still carries the initializing label. fn open_encrypted_volume( &self, @@ -2964,7 +3037,7 @@ impl<'a> Stage0<'a> { .context("Failed to load LUKS2 header")?; let hdr_file = fs::File::open(&in_mem_hdr).context("Failed to open LUKS2 header")?; - let initializing = + let LuksHeaderState { initializing, .. } = validate_luks2_headers(hdr_file).context("Failed to validate LUKS2 header")?; info!("Opening the device"); @@ -3194,8 +3267,7 @@ impl<'a> Stage0<'a> { opts.storage_encrypted, opts.storage_fs ); - self.mount_data_disk(&hex::encode(&app_keys.disk_crypt_key), &opts) - .await?; + self.mount_data_disk(&app_keys, &opts).await?; self.setup_swap(self.shared.app_compose.swap_size, &opts) .await?; self.vmm @@ -3388,18 +3460,148 @@ macro_rules! const_pad { const PAYLOAD_OFFSET: u64 = 16777216; +fn kms_key_epoch(keys: &AppKeys) -> Option { + match keys.key_provider { + KeyProvider::Kms { key_epoch, .. } => Some(key_epoch), + _ => None, + } +} + +fn cryptsetup(args: &[&str], device: &Path) -> Command { + let mut command = Command::new("cryptsetup"); + command.args(args).arg(device); + command +} + +/// Run `command` with `input` on stdin and return its exit code. +fn run_with_input(mut command: Command, input: &[u8]) -> Result> { + let mut child = command + .stdin(Stdio::piped()) + .spawn() + .context("failed to start cryptsetup")?; + child + .stdin + .take() + .context("cryptsetup stdin is unavailable")? + .write_all(input) + .context("failed to write to cryptsetup")?; + Ok(child + .wait() + .context("failed to wait for cryptsetup")? + .code()) +} + +/// Whether `key` unlocks keyslot 0. cryptsetup exits with 2 on a wrong key. +fn luks_unlocks(device: &Path, key: &str) -> Result { + let command = cryptsetup( + &["open", "--test-passphrase", "--key-slot", "0", "-d-"], + device, + ); + match run_with_input(command, key.as_bytes())? { + Some(0) => Ok(true), + Some(2) => Ok(false), + code => bail!("cryptsetup failed to test the data disk key: {code:?}"), + } +} + +fn read_luks_key_epoch(device: &Path) -> Result> { + #[derive(Deserialize)] + struct Metadata { + tokens: BTreeMap, + } + let output = cryptsetup(&["luksDump", "--dump-json-metadata"], device) + .output() + .context("failed to run cryptsetup luksDump")?; + if !output.status.success() { + bail!("cryptsetup luksDump failed"); + } + let metadata: Metadata = + serde_json::from_slice(&output.stdout).context("failed to parse LUKS metadata")?; + Ok(metadata + .tokens + .values() + .find(|token| token.kind == KEY_EPOCH_TOKEN_TYPE) + .and_then(|token| token.key_epoch)) +} + +fn write_luks_key_epoch(device: &Path, epoch: u32, replace: bool) -> Result<()> { + let token = serde_json::json!({ + "type": KEY_EPOCH_TOKEN_TYPE, + "keyslots": ["0"], + "key_epoch": epoch, + }); + let args: &[&str] = if replace { + &["token", "import", "--token-id", "0", "--token-replace"] + } else { + &["token", "import", "--token-id", "0"] + }; + let code = run_with_input(cryptsetup(args, device), token.to_string().as_bytes())?; + if code != Some(0) { + bail!("failed to record key epoch {epoch} in the LUKS header: {code:?}"); + } + Ok(()) +} + +/// Directory on the /run tmpfs for the old key handed to cryptsetup as a file. +const REKEY_DIR: &str = "/run/dstack-rekey"; + +/// Replace the keyslot 0 passphrase in place. The new key goes over stdin; +/// the old one is written to a 0700 directory in `secret_dir`, which is +/// removed before returning. +fn luks_change_key(device: &Path, old_key: &str, new_key: &str, secret_dir: &Path) -> Result<()> { + fs::create_dir_all(secret_dir)?; + defer! { + fs::remove_dir_all(secret_dir).ok(); + } + fs::set_permissions( + secret_dir, + std::os::unix::fs::PermissionsExt::from_mode(0o700), + )?; + let old_file = secret_dir.join("old-key"); + safe_write_with_mode(&old_file, old_key, 0o600)?; + let mut command = cryptsetup( + &[ + "luksChangeKey", + "--batch-mode", + "--pbkdf", + "pbkdf2", + "--key-slot", + "0", + ], + device, + ); + command.arg("--key-file").arg(&old_file); + if run_with_input(command, new_key.as_bytes())? != Some(0) { + bail!("cryptsetup luksChangeKey failed"); + } + Ok(()) +} + /// Set by luksFormat and cleared once the filesystem exists, so a disk still /// carrying it was interrupted mid-initialization and holds no data. const LUKS_INITIALIZING_LABEL: &str = "dstack-initializing"; -/// Returns whether the primary header carries [`LUKS_INITIALIZING_LABEL`]. -fn validate_luks2_headers(mut reader: impl std::io::Read) -> Result { - let initializing = validate_single_luks2_header(&mut reader, 0)?; - validate_single_luks2_header(&mut reader, 1)?; - Ok(initializing) +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct LuksHeaderState { + /// The header carries [`LUKS_INITIALIZING_LABEL`]. + initializing: bool, + /// KMS key epoch recorded in the dstack token; `None` predates epochs. + key_epoch: Option, } -fn validate_single_luks2_header(mut reader: impl std::io::Read, hdr_ind: u64) -> Result { +fn validate_luks2_headers(mut reader: impl std::io::Read) -> Result { + let primary = validate_single_luks2_header(&mut reader, 0)?; + let secondary = validate_single_luks2_header(&mut reader, 1)?; + if primary.key_epoch != secondary.key_epoch { + bail!("LUKS header copies disagree on the key epoch"); + } + Ok(primary) +} + +fn validate_single_luks2_header( + mut reader: impl std::io::Read, + hdr_ind: u64, +) -> Result { let mut hdr_data = vec![0u8; BINARY_HEADER_SIZE as usize]; reader .read_exact(&mut hdr_data) @@ -3470,9 +3672,17 @@ fn validate_single_luks2_header(mut reader: impl std::io::Read, hdr_ind: u64) -> if keyslots.len() != 1 { bail!("Invalid LUKS keyslots"); } - if !tokens.is_empty() { - bail!("Invalid LUKS tokens"); - } + let key_epoch = match tokens.len() { + 0 => None, + 1 => { + let token = tokens.get(&0).context("Invalid LUKS token id")?; + if token.kind != KEY_EPOCH_TOKEN_TYPE || token.keyslots != ["0"] { + bail!("Invalid LUKS token"); + } + Some(token.key_epoch.context("LUKS token has no key epoch")?) + } + _ => bail!("Invalid LUKS tokens"), + }; if segments.len() != 1 { bail!("Invalid LUKS segments"); } @@ -3588,18 +3798,31 @@ fn validate_single_luks2_header(mut reader: impl std::io::Read, hdr_ind: u64) -> bail!("Invalid LUKS digest segments: {segments:?}"); } } - Ok(initializing) + Ok(LuksHeaderState { + initializing, + key_epoch, + }) } #[test] fn test_validate_luks2_header() { let mut header_data = include_bytes!("../tests/fixtures/luks_header_good").to_vec(); - assert!( - !validate_luks2_headers(&mut &header_data[..]).expect("Failed to validate LUKS2 header") + let state = + validate_luks2_headers(&mut &header_data[..]).expect("Failed to validate LUKS2 header"); + assert_eq!( + state, + LuksHeaderState { + initializing: false, + key_epoch: None + } ); header_data[24..24 + LUKS_INITIALIZING_LABEL.len()] .copy_from_slice(LUKS_INITIALIZING_LABEL.as_bytes()); - assert!(validate_luks2_headers(&mut &header_data[..]).expect("Failed to validate LUKS2 header")); + assert!( + validate_luks2_headers(&mut &header_data[..]) + .expect("Failed to validate LUKS2 header") + .initializing + ); let header_data = include_bytes!("../tests/fixtures/luks_header_cipher_null").to_vec(); let error = validate_luks2_headers(&mut &header_data[..]).unwrap_err(); assert!(error @@ -3643,6 +3866,89 @@ fn patch_luks_json(header: &mut [u8], needle: &[u8], replacement: &[u8]) -> usiz patched } +#[test] +fn test_disk_key_migration_with_cryptsetup() { + if Command::new("cryptsetup") + .arg("--version") + .output() + .is_err() + { + eprintln!("cryptsetup is unavailable; skipping"); + return; + } + let dir = tempfile::tempdir().unwrap(); + let image = dir.path().join("disk.img"); + fs::File::create(&image) + .unwrap() + .set_len(2 * PAYLOAD_OFFSET) + .unwrap(); + let offset = (PAYLOAD_OFFSET / 512).to_string(); + let format = cryptsetup( + &[ + "luksFormat", + "--batch-mode", + "--type", + "luks2", + "--offset", + &offset, + "--cipher", + "aes-xts-plain64", + "--pbkdf", + "pbkdf2", + "--sector-size", + "512", + "-d-", + ], + &image, + ); + assert_eq!(run_with_input(format, b"old").unwrap(), Some(0)); + + assert!(luks_unlocks(&image, "old").unwrap()); + assert!(!luks_unlocks(&image, "new").unwrap()); + assert_eq!(read_luks_key_epoch(&image).unwrap(), None); + + let secret_dir = dir.path().join("rekey"); + luks_change_key(&image, "old", "new", &secret_dir).unwrap(); + assert!(!secret_dir.exists()); + assert!(luks_unlocks(&image, "new").unwrap()); + assert!(!luks_unlocks(&image, "old").unwrap()); + + write_luks_key_epoch(&image, 2, false).unwrap(); + assert_eq!(read_luks_key_epoch(&image).unwrap(), Some(2)); + write_luks_key_epoch(&image, 3, true).unwrap(); + assert_eq!(read_luks_key_epoch(&image).unwrap(), Some(3)); + + // The re-keyed, tagged header still passes the boot-time validation. + let header = dir.path().join("header"); + let status = cryptsetup(&["luksHeaderBackup", "--header-backup-file"], &header) + .arg(&image) + .status() + .unwrap(); + assert!(status.success()); + let state = validate_luks2_headers(fs::File::open(&header).unwrap()).unwrap(); + assert_eq!(state.key_epoch, Some(3)); +} + +#[test] +fn test_validate_luks2_header_key_epoch_token() { + let with_token = |token: &str| { + let mut header = include_bytes!("../tests/fixtures/luks_header_good").to_vec(); + let replacement = format!(r#""tokens":{{"0":{token}}}"#); + let patched = patch_luks_json(&mut header, br#""tokens":{}"#, replacement.as_bytes()); + assert_eq!(patched, 2, "expected to patch both header copies"); + validate_luks2_headers(&mut &header[..]) + }; + let state = with_token(r#"{"type":"dstack-kms","keyslots":["0"],"key_epoch":2}"#).unwrap(); + assert_eq!(state.key_epoch, Some(2)); + for token in [ + r#"{"type":"luks2-keyring","keyslots":["0"],"key_epoch":2}"#, + r#"{"type":"dstack-kms","keyslots":["1"],"key_epoch":2}"#, + r#"{"type":"dstack-kms","keyslots":["0"]}"#, + ] { + assert!(with_token(token).is_err(), "{token}"); + } +} + #[test] fn test_validate_luks2_header_rejects_keyslot_area_that_overflows() { let mut header = include_bytes!("../tests/fixtures/luks_header_good").to_vec(); diff --git a/dstack/dstack-util/src/system_setup/luks.rs b/dstack/dstack-util/src/system_setup/luks.rs index 6d0381b9e..254548921 100644 --- a/dstack/dstack-util/src/system_setup/luks.rs +++ b/dstack/dstack-util/src/system_setup/luks.rs @@ -49,12 +49,24 @@ impl LuksHeader { #[derive(Deserialize)] pub(super) struct LuksJson { pub keyslots: HashMap, - pub tokens: HashMap, + pub tokens: HashMap, pub segments: HashMap, pub digests: HashMap, pub config: LuksConfig, } +/// Type of the only token dstack writes: the KMS key epoch of keyslot 0. +pub(super) const KEY_EPOCH_TOKEN_TYPE: &str = "dstack-kms"; + +#[derive(Deserialize)] +pub(super) struct LuksToken { + #[serde(rename = "type")] + pub kind: String, + pub keyslots: Vec, + #[serde(default)] + pub key_epoch: Option, +} + #[derive(Deserialize)] #[serde(tag = "type", rename_all = "lowercase")] pub(super) enum LuksKeyslot { From 1d502def9939a01f87f2bfb766175a309e5ebfe2 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 5 Nov 2026 22:38:35 -0800 Subject: [PATCH 04/17] guest: re-encrypt envs left on an older key epoch and let the VMM swap them in When the encrypted env does not decrypt with the current key, boot tries the older epochs from newest to oldest. Once one decrypts it, the env is re-encrypted to the current env public key inside the CVM, checked to round-trip, and recorded with the SHA-256 of the env it came from. GuestApi.GetReencryptedEnv returns that copy; it reaches only the VMM, which talks to the guest over its own vsock, so its origin is the VM. The new Vmm.ApplyReencryptedEnv RPC (vmm-cli apply-reencrypted-env) replaces the stored env with it only if the stored env still matches the recorded hash, so a concurrent update is never overwritten. Signed-off-by: Kevin Wang --- docs/vmm-cli-user-guide.md | 11 +++ dstack/dstack-types/src/lib.rs | 13 ++++ dstack/dstack-types/src/shared_filenames.rs | 1 + dstack/dstack-util/src/crypto.rs | 37 +++++++++ dstack/dstack-util/src/system_setup.rs | 86 ++++++++++++++------- dstack/guest-agent/src/guest_api_service.rs | 18 ++++- dstack/guest-api/proto/guest_api.proto | 14 ++++ dstack/vmm/rpc/proto/vmm_rpc.proto | 11 +++ dstack/vmm/src/guest_api_service.rs | 9 ++- dstack/vmm/src/main_service.rs | 32 +++++++- dstack/vmm/src/vmm-cli.py | 16 ++++ 11 files changed, 212 insertions(+), 36 deletions(-) diff --git a/docs/vmm-cli-user-guide.md b/docs/vmm-cli-user-guide.md index c2953d42f..94e335c1e 100644 --- a/docs/vmm-cli-user-guide.md +++ b/docs/vmm-cli-user-guide.md @@ -534,6 +534,17 @@ The whitelist is stored in `~/.dstack-vmm/kms-whitelist.json`. ./vmm-cli.py update-env --env-file ./new-secrets.env ``` +After a KMS key rotation, a VM whose env is still encrypted to an older key +epoch re-encrypts it to the current epoch at boot. Replace the stored env with +that copy so later boots no longer need the old key: + +```bash +./vmm-cli.py apply-reencrypted-env +``` + +The env is only replaced if it has not changed since the VM booted. Rotate +any secrets the old key could expose separately. + #### Update Application Compose ```bash diff --git a/dstack/dstack-types/src/lib.rs b/dstack/dstack-types/src/lib.rs index 1d7070313..3b9edfbed 100644 --- a/dstack/dstack-types/src/lib.rs +++ b/dstack/dstack-types/src/lib.rs @@ -2392,6 +2392,19 @@ impl OsImageMeasurementDocument { } } +/// Encrypted env that boot re-encrypted from an older KMS key epoch to the +/// current one, for the host to swap in place of the env it supplied. +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq)] +pub struct ReencryptedEnv { + /// SHA-256 of the encrypted env the host supplied for this boot. + #[serde(with = "hex_bytes")] + pub source_sha256: Vec, + /// The same env encrypted to the current env public key. + #[serde(with = "hex_bytes")] + pub encrypted_env: Vec, + pub key_epoch: u32, +} + #[derive(Serialize, Deserialize, Debug, Clone)] pub struct AppKeys { #[serde(with = "hex_bytes")] diff --git a/dstack/dstack-types/src/shared_filenames.rs b/dstack/dstack-types/src/shared_filenames.rs index 79944cff6..a653490a1 100644 --- a/dstack/dstack-types/src/shared_filenames.rs +++ b/dstack/dstack-types/src/shared_filenames.rs @@ -10,6 +10,7 @@ pub const USER_CONFIG: &str = ".user-config"; pub const ENCRYPTED_ENV: &str = ".encrypted-env"; pub const DECRYPTED_ENV: &str = ".decrypted-env"; pub const DECRYPTED_ENV_JSON: &str = ".decrypted-env.json"; +pub const REENCRYPTED_ENV: &str = ".reencrypted-env.json"; pub const INSTANCE_INFO: &str = ".instance_info"; pub const HOST_SHARED_DIR: &str = "/dstack/.host-shared"; pub const HOST_SHARED_DIR_NAME: &str = ".host-shared"; diff --git a/dstack/dstack-util/src/crypto.rs b/dstack/dstack-util/src/crypto.rs index 1f8227f5b..dc8527620 100644 --- a/dstack/dstack-util/src/crypto.rs +++ b/dstack/dstack-util/src/crypto.rs @@ -68,6 +68,34 @@ pub fn dh_decrypt(secret: [u8; 32], ciphertext: &[u8]) -> Result> { .map_err(|e| anyhow!("Decryption failed: {}", e)) } +/// Encrypt to `remote_public_key` in the format [`dh_decrypt`] reads: +/// ephemeral public key, IV, then the AES-GCM ciphertext. +pub fn dh_encrypt(remote_public_key: [u8; 32], plaintext: &[u8]) -> Result> { + let mut ephemeral_secret = [0u8; 32]; + getrandom::fill(&mut ephemeral_secret).context("failed to generate ephemeral secret")?; + let ephemeral_secret = StaticSecret::from(ephemeral_secret); + let shared_secret = ephemeral_secret + .diffie_hellman(&PublicKey::from(remote_public_key)) + .to_bytes(); + ensure!( + !shared_secret.iter().all(|byte| *byte == 0), + "invalid X25519 shared secret" + ); + let cipher = Aes256Gcm::new_from_slice(&shared_secret) + .map_err(|e| anyhow!("failed to create cipher: {e}"))?; + let mut iv = [0u8; 12]; + getrandom::fill(&mut iv).context("failed to generate IV")?; + let ciphertext = cipher + .encrypt(Nonce::::from_slice(&iv), plaintext) + .map_err(|e| anyhow!("encryption failed: {e}"))?; + Ok([ + PublicKey::from(&ephemeral_secret).as_bytes(), + &iv[..], + &ciphertext, + ] + .concat()) +} + fn stream_nonce(prefix: &[u8; 8], index: u32) -> [u8; 12] { let mut nonce = [0u8; 12]; nonce[..8].copy_from_slice(prefix); @@ -266,6 +294,15 @@ pub fn dh_decrypt_stream( mod tests { use super::*; + #[test] + fn dh_encrypt_round_trips_through_dh_decrypt() { + let secret = [7u8; 32]; + let public = PublicKey::from(&StaticSecret::from(secret)).to_bytes(); + let ciphertext = dh_encrypt(public, b"env").unwrap(); + assert_eq!(dh_decrypt(secret, &ciphertext).unwrap(), b"env"); + assert!(dh_decrypt([8u8; 32], &ciphertext).is_err()); + } + #[test] fn test_dh_agree() { use rand::Rng; diff --git a/dstack/dstack-util/src/system_setup.rs b/dstack/dstack-util/src/system_setup.rs index b5374b3eb..986bd8698 100644 --- a/dstack/dstack-util/src/system_setup.rs +++ b/dstack/dstack-util/src/system_setup.rs @@ -22,9 +22,9 @@ use dstack_types::{ gpu_policy_hash, shared_filenames::{ APP_COMPOSE, APP_KEYS, DECRYPTED_ENV, DECRYPTED_ENV_JSON, ENCRYPTED_ENV, - HOST_SHARED_DIR_NAME, INSTANCE_INFO, SYS_CONFIG, USER_CONFIG, + HOST_SHARED_DIR_NAME, INSTANCE_INFO, REENCRYPTED_ENV, SYS_CONFIG, USER_CONFIG, }, - GpuPolicy, KeyProvider, KeyProviderInfo, GPU_ATTESTATION_OUTPUT, + GpuPolicy, KeyProvider, KeyProviderInfo, ReencryptedEnv, GPU_ATTESTATION_OUTPUT, }; use fs_err as fs; use luks::{ @@ -47,7 +47,7 @@ use tracing::{info, warn}; use crate::{ cmd_show_mrs, - crypto::dh_decrypt, + crypto::{dh_decrypt, dh_encrypt}, gen_app_keys_from_seed, host_api::HostApi, host_shared::{mount_host_shared, unmount_host_shared}, @@ -3292,25 +3292,45 @@ impl<'a> Stage0<'a> { } impl Stage1<'_> { - fn decrypt_env_vars( - &self, - key: &[u8], - ciphertext: &[u8], - allowed: &BTreeSet, - ) -> Result> { - let vars = if !ciphertext.is_empty() { - info!("Processing encrypted env"); - let env_crypt_key: [u8; 32] = key.try_into().ok().context( - "encrypted env is present but the key provider gave no valid env crypt key", - )?; - let decrypted_json = - dh_decrypt(env_crypt_key, ciphertext).context("Failed to decrypt env file")?; - crate::parse_env_file::parse_env(&decrypted_json, allowed)? - } else { - info!("No encrypted env, using default"); - Default::default() + /// Decrypt the host-supplied env. An env still encrypted to an older KMS + /// key epoch is re-encrypted to the current one for the host to swap in. + async fn decrypt_env(&self) -> Result> { + let ciphertext = &self.shared.encrypted_env; + let current_key = env_crypt_key(&self.keys.env_crypt_key)?; + let error = match dh_decrypt(current_key, ciphertext) { + Ok(plaintext) => return Ok(plaintext), + Err(error) => error, }; - Ok(vars) + let current = kms_key_epoch(&self.keys).unwrap_or(0); + for epoch in (0..current).rev() { + let old = self + .shared + .request_historical_keys(&self.keys, epoch) + .await?; + let Ok(plaintext) = dh_decrypt(env_crypt_key(&old.env_crypt_key)?, ciphertext) else { + continue; + }; + info!("re-encrypting the env from key epoch {epoch} to {current}"); + let public_key = + x25519_dalek::PublicKey::from(&x25519_dalek::StaticSecret::from(current_key)); + let encrypted_env = dh_encrypt(public_key.to_bytes(), &plaintext)?; + if dh_decrypt(current_key, &encrypted_env)? != plaintext { + bail!("re-encrypted env does not round-trip"); + } + let migrated = ReencryptedEnv { + source_sha256: sha256(ciphertext).to_vec(), + encrypted_env, + key_epoch: current, + }; + safe_write_with_mode( + self.shared.dir.join(REENCRYPTED_ENV), + serde_json::to_vec(&migrated)?, + 0o600, + ) + .context("failed to write the re-encrypted env")?; + return Ok(plaintext); + } + Err(error).context("failed to decrypt env file") } fn write_env_file(&self, env_vars: &BTreeMap) -> Result<()> { @@ -3327,7 +3347,7 @@ impl Stage1<'_> { Ok(()) } - fn unseal_env_vars(&self) -> Result> { + async fn unseal_env_vars(&self) -> Result> { let allowed_envs: BTreeSet = self .shared .app_compose @@ -3335,18 +3355,20 @@ impl Stage1<'_> { .iter() .cloned() .collect(); - // Decrypt env file - let decrypted_env = self.decrypt_env_vars( - &self.keys.env_crypt_key, - &self.shared.encrypted_env, - &allowed_envs, - )?; + let decrypted_env = if self.shared.encrypted_env.is_empty() { + info!("No encrypted env, using default"); + Default::default() + } else { + info!("Processing encrypted env"); + let plaintext = self.decrypt_env().await?; + crate::parse_env_file::parse_env(&plaintext, &allowed_envs)? + }; self.write_env_file(&decrypted_env)?; Ok(decrypted_env) } async fn setup(&self) -> Result<()> { - let _envs = self.unseal_env_vars()?; + let _envs = self.unseal_env_vars().await?; self.link_files()?; self.setup_socket_dir()?; self.setup_guest_agent_config()?; @@ -3460,6 +3482,12 @@ macro_rules! const_pad { const PAYLOAD_OFFSET: u64 = 16777216; +fn env_crypt_key(key: &[u8]) -> Result<[u8; 32]> { + key.try_into() + .ok() + .context("encrypted env is present but the key provider gave no valid env crypt key") +} + fn kms_key_epoch(keys: &AppKeys) -> Option { match keys.key_provider { KeyProvider::Kms { key_epoch, .. } => Some(key_epoch), diff --git a/dstack/guest-agent/src/guest_api_service.rs b/dstack/guest-agent/src/guest_api_service.rs index cf3e96163..aec2f76b2 100644 --- a/dstack/guest-agent/src/guest_api_service.rs +++ b/dstack/guest-agent/src/guest_api_service.rs @@ -7,13 +7,13 @@ use std::{collections::HashSet, fmt::Debug, path::PathBuf, process::ExitStatus, use anyhow::{Context, Result}; use bollard::{container::ListContainersOptions, Docker}; use dstack_guest_agent_rpc::v0::worker_server::WorkerRpc as _; -use dstack_types::shared_filenames::{HOST_SHARED_DIR, SYS_CONFIG}; +use dstack_types::shared_filenames::{HOST_SHARED_DIR, REENCRYPTED_ENV, SYS_CONFIG}; use dstack_types::SysConfig; use fs_err as fs; use guest_api::{ guest_api_server::{GuestApiRpc, GuestApiServer}, Container, DiskInfo, Gateway, GpuInfoResponse, GuestInfo, Interface, IpAddress, - ListContainersResponse, NetworkInformation, SystemInfo, + ListContainersResponse, NetworkInformation, ReencryptedEnv, SystemInfo, }; use host_api::Notification; use ra_rpc::{CallContext, RpcCall}; @@ -44,6 +44,20 @@ impl RpcCall for GuestApiHandler { } impl GuestApiRpc for GuestApiHandler { + async fn get_reencrypted_env(self) -> Result { + let path = PathBuf::from(HOST_SHARED_DIR).join(REENCRYPTED_ENV); + let migrated: dstack_types::ReencryptedEnv = match fs::read(&path) { + Ok(json) => serde_json::from_slice(&json).context("invalid re-encrypted env")?, + Err(err) if err.kind() == std::io::ErrorKind::NotFound => return Ok(Default::default()), + Err(err) => return Err(err).context("failed to read the re-encrypted env"), + }; + Ok(ReencryptedEnv { + source_sha256: migrated.source_sha256, + encrypted_env: migrated.encrypted_env, + key_epoch: migrated.key_epoch, + }) + } + async fn info(self) -> Result { let ext_rpc = ExternalRpcHandler::new(self.state); let info = timeout(GUEST_API_HANDLER_TIMEOUT, ext_rpc.info()) diff --git a/dstack/guest-api/proto/guest_api.proto b/dstack/guest-api/proto/guest_api.proto index 9e2a9f990..98a31616e 100644 --- a/dstack/guest-api/proto/guest_api.proto +++ b/dstack/guest-api/proto/guest_api.proto @@ -185,6 +185,17 @@ message GpuInfoResponse { optional uint64 sample_age_ms = 5; } +// Env that boot re-encrypted from an older KMS key epoch to the current one. +// Everything is empty when the env this boot started from needed no change. +message ReencryptedEnv { + // SHA-256 of the encrypted env the host supplied for this boot. + bytes source_sha256 = 1; + // The same env encrypted to the current env public key. + bytes encrypted_env = 2; + // KMS key epoch `encrypted_env` is encrypted for. + uint32 key_epoch = 3; +} + // Direct gRPC surface exposed by the in-guest agent. service GuestApi { // Returns attestation material and identifiers for the calling guest. @@ -199,6 +210,8 @@ service GuestApi { rpc ListContainers(google.protobuf.Empty) returns (ListContainersResponse); // Initiates a graceful shutdown inside the guest VM. rpc Shutdown(google.protobuf.Empty) returns (google.protobuf.Empty); + // Returns the env re-encrypted to the current KMS key epoch during boot. + rpc GetReencryptedEnv(google.protobuf.Empty) returns (ReencryptedEnv); } // Same API surface as GuestApi but multiplexed by VM ID through VMM. @@ -209,4 +222,5 @@ service ProxiedGuestApi { rpc NetworkInfo(Id) returns (NetworkInformation); rpc ListContainers(Id) returns (ListContainersResponse); rpc Shutdown(Id) returns (google.protobuf.Empty); + rpc GetReencryptedEnv(Id) returns (ReencryptedEnv); } diff --git a/dstack/vmm/rpc/proto/vmm_rpc.proto b/dstack/vmm/rpc/proto/vmm_rpc.proto index 84b66b2d6..7b2383cd1 100644 --- a/dstack/vmm/rpc/proto/vmm_rpc.proto +++ b/dstack/vmm/rpc/proto/vmm_rpc.proto @@ -448,6 +448,17 @@ service Vmm { // Delete a local guest image by name. rpc DeleteImage(Id) returns (google.protobuf.Empty); + + // Replace the VM's encrypted env with the copy the guest re-encrypted to + // the current KMS key epoch at boot, if the env has not changed since. + rpc ApplyReencryptedEnv(Id) returns (ApplyReencryptedEnvResponse); +} + +message ApplyReencryptedEnvResponse { + // Whether the stored encrypted env was replaced. + bool replaced = 1; + // KMS key epoch the stored env is now encrypted for, when replaced. + uint32 key_epoch = 2; } // Response containing a list of supervisor processes. diff --git a/dstack/vmm/src/guest_api_service.rs b/dstack/vmm/src/guest_api_service.rs index 439df00d1..dbf945cb1 100644 --- a/dstack/vmm/src/guest_api_service.rs +++ b/dstack/vmm/src/guest_api_service.rs @@ -6,7 +6,8 @@ use crate::App as AppState; use anyhow::Result; use guest_api::{ proxied_guest_api_server::{ProxiedGuestApiRpc, ProxiedGuestApiServer}, - GpuInfoResponse, GuestInfo, Id, ListContainersResponse, NetworkInformation, SystemInfo, + GpuInfoResponse, GuestInfo, Id, ListContainersResponse, NetworkInformation, ReencryptedEnv, + SystemInfo, }; use ra_rpc::{CallContext, RpcCall}; use std::ops::Deref; @@ -59,4 +60,10 @@ impl ProxiedGuestApiRpc for GuestApiHandler { async fn shutdown(self, request: Id) -> Result<()> { self.guest_agent_client(&request.id)?.shutdown().await } + + async fn get_reencrypted_env(self, request: Id) -> Result { + self.guest_agent_client(&request.id)? + .get_reencrypted_env() + .await + } } diff --git a/dstack/vmm/src/main_service.rs b/dstack/vmm/src/main_service.rs index ffd06ad61..23af1cdf0 100644 --- a/dstack/vmm/src/main_service.rs +++ b/dstack/vmm/src/main_service.rs @@ -12,10 +12,11 @@ use dstack_types::AppCompose; use dstack_vmm_rpc as rpc; use dstack_vmm_rpc::vmm_server::{VmmRpc, VmmServer}; use dstack_vmm_rpc::{ - AppId, ComposeHash as RpcComposeHash, GatewaySettings, GetInfoResponse, GetMetaResponse, Id, - ImageInfo as RpcImageInfo, ImageListResponse, KmsSettings, ListGpusResponse, PublicKeyResponse, - ReloadVmsResponse, ResizeVmRequest, ResourcesSettings, StatusRequest, StatusResponse, - SvListResponse, SvProcessInfo, UpdateVmRequest, VersionResponse, VmConfiguration, + AppId, ApplyReencryptedEnvResponse, ComposeHash as RpcComposeHash, GatewaySettings, + GetInfoResponse, GetMetaResponse, Id, ImageInfo as RpcImageInfo, ImageListResponse, + KmsSettings, ListGpusResponse, PublicKeyResponse, ReloadVmsResponse, ResizeVmRequest, + ResourcesSettings, StatusRequest, StatusResponse, SvListResponse, SvProcessInfo, + UpdateVmRequest, VersionResponse, VmConfiguration, }; use fs_err as fs; use path_absolutize::Absolutize; @@ -1297,6 +1298,29 @@ impl VmmRpc for RpcHandler { Ok(()) } + async fn apply_reencrypted_env(self, request: Id) -> Result { + use sha2::Digest; + let migrated = self + .guest_agent_client(&request.id)? + .get_reencrypted_env() + .await?; + if migrated.encrypted_env.is_empty() { + return Ok(ApplyReencryptedEnvResponse::default()); + } + let path = self.encrypted_env_path(&request.id)?; + let stored = fs::read(&path).context("Failed to read encrypted env")?; + if sha2::Sha256::digest(&stored)[..] != migrated.source_sha256[..] { + bail!("the encrypted env changed since the VM booted; not replacing it"); + } + safe_write::safe_write(&path, &migrated.encrypted_env) + .context("Failed to write encrypted env")?; + info!(vm_id = %request.id, key_epoch = migrated.key_epoch, "applied re-encrypted env"); + Ok(ApplyReencryptedEnvResponse { + replaced: true, + key_epoch: migrated.key_epoch, + }) + } + async fn shutdown_vm(self, request: Id) -> Result<()> { info!(vm_id = %request.id, "shutdown_vm RPC called"); self.guest_agent_client(&request.id)?.shutdown().await?; diff --git a/dstack/vmm/src/vmm-cli.py b/dstack/vmm/src/vmm-cli.py index ad469f981..39dd50547 100755 --- a/dstack/vmm/src/vmm-cli.py +++ b/dstack/vmm/src/vmm-cli.py @@ -692,6 +692,14 @@ def stop_vm(self, vm_id: str, force: bool = False) -> None: self.rpc_call("ShutdownVm", {"id": vm_id}) print(f"Gracefully shutting down VM {vm_id}") + def apply_reencrypted_env(self, vm_id: str) -> None: + """Swap in the env the guest re-encrypted to the current KMS key epoch.""" + result = self.rpc_call("ApplyReencryptedEnv", {"id": vm_id}) + if result.get("replaced"): + print(f"Replaced env of VM {vm_id} (key epoch {result.get('key_epoch', 0)})") + else: + print(f"VM {vm_id} has no re-encrypted env to apply") + def remove_vm(self, vm_id: str) -> None: """Remove a VM.""" self.rpc_call("RemoveVm", {"id": vm_id}) @@ -2159,6 +2167,12 @@ def _patched_format_help(): "--kms-url", action="append", type=str, help="KMS URL" ) + apply_env_parser = subparsers.add_parser( + "apply-reencrypted-env", + help="Replace a VM's env with the copy it re-encrypted to the current KMS key epoch", + ) + apply_env_parser.add_argument("vm_id", help="VM ID") + # Whitelist command kms_parser = subparsers.add_parser("kms", help="Manage trusted KMS whitelist") kms_subparsers = kms_parser.add_subparsers(dest="kms_action", help="KMS actions") @@ -2395,6 +2409,8 @@ def _patched_format_help(): cli.update_vm_env( args.vm_id, parse_env_file(args.env_file), kms_urls=args.kms_url ) + elif args.command == "apply-reencrypted-env": + cli.apply_reencrypted_env(args.vm_id) elif args.command == "update-user-config": cli.update_vm_user_config(args.vm_id, open(args.user_config, "r").read()) elif args.command == "update-app-compose": From a87f8ed4ae8a0a21c5aafa11f7954af65358881c Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 5 Nov 2026 22:39:23 -0800 Subject: [PATCH 05/17] docs: describe KMS root key rotation Signed-off-by: Kevin Wang --- docs/deployment.md | 22 ++++++++++++++++++++++ dstack/kms/auth-simple/README.md | 1 + 2 files changed, 23 insertions(+) diff --git a/docs/deployment.md b/docs/deployment.md index 4f6656668..003992ef5 100644 --- a/docs/deployment.md +++ b/docs/deployment.md @@ -513,6 +513,28 @@ curl http://:9203/finish > > **Key handover.** It is served as `Admin.GetKmsKey` on the `[core.admin]` listener, which needs `[core.admin.tls]` so it can verify the onboarding KMS's RA-TLS certificate. Pass the source admin URL and its token (`source_token`) in the onboarding request. `core.onboard.public_key_handover = true` keeps the legacy public `KMS.GetKmsKey`. +#### Rotating KMS Root Keys + +KMS root keys are versioned by epoch, and every KMS keeps all epochs so CVMs +can migrate from any of them. See the [rotation design](./kms-emergency-key-rotation-protocol.md) +for the full procedure; the KMS steps are: + +1. If epochs were never registered, register the current root public keys as + epoch 0: `DstackKms.registerKeyEpoch(caPubkey, k256Pubkey)` on chain, or a + `keyEpochs` entry in the auth-simple config. +2. Start a fresh KMS with `[core.key_rotation] current_epoch = N` and onboard it + with `"rotate": true` (the **Rotate** checkbox in the onboard UI) from a KMS + at epoch N-1. It imports epochs 0..N-1, generates epoch N, and returns the new + public keys with an attestation whose report data is + `keccak256("dstack-kms-generated-keys-v2:{N};{ca_pubkey_hex};{k256_pubkey_hex};")`. +3. Verify the attestation and register epoch N. A KMS refuses to start until its + keyring matches the registered epochs, and a KMS still at an older epoch stops + serving once a newer one is registered. +4. Onboard the other KMS nodes from the rotated one with `current_epoch = N`. +5. Boot CVMs on an OS that supports key epochs. They re-key their data disk and + re-encrypt their env at boot; apply the new env with + `vmm-cli.py apply-reencrypted-env `. + --- ## Deploying Apps diff --git a/dstack/kms/auth-simple/README.md b/dstack/kms/auth-simple/README.md index efd2efcb5..fc40f7314 100644 --- a/dstack/kms/auth-simple/README.md +++ b/dstack/kms/auth-simple/README.md @@ -107,6 +107,7 @@ raw NitroTPM evidence by itself. | `apps..composeHashes` | No | Allowed compose hashes for this app | | `apps..devices` | No | Allowed device IDs for this app. Empty denies all unless `allowAnyDevice` is true. | | `apps..allowAnyDevice` | No | If true, skip device ID check for this app | +| `keyEpochs` | No | KMS root public keys per key epoch, `[{"caPubkey": "0x...", "k256Pubkey": "0x..."}]`, indexed by epoch. Append only. When set, the KMS checks its keyring against it and boot responses report the latest epoch as `keyEpoch`. | For experimental AMD SEV-SNP dry-run authorization, keep the default fail-closed TCB policy unless you intentionally want the auth webhook to accept non-up-to-date verifier-derived SNP `BootInfo`. To exercise the dry-run path without enabling key release, allowlist the recomputed SNP `mrAggregated`, `osImageHash`, app/compose identity, device/chip identity, and any non-default `allowedTcbStatuses`/`allowedAdvisoryIds` values explicitly. KMS still rejects SNP before returning app keys, KMS keys, or app certificates. From e42c616fff31862f8a889089e5586e09e194be21 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 5 Nov 2026 23:02:12 -0800 Subject: [PATCH 06/17] Count root epochs from 1, use 0 for latest, and select the epoch in GetKey Epoch 1 is now the original root keys and 0 means the current epoch in every selector, so an absent field asks for the latest keys. Before any rotation that is epoch 1, so older OS images see no change; after one they are cut off by the OS allowlist anyway. A peer that predates root epochs reports 0, which is read as epoch 1. The key_epoch fields are renamed root_epoch to leave room for a per-app epoch later. At boot the guest fetches the app keys of every older root epoch once. Disk and env migration reuse them, and their app root keys go into the keys file. Guest API v1 GetKey takes root_epoch (0 = current) and reports the epoch it derived from; the SDKs expose it as an optional argument (Rust get_key_at_root_epoch, Go GetKeyAtRootEpoch). Signed-off-by: Kevin Wang --- docs/deployment.md | 6 +- dstack/dstack-types/src/lib.rs | 39 +++- dstack/dstack-util/src/main.rs | 6 +- dstack/dstack-util/src/system_setup.rs | 184 ++++++++++-------- dstack/dstack-util/src/system_setup/luks.rs | 4 +- .../guest-agent/rpc/proto/agent_rpc_v1.proto | 8 + dstack/guest-agent/src/guest_api_service.rs | 2 +- dstack/guest-agent/src/rpc_service.rs | 48 +++++ dstack/guest-agent/src/rpc_service_v1.rs | 59 +++++- dstack/guest-api/proto/guest_api.proto | 2 +- dstack/kms/auth-eth/contracts/DstackKms.sol | 33 ++-- dstack/kms/auth-eth/src/ethereum.ts | 22 +-- dstack/kms/auth-eth/src/main.test.ts | 10 +- dstack/kms/auth-eth/src/server.ts | 6 +- dstack/kms/auth-eth/src/types.ts | 6 +- dstack/kms/auth-eth/test/DstackKms.t.sol | 32 +-- dstack/kms/auth-simple/README.md | 2 +- dstack/kms/auth-simple/index.test.ts | 10 +- dstack/kms/auth-simple/index.ts | 18 +- dstack/kms/kms.toml | 4 +- dstack/kms/rpc/proto/kms_rpc.proto | 17 +- dstack/kms/src/config.rs | 47 +++-- dstack/kms/src/main_service.rs | 48 +++-- .../kms/src/main_service/upgrade_authority.rs | 28 +-- dstack/kms/src/onboard_service.rs | 91 +++++---- dstack/vmm/rpc/proto/vmm_rpc.proto | 2 +- dstack/vmm/src/main_service.rs | 4 +- dstack/vmm/src/vmm-cli.py | 2 +- sdk/go/dstack/client_v1.go | 23 ++- sdk/js/src/client-v1.ts | 13 +- sdk/python/src/dstack_sdk/dstack_client_v1.py | 23 ++- sdk/rust/src/dstack_client_v1.rs | 14 +- sdk/rust/types/src/dstack_v1.rs | 4 + 33 files changed, 536 insertions(+), 281 deletions(-) diff --git a/docs/deployment.md b/docs/deployment.md index 003992ef5..e47bc5c17 100644 --- a/docs/deployment.md +++ b/docs/deployment.md @@ -520,11 +520,11 @@ can migrate from any of them. See the [rotation design](./kms-emergency-key-rota for the full procedure; the KMS steps are: 1. If epochs were never registered, register the current root public keys as - epoch 0: `DstackKms.registerKeyEpoch(caPubkey, k256Pubkey)` on chain, or a - `keyEpochs` entry in the auth-simple config. + epoch 1: `DstackKms.registerRootEpoch(caPubkey, k256Pubkey)` on chain, or a + `rootEpochs` entry in the auth-simple config. 2. Start a fresh KMS with `[core.key_rotation] current_epoch = N` and onboard it with `"rotate": true` (the **Rotate** checkbox in the onboard UI) from a KMS - at epoch N-1. It imports epochs 0..N-1, generates epoch N, and returns the new + at epoch N-1. It imports epochs 1..N-1, generates epoch N, and returns the new public keys with an attestation whose report data is `keccak256("dstack-kms-generated-keys-v2:{N};{ca_pubkey_hex};{k256_pubkey_hex};")`. 3. Verify the attestation and register epoch N. A KMS refuses to start until its diff --git a/dstack/dstack-types/src/lib.rs b/dstack/dstack-types/src/lib.rs index 3b9edfbed..e5bfa1b7e 100644 --- a/dstack/dstack-types/src/lib.rs +++ b/dstack/dstack-types/src/lib.rs @@ -2402,7 +2402,17 @@ pub struct ReencryptedEnv { /// The same env encrypted to the current env public key. #[serde(with = "hex_bytes")] pub encrypted_env: Vec, - pub key_epoch: u32, + pub root_epoch: u32, +} + +/// The app root key of an older KMS root epoch. +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq)] +pub struct HistoricalAppRoot { + pub root_epoch: u32, + #[serde(with = "hex_bytes")] + pub k256_key: Vec, + #[serde(with = "hex_bytes")] + pub k256_signature: Vec, } #[derive(Serialize, Deserialize, Debug, Clone)] @@ -2418,6 +2428,21 @@ pub struct AppKeys { pub gateway_app_id: String, pub ca_cert: String, pub key_provider: KeyProvider, + /// App root keys of the KMS root epochs before the current one, from + /// epoch 1, for key requests that select one explicitly. + #[serde(default)] + pub historical_roots: Vec, +} + +impl AppKeys { + /// KMS root epoch the keys were derived from; 1 for other key providers, + /// which have a single root. + pub fn root_epoch(&self) -> u32 { + match self.key_provider { + KeyProvider::Kms { root_epoch, .. } => root_epoch, + _ => first_root_epoch(), + } + } } /// Maximum length of a framed message exchanged with the local key provider. @@ -2444,12 +2469,16 @@ pub enum KeyProvider { pubkey: Vec, tmp_ca_key: String, tmp_ca_cert: String, - /// KMS root key epoch the keys were derived from. - #[serde(default)] - key_epoch: u32, + /// KMS root key epoch the keys were derived from, counting from 1. + #[serde(default = "first_root_epoch")] + root_epoch: u32, }, } +fn first_root_epoch() -> u32 { + 1 +} + impl KeyProvider { pub fn kind(&self) -> KeyProviderKind { match self { @@ -2510,7 +2539,7 @@ mod key_provider_tests { pubkey: vec![0xab; 32], tmp_ca_key: String::new(), tmp_ca_cert: String::new(), - key_epoch: 0, + root_epoch: 1, }; assert_eq!(kms.id(), &[0xab; 32]); } diff --git a/dstack/dstack-util/src/main.rs b/dstack/dstack-util/src/main.rs index f78a77216..fd68d76e0 100644 --- a/dstack/dstack-util/src/main.rs +++ b/dstack/dstack-util/src/main.rs @@ -701,7 +701,7 @@ async fn cmd_get_keys(args: GetKeysArgs) -> Result<()> { .get_app_key(dstack_kms_rpc::GetAppKeyRequest { api_version: 1, vm_config: "".to_string(), - key_epoch: -1, + root_epoch: 0, }) .await .context("Failed to get app key")?; @@ -719,12 +719,13 @@ async fn cmd_get_keys(args: GetKeysArgs) -> Result<()> { k256_key: response.k256_key, k256_signature: response.k256_signature, gateway_app_id: response.gateway_app_id, + historical_roots: Vec::new(), key_provider: KeyProvider::Kms { url: kms_url, pubkey: root_pubkey, tmp_ca_key: tmp_ca.temp_ca_key, tmp_ca_cert: tmp_ca.temp_ca_cert, - key_epoch: response.key_epoch, + root_epoch: response.root_epoch.max(1), }, }; @@ -1236,6 +1237,7 @@ fn make_app_keys( k256_signature: vec![], gateway_app_id: "".to_string(), ca_cert: cert.pem(), + historical_roots: Vec::new(), key_provider, }) } diff --git a/dstack/dstack-util/src/system_setup.rs b/dstack/dstack-util/src/system_setup.rs index 986bd8698..3619a6293 100644 --- a/dstack/dstack-util/src/system_setup.rs +++ b/dstack/dstack-util/src/system_setup.rs @@ -24,12 +24,13 @@ use dstack_types::{ APP_COMPOSE, APP_KEYS, DECRYPTED_ENV, DECRYPTED_ENV_JSON, ENCRYPTED_ENV, HOST_SHARED_DIR_NAME, INSTANCE_INFO, REENCRYPTED_ENV, SYS_CONFIG, USER_CONFIG, }, - GpuPolicy, KeyProvider, KeyProviderInfo, ReencryptedEnv, GPU_ATTESTATION_OUTPUT, + GpuPolicy, HistoricalAppRoot, KeyProvider, KeyProviderInfo, ReencryptedEnv, + GPU_ATTESTATION_OUTPUT, }; use fs_err as fs; use luks::{ LuksAf, LuksArea, LuksConfig, LuksDigest, LuksHeader, LuksJson, LuksKdf, LuksKeyslot, - LuksSegment, BINARY_HEADER_SIZE, KEY_EPOCH_TOKEN_TYPE, + LuksSegment, BINARY_HEADER_SIZE, ROOT_EPOCH_TOKEN_TYPE, }; use ra_rpc::{ client::{CertInfo, RaClient, RaClientConfig}, @@ -2354,13 +2355,13 @@ struct AppInfo { } impl HostShared { - /// Request app keys of `key_epoch` (`-1` = current) from one KMS. + /// Request app keys of `root_epoch` (0 = current) from one KMS. async fn request_kms_keys( &self, kms_url: String, - key_epoch: i32, + root_epoch: u32, ) -> Result<(AppKeys, Vec)> { - info!("Requesting app keys of epoch {key_epoch} from KMS: {kms_url}"); + info!("Requesting app keys of epoch {root_epoch} from KMS: {kms_url}"); let tmp_ca = { info!("Getting temp ca cert"); let client = RaClient::new(kms_url.clone(), true)?; @@ -2389,16 +2390,14 @@ impl HostShared { .get_app_key(rpc::GetAppKeyRequest { api_version: 1, vm_config: self.sys_config.vm_config.clone(), - key_epoch, + root_epoch, }) .await .context("Failed to get app key")?; - // A KMS predating key epochs ignores the selector and serves epoch 0. - if key_epoch >= 0 && response.key_epoch as i32 != key_epoch { - bail!( - "KMS returned keys of epoch {} instead of {key_epoch}", - response.key_epoch - ); + // A KMS predating root epochs ignores the selector and reports 0: epoch 1. + let served = response.root_epoch.max(1); + if root_epoch != 0 && served != root_epoch { + bail!("KMS returned keys of epoch {served} instead of {root_epoch}"); } let (_, ca_pem) = x509_parser::pem::parse_x509_pem(tmp_ca.ca_cert.as_bytes()) @@ -2413,24 +2412,32 @@ impl HostShared { k256_key: response.k256_key, k256_signature: response.k256_signature, gateway_app_id: response.gateway_app_id, + historical_roots: Vec::new(), key_provider: KeyProvider::Kms { url: kms_url, pubkey: root_pubkey, tmp_ca_key: tmp_ca.temp_ca_key, tmp_ca_cert: tmp_ca.temp_ca_cert, - key_epoch: response.key_epoch, + root_epoch: served, }, }; Ok((keys, response.os_image_hash)) } - /// Request keys of an older epoch from the KMS that served `keys`. - async fn request_historical_keys(&self, keys: &AppKeys, epoch: u32) -> Result { - let KeyProvider::Kms { url, .. } = &keys.key_provider else { - bail!("only KMS-provided keys have historical epochs"); + /// Request the keys of every root epoch before the one `keys` came from, + /// oldest first, from the same KMS. Empty for other key providers. + async fn request_key_history(&self, keys: &AppKeys) -> Result> { + let KeyProvider::Kms { + url, root_epoch, .. + } = &keys.key_provider + else { + return Ok(Vec::new()); }; - let epoch = i32::try_from(epoch).context("key epoch out of range")?; - Ok(self.request_kms_keys(url.clone(), epoch).await?.0) + let mut history = Vec::new(); + for epoch in 1..*root_epoch { + history.push(self.request_kms_keys(url.clone(), epoch).await?.0); + } + Ok(history) } } @@ -2445,6 +2452,8 @@ struct Stage1<'a> { vmm: HostApi, shared: HostShared, keys: AppKeys, + /// Keys of root epochs 1 up to, but not including, `keys.root_epoch()`. + key_history: Vec, } fn validate_key_provider_inputs(kind: KeyProviderKind, kms_urls: &[String]) -> Result<()> { @@ -2503,7 +2512,7 @@ impl<'a> Stage0<'a> { let mut errors = vec![]; for kms_url in &self.shared.sys_config.kms_urls { let kms_url = kms_rpc_url(kms_url); - match self.shared.request_kms_keys(kms_url.clone(), -1).await { + match self.shared.request_kms_keys(kms_url.clone(), 0).await { Ok((keys, os_image_hash)) => { // Measured here, outside the per-URL request, so that a // failover can never extend os-image-hash twice. @@ -2746,9 +2755,14 @@ impl<'a> Stage0<'a> { initialized } - async fn mount_data_disk(&self, keys: &AppKeys, opts: &DstackOptions) -> Result<()> { + async fn mount_data_disk( + &self, + keys: &AppKeys, + key_history: &[AppKeys], + opts: &DstackOptions, + ) -> Result<()> { let disk_crypt_key = &hex::encode(&keys.disk_crypt_key); - let key_epoch = kms_key_epoch(keys); + let root_epoch = kms_root_epoch(keys); let name = "dstack_data_disk"; let mount_point = &self.args.mount_point; @@ -2769,9 +2783,8 @@ impl<'a> Stage0<'a> { let has_filesystem = if opts.storage_encrypted && disk_initialized { info!("Mounting encrypted data disk"); - if let Some(epoch) = key_epoch { - self.migrate_disk_key(keys, disk_crypt_key, epoch) - .await + if let Some(epoch) = root_epoch { + self.migrate_disk_key(key_history, disk_crypt_key, epoch) .context("failed to migrate the data disk key")?; } let initializing = @@ -2793,7 +2806,7 @@ impl<'a> Stage0<'a> { info!("Skipping disk encryption as requested by kernel cmdline"); } else if !disk_initialized { info!("Setting up disk encryption"); - self.luks_setup(disk_crypt_key, key_epoch, name, opts.storage_discard)?; + self.luks_setup(disk_crypt_key, root_epoch, name, opts.storage_discard)?; } match opts.storage_fs { @@ -2925,7 +2938,7 @@ impl<'a> Stage0<'a> { fn luks_setup( &self, disk_crypt_key: &str, - key_epoch: Option, + root_epoch: Option, name: &str, discard: bool, ) -> Result<()> { @@ -2966,8 +2979,8 @@ impl<'a> Stage0<'a> { { bail!("Failed to setup luks volume"); } - if let Some(epoch) = key_epoch.filter(|epoch| *epoch > 0) { - write_luks_key_epoch(root_hd, epoch, false)?; + if let Some(epoch) = root_epoch.filter(|epoch| *epoch > 1) { + write_luks_root_epoch(root_hd, epoch, false)?; } self.open_encrypted_volume(disk_crypt_key, name, discard)?; Ok(()) @@ -2976,32 +2989,34 @@ impl<'a> Stage0<'a> { /// Re-key a data disk still protected by an older KMS epoch. /// /// Only the keyslot passphrase changes; the volume key stays. The epoch - /// is recorded in the header so the old key can be requested next time - /// the KMS rotates. An unrotated epoch-0 disk is left untouched, so it - /// stays readable by older OS images. - async fn migrate_disk_key( + /// is recorded in the header so the right old key is used after the next + /// rotation. An unrotated epoch-1 disk is left untouched, so it stays + /// readable by older OS images. + fn migrate_disk_key( &self, - keys: &AppKeys, + key_history: &[AppKeys], current_key: &str, current: u32, ) -> Result<()> { let device = &self.args.device; - let stored = read_luks_key_epoch(device)?; + let stored = read_luks_root_epoch(device)?; if !luks_unlocks(device, current_key)? { - let epoch = stored.unwrap_or(0); + let epoch = stored.unwrap_or(1); if epoch >= current { bail!("the data disk does not unlock with the key of epoch {current}"); } info!("re-keying the data disk from key epoch {epoch} to {current}"); - let old = self.shared.request_historical_keys(keys, epoch).await?; + let old = key_history + .get(epoch as usize - 1) + .with_context(|| format!("no keys for root epoch {epoch}"))?; let old_key = hex::encode(&old.disk_crypt_key); luks_change_key(device, &old_key, current_key, Path::new(REKEY_DIR))?; if !luks_unlocks(device, current_key)? { bail!("the data disk does not unlock after re-keying"); } } - if stored.unwrap_or(0) != current { - write_luks_key_epoch(device, current, stored.is_some())?; + if stored.unwrap_or(1) != current { + write_luks_root_epoch(device, current, stored.is_some())?; } Ok(()) } @@ -3229,10 +3244,23 @@ impl<'a> Stage0<'a> { self.vmm .notify_q("boot.progress", "requesting app keys") .await; - let app_keys = self + let mut app_keys = self .request_app_keys() .await .context("Failed to request app keys")?; + let key_history = self + .shared + .request_key_history(&app_keys) + .await + .context("Failed to request keys of older root epochs")?; + app_keys.historical_roots = key_history + .iter() + .map(|keys| HistoricalAppRoot { + root_epoch: keys.root_epoch(), + k256_key: keys.k256_key.clone(), + k256_signature: keys.k256_signature.clone(), + }) + .collect(); if app_keys.disk_crypt_key.is_empty() { bail!("Failed to get valid key phrase from KMS"); } @@ -3267,7 +3295,7 @@ impl<'a> Stage0<'a> { opts.storage_encrypted, opts.storage_fs ); - self.mount_data_disk(&app_keys, &opts).await?; + self.mount_data_disk(&app_keys, &key_history, &opts).await?; self.setup_swap(self.shared.app_compose.swap_size, &opts) .await?; self.vmm @@ -3287,6 +3315,7 @@ impl<'a> Stage0<'a> { shared: self.shared, vmm: self.vmm, keys: app_keys, + key_history, }) } } @@ -3294,19 +3323,16 @@ impl<'a> Stage0<'a> { impl Stage1<'_> { /// Decrypt the host-supplied env. An env still encrypted to an older KMS /// key epoch is re-encrypted to the current one for the host to swap in. - async fn decrypt_env(&self) -> Result> { + fn decrypt_env(&self) -> Result> { let ciphertext = &self.shared.encrypted_env; let current_key = env_crypt_key(&self.keys.env_crypt_key)?; let error = match dh_decrypt(current_key, ciphertext) { Ok(plaintext) => return Ok(plaintext), Err(error) => error, }; - let current = kms_key_epoch(&self.keys).unwrap_or(0); - for epoch in (0..current).rev() { - let old = self - .shared - .request_historical_keys(&self.keys, epoch) - .await?; + let current = self.keys.root_epoch(); + for (index, old) in self.key_history.iter().enumerate().rev() { + let epoch = index as u32 + 1; let Ok(plaintext) = dh_decrypt(env_crypt_key(&old.env_crypt_key)?, ciphertext) else { continue; }; @@ -3320,7 +3346,7 @@ impl Stage1<'_> { let migrated = ReencryptedEnv { source_sha256: sha256(ciphertext).to_vec(), encrypted_env, - key_epoch: current, + root_epoch: current, }; safe_write_with_mode( self.shared.dir.join(REENCRYPTED_ENV), @@ -3360,7 +3386,7 @@ impl Stage1<'_> { Default::default() } else { info!("Processing encrypted env"); - let plaintext = self.decrypt_env().await?; + let plaintext = self.decrypt_env()?; crate::parse_env_file::parse_env(&plaintext, &allowed_envs)? }; self.write_env_file(&decrypted_env)?; @@ -3488,9 +3514,9 @@ fn env_crypt_key(key: &[u8]) -> Result<[u8; 32]> { .context("encrypted env is present but the key provider gave no valid env crypt key") } -fn kms_key_epoch(keys: &AppKeys) -> Option { +fn kms_root_epoch(keys: &AppKeys) -> Option { match keys.key_provider { - KeyProvider::Kms { key_epoch, .. } => Some(key_epoch), + KeyProvider::Kms { root_epoch, .. } => Some(root_epoch), _ => None, } } @@ -3532,7 +3558,7 @@ fn luks_unlocks(device: &Path, key: &str) -> Result { } } -fn read_luks_key_epoch(device: &Path) -> Result> { +fn read_luks_root_epoch(device: &Path) -> Result> { #[derive(Deserialize)] struct Metadata { tokens: BTreeMap, @@ -3548,15 +3574,15 @@ fn read_luks_key_epoch(device: &Path) -> Result> { Ok(metadata .tokens .values() - .find(|token| token.kind == KEY_EPOCH_TOKEN_TYPE) - .and_then(|token| token.key_epoch)) + .find(|token| token.kind == ROOT_EPOCH_TOKEN_TYPE) + .and_then(|token| token.root_epoch)) } -fn write_luks_key_epoch(device: &Path, epoch: u32, replace: bool) -> Result<()> { +fn write_luks_root_epoch(device: &Path, epoch: u32, replace: bool) -> Result<()> { let token = serde_json::json!({ - "type": KEY_EPOCH_TOKEN_TYPE, + "type": ROOT_EPOCH_TOKEN_TYPE, "keyslots": ["0"], - "key_epoch": epoch, + "root_epoch": epoch, }); let args: &[&str] = if replace { &["token", "import", "--token-id", "0", "--token-replace"] @@ -3613,14 +3639,14 @@ const LUKS_INITIALIZING_LABEL: &str = "dstack-initializing"; struct LuksHeaderState { /// The header carries [`LUKS_INITIALIZING_LABEL`]. initializing: bool, - /// KMS key epoch recorded in the dstack token; `None` predates epochs. - key_epoch: Option, + /// Root epoch recorded in the dstack token; `None` means epoch 1. + root_epoch: Option, } fn validate_luks2_headers(mut reader: impl std::io::Read) -> Result { let primary = validate_single_luks2_header(&mut reader, 0)?; let secondary = validate_single_luks2_header(&mut reader, 1)?; - if primary.key_epoch != secondary.key_epoch { + if primary.root_epoch != secondary.root_epoch { bail!("LUKS header copies disagree on the key epoch"); } Ok(primary) @@ -3700,14 +3726,18 @@ fn validate_single_luks2_header( if keyslots.len() != 1 { bail!("Invalid LUKS keyslots"); } - let key_epoch = match tokens.len() { + let root_epoch = match tokens.len() { 0 => None, 1 => { let token = tokens.get(&0).context("Invalid LUKS token id")?; - if token.kind != KEY_EPOCH_TOKEN_TYPE || token.keyslots != ["0"] { + if token.kind != ROOT_EPOCH_TOKEN_TYPE || token.keyslots != ["0"] { bail!("Invalid LUKS token"); } - Some(token.key_epoch.context("LUKS token has no key epoch")?) + let epoch = token.root_epoch.context("LUKS token has no root epoch")?; + if epoch == 0 { + bail!("Invalid LUKS token root epoch"); + } + Some(epoch) } _ => bail!("Invalid LUKS tokens"), }; @@ -3828,7 +3858,7 @@ fn validate_single_luks2_header( } Ok(LuksHeaderState { initializing, - key_epoch, + root_epoch, }) } @@ -3841,7 +3871,7 @@ fn test_validate_luks2_header() { state, LuksHeaderState { initializing: false, - key_epoch: None + root_epoch: None } ); header_data[24..24 + LUKS_INITIALIZING_LABEL.len()] @@ -3933,7 +3963,7 @@ fn test_disk_key_migration_with_cryptsetup() { assert!(luks_unlocks(&image, "old").unwrap()); assert!(!luks_unlocks(&image, "new").unwrap()); - assert_eq!(read_luks_key_epoch(&image).unwrap(), None); + assert_eq!(read_luks_root_epoch(&image).unwrap(), None); let secret_dir = dir.path().join("rekey"); luks_change_key(&image, "old", "new", &secret_dir).unwrap(); @@ -3941,10 +3971,10 @@ fn test_disk_key_migration_with_cryptsetup() { assert!(luks_unlocks(&image, "new").unwrap()); assert!(!luks_unlocks(&image, "old").unwrap()); - write_luks_key_epoch(&image, 2, false).unwrap(); - assert_eq!(read_luks_key_epoch(&image).unwrap(), Some(2)); - write_luks_key_epoch(&image, 3, true).unwrap(); - assert_eq!(read_luks_key_epoch(&image).unwrap(), Some(3)); + write_luks_root_epoch(&image, 2, false).unwrap(); + assert_eq!(read_luks_root_epoch(&image).unwrap(), Some(2)); + write_luks_root_epoch(&image, 3, true).unwrap(); + assert_eq!(read_luks_root_epoch(&image).unwrap(), Some(3)); // The re-keyed, tagged header still passes the boot-time validation. let header = dir.path().join("header"); @@ -3954,11 +3984,11 @@ fn test_disk_key_migration_with_cryptsetup() { .unwrap(); assert!(status.success()); let state = validate_luks2_headers(fs::File::open(&header).unwrap()).unwrap(); - assert_eq!(state.key_epoch, Some(3)); + assert_eq!(state.root_epoch, Some(3)); } #[test] -fn test_validate_luks2_header_key_epoch_token() { +fn test_validate_luks2_header_root_epoch_token() { let with_token = |token: &str| { let mut header = include_bytes!("../tests/fixtures/luks_header_good").to_vec(); let replacement = format!(r#""tokens":{{"0":{token}}}"#); @@ -3966,11 +3996,11 @@ fn test_validate_luks2_header_key_epoch_token() { assert_eq!(patched, 2, "expected to patch both header copies"); validate_luks2_headers(&mut &header[..]) }; - let state = with_token(r#"{"type":"dstack-kms","keyslots":["0"],"key_epoch":2}"#).unwrap(); - assert_eq!(state.key_epoch, Some(2)); + let state = with_token(r#"{"type":"dstack-kms","keyslots":["0"],"root_epoch":2}"#).unwrap(); + assert_eq!(state.root_epoch, Some(2)); for token in [ - r#"{"type":"luks2-keyring","keyslots":["0"],"key_epoch":2}"#, - r#"{"type":"dstack-kms","keyslots":["1"],"key_epoch":2}"#, + r#"{"type":"luks2-keyring","keyslots":["0"],"root_epoch":2}"#, + r#"{"type":"dstack-kms","keyslots":["1"],"root_epoch":2}"#, r#"{"type":"dstack-kms","keyslots":["0"]}"#, ] { assert!(with_token(token).is_err(), "{token}"); diff --git a/dstack/dstack-util/src/system_setup/luks.rs b/dstack/dstack-util/src/system_setup/luks.rs index 254548921..aaa184a1c 100644 --- a/dstack/dstack-util/src/system_setup/luks.rs +++ b/dstack/dstack-util/src/system_setup/luks.rs @@ -56,7 +56,7 @@ pub(super) struct LuksJson { } /// Type of the only token dstack writes: the KMS key epoch of keyslot 0. -pub(super) const KEY_EPOCH_TOKEN_TYPE: &str = "dstack-kms"; +pub(super) const ROOT_EPOCH_TOKEN_TYPE: &str = "dstack-kms"; #[derive(Deserialize)] pub(super) struct LuksToken { @@ -64,7 +64,7 @@ pub(super) struct LuksToken { pub kind: String, pub keyslots: Vec, #[serde(default)] - pub key_epoch: Option, + pub root_epoch: Option, } #[derive(Deserialize)] diff --git a/dstack/guest-agent/rpc/proto/agent_rpc_v1.proto b/dstack/guest-agent/rpc/proto/agent_rpc_v1.proto index 3df408ad8..ad68adde6 100644 --- a/dstack/guest-agent/rpc/proto/agent_rpc_v1.proto +++ b/dstack/guest-agent/rpc/proto/agent_rpc_v1.proto @@ -189,6 +189,11 @@ message GetKeyRequest { // signing method: v0's `secp256k1_prehashed` was a mode wearing an // algorithm's name, and it does not appear here. string algorithm = 2; + // KMS root key epoch to derive from: 0 (default) = the current epoch, which is + // what new keys should use; N = epoch N, counting from 1. Name an older epoch + // only to reach a key from before a KMS root key rotation, for example to + // migrate data or assets off it. + uint32 root_epoch = 3; } message GetKeyResponse { @@ -209,6 +214,9 @@ message GetKeyResponse { // unversioned surface returns. Link [0] is not: its claim encoding is new, // and it is deliberately not one the v0 claim format can produce. repeated bytes signature_chain = 3; + // KMS root key epoch the key was derived from. Link [1] is signed by that + // epoch's KMS root key. + uint32 root_epoch = 4; } // --------------------------------------------------------------------------- diff --git a/dstack/guest-agent/src/guest_api_service.rs b/dstack/guest-agent/src/guest_api_service.rs index aec2f76b2..54af121f8 100644 --- a/dstack/guest-agent/src/guest_api_service.rs +++ b/dstack/guest-agent/src/guest_api_service.rs @@ -54,7 +54,7 @@ impl GuestApiRpc for GuestApiHandler { Ok(ReencryptedEnv { source_sha256: migrated.source_sha256, encrypted_env: migrated.encrypted_env, - key_epoch: migrated.key_epoch, + root_epoch: migrated.root_epoch, }) } diff --git a/dstack/guest-agent/src/rpc_service.rs b/dstack/guest-agent/src/rpc_service.rs index a6dcd8e5f..cacfce569 100644 --- a/dstack/guest-agent/src/rpc_service.rs +++ b/dstack/guest-agent/src/rpc_service.rs @@ -92,6 +92,14 @@ pub struct AppState { inner: Arc, } +/// The app root key of one KMS root epoch. +pub(crate) struct AppRoot<'a> { + pub epoch: u32, + pub k256_key: &'a [u8], + pub signing_key: std::borrow::Cow<'a, SigningKey>, + pub kms_signature: &'a [u8], +} + struct AppStateInner { config: Config, keys: AppKeys, @@ -377,6 +385,37 @@ impl AppState { Ok(key) } + /// The app root key and its KMS signature of `root_epoch`; the current + /// ones for 0. The key is parsed for older epochs only, which are + /// requested rarely. + pub(crate) fn app_root(&self, root_epoch: u32) -> Result> { + let keys = &self.inner.keys; + let current = keys.root_epoch(); + let epoch = if root_epoch == 0 { current } else { root_epoch }; + if epoch == current { + return Ok(AppRoot { + epoch, + k256_key: self.app_root_k256_key(), + signing_key: std::borrow::Cow::Borrowed(self.app_root_signing_key()?), + kms_signature: self.kms_k256_signature(), + }); + } + let root = keys + .historical_roots + .iter() + .find(|root| root.root_epoch == epoch) + .with_context(|| format!("root key epoch {epoch} is not available"))?; + Ok(AppRoot { + epoch, + k256_key: &root.k256_key, + signing_key: std::borrow::Cow::Owned( + SigningKey::from_slice(&root.k256_key) + .context("the app root k256 key is not a valid secp256k1 scalar")?, + ), + kms_signature: &root.k256_signature, + }) + } + /// The KMS root key's signature over the app root public key: the second /// link of every signature chain, produced outside this agent and passed /// through byte-for-byte on both API surfaces. @@ -1152,6 +1191,8 @@ pub(crate) mod tests { .expect("certificate carries an attestation") } + pub(crate) const HISTORICAL_K256_KEY: [u8; 32] = [0x11; 32]; + pub(crate) async fn setup_test_state() -> (AppState, tempfile::NamedTempFile) { setup_test_state_with_platform(None).await } @@ -1355,6 +1396,13 @@ pNs85uhOZE8z2jr8Pg== k256_signature: Vec::new(), gateway_app_id: String::new(), ca_cert: DUMMY_PEM_CERT.to_string(), + // Only selected by tests that name root epoch 2 explicitly; the + // fixture's current root is epoch 1. + historical_roots: vec![dstack_types::HistoricalAppRoot { + root_epoch: 2, + k256_key: HISTORICAL_K256_KEY.to_vec(), + k256_signature: vec![0xab; 65], + }], key_provider: KeyProvider::None { key: DUMMY_PEM_KEY.to_string(), }, diff --git a/dstack/guest-agent/src/rpc_service_v1.rs b/dstack/guest-agent/src/rpc_service_v1.rs index 354e94618..95fd6e5fa 100644 --- a/dstack/guest-agent/src/rpc_service_v1.rs +++ b/dstack/guest-agent/src/rpc_service_v1.rs @@ -159,16 +159,22 @@ impl V1RpcHandler { Self { state } } - /// Derive the key named by `(domain, algorithm)` together with its - /// signature chain. - fn derive(&self, domain: &str, algorithm: &str) -> Result<(AppKey, Vec>)> { + /// Derive the key named by `(domain, algorithm)` under the app root of + /// `root_epoch`, together with its signature chain and that epoch. + fn derive( + &self, + domain: &str, + algorithm: &str, + root_epoch: u32, + ) -> Result<(AppKey, Vec>, u32)> { let algorithm = Algorithm::parse(algorithm)?; - let key = AppKey::derive(self.state.app_root_k256_key(), domain, algorithm)?; + let root = self.state.app_root(root_epoch)?; + let key = AppKey::derive(root.k256_key, domain, algorithm)?; let chain = vec![ - key.claim_signature(self.state.app_root_signing_key()?)?, - self.state.kms_k256_signature().to_vec(), + key.claim_signature(&root.signing_key)?, + root.kms_signature.to_vec(), ]; - Ok((key, chain)) + Ok((key, chain, root.epoch)) } } @@ -196,11 +202,13 @@ impl DstackGuestRpc for V1RpcHandler { } async fn get_key(self, request: GetKeyRequest) -> Result { - let (key, signature_chain) = self.derive(&request.domain, &request.algorithm)?; + let (key, signature_chain, root_epoch) = + self.derive(&request.domain, &request.algorithm, request.root_epoch)?; Ok(GetKeyResponse { key: key.secret(), public_key: key.public_key(), signature_chain, + root_epoch, }) } @@ -342,9 +350,44 @@ mod tests { GetKeyRequest { domain: domain.to_string(), algorithm: algorithm.to_string(), + root_epoch: 0, } } + #[tokio::test] + async fn get_key_selects_the_root_epoch() { + let (state, _guard) = state().await; + let handler = || V1RpcHandler::new(state.clone()); + let request = |root_epoch| GetKeyRequest { + root_epoch, + ..get_key_request("storage-encryption", "secp256k1") + }; + + let current = handler().get_key(request(0)).await.unwrap(); + assert_eq!(current.root_epoch, 1); + assert_eq!( + handler().get_key(request(1)).await.unwrap().key, + current.key + ); + + let old = handler().get_key(request(2)).await.unwrap(); + assert_eq!(old.root_epoch, 2); + assert_ne!(old.key, current.key); + assert_eq!( + old.key, + AppKey::derive( + &crate::rpc_service::tests::HISTORICAL_K256_KEY, + "storage-encryption", + Algorithm::Secp256k1 + ) + .unwrap() + .secret() + ); + assert_eq!(old.signature_chain[1], vec![0xab; 65]); + + assert!(handler().get_key(request(3)).await.is_err()); + } + /// The fixture's app root key is what `keys::tests` committed its vectors /// against. If the two drift, the vectors stop describing what this /// handler serves and every assertion built on them goes quiet. diff --git a/dstack/guest-api/proto/guest_api.proto b/dstack/guest-api/proto/guest_api.proto index 98a31616e..3ef86458d 100644 --- a/dstack/guest-api/proto/guest_api.proto +++ b/dstack/guest-api/proto/guest_api.proto @@ -193,7 +193,7 @@ message ReencryptedEnv { // The same env encrypted to the current env public key. bytes encrypted_env = 2; // KMS key epoch `encrypted_env` is encrypted for. - uint32 key_epoch = 3; + uint32 root_epoch = 3; } // Direct gRPC surface exposed by the in-guest agent. diff --git a/dstack/kms/auth-eth/contracts/DstackKms.sol b/dstack/kms/auth-eth/contracts/DstackKms.sol index ae44664b1..c3699d7c2 100644 --- a/dstack/kms/auth-eth/contracts/DstackKms.sol +++ b/dstack/kms/auth-eth/contracts/DstackKms.sol @@ -44,14 +44,14 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E // DstackApp implementation contract address for factory deployment address public appImplementation; - struct KeyEpoch { + struct RootEpoch { bytes caPubkey; bytes k256Pubkey; } - // Root public keys of each KMS key epoch, indexed by epoch. Append-only: - // a registered epoch is never changed, so the last entry is current. - KeyEpoch[] private _keyEpochs; + // Root public keys of KMS root epochs 1, 2, ... in order. Append-only: a + // registered epoch is never changed, so the last entry is current. + RootEpoch[] private _rootEpochs; // Events // appId is unindexed for backward compatibility with deployed log indexers. @@ -67,7 +67,7 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E event GatewayAppIdSet(string gatewayAppId); // slither-disable-next-line unindexed-event-address event AppImplementationSet(address implementation); - event KeyEpochRegistered(uint256 indexed epoch, bytes caPubkey, bytes k256Pubkey); + event RootEpochRegistered(uint256 indexed epoch, bytes caPubkey, bytes k256Pubkey); event AppDeployedViaFactory(address indexed appId, address indexed deployer); /// @notice Additive audit event for reconstructing authorization policy. /// @dev `value` is the affected bytes32 value, address, or hash of dynamic public data. @@ -125,8 +125,9 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E _emitPolicy("kms-info", keccak256(info.k256Pubkey), true); } - // Register the root public keys of the next KMS key epoch - function registerKeyEpoch( + // Register the root public keys of the next KMS root epoch; epochs count + // from 1, the original roots + function registerRootEpoch( bytes calldata caPubkey, bytes calldata k256Pubkey ) @@ -135,18 +136,20 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E returns (uint256 epoch) { require(caPubkey.length > 0 && k256Pubkey.length > 0, "Empty epoch public key"); - epoch = _keyEpochs.length; - _keyEpochs.push(KeyEpoch(caPubkey, k256Pubkey)); - emit KeyEpochRegistered(epoch, caPubkey, k256Pubkey); - _emitPolicy("key-epoch", keccak256(abi.encode(caPubkey, k256Pubkey)), true); + _rootEpochs.push(RootEpoch(caPubkey, k256Pubkey)); + epoch = _rootEpochs.length; + emit RootEpochRegistered(epoch, caPubkey, k256Pubkey); + _emitPolicy("root-epoch", keccak256(abi.encode(caPubkey, k256Pubkey)), true); } - function keyEpochCount() external view returns (uint256) { - return _keyEpochs.length; + // The latest registered root epoch; 0 when none is registered + function latestRootEpoch() external view returns (uint256) { + return _rootEpochs.length; } - function keyEpoch(uint256 epoch) external view returns (bytes memory caPubkey, bytes memory k256Pubkey) { - KeyEpoch storage keys = _keyEpochs[epoch]; + function rootEpoch(uint256 epoch) external view returns (bytes memory caPubkey, bytes memory k256Pubkey) { + require(epoch > 0 && epoch <= _rootEpochs.length, "Unknown root epoch"); + RootEpoch storage keys = _rootEpochs[epoch - 1]; return (keys.caPubkey, keys.k256Pubkey); } diff --git a/dstack/kms/auth-eth/src/ethereum.ts b/dstack/kms/auth-eth/src/ethereum.ts index 3c4a696ad..4f57b54b6 100644 --- a/dstack/kms/auth-eth/src/ethereum.ts +++ b/dstack/kms/auth-eth/src/ethereum.ts @@ -3,7 +3,7 @@ // SPDX-License-Identifier: Apache-2.0 import { ethers } from 'ethers'; -import { BootInfo, BootResponse, KeyEpoch } from './types'; +import { BootInfo, BootResponse, RootEpoch } from './types'; // Minimal ABI for DstackKms contract const DSTACK_KMS_ABI = [ @@ -11,8 +11,8 @@ const DSTACK_KMS_ABI = [ "function isKmsAllowed((address appId,bytes32 composeHash,address instanceId,bytes32 deviceId,bytes32 mrAggregated,bytes32 mrSystem,bytes32 osImageHash,string tcbStatus,string[] advisoryIds) bootInfo) view returns (bool, string)", "function gatewayAppId() view returns (string)", "function appImplementation() view returns (address)", - "function keyEpochCount() view returns (uint256)", - "function keyEpoch(uint256 epoch) view returns (bytes caPubkey, bytes k256Pubkey)" + "function latestRootEpoch() view returns (uint256)", + "function rootEpoch(uint256 epoch) view returns (bytes caPubkey, bytes k256Pubkey)" ]; // Contracts deployed before key epochs existed lack the registry getters. @@ -67,30 +67,30 @@ export class EthereumBackend { } const [isAllowed, reason] = response; const gatewayAppId = await this.kmsContract.gatewayAppId({ blockTag }); - const epochCount = await this.getKeyEpochCount(blockTag); + const latestEpoch = await this.getLatestRootEpoch(blockTag); return { isAllowed, reason, gatewayAppId, - ...(epochCount > 0 && { keyEpoch: epochCount - 1 }), + ...(latestEpoch > 0 && { rootEpoch: latestEpoch }), } } - private async getKeyEpochCount(blockTag?: number): Promise { + private async getLatestRootEpoch(blockTag?: number): Promise { try { - return Number(await this.kmsContract.keyEpochCount({ blockTag })); + return Number(await this.kmsContract.latestRootEpoch({ blockTag })); } catch (error) { if (isMissingFunction(error)) return 0; throw error; } } - async getKeyEpochs(): Promise { + async getRootEpochs(): Promise { const blockTag = await this.provider.getBlockNumber(); - const count = await this.getKeyEpochCount(blockTag); + const latest = await this.getLatestRootEpoch(blockTag); return Promise.all( - Array.from({ length: count }, async (_, epoch) => { - const [caPubkey, k256Pubkey] = await this.kmsContract.keyEpoch(epoch, { blockTag }); + Array.from({ length: latest }, async (_, index) => { + const [caPubkey, k256Pubkey] = await this.kmsContract.rootEpoch(index + 1, { blockTag }); return { caPubkey, k256Pubkey }; }) ); diff --git a/dstack/kms/auth-eth/src/main.test.ts b/dstack/kms/auth-eth/src/main.test.ts index 43605eb7b..1a6beff9d 100644 --- a/dstack/kms/auth-eth/src/main.test.ts +++ b/dstack/kms/auth-eth/src/main.test.ts @@ -94,7 +94,7 @@ describe('Server', () => { isAllowed: true, reason: '', gatewayAppId: '', - keyEpoch: 2, + rootEpoch: 2, }); const response = await app.inject({ @@ -103,22 +103,22 @@ describe('Server', () => { payload: mockBootInfo }); - expect(JSON.parse(response.payload).keyEpoch).toBe(2); + expect(JSON.parse(response.payload).rootEpoch).toBe(2); }); it('should list registered key epochs in the info response', async () => { - const keyEpochs = [{ caPubkey: '0xa0', k256Pubkey: '0xb0' }]; + const rootEpochs = [{ caPubkey: '0xa0', k256Pubkey: '0xb0' }]; Object.assign(app.ethereum, { getGatewayAppId: jest.fn().mockResolvedValue('gw'), getChainId: jest.fn().mockResolvedValue(1), getAppImplementation: jest.fn().mockResolvedValue('0x00'), - getKeyEpochs: jest.fn().mockResolvedValue(keyEpochs), + getRootEpochs: jest.fn().mockResolvedValue(rootEpochs), }); const response = await app.inject({ method: 'GET', url: '/' }); expect(response.statusCode).toBe(200); - expect(JSON.parse(response.payload).keyEpochs).toEqual(keyEpochs); + expect(JSON.parse(response.payload).rootEpochs).toEqual(rootEpochs); }); }); diff --git a/dstack/kms/auth-eth/src/server.ts b/dstack/kms/auth-eth/src/server.ts index 08a7e913c..747c4f3d5 100644 --- a/dstack/kms/auth-eth/src/server.ts +++ b/dstack/kms/auth-eth/src/server.ts @@ -50,7 +50,7 @@ export async function build(): Promise { isAllowed: { type: 'boolean' }, reason: { type: 'string' }, gatewayAppId: { type: 'string' }, - keyEpoch: { type: 'integer' }, + rootEpoch: { type: 'integer' }, } }); @@ -85,7 +85,7 @@ export async function build(): Promise { server.ethereum.getGatewayAppId(), server.ethereum.getChainId(), server.ethereum.getAppImplementation(), - server.ethereum.getKeyEpochs(), + server.ethereum.getRootEpochs(), ]); return { status: 'ok', @@ -95,7 +95,7 @@ export async function build(): Promise { chainId: batch[1], appAuthImplementation: batch[2], // NOTE: for backward compatibility appImplementation: batch[2], - keyEpochs: batch[3], + rootEpochs: batch[3], }; } catch { _request.log.error('authorization backend health check failed'); diff --git a/dstack/kms/auth-eth/src/types.ts b/dstack/kms/auth-eth/src/types.ts index 5442a7e69..f74defa2e 100644 --- a/dstack/kms/auth-eth/src/types.ts +++ b/dstack/kms/auth-eth/src/types.ts @@ -18,11 +18,11 @@ export interface BootResponse { isAllowed: boolean; gatewayAppId: string; reason: string; - // Latest registered KMS key epoch; absent when none is registered. - keyEpoch?: number; + // Latest registered KMS root epoch (from 1); absent when none is registered. + rootEpoch?: number; } -export interface KeyEpoch { +export interface RootEpoch { caPubkey: string; k256Pubkey: string; } diff --git a/dstack/kms/auth-eth/test/DstackKms.t.sol b/dstack/kms/auth-eth/test/DstackKms.t.sol index ab99f12ff..cfafce9ce 100644 --- a/dstack/kms/auth-eth/test/DstackKms.t.sol +++ b/dstack/kms/auth-eth/test/DstackKms.t.sol @@ -19,7 +19,7 @@ contract DstackKmsTest is Test { address public user; event KmsInfoSet(bytes k256Pubkey); - event KeyEpochRegistered(uint256 indexed epoch, bytes caPubkey, bytes k256Pubkey); + event RootEpochRegistered(uint256 indexed epoch, bytes caPubkey, bytes k256Pubkey); event AppRegistered(address appId); @@ -67,37 +67,41 @@ contract DstackKmsTest is Test { assertEq(storedEventlog, eventlog); } - function test_RegisterKeyEpochAppends() public { - assertEq(kms.keyEpochCount(), 0); + function test_RegisterRootEpochAppends() public { + assertEq(kms.latestRootEpoch(), 0); vm.startPrank(owner); vm.expectEmit(true, true, true, true); - emit KeyEpochRegistered(0, hex"a0", hex"b0"); - assertEq(kms.registerKeyEpoch(hex"a0", hex"b0"), 0); - assertEq(kms.registerKeyEpoch(hex"a1", hex"b1"), 1); + emit RootEpochRegistered(1, hex"a0", hex"b0"); + assertEq(kms.registerRootEpoch(hex"a0", hex"b0"), 1); + assertEq(kms.registerRootEpoch(hex"a1", hex"b1"), 2); vm.stopPrank(); - assertEq(kms.keyEpochCount(), 2); - (bytes memory ca, bytes memory k256) = kms.keyEpoch(0); + assertEq(kms.latestRootEpoch(), 2); + (bytes memory ca, bytes memory k256) = kms.rootEpoch(1); assertEq(ca, hex"a0"); assertEq(k256, hex"b0"); - (ca, k256) = kms.keyEpoch(1); + (ca, k256) = kms.rootEpoch(2); assertEq(ca, hex"a1"); assertEq(k256, hex"b1"); + vm.expectRevert("Unknown root epoch"); + kms.rootEpoch(0); + vm.expectRevert("Unknown root epoch"); + kms.rootEpoch(3); } - function test_RegisterKeyEpochOnlyOwner() public { + function test_RegisterRootEpochOnlyOwner() public { vm.prank(user); vm.expectRevert(abi.encodeWithSelector(OwnableUpgradeable.OwnableUnauthorizedAccount.selector, user)); - kms.registerKeyEpoch(hex"a0", hex"b0"); + kms.registerRootEpoch(hex"a0", hex"b0"); } - function test_RegisterKeyEpochRejectsEmptyKeys() public { + function test_RegisterRootEpochRejectsEmptyKeys() public { vm.startPrank(owner); vm.expectRevert("Empty epoch public key"); - kms.registerKeyEpoch("", hex"b0"); + kms.registerRootEpoch("", hex"b0"); vm.expectRevert("Empty epoch public key"); - kms.registerKeyEpoch(hex"a0", ""); + kms.registerRootEpoch(hex"a0", ""); vm.stopPrank(); } diff --git a/dstack/kms/auth-simple/README.md b/dstack/kms/auth-simple/README.md index fc40f7314..a7937e8d8 100644 --- a/dstack/kms/auth-simple/README.md +++ b/dstack/kms/auth-simple/README.md @@ -107,7 +107,7 @@ raw NitroTPM evidence by itself. | `apps..composeHashes` | No | Allowed compose hashes for this app | | `apps..devices` | No | Allowed device IDs for this app. Empty denies all unless `allowAnyDevice` is true. | | `apps..allowAnyDevice` | No | If true, skip device ID check for this app | -| `keyEpochs` | No | KMS root public keys per key epoch, `[{"caPubkey": "0x...", "k256Pubkey": "0x..."}]`, indexed by epoch. Append only. When set, the KMS checks its keyring against it and boot responses report the latest epoch as `keyEpoch`. | +| `rootEpochs` | No | KMS root public keys of root epochs 1, 2, ... in order, `[{"caPubkey": "0x...", "k256Pubkey": "0x..."}]`. Epoch 1 is the original roots. Append only. When set, the KMS checks its keyring against it and boot responses report the latest epoch as `rootEpoch`. | For experimental AMD SEV-SNP dry-run authorization, keep the default fail-closed TCB policy unless you intentionally want the auth webhook to accept non-up-to-date verifier-derived SNP `BootInfo`. To exercise the dry-run path without enabling key release, allowlist the recomputed SNP `mrAggregated`, `osImageHash`, app/compose identity, device/chip identity, and any non-default `allowedTcbStatuses`/`allowedAdvisoryIds` values explicitly. KMS still rejects SNP before returning app keys, KMS keys, or app certificates. diff --git a/dstack/kms/auth-simple/index.test.ts b/dstack/kms/auth-simple/index.test.ts index a75120fca..043ca22d4 100644 --- a/dstack/kms/auth-simple/index.test.ts +++ b/dstack/kms/auth-simple/index.test.ts @@ -447,23 +447,23 @@ describe('auth-simple', () => { }); describe('key epochs', () => { - const keyEpochs = [ + const rootEpochs = [ { caPubkey: '0xa0', k256Pubkey: '0xb0' }, { caPubkey: '0xa1', k256Pubkey: '0xb1' } ]; it('reports the latest key epoch in boot responses and lists all epochs', async () => { - writeTestConfig({ gatewayAppId: '0xgateway', keyEpochs }); + writeTestConfig({ gatewayAppId: '0xgateway', rootEpochs }); const boot = await app.fetch(new Request('http://localhost/bootAuth/app', { method: 'POST', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify(baseBootInfo) })); - expect((await boot.json()).keyEpoch).toBe(1); + expect((await boot.json()).rootEpoch).toBe(2); const info = await app.fetch(new Request('http://localhost/', { method: 'GET' })); - expect((await info.json()).keyEpochs).toEqual(keyEpochs); + expect((await info.json()).rootEpochs).toEqual(rootEpochs); }); it('omits the key epoch when none is configured', async () => { @@ -474,7 +474,7 @@ describe('auth-simple', () => { headers: { 'Content-Type': 'application/json' }, body: JSON.stringify(baseBootInfo) })); - expect((await boot.json()).keyEpoch).toBeUndefined(); + expect((await boot.json()).rootEpoch).toBeUndefined(); }); }); diff --git a/dstack/kms/auth-simple/index.ts b/dstack/kms/auth-simple/index.ts index d2be75515..a7197eea1 100644 --- a/dstack/kms/auth-simple/index.ts +++ b/dstack/kms/auth-simple/index.ts @@ -25,10 +25,10 @@ const BootResponseSchema = z.object({ isAllowed: z.boolean(), reason: z.string(), gatewayAppId: z.string(), - keyEpoch: z.number().optional() + rootEpoch: z.number().optional() }); -const KeyEpochSchema = z.object({ +const RootEpochSchema = z.object({ caPubkey: z.string(), k256Pubkey: z.string() }); @@ -59,8 +59,8 @@ const AuthConfigSchema = z.object({ allowedAdvisoryIds: z.array(z.string()).default([]), kms: KmsConfigSchema.default({}), apps: z.record(z.string(), AppConfigSchema).default({}), - // Root public keys of each KMS key epoch, indexed by epoch; append only. - keyEpochs: z.array(KeyEpochSchema).default([]) + // Root public keys of KMS root epochs 1, 2, ... in order; append only. + rootEpochs: z.array(RootEpochSchema).default([]) }); type BootInfo = z.infer; @@ -102,8 +102,8 @@ class ConfigBackend { async checkBoot(bootInfo: BootInfo, isKms: boolean): Promise { const config = this.loadConfig(); const response = this.evaluateBoot(bootInfo, isKms, config); - const epochCount = config.keyEpochs.length; - return epochCount > 0 ? { ...response, keyEpoch: epochCount - 1 } : response; + const latestEpoch = config.rootEpochs.length; + return latestEpoch > 0 ? { ...response, rootEpoch: latestEpoch } : response; } private evaluateBoot(bootInfo: BootInfo, isKms: boolean, config: AuthConfig): BootResponse { @@ -234,7 +234,7 @@ class ConfigBackend { kmsContractAddr: string; chainId: number; appImplementation: string; - keyEpochs: AuthConfig['keyEpochs']; + rootEpochs: AuthConfig['rootEpochs']; }> { const config = this.loadConfig(); return { @@ -242,7 +242,7 @@ class ConfigBackend { kmsContractAddr: config.kmsContractAddr, chainId: config.chainId, appImplementation: config.appImplementation, - keyEpochs: config.keyEpochs, + rootEpochs: config.rootEpochs, }; } } @@ -264,7 +264,7 @@ app.get('/', async (c) => { chainId: info.chainId, appAuthImplementation: info.appImplementation, // backward compat appImplementation: info.appImplementation, - keyEpochs: info.keyEpochs, + rootEpochs: info.rootEpochs, }); } catch (error) { console.error('error in health check:', error); diff --git a/dstack/kms/kms.toml b/dstack/kms/kms.toml index 233fda094..48efa121f 100644 --- a/dstack/kms/kms.toml +++ b/dstack/kms/kms.toml @@ -125,10 +125,10 @@ timeout = "60s" gateway_app_id = "any" [core.key_rotation] -# Root key epoch served as current. Raising it requires a KMS onboarded with +# Root key epoch served as current; the original roots are epoch 1. Raising it requires a KMS onboarded with # `rotate` (or from a source already at this epoch); editing it alone fails. # Older epochs stay available so CVMs can migrate their data at any time. -current_epoch = 0 +current_epoch = 1 [core.onboard] enabled = true diff --git a/dstack/kms/rpc/proto/kms_rpc.proto b/dstack/kms/rpc/proto/kms_rpc.proto index 0405d2606..ca2031735 100644 --- a/dstack/kms/rpc/proto/kms_rpc.proto +++ b/dstack/kms/rpc/proto/kms_rpc.proto @@ -11,9 +11,10 @@ package kms; message GetAppKeyRequest { uint32 api_version = 1; string vm_config = 2; - // Root key epoch to derive from: 0 (default) = epoch 0, N > 0 = epoch N, - // -1 = the current epoch. Older epochs remain available for migration. - sint32 key_epoch = 3; + // Root key epoch to derive from: 0 (default) = the current epoch, N = epoch N. + // Epoch 1 is the original root keys. Older epochs remain available so CVMs + // can migrate. + uint32 root_epoch = 3; } message AppId { @@ -49,8 +50,9 @@ message AppKeyResponse { string gateway_app_id = 7; // OS Image hash bytes os_image_hash = 8; - // Root key epoch all epoch-dependent fields above were derived from. - uint32 key_epoch = 9; + // Root key epoch all epoch-dependent fields above were derived from. 0 from a + // KMS that predates root epochs, which means epoch 1. + uint32 root_epoch = 9; } message GetMetaResponse { @@ -75,7 +77,8 @@ message GetKmsKeyRequest { message KmsKeys { string ca_key = 1; bytes k256_key = 2; - // Root key epoch of this pair. + // Root key epoch of this pair. 0 from a KMS that predates root epochs, which + // means epoch 1. uint32 epoch = 3; } @@ -165,7 +168,7 @@ message OnboardRequest { string domain = 2; // Optional bearer token for the source KMS admin listener. string source_token = 3; - // Import epochs 0..N-1 from the source and generate epoch N locally, where N + // Import epochs 1..N-1 from the source and generate epoch N locally, where N // is this KMS's configured current epoch. bool rotate = 4; } diff --git a/dstack/kms/src/config.rs b/dstack/kms/src/config.rs index f3869ca23..1476c98ff 100644 --- a/dstack/kms/src/config.rs +++ b/dstack/kms/src/config.rs @@ -2,7 +2,7 @@ // // SPDX-License-Identifier: Apache-2.0 -use anyhow::{anyhow, bail, Result}; +use anyhow::{bail, Result}; use dstack_attest::attestation::AttestationVerifierConfig; use load_config::load_config; use rocket::figment::Figment; @@ -176,7 +176,7 @@ impl KmsConfig { self.cert_dir.join(K256_KEY) } - /// Records which epoch the top-level root keys belong to. Absent means 0. + /// Records which epoch the top-level root keys belong to. Absent means 1. pub fn root_epoch(&self) -> PathBuf { self.cert_dir.join(ROOT_EPOCH) } @@ -241,25 +241,35 @@ pub(crate) struct RootKeyPaths { pub k256_key: PathBuf, } -#[derive(Debug, Clone, Default, Deserialize)] +#[derive(Debug, Clone, Deserialize)] pub(crate) struct KeyRotationConfig { - /// Root key epoch served as current. Every KMS holds all epochs up to it. - #[serde(default)] + /// Root key epoch served as current; epochs count from 1, the original + /// roots. Every KMS holds all epochs up to it. + #[serde(default = "first_epoch")] pub current_epoch: u32, } +impl Default for KeyRotationConfig { + fn default() -> Self { + Self { + current_epoch: first_epoch(), + } + } +} + +fn first_epoch() -> u32 { + 1 +} + impl KeyRotationConfig { - /// Resolve a GetAppKey epoch selector: `-1` is the current epoch. - pub fn resolve_epoch(&self, requested: i32) -> Result { + /// Resolve a root epoch selector: 0 is the current epoch. + pub fn resolve_epoch(&self, requested: u32) -> Result { let current = self.current_epoch; - let epoch = match requested { - -1 => current, - epoch => u32::try_from(epoch).map_err(|_| anyhow!("invalid key epoch {requested}"))?, - }; - if epoch > current { - bail!("unknown key epoch {epoch}; the current epoch is {current}"); + match requested { + 0 => Ok(current), + epoch if epoch <= current => Ok(epoch), + epoch => bail!("unknown root epoch {epoch}; the current epoch is {current}"), } - Ok(epoch) } } @@ -293,14 +303,13 @@ mod tests { } #[test] - fn key_epoch_selector_resolution() { + fn root_epoch_selector_resolution() { let cfg = KeyRotationConfig { current_epoch: 2 }; - assert_eq!(cfg.resolve_epoch(-1).unwrap(), 2); - assert_eq!(cfg.resolve_epoch(0).unwrap(), 0); + assert_eq!(cfg.resolve_epoch(0).unwrap(), 2); + assert_eq!(cfg.resolve_epoch(1).unwrap(), 1); assert_eq!(cfg.resolve_epoch(2).unwrap(), 2); assert!(cfg.resolve_epoch(3).is_err()); - assert!(cfg.resolve_epoch(-2).is_err()); - assert_eq!(KeyRotationConfig::default().resolve_epoch(-1).unwrap(), 0); + assert_eq!(KeyRotationConfig::default().resolve_epoch(0).unwrap(), 1); } #[test] diff --git a/dstack/kms/src/main_service.rs b/dstack/kms/src/main_service.rs index f7187e288..bd21b3236 100644 --- a/dstack/kms/src/main_service.rs +++ b/dstack/kms/src/main_service.rs @@ -31,7 +31,7 @@ use scale::Decode; use tracing::warn; use upgrade_authority::{ build_boot_info, ensure_app_id_len, local_kms_boot_info, BootInfo, BootResponse, - GetInfoResponse, KeyEpochPubkeys, + GetInfoResponse, RootEpochPubkeys, }; use crate::{ @@ -83,8 +83,8 @@ impl RootKeys { Ok(Self { ca, k256 }) } - fn pubkeys(&self) -> KeyEpochPubkeys { - KeyEpochPubkeys { + fn pubkeys(&self) -> RootEpochPubkeys { + RootEpochPubkeys { ca_pubkey: self.ca.key.public_key_der(), k256_pubkey: self.k256.verifying_key().to_sec1_bytes().to_vec(), } @@ -93,7 +93,7 @@ impl RootKeys { pub struct KmsStateInner { config: KmsConfig, - /// Root keys of every epoch, indexed by epoch; the last is current. + /// Root keys of epochs 1, 2, ... in order; the last is current. roots: Vec, temp_ca_cert: String, temp_ca_key: String, @@ -153,6 +153,9 @@ impl KmsState { pub fn new(config: KmsConfig) -> Result { let current_epoch = config.key_rotation.current_epoch; + if current_epoch == 0 { + bail!("key_rotation.current_epoch counts from 1"); + } let stored_epoch = read_root_epoch(&config)?; if stored_epoch != current_epoch { bail!( @@ -160,7 +163,7 @@ impl KmsState { is {current_epoch}; a new epoch must be imported or generated by onboarding" ); } - let roots = (0..=current_epoch) + let roots = (1..=current_epoch) .map(|epoch| { RootKeys::load(&config.root_key_paths(epoch)) .with_context(|| format!("failed to load root keys of epoch {epoch}")) @@ -214,34 +217,39 @@ impl KmsState { /// The current epoch's root keys. fn root(&self) -> &RootKeys { - self.roots.last().expect("the keyring always holds epoch 0") + self.roots.last().expect("the keyring always holds epoch 1") + } + + /// The root keys of `epoch`, which must be one this KMS serves. + fn epoch_root(&self, epoch: u32) -> &RootKeys { + &self.roots[epoch as usize - 1] } /// Refuse to serve a keyring governance has not registered. Deployments - /// that never registered epochs keep working while they stay at epoch 0. + /// that never registered epochs keep working while they stay at epoch 1. pub async fn ensure_keyring_registered(&self) -> Result<()> { let current = self.config.key_rotation.current_epoch; let info = match self.config.auth_api.get_info().await { Ok(info) => info, - Err(err) if current == 0 => { + Err(err) if current == 1 => { warn!("skipping key epoch registration check: {err:#}"); return Ok(()); } Err(err) => return Err(err).context("failed to read registered key epochs"), }; - let Some(registered) = info.key_epochs else { + let Some(registered) = info.root_epochs else { return Ok(()); }; - if registered.is_empty() && current == 0 { + if registered.is_empty() && current == 1 { return Ok(()); } - if registered.len() != current as usize + 1 { + if registered.len() != current as usize { bail!( "this KMS serves key epoch {current}, but {} epochs are registered", registered.len() ); } - for (epoch, (root, registered)) in self.roots.iter().zip(®istered).enumerate() { + for (epoch, (root, registered)) in (1..).zip(self.roots.iter().zip(®istered)) { if root.pubkeys() != *registered { bail!("root keys of epoch {epoch} do not match the registered public keys"); } @@ -252,11 +260,11 @@ impl KmsState { /// Stop serving once governance registers a newer epoch than this KMS has. pub(crate) fn ensure_epoch_current(&self, response: &BootResponse) -> Result<()> { let current = self.config.key_rotation.current_epoch; - match response.key_epoch { + match response.root_epoch { Some(epoch) if epoch != current => { bail!("the registered key epoch is {epoch}, but this KMS serves epoch {current}") } - None if current > 0 && !self.config.auth_api.is_dev() => { + None if current > 1 && !self.config.auth_api.is_dev() => { bail!("the auth API reports no key epoch, but this KMS serves epoch {current}") } _ => Ok(()), @@ -267,7 +275,7 @@ impl KmsState { pub(crate) fn read_root_epoch(config: &KmsConfig) -> Result { match fs::read_to_string(config.root_epoch()) { Ok(epoch) => epoch.trim().parse().context("invalid root epoch file"), - Err(err) if err.kind() == std::io::ErrorKind::NotFound => Ok(0), + Err(err) if err.kind() == std::io::ErrorKind::NotFound => Ok(1), Err(err) => Err(err).context("failed to read root epoch file"), } } @@ -293,7 +301,7 @@ impl RpcHandler { ensure_key_release_allowed(&info, (&self.state.config).into())?; Ok(KmsKeyResponse { temp_ca_key: self.state.inner.temp_ca_key.clone(), - keys: (0..) + keys: (1..) .zip(&self.state.roots) .map(|(epoch, root)| KmsKeys { ca_key: root.ca.key.serialize_pem(), @@ -501,12 +509,12 @@ impl KmsRpc for RpcHandler { .await .context("App not allowed")?; ensure_key_release_allowed(&boot_info, (&self.state.config).into())?; - let key_epoch = self + let root_epoch = self .state .config .key_rotation - .resolve_epoch(request.key_epoch)?; - let root = &self.state.roots[key_epoch as usize]; + .resolve_epoch(request.root_epoch)?; + let root = self.state.epoch_root(root_epoch); let app_id = boot_info.app_id; let instance_id = boot_info.instance_id; let os_image_hash = boot_info.os_image_hash; @@ -536,7 +544,7 @@ impl KmsRpc for RpcHandler { tproxy_app_id: gateway_app_id.clone(), gateway_app_id, os_image_hash, - key_epoch, + root_epoch, }) } diff --git a/dstack/kms/src/main_service/upgrade_authority.rs b/dstack/kms/src/main_service/upgrade_authority.rs index e9d7e97e0..7d37dcd0e 100644 --- a/dstack/kms/src/main_service/upgrade_authority.rs +++ b/dstack/kms/src/main_service/upgrade_authority.rs @@ -86,13 +86,13 @@ pub(crate) struct BootResponse { pub reason: String, /// Latest KMS key epoch registered by governance, if any. #[serde(default)] - pub key_epoch: Option, + pub root_epoch: Option, } /// Root public keys governance registered for one KMS key epoch. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] -pub(crate) struct KeyEpochPubkeys { +pub(crate) struct RootEpochPubkeys { #[serde(deserialize_with = "deserialize_hex")] pub ca_pubkey: Vec, #[serde(deserialize_with = "deserialize_hex")] @@ -115,7 +115,7 @@ pub(crate) struct AuthApiInfoResponse { pub chain_id: u64, pub app_implementation: String, #[serde(default)] - pub key_epochs: Vec, + pub root_epochs: Vec, } #[derive(Debug, Clone, Serialize, Deserialize)] @@ -128,7 +128,7 @@ pub(crate) struct GetInfoResponse { pub chain_id: Option, pub app_implementation: Option, /// Registered key epochs, indexed by epoch. `None` means no registry. - pub key_epochs: Option>, + pub root_epochs: Option>, } const AUTH_API_CONNECT_TIMEOUT: Duration = Duration::from_secs(5); @@ -196,7 +196,7 @@ impl AuthApi { is_allowed: true, reason: "".to_string(), gateway_app_id: dev.gateway_app_id.clone(), - key_epoch: None, + root_epoch: None, }), AuthApi::Webhook { webhook } => { let path = if is_kms { @@ -219,7 +219,7 @@ impl AuthApi { gateway_app_id: Some(dev.gateway_app_id.clone()), chain_id: None, app_implementation: None, - key_epochs: None, + root_epochs: None, }), AuthApi::Webhook { webhook } => { let info: AuthApiInfoResponse = http_get(&webhook.url, webhook.timeout).await?; @@ -235,7 +235,7 @@ impl AuthApi { chain_id: Some(info.chain_id), gateway_app_id: Some(info.gateway_app_id.clone()), app_implementation: Some(info.app_implementation.clone()), - key_epochs: Some(info.key_epochs), + root_epochs: Some(info.root_epochs), }) } } @@ -538,30 +538,30 @@ pub(crate) mod tests { } #[rocket::async_test] - async fn key_epochs_are_parsed_from_info_and_boot_responses() { + async fn root_epochs_are_parsed_from_info_and_boot_responses() { let (url, server) = serve(vec![ r#"{"status":"ok","kmsContractAddr":"0x00","gatewayAppId":"gw","chainId":1, - "appImplementation":"0x00","keyEpochs":[{"caPubkey":"0xa0","k256Pubkey":"b0"}]}"#, - r#"{"isAllowed":true,"gatewayAppId":"gw","reason":"","keyEpoch":3}"#, + "appImplementation":"0x00","rootEpochs":[{"caPubkey":"0xa0","k256Pubkey":"b0"}]}"#, + r#"{"isAllowed":true,"gatewayAppId":"gw","reason":"","rootEpoch":3}"#, r#"{"isAllowed":true,"gatewayAppId":"gw","reason":""}"#, ]); let auth = webhook(url); let info = auth.get_info().await.unwrap(); assert_eq!( - info.key_epochs, - Some(vec![KeyEpochPubkeys { + info.root_epochs, + Some(vec![RootEpochPubkeys { ca_pubkey: vec![0xa0], k256_pubkey: vec![0xb0], }]) ); let info = boot_info(50); assert_eq!( - auth.is_app_allowed(&info, false).await.unwrap().key_epoch, + auth.is_app_allowed(&info, false).await.unwrap().root_epoch, Some(3) ); assert_eq!( - auth.is_app_allowed(&info, false).await.unwrap().key_epoch, + auth.is_app_allowed(&info, false).await.unwrap().root_epoch, None ); server.join().unwrap(); diff --git a/dstack/kms/src/onboard_service.rs b/dstack/kms/src/onboard_service.rs index 764023562..3e1a96eac 100644 --- a/dstack/kms/src/onboard_service.rs +++ b/dstack/kms/src/onboard_service.rs @@ -295,7 +295,7 @@ mod tests { use crate::main_service::amd_attest::{ compute_expected_measurement, MeasurementInput, OvmfSectionParam, }; - use crate::main_service::upgrade_authority::{BootResponse, KeyEpochPubkeys}; + use crate::main_service::upgrade_authority::{BootResponse, RootEpochPubkeys}; use dstack_kms_rpc::KmsKeys; use sha2::Digest; @@ -551,7 +551,7 @@ mod tests { temp_ca_key: KeyPair::generate_for(&PKCS_ECDSA_P256_SHA256) .unwrap() .serialize_pem(), - keys: (0..epochs) + keys: (1..=epochs) .map(|epoch| KmsKeys { ca_key: KeyPair::generate_for(&PKCS_ECDSA_P256_SHA256) .unwrap() @@ -574,7 +574,7 @@ mod tests { #[rocket::async_test] async fn join_imports_and_stores_every_epoch() { let cert_dir = tempfile::tempdir().unwrap(); - let cfg = epoch_config(cert_dir.path(), 2); + let cfg = epoch_config(cert_dir.path(), 3); let res = handover(3); let sent: Vec<_> = res.keys.iter().map(|k| k.k256_key.clone()).collect(); @@ -583,24 +583,36 @@ mod tests { .unwrap(); keys.store(&cfg).unwrap(); - assert_eq!(crate::main_service::read_root_epoch(&cfg).unwrap(), 2); - for (epoch, k256) in (0..).zip(&sent) { + assert_eq!(crate::main_service::read_root_epoch(&cfg).unwrap(), 3); + for (epoch, k256) in (1..).zip(&sent) { assert_eq!(&fs::read(cfg.root_key_paths(epoch).k256_key).unwrap(), k256); } assert_eq!(&fs::read(cfg.k256_key()).unwrap(), &sent[2]); crate::main_service::KmsState::new(cfg.clone()).unwrap(); // Editing the configured epoch cannot relabel the stored roots. - let err = crate::main_service::KmsState::new(epoch_config(cert_dir.path(), 3)) + let err = crate::main_service::KmsState::new(epoch_config(cert_dir.path(), 4)) .err() .unwrap(); - assert!(format!("{err:#}").contains("belong to epoch 2"), "{err:#}"); + assert!(format!("{err:#}").contains("belong to epoch 3"), "{err:#}"); + } + + #[rocket::async_test] + async fn join_accepts_a_source_predating_root_epochs() { + let cert_dir = tempfile::tempdir().unwrap(); + let cfg = epoch_config(cert_dir.path(), 1); + let mut res = handover(1); + res.keys[0].epoch = 0; + let keys = Keys::from_key_response(&cfg, res, "kms.example.com", false) + .await + .unwrap(); + assert_eq!(keys.epoch, 1); } - fn pubkeys_of(res: &KmsKeyResponse) -> Vec { + fn pubkeys_of(res: &KmsKeyResponse) -> Vec { res.keys .iter() - .map(|k| KeyEpochPubkeys { + .map(|k| RootEpochPubkeys { ca_pubkey: KeyPair::from_pem(&k.ca_key).unwrap().public_key_der(), k256_pubkey: SigningKey::from_slice(&k.k256_key) .unwrap() @@ -613,7 +625,7 @@ mod tests { async fn joined_state(res: KmsKeyResponse, url: String) -> crate::main_service::KmsState { let cert_dir = tempfile::tempdir().unwrap(); - let mut cfg = epoch_config(cert_dir.path(), res.keys.len() as u32 - 1); + let mut cfg = epoch_config(cert_dir.path(), res.keys.len() as u32); Keys::from_key_response(&cfg, res, "kms.example.com", false) .await .unwrap() @@ -623,7 +635,7 @@ mod tests { crate::main_service::KmsState::new(cfg).unwrap() } - fn info_with_epochs(epochs: &[KeyEpochPubkeys]) -> String { + fn info_with_epochs(epochs: &[RootEpochPubkeys]) -> String { let epochs: Vec<_> = epochs .iter() .map(|e| { @@ -635,7 +647,7 @@ mod tests { .collect(); serde_json::json!({ "status": "ok", "kmsContractAddr": "0x00", "gatewayAppId": "", "chainId": 1, - "appImplementation": "0x00", "keyEpochs": epochs, + "appImplementation": "0x00", "rootEpochs": epochs, }) .to_string() } @@ -644,7 +656,7 @@ mod tests { async fn keyring_must_match_the_registered_epochs() { use crate::main_service::upgrade_authority::tests::serve; - type Registered = fn(Vec) -> Vec; + type Registered = fn(Vec) -> Vec; let cases: [(Registered, bool); 4] = [ (|keys| keys, true), (|keys| keys[..1].to_vec(), false), @@ -668,21 +680,21 @@ mod tests { #[rocket::async_test] async fn a_newer_registered_epoch_stops_a_stale_kms() { let state = joined_state(handover(2), "http://127.0.0.1:1".into()).await; - let response = |key_epoch| BootResponse { + let response = |root_epoch| BootResponse { is_allowed: true, gateway_app_id: String::new(), reason: String::new(), - key_epoch, + root_epoch, }; - assert!(state.ensure_epoch_current(&response(Some(1))).is_ok()); - assert!(state.ensure_epoch_current(&response(Some(2))).is_err()); + assert!(state.ensure_epoch_current(&response(Some(2))).is_ok()); + assert!(state.ensure_epoch_current(&response(Some(3))).is_err()); assert!(state.ensure_epoch_current(&response(None)).is_err()); } #[rocket::async_test] async fn rotation_imports_history_and_generates_the_new_epoch() { let cert_dir = tempfile::tempdir().unwrap(); - let cfg = epoch_config(cert_dir.path(), 2); + let cfg = epoch_config(cert_dir.path(), 3); let res = handover(2); let sent: Vec<_> = res.keys.iter().map(|k| k.k256_key.clone()).collect(); @@ -691,7 +703,7 @@ mod tests { .unwrap(); keys.store(&cfg).unwrap(); - assert_eq!(fs::read(cfg.root_key_paths(1).k256_key).unwrap(), sent[1]); + assert_eq!(fs::read(cfg.root_key_paths(2).k256_key).unwrap(), sent[1]); let current = fs::read(cfg.k256_key()).unwrap(); assert!(!sent.contains(¤t)); } @@ -699,7 +711,7 @@ mod tests { #[rocket::async_test] async fn onboarding_rejects_incomplete_or_extra_history() { let cert_dir = tempfile::tempdir().unwrap(); - let cfg = epoch_config(cert_dir.path(), 1); + let cfg = epoch_config(cert_dir.path(), 2); for (epochs, rotate) in [(1, false), (3, false), (2, true), (0, true)] { let res = handover(epochs); assert!( @@ -709,7 +721,7 @@ mod tests { "epochs={epochs} rotate={rotate}" ); } - let cfg = epoch_config(cert_dir.path(), 0); + let cfg = epoch_config(cert_dir.path(), 1); assert!( Keys::from_key_response(&cfg, handover(0), "kms.example.com", true) .await @@ -720,7 +732,7 @@ mod tests { #[rocket::async_test] async fn bootstrap_refuses_a_later_epoch() { let cert_dir = tempfile::tempdir().unwrap(); - let cfg = epoch_config(cert_dir.path(), 1); + let cfg = epoch_config(cert_dir.path(), 2); let verifier = AttestationVerifier::load(&cfg.attestation).unwrap(); assert!(bootstrap_keys(&cfg, &verifier).await.is_err()); assert!(!cfg.root_keys_exist()); @@ -816,7 +828,7 @@ impl HistoricalRoot { } struct Keys { - /// Epoch of the root keys below; `history` holds epochs `0..epoch`. + /// Epoch of the root keys below; `history` holds epochs `1..epoch`. epoch: u32, history: Vec, k256_key: SigningKey, @@ -846,7 +858,7 @@ impl Keys { .await } - /// Attest the current root public keys. Epoch 0 keeps the original + /// Attest the current root public keys. Epoch 1 keeps the original /// bootstrap encoding; later epochs also bind the epoch number. async fn attest(&self) -> Result { let ca_pubkey = self.ca_key.public_key_der(); @@ -906,7 +918,7 @@ impl Keys { .build() .signed_by(&ca_cert, &ca_key)?; Ok(Keys { - epoch: 0, + epoch: 1, history: Vec::new(), k256_key, tmp_ca_key, @@ -1020,19 +1032,18 @@ impl Keys { ) -> Result { let epoch = cfg.key_rotation.current_epoch; // A rotation imports every epoch before the new one; a join imports all. - let expected = if rotate { - if epoch == 0 { - bail!("rotation requires key_rotation.current_epoch > 0"); + let last = if rotate { + if epoch < 2 { + bail!("rotation requires key_rotation.current_epoch > 1"); } - epoch + epoch - 1 } else { - epoch + 1 + epoch }; - if !keys_res.keys.iter().map(|k| k.epoch).eq(0..expected) { - bail!( - "source KMS must hand over root key epochs 0..{expected}, got {:?}", - keys_res.keys.iter().map(|k| k.epoch).collect::>() - ); + // A source predating root epochs sends its single pair as epoch 0. + let received: Vec = keys_res.keys.iter().map(|k| k.epoch.max(1)).collect(); + if !received.iter().copied().eq(1..=last) { + bail!("source KMS must hand over root key epochs 1..={last}, got {received:?}"); } let mut roots = keys_res .keys @@ -1078,7 +1089,7 @@ impl Keys { /// The root keys go last: see [`KmsConfig::root_keys_exist`]. fn store(&self, cfg: &KmsConfig) -> Result<()> { - for (epoch, root) in (0..).zip(&self.history) { + for (epoch, root) in (1..).zip(&self.history) { let paths = cfg.root_key_paths(epoch); if let Some(dir) = paths.ca_key.parent() { fs::create_dir_all(dir)?; @@ -1210,10 +1221,10 @@ pub(crate) async fn bootstrap_keys(cfg: &KmsConfig, verifier: &AttestationVerifi Ok(()) } -/// Bootstrapping creates epoch 0 only; later epochs come from rotation onboarding. +/// Bootstrapping creates epoch 1 only; later epochs come from rotation onboarding. fn ensure_bootstrap_epoch(cfg: &KmsConfig) -> Result<()> { - if cfg.key_rotation.current_epoch != 0 { - bail!("bootstrap only creates epoch 0; onboard with rotate to create a later epoch"); + if cfg.key_rotation.current_epoch != 1 { + bail!("bootstrap only creates epoch 1; onboard with rotate to create a later epoch"); } Ok(()) } @@ -1231,7 +1242,7 @@ fn self_signed_root_ca(ca_key: &KeyPair) -> Result { async fn attest_keys(epoch: u32, p256_pubkey: &[u8], k256_pubkey: &[u8]) -> Result> { let p256_hex = hex::encode(p256_pubkey); let k256_hex = hex::encode(k256_pubkey); - let content_to_quote = if epoch == 0 { + let content_to_quote = if epoch == 1 { format!("dstack-kms-genereted-keys-v1:{p256_hex};{k256_hex};") } else { format!("dstack-kms-generated-keys-v2:{epoch};{p256_hex};{k256_hex};") diff --git a/dstack/vmm/rpc/proto/vmm_rpc.proto b/dstack/vmm/rpc/proto/vmm_rpc.proto index 7b2383cd1..e71743b87 100644 --- a/dstack/vmm/rpc/proto/vmm_rpc.proto +++ b/dstack/vmm/rpc/proto/vmm_rpc.proto @@ -458,7 +458,7 @@ message ApplyReencryptedEnvResponse { // Whether the stored encrypted env was replaced. bool replaced = 1; // KMS key epoch the stored env is now encrypted for, when replaced. - uint32 key_epoch = 2; + uint32 root_epoch = 2; } // Response containing a list of supervisor processes. diff --git a/dstack/vmm/src/main_service.rs b/dstack/vmm/src/main_service.rs index 23af1cdf0..8ebda3497 100644 --- a/dstack/vmm/src/main_service.rs +++ b/dstack/vmm/src/main_service.rs @@ -1314,10 +1314,10 @@ impl VmmRpc for RpcHandler { } safe_write::safe_write(&path, &migrated.encrypted_env) .context("Failed to write encrypted env")?; - info!(vm_id = %request.id, key_epoch = migrated.key_epoch, "applied re-encrypted env"); + info!(vm_id = %request.id, root_epoch = migrated.root_epoch, "applied re-encrypted env"); Ok(ApplyReencryptedEnvResponse { replaced: true, - key_epoch: migrated.key_epoch, + root_epoch: migrated.root_epoch, }) } diff --git a/dstack/vmm/src/vmm-cli.py b/dstack/vmm/src/vmm-cli.py index 39dd50547..f43c58114 100755 --- a/dstack/vmm/src/vmm-cli.py +++ b/dstack/vmm/src/vmm-cli.py @@ -696,7 +696,7 @@ def apply_reencrypted_env(self, vm_id: str) -> None: """Swap in the env the guest re-encrypted to the current KMS key epoch.""" result = self.rpc_call("ApplyReencryptedEnv", {"id": vm_id}) if result.get("replaced"): - print(f"Replaced env of VM {vm_id} (key epoch {result.get('key_epoch', 0)})") + print(f"Replaced env of VM {vm_id} (key epoch {result.get('root_epoch', 0)})") else: print(f"VM {vm_id} has no re-encrypted env to apply") diff --git a/sdk/go/dstack/client_v1.go b/sdk/go/dstack/client_v1.go index fd0f98ace..eb718aa74 100644 --- a/sdk/go/dstack/client_v1.go +++ b/sdk/go/dstack/client_v1.go @@ -42,6 +42,9 @@ type GetKeyV1Response struct { // key over the app root public key. See `docs/guest-api-v1.md` for the claim // encoding and the verification steps. SignatureChain [][]byte + // KMS root key epoch the key was derived from; 0 from an agent that + // predates root epochs, which means epoch 1. + RootEpoch uint32 } // Response from a v1 attestation request. @@ -442,13 +445,21 @@ func (c *DstackClientV1) IssueCert(ctx context.Context, options ...IssueCertV1Op // Derivation is flat -- `a/b` is not a child of `a` -- and binds both arguments, // so these keys differ from the v0 keys of the same name. func (c *DstackClientV1) GetKey(ctx context.Context, domain string, algorithm string) (*GetKeyV1Response, error) { + return c.GetKeyAtRootEpoch(ctx, domain, algorithm, 0) +} + +// GetKeyAtRootEpoch is GetKey under the KMS root key epoch rootEpoch, counting +// from 1; 0 is the current epoch. Name an older epoch only to reach a key from +// before a KMS root key rotation, for example to migrate off it. +func (c *DstackClientV1) GetKeyAtRootEpoch(ctx context.Context, domain string, algorithm string, rootEpoch uint32) (*GetKeyV1Response, error) { if algorithm == "" { return nil, fmt.Errorf("algorithm is required, use `secp256k1` or `ed25519`") } payload := map[string]interface{}{ - "domain": domain, - "algorithm": algorithm, + "domain": domain, + "algorithm": algorithm, + "root_epoch": rootEpoch, } data, err := c.sendRPCRequest(ctx, "/v1/GetKey", payload) @@ -464,6 +475,7 @@ func (c *DstackClientV1) GetKey(ctx context.Context, domain string, algorithm st Key *string `json:"key"` PublicKey *string `json:"public_key"` SignatureChain *[]*string `json:"signature_chain"` + RootEpoch uint32 `json:"root_epoch"` } if err := json.Unmarshal(data, &response); err != nil { return nil, err @@ -489,7 +501,12 @@ func (c *DstackClientV1) GetKey(ctx context.Context, domain string, algorithm st } } - return &GetKeyV1Response{Key: key, PublicKey: publicKey, SignatureChain: chain}, nil + return &GetKeyV1Response{ + Key: key, + PublicKey: publicKey, + SignatureChain: chain, + RootEpoch: response.RootEpoch, + }, nil } // Attest produces a versioned attestation over reportData, which must be diff --git a/sdk/js/src/client-v1.ts b/sdk/js/src/client-v1.ts index d35de8590..f743aebd5 100644 --- a/sdk/js/src/client-v1.ts +++ b/sdk/js/src/client-v1.ts @@ -68,6 +68,8 @@ export interface GetKeyResponseV1 { public_key: Uint8Array /** Two links: the app root key over the v1 key claim, then the KMS root key. */ signature_chain: Uint8Array[] + /** KMS root key epoch the key was derived from; 0 from an agent predating root epochs, meaning 1. */ + root_epoch: number } export interface AttestResponseV1 { @@ -333,14 +335,18 @@ export class DstackClientV1 { * @param domain Caller-chosen domain-separation string. May be empty. * @param algorithm Exactly `secp256k1` or `ed25519`. No default, no `k256` alias. */ - async getKey(domain: string, algorithm: string): Promise { + async getKey(domain: string, algorithm: string, rootEpoch = 0): Promise { // v0 defaulted an empty algorithm to secp256k1, which let a typo hand back a // key of the wrong type under a name the caller thought meant something else. if (!algorithm) { throw new Error('algorithm is required, use "secp256k1" or "ed25519"') } - const payload = JSON.stringify({ domain, algorithm }) - const result = await send_rpc_request<{ key: string, public_key: string, signature_chain: string[] }>( + // 0 is the current KMS root key epoch; name an older one only to migrate off it. + if (!Number.isInteger(rootEpoch) || rootEpoch < 0) { + throw new Error('rootEpoch must be 0 (current) or an epoch number') + } + const payload = JSON.stringify({ domain, algorithm, root_epoch: rootEpoch }) + const result = await send_rpc_request<{ key: string, public_key: string, signature_chain: string[], root_epoch?: number }>( this.endpoint, '/v1/GetKey', payload) throwOnRpcError(result) return Object.freeze({ @@ -348,6 +354,7 @@ export class DstackClientV1 { public_key: from_hex(result.public_key, 'public_key'), signature_chain: to_list(result.signature_chain, 'signature_chain', 'error') .map((link, i) => from_hex(link, `signature_chain[${i}]`)), + root_epoch: result.root_epoch ?? 0, __name__: 'GetKeyResponseV1' as const, }) } diff --git a/sdk/python/src/dstack_sdk/dstack_client_v1.py b/sdk/python/src/dstack_sdk/dstack_client_v1.py index e57c8bba5..a1ebe2c4f 100644 --- a/sdk/python/src/dstack_sdk/dstack_client_v1.py +++ b/sdk/python/src/dstack_sdk/dstack_client_v1.py @@ -107,6 +107,9 @@ class GetKeyResponseV1(BaseModel): # Two links: the app root key's signature over the v1 key claim, then the # KMS root key's signature over the app root public key. signature_chain: List[HexBytes] + # KMS root key epoch the key was derived from. 0 from an agent that + # predates root epochs, which means epoch 1. + root_epoch: int = 0 class GpuEvidenceBundleV1(BaseModel): @@ -237,7 +240,9 @@ async def issue_cert( result = await self._send_rpc_request("IssueCert", data) return IssueCertResponseV1(**result) - async def get_key(self, domain: str, algorithm: str) -> GetKeyResponseV1: + async def get_key( + self, domain: str, algorithm: str, root_epoch: int = 0 + ) -> GetKeyResponseV1: """Derive an application key from ``(domain, algorithm)``. Both arguments are required. Derivation is flat: two domains yield @@ -245,10 +250,20 @@ async def get_key(self, domain: str, algorithm: str) -> GetKeyResponseV1: exactly ``secp256k1`` or ``ed25519`` -- there is no default and no ``k256`` alias, because in v0 a typo silently produced a key of the wrong type under a name the caller thought meant something else. + + ``root_epoch`` selects the KMS root key epoch, counting from 1; the + default 0 is the current epoch. Name an older epoch only to reach a key + from before a KMS root key rotation, for example to migrate off it. """ if not algorithm: raise ValueError("algorithm is required, use `secp256k1` or `ed25519`") - data: Dict[str, Any] = {"domain": domain, "algorithm": algorithm} + if root_epoch < 0: + raise ValueError("root_epoch must be 0 (current) or an epoch number") + data: Dict[str, Any] = { + "domain": domain, + "algorithm": algorithm, + "root_epoch": root_epoch, + } result = await self._send_rpc_request("GetKey", data) return GetKeyResponseV1(**result) @@ -350,7 +365,9 @@ def issue_cert( raise NotImplementedError @call_async - def get_key(self, domain: str, algorithm: str) -> GetKeyResponseV1: + def get_key( + self, domain: str, algorithm: str, root_epoch: int = 0 + ) -> GetKeyResponseV1: """Derive an application key from ``(domain, algorithm)``.""" raise NotImplementedError diff --git a/sdk/rust/src/dstack_client_v1.rs b/sdk/rust/src/dstack_client_v1.rs index efcc78236..3ab11e83d 100644 --- a/sdk/rust/src/dstack_client_v1.rs +++ b/sdk/rust/src/dstack_client_v1.rs @@ -95,10 +95,22 @@ impl DstackClientV1 { /// default and no `k256` alias -- a typo is an error rather than a key of /// the wrong type under a name the caller misread. pub async fn get_key(&self, domain: &str, algorithm: &str) -> Result { + self.get_key_at_root_epoch(domain, algorithm, 0).await + } + + /// [`Self::get_key`] under the KMS root key epoch `root_epoch` (counting + /// from 1; 0 is the current epoch). Name an older epoch only to reach a key + /// from before a KMS root key rotation, for example to migrate off it. + pub async fn get_key_at_root_epoch( + &self, + domain: &str, + algorithm: &str, + root_epoch: u32, + ) -> Result { if algorithm.is_empty() { anyhow::bail!("algorithm is required, use `secp256k1` or `ed25519`") } - let data = json!({ "domain": domain, "algorithm": algorithm }); + let data = json!({ "domain": domain, "algorithm": algorithm, "root_epoch": root_epoch }); let response = self.send_rpc_request("GetKey", &data).await?; Ok(serde_json::from_value::(response)?) } diff --git a/sdk/rust/types/src/dstack_v1.rs b/sdk/rust/types/src/dstack_v1.rs index 917db972e..02a11a3e9 100644 --- a/sdk/rust/types/src/dstack_v1.rs +++ b/sdk/rust/types/src/dstack_v1.rs @@ -125,6 +125,10 @@ pub struct GetKeyResponse { /// steps. Verifying is the relying party's job; this SDK does not do it. #[serde(with = "hex_vec")] pub signature_chain: Vec>, + /// KMS root key epoch the key was derived from. 0 from an agent that + /// predates root epochs, which means epoch 1. + #[serde(default)] + pub root_epoch: u32, } /// The value an attestation binds to: bytes, and only bytes. From 311224b6b3b861531d6c3e3daac0ea22d8188301 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 5 Nov 2026 23:46:14 -0800 Subject: [PATCH 07/17] kms: record each root epoch's attestation on chain registerRootEpoch now takes the attestation the generating KMS produced for the new keys, whose report data binds the epoch and both public keys, and rootEpoch returns it, so anyone can check that a registered epoch came from an approved KMS in a TEE. Add RegisterRootEpoch to Manage.s.sol and list the registered epochs in GetKmsSettings. Signed-off-by: Kevin Wang --- docs/deployment.md | 9 +++++--- dstack/kms/auth-eth/contracts/DstackKms.sol | 23 +++++++++++++++------ dstack/kms/auth-eth/script/Manage.s.sol | 14 +++++++++++++ dstack/kms/auth-eth/script/Query.s.sol | 10 +++++++++ dstack/kms/auth-eth/script/README.md | 9 ++++++++ dstack/kms/auth-eth/src/ethereum.ts | 2 +- dstack/kms/auth-eth/test/DstackKms.t.sol | 18 +++++++++------- 7 files changed, 68 insertions(+), 17 deletions(-) diff --git a/docs/deployment.md b/docs/deployment.md index e47bc5c17..133a234c0 100644 --- a/docs/deployment.md +++ b/docs/deployment.md @@ -520,14 +520,17 @@ can migrate from any of them. See the [rotation design](./kms-emergency-key-rota for the full procedure; the KMS steps are: 1. If epochs were never registered, register the current root public keys as - epoch 1: `DstackKms.registerRootEpoch(caPubkey, k256Pubkey)` on chain, or a - `rootEpochs` entry in the auth-simple config. + epoch 1 with the attestation the KMS produced at bootstrap (kept in + `bootstrap-info.json` on the bootstrapped node): + `DstackKms.registerRootEpoch(caPubkey, k256Pubkey, attestation)`, via + `script/Manage.s.sol:RegisterRootEpoch`, or a `rootEpochs` entry in the + auth-simple config. 2. Start a fresh KMS with `[core.key_rotation] current_epoch = N` and onboard it with `"rotate": true` (the **Rotate** checkbox in the onboard UI) from a KMS at epoch N-1. It imports epochs 1..N-1, generates epoch N, and returns the new public keys with an attestation whose report data is `keccak256("dstack-kms-generated-keys-v2:{N};{ca_pubkey_hex};{k256_pubkey_hex};")`. -3. Verify the attestation and register epoch N. A KMS refuses to start until its +3. Verify the attestation and register epoch N together with it. A KMS refuses to start until its keyring matches the registered epochs, and a KMS still at an older epoch stops serving once a newer one is registered. 4. Onboard the other KMS nodes from the rotated one with `current_epoch = N`. diff --git a/dstack/kms/auth-eth/contracts/DstackKms.sol b/dstack/kms/auth-eth/contracts/DstackKms.sol index c3699d7c2..d32f32646 100644 --- a/dstack/kms/auth-eth/contracts/DstackKms.sol +++ b/dstack/kms/auth-eth/contracts/DstackKms.sol @@ -47,6 +47,10 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E struct RootEpoch { bytes caPubkey; bytes k256Pubkey; + // The generating KMS's attestation; its report data binds the epoch + // and both public keys, so anyone can check the keys came from an + // approved KMS in a TEE. + bytes attestation; } // Root public keys of KMS root epochs 1, 2, ... in order. Append-only: a @@ -125,18 +129,21 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E _emitPolicy("kms-info", keccak256(info.k256Pubkey), true); } - // Register the root public keys of the next KMS root epoch; epochs count - // from 1, the original roots + // Register the root public keys of the next KMS root epoch, with the + // attestation proving where they were generated; epochs count from 1, the + // original roots function registerRootEpoch( bytes calldata caPubkey, - bytes calldata k256Pubkey + bytes calldata k256Pubkey, + bytes calldata attestation ) external onlyOwner returns (uint256 epoch) { require(caPubkey.length > 0 && k256Pubkey.length > 0, "Empty epoch public key"); - _rootEpochs.push(RootEpoch(caPubkey, k256Pubkey)); + require(attestation.length > 0, "Empty epoch attestation"); + _rootEpochs.push(RootEpoch(caPubkey, k256Pubkey, attestation)); epoch = _rootEpochs.length; emit RootEpochRegistered(epoch, caPubkey, k256Pubkey); _emitPolicy("root-epoch", keccak256(abi.encode(caPubkey, k256Pubkey)), true); @@ -147,10 +154,14 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E return _rootEpochs.length; } - function rootEpoch(uint256 epoch) external view returns (bytes memory caPubkey, bytes memory k256Pubkey) { + function rootEpoch(uint256 epoch) + external + view + returns (bytes memory caPubkey, bytes memory k256Pubkey, bytes memory attestation) + { require(epoch > 0 && epoch <= _rootEpochs.length, "Unknown root epoch"); RootEpoch storage keys = _rootEpochs[epoch - 1]; - return (keys.caPubkey, keys.k256Pubkey); + return (keys.caPubkey, keys.k256Pubkey, keys.attestation); } // Function to set KMS quote diff --git a/dstack/kms/auth-eth/script/Manage.s.sol b/dstack/kms/auth-eth/script/Manage.s.sol index 110612b7e..c0a84ed7a 100644 --- a/dstack/kms/auth-eth/script/Manage.s.sol +++ b/dstack/kms/auth-eth/script/Manage.s.sol @@ -109,6 +109,20 @@ contract SetKmsInfo is BaseScript { } } +contract RegisterRootEpoch is BaseScript { + function run() external { + bytes memory caPubkey = vm.envBytes("CA_PUBKEY"); + bytes memory k256Pubkey = vm.envBytes("K256_PUBKEY"); + bytes memory attestation = vm.envBytes("ATTESTATION"); + + vm.startBroadcast(vm.envUint("PRIVATE_KEY")); + uint256 epoch = kms.registerRootEpoch(caPubkey, k256Pubkey, attestation); + vm.stopBroadcast(); + + console.log("Registered KMS root epoch", epoch); + } +} + // App Management Scripts contract AddComposeHash is BaseScript { function run() external { diff --git a/dstack/kms/auth-eth/script/Query.s.sol b/dstack/kms/auth-eth/script/Query.s.sol index 4a9a01496..e7f4404c4 100644 --- a/dstack/kms/auth-eth/script/Query.s.sol +++ b/dstack/kms/auth-eth/script/Query.s.sol @@ -211,6 +211,16 @@ contract GetKmsSettings is BaseQueryScript { console.log(" Quote length:", quote.length); console.log(" Eventlog length:", eventlog.length); + uint256 latestEpoch = kms.latestRootEpoch(); + console.log("\nLatest root epoch:", latestEpoch); + for (uint256 epoch = 1; epoch <= latestEpoch; epoch++) { + (bytes memory epochCa, bytes memory epochK256, bytes memory attestation) = kms.rootEpoch(epoch); + console.log(" Epoch", epoch); + console.log(" CA Pubkey:", vm.toString(epochCa)); + console.log(" K256 Pubkey:", vm.toString(epochK256)); + console.log(" Attestation length:", attestation.length); + } + // Check implementation via storage bytes32 implSlot = 0x360894a13ba1a3210667c828492db98dca3e2076cc3735a920a3ca505d382bbc; address implementation = address(uint160(uint256(vm.load(address(kms), implSlot)))); diff --git a/dstack/kms/auth-eth/script/README.md b/dstack/kms/auth-eth/script/README.md index 8ff68a00e..9aaae12a1 100644 --- a/dstack/kms/auth-eth/script/README.md +++ b/dstack/kms/auth-eth/script/README.md @@ -58,6 +58,15 @@ KMS_CONTRACT_ADDR=0x... K256_PUBKEY=0x... CA_PUBKEY=0x... QUOTE=0x... EVENTLOG=0 forge script script/Manage.s.sol:SetKmsInfo --broadcast --rpc-url $RPC_URL ``` +#### Register a KMS Root Epoch +Appends the next root epoch (the original roots are epoch 1). `ATTESTATION` is +the one the KMS returned when it generated the keys (Bootstrap, or Onboard with +`rotate`); verify it before registering. Registered epochs cannot be changed. +```bash +KMS_CONTRACT_ADDR=0x... CA_PUBKEY=0x... K256_PUBKEY=0x... ATTESTATION=0x... \ +forge script script/Manage.s.sol:RegisterRootEpoch --broadcast --rpc-url $RPC_URL +``` + #### Remove KMS Aggregated MR ```bash KMS_CONTRACT_ADDR=0x... MR_AGGREGATED=0x1234... \ diff --git a/dstack/kms/auth-eth/src/ethereum.ts b/dstack/kms/auth-eth/src/ethereum.ts index 4f57b54b6..fdb0e09c5 100644 --- a/dstack/kms/auth-eth/src/ethereum.ts +++ b/dstack/kms/auth-eth/src/ethereum.ts @@ -12,7 +12,7 @@ const DSTACK_KMS_ABI = [ "function gatewayAppId() view returns (string)", "function appImplementation() view returns (address)", "function latestRootEpoch() view returns (uint256)", - "function rootEpoch(uint256 epoch) view returns (bytes caPubkey, bytes k256Pubkey)" + "function rootEpoch(uint256 epoch) view returns (bytes caPubkey, bytes k256Pubkey, bytes attestation)" ]; // Contracts deployed before key epochs existed lack the registry getters. diff --git a/dstack/kms/auth-eth/test/DstackKms.t.sol b/dstack/kms/auth-eth/test/DstackKms.t.sol index cfafce9ce..c32250f3e 100644 --- a/dstack/kms/auth-eth/test/DstackKms.t.sol +++ b/dstack/kms/auth-eth/test/DstackKms.t.sol @@ -73,17 +73,19 @@ contract DstackKmsTest is Test { vm.startPrank(owner); vm.expectEmit(true, true, true, true); emit RootEpochRegistered(1, hex"a0", hex"b0"); - assertEq(kms.registerRootEpoch(hex"a0", hex"b0"), 1); - assertEq(kms.registerRootEpoch(hex"a1", hex"b1"), 2); + assertEq(kms.registerRootEpoch(hex"a0", hex"b0", hex"c0"), 1); + assertEq(kms.registerRootEpoch(hex"a1", hex"b1", hex"c1"), 2); vm.stopPrank(); assertEq(kms.latestRootEpoch(), 2); - (bytes memory ca, bytes memory k256) = kms.rootEpoch(1); + (bytes memory ca, bytes memory k256, bytes memory attestation) = kms.rootEpoch(1); assertEq(ca, hex"a0"); assertEq(k256, hex"b0"); - (ca, k256) = kms.rootEpoch(2); + assertEq(attestation, hex"c0"); + (ca, k256, attestation) = kms.rootEpoch(2); assertEq(ca, hex"a1"); assertEq(k256, hex"b1"); + assertEq(attestation, hex"c1"); vm.expectRevert("Unknown root epoch"); kms.rootEpoch(0); vm.expectRevert("Unknown root epoch"); @@ -93,15 +95,17 @@ contract DstackKmsTest is Test { function test_RegisterRootEpochOnlyOwner() public { vm.prank(user); vm.expectRevert(abi.encodeWithSelector(OwnableUpgradeable.OwnableUnauthorizedAccount.selector, user)); - kms.registerRootEpoch(hex"a0", hex"b0"); + kms.registerRootEpoch(hex"a0", hex"b0", hex"c0"); } function test_RegisterRootEpochRejectsEmptyKeys() public { vm.startPrank(owner); vm.expectRevert("Empty epoch public key"); - kms.registerRootEpoch("", hex"b0"); + kms.registerRootEpoch("", hex"b0", hex"c0"); vm.expectRevert("Empty epoch public key"); - kms.registerRootEpoch(hex"a0", ""); + kms.registerRootEpoch(hex"a0", "", hex"c0"); + vm.expectRevert("Empty epoch attestation"); + kms.registerRootEpoch(hex"a0", hex"b0", ""); vm.stopPrank(); } From 9355df0c9770a77dbc24f738b2acd57e1205b47a Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 5 Nov 2026 23:48:48 -0800 Subject: [PATCH 08/17] kms: refuse to onboard from a source whose keyring is not the registered one The source KMS was only checked for running an allowed image, so a legitimate image that had bootstrapped its own keys could hand over an unrelated keyring, and a rotation would build the new epoch on top of it. Onboarding now reads the registered root epochs and requires the handed- over pairs to be exactly epochs 1..=N (a join) or 1..=N-1 (a rotation), key for key, which also proves the source is at the expected epoch. Rotating past epoch 1 therefore needs the earlier epochs registered. Deployments that never registered epochs can still join at epoch 1. Signed-off-by: Kevin Wang --- dstack/kms/src/onboard_service.rs | 111 +++++++++++++++++++++++++++--- 1 file changed, 102 insertions(+), 9 deletions(-) diff --git a/dstack/kms/src/onboard_service.rs b/dstack/kms/src/onboard_service.rs index 3e1a96eac..d547285ba 100644 --- a/dstack/kms/src/onboard_service.rs +++ b/dstack/kms/src/onboard_service.rs @@ -39,6 +39,7 @@ use crate::{ build_boot_info_for_attestation, upgrade_authority::{ app_attest, dstack_client, ensure_kms_allowed, ensure_self_kms_allowed, pad64, + RootEpochPubkeys, }, }, }; @@ -295,7 +296,7 @@ mod tests { use crate::main_service::amd_attest::{ compute_expected_measurement, MeasurementInput, OvmfSectionParam, }; - use crate::main_service::upgrade_authority::{BootResponse, RootEpochPubkeys}; + use crate::main_service::upgrade_authority::BootResponse; use dstack_kms_rpc::KmsKeys; use sha2::Digest; @@ -578,7 +579,7 @@ mod tests { let res = handover(3); let sent: Vec<_> = res.keys.iter().map(|k| k.k256_key.clone()).collect(); - let keys = Keys::from_key_response(&cfg, res, "kms.example.com", false) + let keys = Keys::from_key_response(&cfg, res, "kms.example.com", false, None) .await .unwrap(); keys.store(&cfg).unwrap(); @@ -603,7 +604,7 @@ mod tests { let cfg = epoch_config(cert_dir.path(), 1); let mut res = handover(1); res.keys[0].epoch = 0; - let keys = Keys::from_key_response(&cfg, res, "kms.example.com", false) + let keys = Keys::from_key_response(&cfg, res, "kms.example.com", false, None) .await .unwrap(); assert_eq!(keys.epoch, 1); @@ -626,7 +627,7 @@ mod tests { async fn joined_state(res: KmsKeyResponse, url: String) -> crate::main_service::KmsState { let cert_dir = tempfile::tempdir().unwrap(); let mut cfg = epoch_config(cert_dir.path(), res.keys.len() as u32); - Keys::from_key_response(&cfg, res, "kms.example.com", false) + Keys::from_key_response(&cfg, res, "kms.example.com", false, None) .await .unwrap() .store(&cfg) @@ -698,7 +699,7 @@ mod tests { let res = handover(2); let sent: Vec<_> = res.keys.iter().map(|k| k.k256_key.clone()).collect(); - let keys = Keys::from_key_response(&cfg, res, "kms.example.com", true) + let keys = Keys::from_key_response(&cfg, res, "kms.example.com", true, None) .await .unwrap(); keys.store(&cfg).unwrap(); @@ -715,7 +716,7 @@ mod tests { for (epochs, rotate) in [(1, false), (3, false), (2, true), (0, true)] { let res = handover(epochs); assert!( - Keys::from_key_response(&cfg, res, "kms.example.com", rotate) + Keys::from_key_response(&cfg, res, "kms.example.com", rotate, None) .await .is_err(), "epochs={epochs} rotate={rotate}" @@ -723,12 +724,53 @@ mod tests { } let cfg = epoch_config(cert_dir.path(), 1); assert!( - Keys::from_key_response(&cfg, handover(0), "kms.example.com", true) + Keys::from_key_response(&cfg, handover(0), "kms.example.com", true, None) .await .is_err() ); } + #[rocket::async_test] + async fn onboarding_requires_the_registered_keyring() { + let cert_dir = tempfile::tempdir().unwrap(); + let join = epoch_config(cert_dir.path(), 2); + let rotate = epoch_config(cert_dir.path(), 3); + let res = handover(2); + let registered = pubkeys_of(&res); + let onboard = |cfg: &KmsConfig, rotate: bool, registered: Vec| { + let (cfg, res) = (cfg.clone(), res.clone()); + async move { + Keys::from_key_response(&cfg, res, "kms.example.com", rotate, Some(®istered)) + .await + } + }; + + assert!(onboard(&join, false, registered.clone()).await.is_ok()); + assert!(onboard(&rotate, true, registered.clone()).await.is_ok()); + // A source at an older epoch than the registry. + let mut newer = registered.clone(); + newer.push(registered[0].clone()); + assert!(onboard(&join, false, newer).await.is_err()); + // A source holding a keyring other than the registered one. + let mut other = registered.clone(); + other[0] = pubkeys_of(&handover(1))[0].clone(); + assert!(onboard(&join, false, other.clone()).await.is_err()); + assert!(onboard(&rotate, true, other).await.is_err()); + // Epochs must be registered before rotating past epoch 1. + assert!(onboard(&rotate, true, Vec::new()).await.is_err()); + } + + #[rocket::async_test] + async fn an_unregistered_deployment_can_still_join_at_epoch_1() { + let cert_dir = tempfile::tempdir().unwrap(); + let cfg = epoch_config(cert_dir.path(), 1); + assert!( + Keys::from_key_response(&cfg, handover(1), "kms.example.com", false, Some(&[])) + .await + .is_ok() + ); + } + #[rocket::async_test] async fn bootstrap_refuses_a_later_epoch() { let cert_dir = tempfile::tempdir().unwrap(); @@ -988,7 +1030,7 @@ impl Keys { ensure_kms_allowed(cfg, &source_attestation()?, &attestation_verifier) .await .context("Source KMS is not allowed for onboarding")?; - return Self::from_key_response(cfg, keys_res, domain, rotate).await; + return Self::from_registered_key_response(cfg, keys_res, domain, rotate).await; } let client = RaClientConfig::builder() @@ -1021,14 +1063,35 @@ impl Keys { vm_config: info.vm_config, }) .await?; - Self::from_key_response(cfg, keys_res, domain, rotate).await + Self::from_registered_key_response(cfg, keys_res, domain, rotate).await + } + + /// [`Self::from_key_response`] after reading the registered epochs, so a + /// source holding any keyring other than the registered one is refused. + async fn from_registered_key_response( + cfg: &KmsConfig, + keys_res: KmsKeyResponse, + domain: &str, + rotate: bool, + ) -> Result { + let registered = cfg + .auth_api + .get_info() + .await + .context("failed to read the registered root epochs")? + .root_epochs; + Self::from_key_response(cfg, keys_res, domain, rotate, registered.as_deref()).await } + /// `registered` is the governance registry; `None` means the auth backend + /// has none (dev mode), and an empty one is accepted only while joining at + /// epoch 1, as deployments that never registered epochs did before. async fn from_key_response( cfg: &KmsConfig, keys_res: KmsKeyResponse, domain: &str, rotate: bool, + registered: Option<&[RootEpochPubkeys]>, ) -> Result { let epoch = cfg.key_rotation.current_epoch; // A rotation imports every epoch before the new one; a join imports all. @@ -1055,6 +1118,7 @@ impl Keys { Ok((ca_key, k256_key)) }) .collect::>>()?; + ensure_handover_registered(&roots, registered, last)?; let (ca_key, k256_key) = if rotate { ( KeyPair::generate_for(&PKCS_ECDSA_P256_SHA256)?, @@ -1221,6 +1285,35 @@ pub(crate) async fn bootstrap_keys(cfg: &KmsConfig, verifier: &AttestationVerifi Ok(()) } +/// Check that the handed-over roots are exactly the registered epochs 1..=last, +/// which also proves the source is at the expected epoch. +fn ensure_handover_registered( + roots: &[(KeyPair, SigningKey)], + registered: Option<&[RootEpochPubkeys]>, + last: u32, +) -> Result<()> { + let Some(registered) = registered else { + return Ok(()); + }; + if registered.is_empty() && last == 1 { + return Ok(()); + } + if registered.len() != last as usize { + bail!( + "the source hands over root epochs 1..={last}, but {} epochs are registered", + registered.len() + ); + } + for (epoch, ((ca_key, k256_key), registered)) in (1..).zip(roots.iter().zip(registered)) { + if ca_key.public_key_der() != registered.ca_pubkey + || k256_key.verifying_key().to_sec1_bytes()[..] != registered.k256_pubkey[..] + { + bail!("the source's root keys of epoch {epoch} are not the registered ones"); + } + } + Ok(()) +} + /// Bootstrapping creates epoch 1 only; later epochs come from rotation onboarding. fn ensure_bootstrap_epoch(cfg: &KmsConfig) -> Result<()> { if cfg.key_rotation.current_epoch != 1 { From e1f33f7e146882c2e2748559cd7ebb913efc0d7d Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Fri, 6 Nov 2026 00:20:12 -0800 Subject: [PATCH 09/17] kms: make rotation onboarding authorize the source strictly A rotation builds the new epoch on the source's keyring, so the source must be authorized on chain. Refuse to rotate with the dev auth backend, which allows any KMS and has no root epoch registry, and do not accept a source certificate that lacks its OS image hash, which would judge its image by mrAggregated alone. A plain join keeps the existing checks as a safeguard; there the source authorizing the receiver is what matters. Signed-off-by: Kevin Wang --- .../kms/src/main_service/upgrade_authority.rs | 6 ++++ dstack/kms/src/onboard_service.rs | 34 +++++++++++++++++-- 2 files changed, 38 insertions(+), 2 deletions(-) diff --git a/dstack/kms/src/main_service/upgrade_authority.rs b/dstack/kms/src/main_service/upgrade_authority.rs index 7d37dcd0e..93df207b3 100644 --- a/dstack/kms/src/main_service/upgrade_authority.rs +++ b/dstack/kms/src/main_service/upgrade_authority.rs @@ -289,10 +289,13 @@ pub(crate) async fn ensure_self_kms_allowed( Ok(()) } +/// `allow_legacy_cert` accepts a source whose certificate predates the unified +/// attestation format, judging its image by mrAggregated alone. pub(crate) async fn ensure_kms_allowed( cfg: &KmsConfig, attestation: &VerifiedAttestation, verifier: &AttestationVerifier, + allow_legacy_cert: bool, ) -> Result<()> { let mut boot_info = build_boot_info_for_attestation(attestation, false, "") .context("failed to build KMS boot info from attestation")?; @@ -302,6 +305,9 @@ pub(crate) async fn ensure_kms_allowed( // A source that presents a config but no os_image_hash is not legacy and gets no fallback. // TODO: remove once all source KMS instances use the unified PHALA_RATLS_ATTESTATION format. if boot_info.os_image_hash.is_empty() && attestation.config.is_empty() { + if !allow_legacy_cert { + bail!("the source KMS certificate does not carry its OS image hash"); + } warn!("source KMS uses the legacy cert format, falling back to the local os_image_hash"); let local_info = local_kms_boot_info(verifier) .await diff --git a/dstack/kms/src/onboard_service.rs b/dstack/kms/src/onboard_service.rs index d547285ba..e1677666b 100644 --- a/dstack/kms/src/onboard_service.rs +++ b/dstack/kms/src/onboard_service.rs @@ -760,6 +760,30 @@ mod tests { assert!(onboard(&rotate, true, Vec::new()).await.is_err()); } + #[rocket::async_test] + async fn rotation_refuses_a_dev_auth_backend() { + let cert_dir = tempfile::tempdir().unwrap(); + let mut cfg = epoch_config(cert_dir.path(), 2); + cfg.auth_api = crate::config::AuthApi::Dev { + dev: crate::config::Dev { + gateway_app_id: String::new(), + }, + }; + let verifier = Arc::new(AttestationVerifier::load(&cfg.attestation).unwrap()); + let err = Keys::onboard( + &cfg, + "http://127.0.0.1:1", + "", + "kms.example.com", + true, + verifier, + ) + .await + .err() + .unwrap(); + assert!(format!("{err:#}").contains("rotation needs"), "{err:#}"); + } + #[rocket::async_test] async fn an_unregistered_deployment_can_still_join_at_epoch_1() { let cert_dir = tempfile::tempdir().unwrap(); @@ -981,6 +1005,12 @@ impl Keys { rotate: bool, attestation_verifier: Arc, ) -> Result { + // A new epoch is built on the source's keyring, so a rotation must be + // able to authorize the source on chain; a plain join is also checked, + // but only as a safeguard, since the source authorizes the receiver. + if rotate && cfg.auth_api.is_dev() { + bail!("rotation needs an auth backend that authorizes KMS instances and registers root epochs"); + } let attestation_slot = Arc::new(Mutex::new(None::)); let attestation_slot_out = attestation_slot.clone(); let capture_source_attestation = Box::new(move |info: Option| { @@ -1027,7 +1057,7 @@ impl Keys { .await?; // The attestation is only captured by that first request; nothing // received is used before the source passes this check. - ensure_kms_allowed(cfg, &source_attestation()?, &attestation_verifier) + ensure_kms_allowed(cfg, &source_attestation()?, &attestation_verifier, !rotate) .await .context("Source KMS is not allowed for onboarding")?; return Self::from_registered_key_response(cfg, keys_res, domain, rotate).await; @@ -1053,7 +1083,7 @@ impl Keys { ) .context("Failed to create client")?; kms_client = KmsClient::new(ra_client); - ensure_kms_allowed(cfg, &source_attestation()?, &attestation_verifier) + ensure_kms_allowed(cfg, &source_attestation()?, &attestation_verifier, !rotate) .await .context("Source KMS is not allowed for onboarding")?; From 584a7f7804ebb5f853a2c8f3f05f79853a82317d Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Fri, 6 Nov 2026 00:24:46 -0800 Subject: [PATCH 10/17] kms: allow rotation onboarding with the dev auth backend The dev backend exists for development and testing, where rotating without a registry is useful; production backends still authorize the source and check its keyring against the registry. Signed-off-by: Kevin Wang --- dstack/kms/src/onboard_service.rs | 30 ------------------------------ 1 file changed, 30 deletions(-) diff --git a/dstack/kms/src/onboard_service.rs b/dstack/kms/src/onboard_service.rs index e1677666b..fe914070b 100644 --- a/dstack/kms/src/onboard_service.rs +++ b/dstack/kms/src/onboard_service.rs @@ -760,30 +760,6 @@ mod tests { assert!(onboard(&rotate, true, Vec::new()).await.is_err()); } - #[rocket::async_test] - async fn rotation_refuses_a_dev_auth_backend() { - let cert_dir = tempfile::tempdir().unwrap(); - let mut cfg = epoch_config(cert_dir.path(), 2); - cfg.auth_api = crate::config::AuthApi::Dev { - dev: crate::config::Dev { - gateway_app_id: String::new(), - }, - }; - let verifier = Arc::new(AttestationVerifier::load(&cfg.attestation).unwrap()); - let err = Keys::onboard( - &cfg, - "http://127.0.0.1:1", - "", - "kms.example.com", - true, - verifier, - ) - .await - .err() - .unwrap(); - assert!(format!("{err:#}").contains("rotation needs"), "{err:#}"); - } - #[rocket::async_test] async fn an_unregistered_deployment_can_still_join_at_epoch_1() { let cert_dir = tempfile::tempdir().unwrap(); @@ -1005,12 +981,6 @@ impl Keys { rotate: bool, attestation_verifier: Arc, ) -> Result { - // A new epoch is built on the source's keyring, so a rotation must be - // able to authorize the source on chain; a plain join is also checked, - // but only as a safeguard, since the source authorizes the receiver. - if rotate && cfg.auth_api.is_dev() { - bail!("rotation needs an auth backend that authorizes KMS instances and registers root epochs"); - } let attestation_slot = Arc::new(Mutex::new(None::)); let attestation_slot_out = attestation_slot.clone(); let capture_source_attestation = Box::new(move |info: Option| { From 009a54f79b33fe750bc7a1f61cd7270d88672b14 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Fri, 6 Nov 2026 00:29:31 -0800 Subject: [PATCH 11/17] auth: report root epochs from auth-eth-bun, dstack-auth, and auth-mock Every webhook backend now lists the registered root epochs in its info response and reports the latest one in boot decisions, like auth-eth and auth-simple. auth-eth-bun reads the DstackKms registry at the same finalized block as the decision and treats a contract without it as having none; dstack-auth reads rootEpochs from its allowlist, which dstack run preserves when it rewrites the file; auth-mock takes them from MOCK_ROOT_EPOCHS. Without this, a KMS past epoch 1 refuses to serve. Signed-off-by: Kevin Wang --- dstack/crates/dstack-auth/src/main.rs | 53 ++++++++++++++++++-- dstack/kms/auth-eth-bun/index.test.ts | 65 ++++++++++++++++++++++++ dstack/kms/auth-eth-bun/index.ts | 72 ++++++++++++++++++++++++++- dstack/kms/auth-eth-bun/openapi.json | 16 ++++++ dstack/kms/auth-mock/README.md | 1 + dstack/kms/auth-mock/index.test.ts | 22 ++++++++ dstack/kms/auth-mock/index.ts | 17 ++++++- dstack/kms/auth-mock/openapi.json | 15 ++++++ 8 files changed, 255 insertions(+), 6 deletions(-) diff --git a/dstack/crates/dstack-auth/src/main.rs b/dstack/crates/dstack-auth/src/main.rs index 2a752e4b0..e28f354d3 100644 --- a/dstack/crates/dstack-auth/src/main.rs +++ b/dstack/crates/dstack-auth/src/main.rs @@ -66,6 +66,9 @@ struct BootResponse { is_allowed: bool, gateway_app_id: String, reason: String, + /// latest registered KMS root epoch (from 1); absent when none is. + #[serde(skip_serializing_if = "Option::is_none")] + root_epoch: Option, } #[derive(Deserialize, Default)] @@ -75,6 +78,21 @@ struct Allowlist { gateway_app_id: String, kms: KmsRules, apps: HashMap, + /// root public keys of KMS root epochs 1, 2, ... in order; append only. + root_epochs: Vec, +} + +#[derive(Deserialize, Serialize, Default, Clone)] +#[serde(rename_all = "camelCase", default)] +struct RootEpoch { + ca_pubkey: String, + k256_pubkey: String, +} + +impl Allowlist { + fn latest_root_epoch(&self) -> Option { + Some(self.root_epochs.len()).filter(|epoch| *epoch > 0) + } } #[derive(Deserialize, Default)] @@ -122,6 +140,7 @@ fn deny(al: &Allowlist, reason: &str) -> BootResponse { is_allowed: false, gateway_app_id: al.gateway_app_id.clone(), reason: reason.to_string(), + root_epoch: al.latest_root_epoch(), } } @@ -130,6 +149,7 @@ fn allow(al: &Allowlist) -> BootResponse { is_allowed: true, gateway_app_id: al.gateway_app_id.clone(), reason: "ok".to_string(), + root_epoch: al.latest_root_epoch(), } } @@ -206,15 +226,20 @@ fn boot_kms(info: Json, cli: &State) -> Json { /// info endpoint the KMS GETs to populate its metadata. Single-node: no chain. #[get("/")] -fn info() -> Json { - Json(json!({ +fn info(cli: &State) -> Json { + Json(info_json(&load(&cli.config))) +} + +fn info_json(al: &Allowlist) -> serde_json::Value { + json!({ "status": "ok", "kmsContractAddr": "", "ethRpcUrl": "", "gatewayAppId": "", "chainId": 0, - "appImplementation": "" - })) + "appImplementation": "", + "rootEpochs": al.root_epochs, + }) } #[rocket::main] @@ -283,6 +308,26 @@ mod tests { assert!(!check_kms(&info, &Allowlist::default()).is_allowed); } + #[test] + fn reports_registered_root_epochs() { + let al: Allowlist = serde_json::from_str( + r#"{"kms":{"mrAggregated":["0xMR"]}, + "rootEpochs":[{"caPubkey":"0xa1","k256Pubkey":"0xb1"}, + {"caPubkey":"0xa2","k256Pubkey":"0xb2"}]}"#, + ) + .unwrap(); + let info = BootInfo { + mr_aggregated: "0xMR".into(), + ..Default::default() + }; + let response = serde_json::to_value(check_kms(&info, &al)).unwrap(); + assert_eq!(response["rootEpoch"], 2); + assert_eq!(info_json(&al)["rootEpochs"][1]["k256Pubkey"], "0xb2"); + + let response = serde_json::to_value(check_kms(&info, &allowlist())).unwrap(); + assert!(response.get("rootEpoch").is_none()); + } + // wire-contract snapshot: BootInfo as the KMS serializes it (camelCase). // Keep these field names in sync with the kms BootInfo. `#[serde(default)]` // means extra fields are ignored AND a renamed field deserializes to "" — diff --git a/dstack/kms/auth-eth-bun/index.test.ts b/dstack/kms/auth-eth-bun/index.test.ts index 4d8a85f45..d024e1ec5 100644 --- a/dstack/kms/auth-eth-bun/index.test.ts +++ b/dstack/kms/auth-eth-bun/index.test.ts @@ -230,6 +230,64 @@ describe('API Compatibility Tests', () => { }); }); + describe('root epochs', () => { + const validBootInfo = { + mrAggregated: '0x' + '12'.repeat(32), + osImageHash: '0x' + 'ab'.repeat(32), + appId: '0x' + '11'.repeat(20), + composeHash: '0x' + '22'.repeat(32), + instanceId: '0x' + '33'.repeat(20), + deviceId: '0x' + '44'.repeat(32), + }; + + it('reports the latest registered root epoch and lists all epochs', async () => { + mockReadContract.mockImplementation((params) => { + switch (params.functionName) { + case 'isAppAllowed': return [true, '']; + case 'gatewayAppId': return '0xgateway'; + case 'appImplementation': return '0x9876543210987654321098765432109876543210'; + case 'latestRootEpoch': return 2n; + case 'rootEpoch': return [`0xa${params.args[0]}`, `0xb${params.args[0]}`, '0xc0']; + } + }); + + const boot = await appFetch(new Request('http://localhost:3001/bootAuth/app', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify(validBootInfo), + })); + expect((await boot.json()).rootEpoch).toBe(2); + + const info = await appFetch(new Request('http://localhost:3001/')); + expect((await info.json()).rootEpochs).toEqual([ + { caPubkey: '0xa1', k256Pubkey: '0xb1' }, + { caPubkey: '0xa2', k256Pubkey: '0xb2' }, + ]); + }); + + it('treats a contract without the registry as having no epochs', async () => { + mockReadContract.mockImplementation((params) => { + switch (params.functionName) { + case 'isAppAllowed': return [true, '']; + case 'gatewayAppId': return '0xgateway'; + case 'latestRootEpoch': + throw Object.assign(new Error('call failed'), { + cause: { name: 'ContractFunctionZeroDataError' }, + }); + } + }); + + const boot = await appFetch(new Request('http://localhost:3001/bootAuth/app', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify(validBootInfo), + })); + const data = await boot.json(); + expect(data.isAllowed).toBe(true); + expect(data.rootEpoch).toBeUndefined(); + }); + }); + describe('POST /bootAuth/kms', () => { const validBootInfo = { mrAggregated: '0x1234567890123456789012345678901234567890123456789012345678901234', @@ -449,6 +507,7 @@ describe('Authorization freshness and domain binding', () => { return decisions === 1 ? [true, 'initial allow'] : [false, 'policy changed']; } if (params.functionName === 'gatewayAppId') return 'gateway-app'; + if (params.functionName === 'latestRootEpoch') return 0n; throw new Error(`unexpected function ${params.functionName}`); }); @@ -468,6 +527,7 @@ describe('Authorization freshness and domain binding', () => { return [true, 'allowed']; } if (params.functionName === 'gatewayAppId') return 'gateway-app'; + if (params.functionName === 'latestRootEpoch') return 0n; throw new Error(`unexpected function ${params.functionName}`); }); @@ -492,6 +552,7 @@ describe('Authorization freshness and domain binding', () => { mockReadContract.mockImplementation((params) => { if (params.functionName === 'isAppAllowed') return [true, 'recovered']; if (params.functionName === 'gatewayAppId') return 'gateway-app'; + if (params.functionName === 'latestRootEpoch') return 0n; throw new Error(`unexpected function ${params.functionName}`); }); const recovered = await postApp(); @@ -522,6 +583,7 @@ describe('Ethereum finalized snapshot authorization', () => { expect(params.blockNumber).toBe(98n); if (params.functionName === 'isAppAllowed') return [true, 'finalized allow']; if (params.functionName === 'gatewayAppId') return 'gateway-app'; + if (params.functionName === 'latestRootEpoch') return 0n; throw new Error(`unexpected function ${params.functionName}`); }); @@ -541,6 +603,7 @@ describe('Ethereum finalized snapshot authorization', () => { return decisions === 1 ? [true, 'old canonical allow'] : [false, 'new canonical deny']; } if (params.functionName === 'gatewayAppId') return 'gateway-app'; + if (params.functionName === 'latestRootEpoch') return 0n; throw new Error(`unexpected function ${params.functionName}`); }); @@ -569,6 +632,7 @@ describe('Ethereum finalized snapshot authorization', () => { mockReadContract.mockImplementation((params) => { if (params.functionName === 'isAppAllowed') return [true, 'recovered']; if (params.functionName === 'gatewayAppId') return 'gateway-app'; + if (params.functionName === 'latestRootEpoch') return 0n; throw new Error(`unexpected function ${params.functionName}`); }); const recovered = await authorize(); @@ -611,6 +675,7 @@ describe('dstack-kms 0.6 bootAuth payload', () => { return [true, '']; } if (params.functionName === 'gatewayAppId') return 'gateway-app'; + if (params.functionName === 'latestRootEpoch') return 0n; throw new Error(`unexpected function ${params.functionName}`); }); }); diff --git a/dstack/kms/auth-eth-bun/index.ts b/dstack/kms/auth-eth-bun/index.ts index b192d9f0f..d1303fdeb 100644 --- a/dstack/kms/auth-eth-bun/index.ts +++ b/dstack/kms/auth-eth-bun/index.ts @@ -35,9 +35,13 @@ const BootInfoSchema = z.object({ const BootResponseSchema = z.object({ isAllowed: z.boolean(), reason: z.string(), - gatewayAppId: z.string() + gatewayAppId: z.string(), + // Latest registered KMS root epoch (from 1); absent when none is registered. + rootEpoch: z.number().optional() }); +type RootEpoch = { caPubkey: Hex; k256Pubkey: Hex }; + type BootInfo = z.infer; type BootResponse = z.infer; @@ -108,9 +112,37 @@ const DSTACK_KMS_ABI = [ stateMutability: 'view', inputs: [], outputs: [{ name: '', type: 'address' }] + }, + { + name: 'latestRootEpoch', + type: 'function', + stateMutability: 'view', + inputs: [], + outputs: [{ name: '', type: 'uint256' }] + }, + { + name: 'rootEpoch', + type: 'function', + stateMutability: 'view', + inputs: [{ name: 'epoch', type: 'uint256' }], + outputs: [ + { name: 'caPubkey', type: 'bytes' }, + { name: 'k256Pubkey', type: 'bytes' }, + { name: 'attestation', type: 'bytes' } + ] } ] as const; +// Contracts deployed before root epochs lack the registry getters. +const isMissingFunction = (error: unknown): boolean => { + for (let cause: any = error; cause; cause = cause.cause) { + if (cause.name === 'ContractFunctionZeroDataError' || cause.name === 'ContractFunctionRevertedError') { + return true; + } + } + return false; +}; + // ethereum backend class class EthereumBackend { private client: ReturnType; @@ -193,13 +225,49 @@ class EthereumBackend { blockNumber }); + const latestEpoch = await this.getLatestRootEpoch(blockNumber); + return { isAllowed, reason, gatewayAppId: gatewayAppId as string, + ...(latestEpoch > 0 && { rootEpoch: latestEpoch }), }; } + private async getLatestRootEpoch(blockNumber?: bigint): Promise { + try { + const latest = await this.client.readContract({ + address: this.kmsContractAddr, + abi: DSTACK_KMS_ABI, + functionName: 'latestRootEpoch', + blockNumber + }); + return Number(latest ?? 0); + } catch (error) { + if (isMissingFunction(error)) return 0; + throw error; + } + } + + // Registered root epochs 1, 2, ... in order. + async getRootEpochs(): Promise { + const blockNumber = await this.finalizedBlockNumber(); + const latest = await this.getLatestRootEpoch(blockNumber); + return Promise.all( + Array.from({ length: latest }, async (_, index) => { + const [caPubkey, k256Pubkey] = await this.client.readContract({ + address: this.kmsContractAddr, + abi: DSTACK_KMS_ABI, + functionName: 'rootEpoch', + args: [BigInt(index + 1)], + blockNumber + }); + return { caPubkey, k256Pubkey }; + }) + ); + } + async getGatewayAppId(): Promise { const result = await this.client.readContract({ address: this.kmsContractAddr, @@ -278,6 +346,7 @@ app.get('/', async (c) => { ethereum.getGatewayAppId(), ethereum.getChainId(), ethereum.getAppImplementation(), + ethereum.getRootEpochs(), ]); console.log('batch', batch); @@ -289,6 +358,7 @@ app.get('/', async (c) => { chainId: batch[1], appAuthImplementation: batch[2], // NOTE: for backward compatibility appImplementation: batch[2], + rootEpochs: batch[3], }); } catch { console.error('authorization backend health check failed'); diff --git a/dstack/kms/auth-eth-bun/openapi.json b/dstack/kms/auth-eth-bun/openapi.json index e87c03611..88947e83c 100644 --- a/dstack/kms/auth-eth-bun/openapi.json +++ b/dstack/kms/auth-eth-bun/openapi.json @@ -173,6 +173,10 @@ "gatewayAppId": { "type": "string", "description": "Gateway application ID" + }, + "rootEpoch": { + "type": "integer", + "description": "Latest registered KMS root epoch, counting from 1; absent when none is registered" } } }, @@ -216,6 +220,18 @@ "appImplementation": { "type": "string", "description": "App implementation address" + }, + "rootEpochs": { + "type": "array", + "description": "Registered KMS root epochs 1, 2, ... in order", + "items": { + "type": "object", + "required": ["caPubkey", "k256Pubkey"], + "properties": { + "caPubkey": { "type": "string", "description": "Root CA public key (hex)" }, + "k256Pubkey": { "type": "string", "description": "Root secp256k1 public key (hex)" } + } + } } } }, diff --git a/dstack/kms/auth-mock/README.md b/dstack/kms/auth-mock/README.md index c26ed0e19..51aff51f6 100644 --- a/dstack/kms/auth-mock/README.md +++ b/dstack/kms/auth-mock/README.md @@ -72,6 +72,7 @@ bun run check - `MOCK_GATEWAY_APP_ID` - mock gateway app ID (default: 0xmockgateway1234567890123456789012345678) - `MOCK_CHAIN_ID` - mock blockchain chain ID (default: 1337) - `MOCK_APP_IMPLEMENTATION` - mock app implementation address (default: 0xmockapp9876543210987654321098765432109) +- `MOCK_ROOT_EPOCHS` - JSON array of `{"caPubkey", "k256Pubkey"}` for KMS root epochs 1, 2, ... (default: none registered) ## API endpoints diff --git a/dstack/kms/auth-mock/index.test.ts b/dstack/kms/auth-mock/index.test.ts index 493745613..872705107 100644 --- a/dstack/kms/auth-mock/index.test.ts +++ b/dstack/kms/auth-mock/index.test.ts @@ -160,6 +160,28 @@ describe('Mock Backend Tests', () => { }); }); + describe('root epochs', () => { + it('reports MOCK_ROOT_EPOCHS', async () => { + const rootEpochs = [{ caPubkey: '0xa1', k256Pubkey: '0xb1' }]; + process.env.MOCK_ROOT_EPOCHS = JSON.stringify(rootEpochs); + try { + const boot = await appFetch(new Request('http://localhost:3002/bootAuth/kms', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + mrAggregated: '0x01', osImageHash: '0x02', appId: '0x03', + composeHash: '0x04', instanceId: '0x05', deviceId: '0x06', + }), + })); + expect((await boot.json()).rootEpoch).toBe(1); + const info = await appFetch(new Request('http://localhost:3002/')); + expect((await info.json()).rootEpochs).toEqual(rootEpochs); + } finally { + delete process.env.MOCK_ROOT_EPOCHS; + } + }); + }); + describe('POST /bootAuth/kms', () => { const validBootInfo = { mrAggregated: '0x1234567890123456789012345678901234567890123456789012345678901234', diff --git a/dstack/kms/auth-mock/index.ts b/dstack/kms/auth-mock/index.ts index 577cae583..23f84944b 100644 --- a/dstack/kms/auth-mock/index.ts +++ b/dstack/kms/auth-mock/index.ts @@ -28,9 +28,17 @@ const BootInfoSchema = z.object({ const BootResponseSchema = z.object({ isAllowed: z.boolean(), reason: z.string(), - gatewayAppId: z.string() + gatewayAppId: z.string(), + rootEpoch: z.number().optional() }); +const RootEpochsSchema = z.array(z.object({ caPubkey: z.string(), k256Pubkey: z.string() })); + +// MOCK_ROOT_EPOCHS: JSON array of {caPubkey, k256Pubkey} for root epochs 1, 2, ... +function getRootEpochs(): z.infer { + return RootEpochsSchema.parse(JSON.parse(process.env.MOCK_ROOT_EPOCHS || '[]')); +} + type BootInfo = z.infer; type BootResponse = z.infer; @@ -70,6 +78,12 @@ class MockBackend { } async checkBoot(bootInfo: BootInfo, isKms: boolean): Promise { + const response = this.evaluateBoot(bootInfo, isKms); + const latestEpoch = getRootEpochs().length; + return latestEpoch > 0 ? { ...response, rootEpoch: latestEpoch } : response; + } + + private evaluateBoot(bootInfo: BootInfo, isKms: boolean): BootResponse { const policy = getPolicy(); const deny = (reason: string): BootResponse => ({ isAllowed: false, @@ -147,6 +161,7 @@ app.get('/', async (c) => { chainId: batch[1], appAuthImplementation: batch[2], // NOTE: for backward compatibility appImplementation: batch[2], + rootEpochs: getRootEpochs(), note: 'this is a mock backend - all authentications will succeed' }); } catch (error) { diff --git a/dstack/kms/auth-mock/openapi.json b/dstack/kms/auth-mock/openapi.json index 1689308d0..66dd817e4 100644 --- a/dstack/kms/auth-mock/openapi.json +++ b/dstack/kms/auth-mock/openapi.json @@ -176,6 +176,10 @@ "type": "string", "description": "Mock gateway application ID", "example": "0xmockgateway1234567890123456789012345678" + }, + "rootEpoch": { + "type": "integer", + "description": "Latest KMS root epoch from MOCK_ROOT_EPOCHS; absent when none" } } }, @@ -221,6 +225,17 @@ "description": "Mock app implementation address", "example": "0xmockapp9876543210987654321098765432109" }, + "rootEpochs": { + "type": "array", + "description": "KMS root epochs 1, 2, ... from MOCK_ROOT_EPOCHS", + "items": { + "type": "object", + "properties": { + "caPubkey": { "type": "string" }, + "k256Pubkey": { "type": "string" } + } + } + }, "note": { "type": "string", "description": "Note about mock backend behavior", From 5bab8483a5bede048339b6acf4a1420b55e7f973 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Fri, 6 Nov 2026 00:34:56 -0800 Subject: [PATCH 12/17] guest-api: rename GetReencryptedEnv to GetEncryptedEnv It now always returns the encrypted env the host should store: the copy boot re-encrypted to the current root epoch, or the boot env itself when it needed no change. ApplyReencryptedEnv replaces nothing in that case. Signed-off-by: Kevin Wang --- dstack/guest-agent/src/guest_api_service.rs | 38 ++++++++++++++------- dstack/guest-agent/src/rpc_service.rs | 5 +++ dstack/guest-api/proto/guest_api.proto | 17 ++++----- dstack/vmm/src/guest_api_service.rs | 6 ++-- dstack/vmm/src/main_service.rs | 6 ++-- 5 files changed, 46 insertions(+), 26 deletions(-) diff --git a/dstack/guest-agent/src/guest_api_service.rs b/dstack/guest-agent/src/guest_api_service.rs index 54af121f8..f26ca4640 100644 --- a/dstack/guest-agent/src/guest_api_service.rs +++ b/dstack/guest-agent/src/guest_api_service.rs @@ -7,16 +7,17 @@ use std::{collections::HashSet, fmt::Debug, path::PathBuf, process::ExitStatus, use anyhow::{Context, Result}; use bollard::{container::ListContainersOptions, Docker}; use dstack_guest_agent_rpc::v0::worker_server::WorkerRpc as _; -use dstack_types::shared_filenames::{HOST_SHARED_DIR, REENCRYPTED_ENV, SYS_CONFIG}; +use dstack_types::shared_filenames::{ENCRYPTED_ENV, HOST_SHARED_DIR, REENCRYPTED_ENV, SYS_CONFIG}; use dstack_types::SysConfig; use fs_err as fs; use guest_api::{ guest_api_server::{GuestApiRpc, GuestApiServer}, - Container, DiskInfo, Gateway, GpuInfoResponse, GuestInfo, Interface, IpAddress, - ListContainersResponse, NetworkInformation, ReencryptedEnv, SystemInfo, + Container, DiskInfo, EncryptedEnv, Gateway, GpuInfoResponse, GuestInfo, Interface, IpAddress, + ListContainersResponse, NetworkInformation, SystemInfo, }; use host_api::Notification; use ra_rpc::{CallContext, RpcCall}; +use sha2::Digest as _; use tokio::{process::Command, task::spawn_blocking, time::timeout}; use tracing::error; @@ -44,17 +45,28 @@ impl RpcCall for GuestApiHandler { } impl GuestApiRpc for GuestApiHandler { - async fn get_reencrypted_env(self) -> Result { - let path = PathBuf::from(HOST_SHARED_DIR).join(REENCRYPTED_ENV); - let migrated: dstack_types::ReencryptedEnv = match fs::read(&path) { - Ok(json) => serde_json::from_slice(&json).context("invalid re-encrypted env")?, - Err(err) if err.kind() == std::io::ErrorKind::NotFound => return Ok(Default::default()), - Err(err) => return Err(err).context("failed to read the re-encrypted env"), + async fn get_encrypted_env(self) -> Result { + let read = |name: &str| match fs::read(PathBuf::from(HOST_SHARED_DIR).join(name)) { + Ok(data) => Ok(Some(data)), + Err(err) if err.kind() == std::io::ErrorKind::NotFound => Ok(None), + Err(err) => Err(err), }; - Ok(ReencryptedEnv { - source_sha256: migrated.source_sha256, - encrypted_env: migrated.encrypted_env, - root_epoch: migrated.root_epoch, + if let Some(json) = read(REENCRYPTED_ENV).context("failed to read the re-encrypted env")? { + let migrated: dstack_types::ReencryptedEnv = + serde_json::from_slice(&json).context("invalid re-encrypted env")?; + return Ok(EncryptedEnv { + source_sha256: migrated.source_sha256, + encrypted_env: migrated.encrypted_env, + root_epoch: migrated.root_epoch, + }); + } + let Some(env) = read(ENCRYPTED_ENV).context("failed to read the encrypted env")? else { + return Ok(Default::default()); + }; + Ok(EncryptedEnv { + source_sha256: sha2::Sha256::digest(&env).to_vec(), + encrypted_env: env, + root_epoch: self.state.root_epoch(), }) } diff --git a/dstack/guest-agent/src/rpc_service.rs b/dstack/guest-agent/src/rpc_service.rs index cacfce569..7eb792341 100644 --- a/dstack/guest-agent/src/rpc_service.rs +++ b/dstack/guest-agent/src/rpc_service.rs @@ -385,6 +385,11 @@ impl AppState { Ok(key) } + /// KMS root epoch of the app keys this agent serves. + pub(crate) fn root_epoch(&self) -> u32 { + self.inner.keys.root_epoch() + } + /// The app root key and its KMS signature of `root_epoch`; the current /// ones for 0. The key is parsed for older epochs only, which are /// requested rarely. diff --git a/dstack/guest-api/proto/guest_api.proto b/dstack/guest-api/proto/guest_api.proto index 3ef86458d..60f214452 100644 --- a/dstack/guest-api/proto/guest_api.proto +++ b/dstack/guest-api/proto/guest_api.proto @@ -185,14 +185,15 @@ message GpuInfoResponse { optional uint64 sample_age_ms = 5; } -// Env that boot re-encrypted from an older KMS key epoch to the current one. -// Everything is empty when the env this boot started from needed no change. -message ReencryptedEnv { +// The encrypted env the host should store for this VM: the env this boot +// started from, or, if it was encrypted to an older KMS root epoch, the copy +// boot re-encrypted to the current one. Empty when the VM has no env. +message EncryptedEnv { // SHA-256 of the encrypted env the host supplied for this boot. bytes source_sha256 = 1; - // The same env encrypted to the current env public key. + // The env encrypted to the current env public key. bytes encrypted_env = 2; - // KMS key epoch `encrypted_env` is encrypted for. + // KMS root epoch `encrypted_env` is encrypted for. uint32 root_epoch = 3; } @@ -210,8 +211,8 @@ service GuestApi { rpc ListContainers(google.protobuf.Empty) returns (ListContainersResponse); // Initiates a graceful shutdown inside the guest VM. rpc Shutdown(google.protobuf.Empty) returns (google.protobuf.Empty); - // Returns the env re-encrypted to the current KMS key epoch during boot. - rpc GetReencryptedEnv(google.protobuf.Empty) returns (ReencryptedEnv); + // Returns the encrypted env the host should store for this VM. + rpc GetEncryptedEnv(google.protobuf.Empty) returns (EncryptedEnv); } // Same API surface as GuestApi but multiplexed by VM ID through VMM. @@ -222,5 +223,5 @@ service ProxiedGuestApi { rpc NetworkInfo(Id) returns (NetworkInformation); rpc ListContainers(Id) returns (ListContainersResponse); rpc Shutdown(Id) returns (google.protobuf.Empty); - rpc GetReencryptedEnv(Id) returns (ReencryptedEnv); + rpc GetEncryptedEnv(Id) returns (EncryptedEnv); } diff --git a/dstack/vmm/src/guest_api_service.rs b/dstack/vmm/src/guest_api_service.rs index dbf945cb1..763df5005 100644 --- a/dstack/vmm/src/guest_api_service.rs +++ b/dstack/vmm/src/guest_api_service.rs @@ -6,7 +6,7 @@ use crate::App as AppState; use anyhow::Result; use guest_api::{ proxied_guest_api_server::{ProxiedGuestApiRpc, ProxiedGuestApiServer}, - GpuInfoResponse, GuestInfo, Id, ListContainersResponse, NetworkInformation, ReencryptedEnv, + EncryptedEnv, GpuInfoResponse, GuestInfo, Id, ListContainersResponse, NetworkInformation, SystemInfo, }; use ra_rpc::{CallContext, RpcCall}; @@ -61,9 +61,9 @@ impl ProxiedGuestApiRpc for GuestApiHandler { self.guest_agent_client(&request.id)?.shutdown().await } - async fn get_reencrypted_env(self, request: Id) -> Result { + async fn get_encrypted_env(self, request: Id) -> Result { self.guest_agent_client(&request.id)? - .get_reencrypted_env() + .get_encrypted_env() .await } } diff --git a/dstack/vmm/src/main_service.rs b/dstack/vmm/src/main_service.rs index 8ebda3497..cca11541a 100644 --- a/dstack/vmm/src/main_service.rs +++ b/dstack/vmm/src/main_service.rs @@ -1302,9 +1302,11 @@ impl VmmRpc for RpcHandler { use sha2::Digest; let migrated = self .guest_agent_client(&request.id)? - .get_reencrypted_env() + .get_encrypted_env() .await?; - if migrated.encrypted_env.is_empty() { + if migrated.encrypted_env.is_empty() + || sha2::Sha256::digest(&migrated.encrypted_env)[..] == migrated.source_sha256[..] + { return Ok(ApplyReencryptedEnvResponse::default()); } let path = self.encrypted_env_path(&request.id)?; From 14eec15ec864da225b8576c4b1cb269549b7ff42 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Fri, 6 Nov 2026 00:54:12 -0800 Subject: [PATCH 13/17] Number root epochs from 0 again and select the latest with an unset field Epoch 0 is the original root keys, as before root epochs existed, so a peer that predates them reports exactly what it serves. Selectors are now optional: an unset root_epoch in GetAppKey and GetKey means the current epoch, and any number names that epoch. Old OS images, which send no selector, get the current keys; before a rotation those are epoch 0. The registry indexes epochs from 0 and exposes rootEpochCount; auth backends report the latest epoch as the count minus one. Signed-off-by: Kevin Wang --- docs/deployment.md | 4 +- dstack/crates/dstack-auth/src/main.rs | 8 +- dstack/dstack-types/src/lib.rs | 18 ++-- dstack/dstack-util/src/main.rs | 4 +- dstack/dstack-util/src/system_setup.rs | 42 ++++----- .../guest-agent/rpc/proto/agent_rpc_v1.proto | 10 +- dstack/guest-agent/src/rpc_service.rs | 12 +-- dstack/guest-agent/src/rpc_service_v1.rs | 16 ++-- dstack/kms/auth-eth-bun/index.test.ts | 22 ++--- dstack/kms/auth-eth-bun/index.ts | 24 ++--- dstack/kms/auth-eth-bun/openapi.json | 4 +- dstack/kms/auth-eth/contracts/DstackKms.sol | 18 ++-- dstack/kms/auth-eth/script/Query.s.sol | 6 +- dstack/kms/auth-eth/script/README.md | 2 +- dstack/kms/auth-eth/src/ethereum.ts | 16 ++-- dstack/kms/auth-eth/src/types.ts | 2 +- dstack/kms/auth-eth/test/DstackKms.t.sol | 18 ++-- dstack/kms/auth-mock/README.md | 2 +- dstack/kms/auth-mock/index.test.ts | 2 +- dstack/kms/auth-mock/index.ts | 6 +- dstack/kms/auth-mock/openapi.json | 2 +- dstack/kms/auth-simple/README.md | 2 +- dstack/kms/auth-simple/index.test.ts | 2 +- dstack/kms/auth-simple/index.ts | 6 +- dstack/kms/kms.toml | 4 +- dstack/kms/rpc/proto/kms_rpc.proto | 15 ++- dstack/kms/src/config.rs | 42 +++------ dstack/kms/src/main_service.rs | 27 +++--- dstack/kms/src/onboard_service.rs | 92 ++++++++----------- sdk/go/dstack/client_v1.go | 24 +++-- sdk/js/src/client-v1.ts | 13 +-- sdk/python/src/dstack_sdk/dstack_client_v1.py | 26 +++--- sdk/rust/src/dstack_client_v1.rs | 23 ++++- sdk/rust/types/src/dstack_v1.rs | 3 +- 34 files changed, 246 insertions(+), 271 deletions(-) diff --git a/docs/deployment.md b/docs/deployment.md index 133a234c0..12bb1efd2 100644 --- a/docs/deployment.md +++ b/docs/deployment.md @@ -520,14 +520,14 @@ can migrate from any of them. See the [rotation design](./kms-emergency-key-rota for the full procedure; the KMS steps are: 1. If epochs were never registered, register the current root public keys as - epoch 1 with the attestation the KMS produced at bootstrap (kept in + epoch 0 with the attestation the KMS produced at bootstrap (kept in `bootstrap-info.json` on the bootstrapped node): `DstackKms.registerRootEpoch(caPubkey, k256Pubkey, attestation)`, via `script/Manage.s.sol:RegisterRootEpoch`, or a `rootEpochs` entry in the auth-simple config. 2. Start a fresh KMS with `[core.key_rotation] current_epoch = N` and onboard it with `"rotate": true` (the **Rotate** checkbox in the onboard UI) from a KMS - at epoch N-1. It imports epochs 1..N-1, generates epoch N, and returns the new + at epoch N-1. It imports epochs 0..N-1, generates epoch N, and returns the new public keys with an attestation whose report data is `keccak256("dstack-kms-generated-keys-v2:{N};{ca_pubkey_hex};{k256_pubkey_hex};")`. 3. Verify the attestation and register epoch N together with it. A KMS refuses to start until its diff --git a/dstack/crates/dstack-auth/src/main.rs b/dstack/crates/dstack-auth/src/main.rs index e28f354d3..35b705b4a 100644 --- a/dstack/crates/dstack-auth/src/main.rs +++ b/dstack/crates/dstack-auth/src/main.rs @@ -66,7 +66,7 @@ struct BootResponse { is_allowed: bool, gateway_app_id: String, reason: String, - /// latest registered KMS root epoch (from 1); absent when none is. + /// latest registered KMS root epoch; absent when none is. #[serde(skip_serializing_if = "Option::is_none")] root_epoch: Option, } @@ -78,7 +78,7 @@ struct Allowlist { gateway_app_id: String, kms: KmsRules, apps: HashMap, - /// root public keys of KMS root epochs 1, 2, ... in order; append only. + /// root public keys of each KMS root epoch, indexed by epoch; append only. root_epochs: Vec, } @@ -91,7 +91,7 @@ struct RootEpoch { impl Allowlist { fn latest_root_epoch(&self) -> Option { - Some(self.root_epochs.len()).filter(|epoch| *epoch > 0) + self.root_epochs.len().checked_sub(1) } } @@ -321,7 +321,7 @@ mod tests { ..Default::default() }; let response = serde_json::to_value(check_kms(&info, &al)).unwrap(); - assert_eq!(response["rootEpoch"], 2); + assert_eq!(response["rootEpoch"], 1); assert_eq!(info_json(&al)["rootEpochs"][1]["k256Pubkey"], "0xb2"); let response = serde_json::to_value(check_kms(&info, &allowlist())).unwrap(); diff --git a/dstack/dstack-types/src/lib.rs b/dstack/dstack-types/src/lib.rs index e5bfa1b7e..532dc9911 100644 --- a/dstack/dstack-types/src/lib.rs +++ b/dstack/dstack-types/src/lib.rs @@ -2428,19 +2428,19 @@ pub struct AppKeys { pub gateway_app_id: String, pub ca_cert: String, pub key_provider: KeyProvider, - /// App root keys of the KMS root epochs before the current one, from - /// epoch 1, for key requests that select one explicitly. + /// App root keys of the KMS root epochs before the current one, indexed by + /// epoch, for key requests that select one explicitly. #[serde(default)] pub historical_roots: Vec, } impl AppKeys { - /// KMS root epoch the keys were derived from; 1 for other key providers, + /// KMS root epoch the keys were derived from; 0 for other key providers, /// which have a single root. pub fn root_epoch(&self) -> u32 { match self.key_provider { KeyProvider::Kms { root_epoch, .. } => root_epoch, - _ => first_root_epoch(), + _ => 0, } } } @@ -2469,16 +2469,12 @@ pub enum KeyProvider { pubkey: Vec, tmp_ca_key: String, tmp_ca_cert: String, - /// KMS root key epoch the keys were derived from, counting from 1. - #[serde(default = "first_root_epoch")] + /// KMS root key epoch the keys were derived from. + #[serde(default)] root_epoch: u32, }, } -fn first_root_epoch() -> u32 { - 1 -} - impl KeyProvider { pub fn kind(&self) -> KeyProviderKind { match self { @@ -2539,7 +2535,7 @@ mod key_provider_tests { pubkey: vec![0xab; 32], tmp_ca_key: String::new(), tmp_ca_cert: String::new(), - root_epoch: 1, + root_epoch: 0, }; assert_eq!(kms.id(), &[0xab; 32]); } diff --git a/dstack/dstack-util/src/main.rs b/dstack/dstack-util/src/main.rs index fd68d76e0..05e8fdc34 100644 --- a/dstack/dstack-util/src/main.rs +++ b/dstack/dstack-util/src/main.rs @@ -701,7 +701,7 @@ async fn cmd_get_keys(args: GetKeysArgs) -> Result<()> { .get_app_key(dstack_kms_rpc::GetAppKeyRequest { api_version: 1, vm_config: "".to_string(), - root_epoch: 0, + root_epoch: None, }) .await .context("Failed to get app key")?; @@ -725,7 +725,7 @@ async fn cmd_get_keys(args: GetKeysArgs) -> Result<()> { pubkey: root_pubkey, tmp_ca_key: tmp_ca.temp_ca_key, tmp_ca_cert: tmp_ca.temp_ca_cert, - root_epoch: response.root_epoch.max(1), + root_epoch: response.root_epoch, }, }; diff --git a/dstack/dstack-util/src/system_setup.rs b/dstack/dstack-util/src/system_setup.rs index 3619a6293..f8a18419b 100644 --- a/dstack/dstack-util/src/system_setup.rs +++ b/dstack/dstack-util/src/system_setup.rs @@ -2355,13 +2355,13 @@ struct AppInfo { } impl HostShared { - /// Request app keys of `root_epoch` (0 = current) from one KMS. + /// Request app keys of `root_epoch` (`None` = current) from one KMS. async fn request_kms_keys( &self, kms_url: String, - root_epoch: u32, + root_epoch: Option, ) -> Result<(AppKeys, Vec)> { - info!("Requesting app keys of epoch {root_epoch} from KMS: {kms_url}"); + info!("Requesting app keys of root epoch {root_epoch:?} from KMS: {kms_url}"); let tmp_ca = { info!("Getting temp ca cert"); let client = RaClient::new(kms_url.clone(), true)?; @@ -2394,10 +2394,10 @@ impl HostShared { }) .await .context("Failed to get app key")?; - // A KMS predating root epochs ignores the selector and reports 0: epoch 1. - let served = response.root_epoch.max(1); - if root_epoch != 0 && served != root_epoch { - bail!("KMS returned keys of epoch {served} instead of {root_epoch}"); + // A KMS predating root epochs ignores the selector and serves epoch 0. + let served = response.root_epoch; + if root_epoch.is_some_and(|epoch| epoch != served) { + bail!("KMS returned keys of epoch {served} instead of {root_epoch:?}"); } let (_, ca_pem) = x509_parser::pem::parse_x509_pem(tmp_ca.ca_cert.as_bytes()) @@ -2434,8 +2434,8 @@ impl HostShared { return Ok(Vec::new()); }; let mut history = Vec::new(); - for epoch in 1..*root_epoch { - history.push(self.request_kms_keys(url.clone(), epoch).await?.0); + for epoch in 0..*root_epoch { + history.push(self.request_kms_keys(url.clone(), Some(epoch)).await?.0); } Ok(history) } @@ -2452,7 +2452,7 @@ struct Stage1<'a> { vmm: HostApi, shared: HostShared, keys: AppKeys, - /// Keys of root epochs 1 up to, but not including, `keys.root_epoch()`. + /// Keys of the root epochs before `keys.root_epoch()`, indexed by epoch. key_history: Vec, } @@ -2512,7 +2512,7 @@ impl<'a> Stage0<'a> { let mut errors = vec![]; for kms_url in &self.shared.sys_config.kms_urls { let kms_url = kms_rpc_url(kms_url); - match self.shared.request_kms_keys(kms_url.clone(), 0).await { + match self.shared.request_kms_keys(kms_url.clone(), None).await { Ok((keys, os_image_hash)) => { // Measured here, outside the per-URL request, so that a // failover can never extend os-image-hash twice. @@ -2979,7 +2979,7 @@ impl<'a> Stage0<'a> { { bail!("Failed to setup luks volume"); } - if let Some(epoch) = root_epoch.filter(|epoch| *epoch > 1) { + if let Some(epoch) = root_epoch.filter(|epoch| *epoch > 0) { write_luks_root_epoch(root_hd, epoch, false)?; } self.open_encrypted_volume(disk_crypt_key, name, discard)?; @@ -2990,7 +2990,7 @@ impl<'a> Stage0<'a> { /// /// Only the keyslot passphrase changes; the volume key stays. The epoch /// is recorded in the header so the right old key is used after the next - /// rotation. An unrotated epoch-1 disk is left untouched, so it stays + /// rotation. An unrotated epoch-0 disk is left untouched, so it stays /// readable by older OS images. fn migrate_disk_key( &self, @@ -3001,13 +3001,13 @@ impl<'a> Stage0<'a> { let device = &self.args.device; let stored = read_luks_root_epoch(device)?; if !luks_unlocks(device, current_key)? { - let epoch = stored.unwrap_or(1); + let epoch = stored.unwrap_or(0); if epoch >= current { bail!("the data disk does not unlock with the key of epoch {current}"); } info!("re-keying the data disk from key epoch {epoch} to {current}"); let old = key_history - .get(epoch as usize - 1) + .get(epoch as usize) .with_context(|| format!("no keys for root epoch {epoch}"))?; let old_key = hex::encode(&old.disk_crypt_key); luks_change_key(device, &old_key, current_key, Path::new(REKEY_DIR))?; @@ -3015,7 +3015,7 @@ impl<'a> Stage0<'a> { bail!("the data disk does not unlock after re-keying"); } } - if stored.unwrap_or(1) != current { + if stored.unwrap_or(0) != current { write_luks_root_epoch(device, current, stored.is_some())?; } Ok(()) @@ -3332,7 +3332,7 @@ impl Stage1<'_> { }; let current = self.keys.root_epoch(); for (index, old) in self.key_history.iter().enumerate().rev() { - let epoch = index as u32 + 1; + let epoch = index as u32; let Ok(plaintext) = dh_decrypt(env_crypt_key(&old.env_crypt_key)?, ciphertext) else { continue; }; @@ -3639,7 +3639,7 @@ const LUKS_INITIALIZING_LABEL: &str = "dstack-initializing"; struct LuksHeaderState { /// The header carries [`LUKS_INITIALIZING_LABEL`]. initializing: bool, - /// Root epoch recorded in the dstack token; `None` means epoch 1. + /// Root epoch recorded in the dstack token; `None` means epoch 0. root_epoch: Option, } @@ -3733,11 +3733,7 @@ fn validate_single_luks2_header( if token.kind != ROOT_EPOCH_TOKEN_TYPE || token.keyslots != ["0"] { bail!("Invalid LUKS token"); } - let epoch = token.root_epoch.context("LUKS token has no root epoch")?; - if epoch == 0 { - bail!("Invalid LUKS token root epoch"); - } - Some(epoch) + Some(token.root_epoch.context("LUKS token has no root epoch")?) } _ => bail!("Invalid LUKS tokens"), }; diff --git a/dstack/guest-agent/rpc/proto/agent_rpc_v1.proto b/dstack/guest-agent/rpc/proto/agent_rpc_v1.proto index ad68adde6..bebc85f99 100644 --- a/dstack/guest-agent/rpc/proto/agent_rpc_v1.proto +++ b/dstack/guest-agent/rpc/proto/agent_rpc_v1.proto @@ -189,11 +189,11 @@ message GetKeyRequest { // signing method: v0's `secp256k1_prehashed` was a mode wearing an // algorithm's name, and it does not appear here. string algorithm = 2; - // KMS root key epoch to derive from: 0 (default) = the current epoch, which is - // what new keys should use; N = epoch N, counting from 1. Name an older epoch - // only to reach a key from before a KMS root key rotation, for example to - // migrate data or assets off it. - uint32 root_epoch = 3; + // KMS root key epoch to derive from. Unset means the current epoch, which is + // what new keys should use; epoch 0 is the original root keys. Name an older + // epoch only to reach a key from before a KMS root key rotation, for example + // to migrate data or assets off it. + optional uint32 root_epoch = 3; } message GetKeyResponse { diff --git a/dstack/guest-agent/src/rpc_service.rs b/dstack/guest-agent/src/rpc_service.rs index 7eb792341..3a3fcbb5c 100644 --- a/dstack/guest-agent/src/rpc_service.rs +++ b/dstack/guest-agent/src/rpc_service.rs @@ -391,12 +391,12 @@ impl AppState { } /// The app root key and its KMS signature of `root_epoch`; the current - /// ones for 0. The key is parsed for older epochs only, which are + /// ones when unset. The key is parsed for older epochs only, which are /// requested rarely. - pub(crate) fn app_root(&self, root_epoch: u32) -> Result> { + pub(crate) fn app_root(&self, root_epoch: Option) -> Result> { let keys = &self.inner.keys; let current = keys.root_epoch(); - let epoch = if root_epoch == 0 { current } else { root_epoch }; + let epoch = root_epoch.unwrap_or(current); if epoch == current { return Ok(AppRoot { epoch, @@ -1401,10 +1401,10 @@ pNs85uhOZE8z2jr8Pg== k256_signature: Vec::new(), gateway_app_id: String::new(), ca_cert: DUMMY_PEM_CERT.to_string(), - // Only selected by tests that name root epoch 2 explicitly; the - // fixture's current root is epoch 1. + // Only selected by tests that name root epoch 1 explicitly; the + // fixture's current root is epoch 0. historical_roots: vec![dstack_types::HistoricalAppRoot { - root_epoch: 2, + root_epoch: 1, k256_key: HISTORICAL_K256_KEY.to_vec(), k256_signature: vec![0xab; 65], }], diff --git a/dstack/guest-agent/src/rpc_service_v1.rs b/dstack/guest-agent/src/rpc_service_v1.rs index 95fd6e5fa..da6ee9765 100644 --- a/dstack/guest-agent/src/rpc_service_v1.rs +++ b/dstack/guest-agent/src/rpc_service_v1.rs @@ -165,7 +165,7 @@ impl V1RpcHandler { &self, domain: &str, algorithm: &str, - root_epoch: u32, + root_epoch: Option, ) -> Result<(AppKey, Vec>, u32)> { let algorithm = Algorithm::parse(algorithm)?; let root = self.state.app_root(root_epoch)?; @@ -350,7 +350,7 @@ mod tests { GetKeyRequest { domain: domain.to_string(), algorithm: algorithm.to_string(), - root_epoch: 0, + root_epoch: None, } } @@ -363,15 +363,15 @@ mod tests { ..get_key_request("storage-encryption", "secp256k1") }; - let current = handler().get_key(request(0)).await.unwrap(); - assert_eq!(current.root_epoch, 1); + let current = handler().get_key(request(None)).await.unwrap(); + assert_eq!(current.root_epoch, 0); assert_eq!( - handler().get_key(request(1)).await.unwrap().key, + handler().get_key(request(Some(0))).await.unwrap().key, current.key ); - let old = handler().get_key(request(2)).await.unwrap(); - assert_eq!(old.root_epoch, 2); + let old = handler().get_key(request(Some(1))).await.unwrap(); + assert_eq!(old.root_epoch, 1); assert_ne!(old.key, current.key); assert_eq!( old.key, @@ -385,7 +385,7 @@ mod tests { ); assert_eq!(old.signature_chain[1], vec![0xab; 65]); - assert!(handler().get_key(request(3)).await.is_err()); + assert!(handler().get_key(request(Some(2))).await.is_err()); } /// The fixture's app root key is what `keys::tests` committed its vectors diff --git a/dstack/kms/auth-eth-bun/index.test.ts b/dstack/kms/auth-eth-bun/index.test.ts index d024e1ec5..658c06fc7 100644 --- a/dstack/kms/auth-eth-bun/index.test.ts +++ b/dstack/kms/auth-eth-bun/index.test.ts @@ -246,7 +246,7 @@ describe('API Compatibility Tests', () => { case 'isAppAllowed': return [true, '']; case 'gatewayAppId': return '0xgateway'; case 'appImplementation': return '0x9876543210987654321098765432109876543210'; - case 'latestRootEpoch': return 2n; + case 'rootEpochCount': return 2n; case 'rootEpoch': return [`0xa${params.args[0]}`, `0xb${params.args[0]}`, '0xc0']; } }); @@ -256,12 +256,12 @@ describe('API Compatibility Tests', () => { headers: { 'Content-Type': 'application/json' }, body: JSON.stringify(validBootInfo), })); - expect((await boot.json()).rootEpoch).toBe(2); + expect((await boot.json()).rootEpoch).toBe(1); const info = await appFetch(new Request('http://localhost:3001/')); expect((await info.json()).rootEpochs).toEqual([ + { caPubkey: '0xa0', k256Pubkey: '0xb0' }, { caPubkey: '0xa1', k256Pubkey: '0xb1' }, - { caPubkey: '0xa2', k256Pubkey: '0xb2' }, ]); }); @@ -270,7 +270,7 @@ describe('API Compatibility Tests', () => { switch (params.functionName) { case 'isAppAllowed': return [true, '']; case 'gatewayAppId': return '0xgateway'; - case 'latestRootEpoch': + case 'rootEpochCount': throw Object.assign(new Error('call failed'), { cause: { name: 'ContractFunctionZeroDataError' }, }); @@ -507,7 +507,7 @@ describe('Authorization freshness and domain binding', () => { return decisions === 1 ? [true, 'initial allow'] : [false, 'policy changed']; } if (params.functionName === 'gatewayAppId') return 'gateway-app'; - if (params.functionName === 'latestRootEpoch') return 0n; + if (params.functionName === 'rootEpochCount') return 0n; throw new Error(`unexpected function ${params.functionName}`); }); @@ -527,7 +527,7 @@ describe('Authorization freshness and domain binding', () => { return [true, 'allowed']; } if (params.functionName === 'gatewayAppId') return 'gateway-app'; - if (params.functionName === 'latestRootEpoch') return 0n; + if (params.functionName === 'rootEpochCount') return 0n; throw new Error(`unexpected function ${params.functionName}`); }); @@ -552,7 +552,7 @@ describe('Authorization freshness and domain binding', () => { mockReadContract.mockImplementation((params) => { if (params.functionName === 'isAppAllowed') return [true, 'recovered']; if (params.functionName === 'gatewayAppId') return 'gateway-app'; - if (params.functionName === 'latestRootEpoch') return 0n; + if (params.functionName === 'rootEpochCount') return 0n; throw new Error(`unexpected function ${params.functionName}`); }); const recovered = await postApp(); @@ -583,7 +583,7 @@ describe('Ethereum finalized snapshot authorization', () => { expect(params.blockNumber).toBe(98n); if (params.functionName === 'isAppAllowed') return [true, 'finalized allow']; if (params.functionName === 'gatewayAppId') return 'gateway-app'; - if (params.functionName === 'latestRootEpoch') return 0n; + if (params.functionName === 'rootEpochCount') return 0n; throw new Error(`unexpected function ${params.functionName}`); }); @@ -603,7 +603,7 @@ describe('Ethereum finalized snapshot authorization', () => { return decisions === 1 ? [true, 'old canonical allow'] : [false, 'new canonical deny']; } if (params.functionName === 'gatewayAppId') return 'gateway-app'; - if (params.functionName === 'latestRootEpoch') return 0n; + if (params.functionName === 'rootEpochCount') return 0n; throw new Error(`unexpected function ${params.functionName}`); }); @@ -632,7 +632,7 @@ describe('Ethereum finalized snapshot authorization', () => { mockReadContract.mockImplementation((params) => { if (params.functionName === 'isAppAllowed') return [true, 'recovered']; if (params.functionName === 'gatewayAppId') return 'gateway-app'; - if (params.functionName === 'latestRootEpoch') return 0n; + if (params.functionName === 'rootEpochCount') return 0n; throw new Error(`unexpected function ${params.functionName}`); }); const recovered = await authorize(); @@ -675,7 +675,7 @@ describe('dstack-kms 0.6 bootAuth payload', () => { return [true, '']; } if (params.functionName === 'gatewayAppId') return 'gateway-app'; - if (params.functionName === 'latestRootEpoch') return 0n; + if (params.functionName === 'rootEpochCount') return 0n; throw new Error(`unexpected function ${params.functionName}`); }); }); diff --git a/dstack/kms/auth-eth-bun/index.ts b/dstack/kms/auth-eth-bun/index.ts index d1303fdeb..e9be9c35a 100644 --- a/dstack/kms/auth-eth-bun/index.ts +++ b/dstack/kms/auth-eth-bun/index.ts @@ -36,7 +36,7 @@ const BootResponseSchema = z.object({ isAllowed: z.boolean(), reason: z.string(), gatewayAppId: z.string(), - // Latest registered KMS root epoch (from 1); absent when none is registered. + // Latest registered KMS root epoch; absent when none is registered. rootEpoch: z.number().optional() }); @@ -114,7 +114,7 @@ const DSTACK_KMS_ABI = [ outputs: [{ name: '', type: 'address' }] }, { - name: 'latestRootEpoch', + name: 'rootEpochCount', type: 'function', stateMutability: 'view', inputs: [], @@ -225,42 +225,42 @@ class EthereumBackend { blockNumber }); - const latestEpoch = await this.getLatestRootEpoch(blockNumber); + const epochCount = await this.getRootEpochCount(blockNumber); return { isAllowed, reason, gatewayAppId: gatewayAppId as string, - ...(latestEpoch > 0 && { rootEpoch: latestEpoch }), + ...(epochCount > 0 && { rootEpoch: epochCount - 1 }), }; } - private async getLatestRootEpoch(blockNumber?: bigint): Promise { + private async getRootEpochCount(blockNumber?: bigint): Promise { try { - const latest = await this.client.readContract({ + const count = await this.client.readContract({ address: this.kmsContractAddr, abi: DSTACK_KMS_ABI, - functionName: 'latestRootEpoch', + functionName: 'rootEpochCount', blockNumber }); - return Number(latest ?? 0); + return Number(count ?? 0); } catch (error) { if (isMissingFunction(error)) return 0; throw error; } } - // Registered root epochs 1, 2, ... in order. + // Registered root epochs, indexed by epoch. async getRootEpochs(): Promise { const blockNumber = await this.finalizedBlockNumber(); - const latest = await this.getLatestRootEpoch(blockNumber); + const count = await this.getRootEpochCount(blockNumber); return Promise.all( - Array.from({ length: latest }, async (_, index) => { + Array.from({ length: count }, async (_, epoch) => { const [caPubkey, k256Pubkey] = await this.client.readContract({ address: this.kmsContractAddr, abi: DSTACK_KMS_ABI, functionName: 'rootEpoch', - args: [BigInt(index + 1)], + args: [BigInt(epoch)], blockNumber }); return { caPubkey, k256Pubkey }; diff --git a/dstack/kms/auth-eth-bun/openapi.json b/dstack/kms/auth-eth-bun/openapi.json index 88947e83c..5ca950e5d 100644 --- a/dstack/kms/auth-eth-bun/openapi.json +++ b/dstack/kms/auth-eth-bun/openapi.json @@ -176,7 +176,7 @@ }, "rootEpoch": { "type": "integer", - "description": "Latest registered KMS root epoch, counting from 1; absent when none is registered" + "description": "Latest registered KMS root epoch (the original roots are 0); absent when none is registered" } } }, @@ -223,7 +223,7 @@ }, "rootEpochs": { "type": "array", - "description": "Registered KMS root epochs 1, 2, ... in order", + "description": "Registered KMS root epochs, indexed by epoch", "items": { "type": "object", "required": ["caPubkey", "k256Pubkey"], diff --git a/dstack/kms/auth-eth/contracts/DstackKms.sol b/dstack/kms/auth-eth/contracts/DstackKms.sol index d32f32646..930512d24 100644 --- a/dstack/kms/auth-eth/contracts/DstackKms.sol +++ b/dstack/kms/auth-eth/contracts/DstackKms.sol @@ -53,8 +53,9 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E bytes attestation; } - // Root public keys of KMS root epochs 1, 2, ... in order. Append-only: a - // registered epoch is never changed, so the last entry is current. + // Root public keys of each KMS root epoch, indexed by epoch; epoch 0 is the + // original roots. Append-only: a registered epoch is never changed, so the + // last entry is current. RootEpoch[] private _rootEpochs; // Events @@ -130,8 +131,7 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E } // Register the root public keys of the next KMS root epoch, with the - // attestation proving where they were generated; epochs count from 1, the - // original roots + // attestation proving where they were generated function registerRootEpoch( bytes calldata caPubkey, bytes calldata k256Pubkey, @@ -143,14 +143,14 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E { require(caPubkey.length > 0 && k256Pubkey.length > 0, "Empty epoch public key"); require(attestation.length > 0, "Empty epoch attestation"); - _rootEpochs.push(RootEpoch(caPubkey, k256Pubkey, attestation)); epoch = _rootEpochs.length; + _rootEpochs.push(RootEpoch(caPubkey, k256Pubkey, attestation)); emit RootEpochRegistered(epoch, caPubkey, k256Pubkey); _emitPolicy("root-epoch", keccak256(abi.encode(caPubkey, k256Pubkey)), true); } - // The latest registered root epoch; 0 when none is registered - function latestRootEpoch() external view returns (uint256) { + // Number of registered root epochs; the latest is this minus one + function rootEpochCount() external view returns (uint256) { return _rootEpochs.length; } @@ -159,8 +159,8 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E view returns (bytes memory caPubkey, bytes memory k256Pubkey, bytes memory attestation) { - require(epoch > 0 && epoch <= _rootEpochs.length, "Unknown root epoch"); - RootEpoch storage keys = _rootEpochs[epoch - 1]; + require(epoch < _rootEpochs.length, "Unknown root epoch"); + RootEpoch storage keys = _rootEpochs[epoch]; return (keys.caPubkey, keys.k256Pubkey, keys.attestation); } diff --git a/dstack/kms/auth-eth/script/Query.s.sol b/dstack/kms/auth-eth/script/Query.s.sol index e7f4404c4..c4bc628be 100644 --- a/dstack/kms/auth-eth/script/Query.s.sol +++ b/dstack/kms/auth-eth/script/Query.s.sol @@ -211,9 +211,9 @@ contract GetKmsSettings is BaseQueryScript { console.log(" Quote length:", quote.length); console.log(" Eventlog length:", eventlog.length); - uint256 latestEpoch = kms.latestRootEpoch(); - console.log("\nLatest root epoch:", latestEpoch); - for (uint256 epoch = 1; epoch <= latestEpoch; epoch++) { + uint256 epochCount = kms.rootEpochCount(); + console.log("\nRegistered root epochs:", epochCount); + for (uint256 epoch = 0; epoch < epochCount; epoch++) { (bytes memory epochCa, bytes memory epochK256, bytes memory attestation) = kms.rootEpoch(epoch); console.log(" Epoch", epoch); console.log(" CA Pubkey:", vm.toString(epochCa)); diff --git a/dstack/kms/auth-eth/script/README.md b/dstack/kms/auth-eth/script/README.md index 9aaae12a1..37ff17ea1 100644 --- a/dstack/kms/auth-eth/script/README.md +++ b/dstack/kms/auth-eth/script/README.md @@ -59,7 +59,7 @@ forge script script/Manage.s.sol:SetKmsInfo --broadcast --rpc-url $RPC_URL ``` #### Register a KMS Root Epoch -Appends the next root epoch (the original roots are epoch 1). `ATTESTATION` is +Appends the next root epoch (the original roots are epoch 0). `ATTESTATION` is the one the KMS returned when it generated the keys (Bootstrap, or Onboard with `rotate`); verify it before registering. Registered epochs cannot be changed. ```bash diff --git a/dstack/kms/auth-eth/src/ethereum.ts b/dstack/kms/auth-eth/src/ethereum.ts index fdb0e09c5..e44b88d73 100644 --- a/dstack/kms/auth-eth/src/ethereum.ts +++ b/dstack/kms/auth-eth/src/ethereum.ts @@ -11,7 +11,7 @@ const DSTACK_KMS_ABI = [ "function isKmsAllowed((address appId,bytes32 composeHash,address instanceId,bytes32 deviceId,bytes32 mrAggregated,bytes32 mrSystem,bytes32 osImageHash,string tcbStatus,string[] advisoryIds) bootInfo) view returns (bool, string)", "function gatewayAppId() view returns (string)", "function appImplementation() view returns (address)", - "function latestRootEpoch() view returns (uint256)", + "function rootEpochCount() view returns (uint256)", "function rootEpoch(uint256 epoch) view returns (bytes caPubkey, bytes k256Pubkey, bytes attestation)" ]; @@ -67,18 +67,18 @@ export class EthereumBackend { } const [isAllowed, reason] = response; const gatewayAppId = await this.kmsContract.gatewayAppId({ blockTag }); - const latestEpoch = await this.getLatestRootEpoch(blockTag); + const epochCount = await this.getRootEpochCount(blockTag); return { isAllowed, reason, gatewayAppId, - ...(latestEpoch > 0 && { rootEpoch: latestEpoch }), + ...(epochCount > 0 && { rootEpoch: epochCount - 1 }), } } - private async getLatestRootEpoch(blockTag?: number): Promise { + private async getRootEpochCount(blockTag?: number): Promise { try { - return Number(await this.kmsContract.latestRootEpoch({ blockTag })); + return Number(await this.kmsContract.rootEpochCount({ blockTag })); } catch (error) { if (isMissingFunction(error)) return 0; throw error; @@ -87,10 +87,10 @@ export class EthereumBackend { async getRootEpochs(): Promise { const blockTag = await this.provider.getBlockNumber(); - const latest = await this.getLatestRootEpoch(blockTag); + const count = await this.getRootEpochCount(blockTag); return Promise.all( - Array.from({ length: latest }, async (_, index) => { - const [caPubkey, k256Pubkey] = await this.kmsContract.rootEpoch(index + 1, { blockTag }); + Array.from({ length: count }, async (_, epoch) => { + const [caPubkey, k256Pubkey] = await this.kmsContract.rootEpoch(epoch, { blockTag }); return { caPubkey, k256Pubkey }; }) ); diff --git a/dstack/kms/auth-eth/src/types.ts b/dstack/kms/auth-eth/src/types.ts index f74defa2e..1ed4ad1a2 100644 --- a/dstack/kms/auth-eth/src/types.ts +++ b/dstack/kms/auth-eth/src/types.ts @@ -18,7 +18,7 @@ export interface BootResponse { isAllowed: boolean; gatewayAppId: string; reason: string; - // Latest registered KMS root epoch (from 1); absent when none is registered. + // Latest registered KMS root epoch; absent when none is registered. rootEpoch?: number; } diff --git a/dstack/kms/auth-eth/test/DstackKms.t.sol b/dstack/kms/auth-eth/test/DstackKms.t.sol index c32250f3e..6278c663e 100644 --- a/dstack/kms/auth-eth/test/DstackKms.t.sol +++ b/dstack/kms/auth-eth/test/DstackKms.t.sol @@ -68,28 +68,26 @@ contract DstackKmsTest is Test { } function test_RegisterRootEpochAppends() public { - assertEq(kms.latestRootEpoch(), 0); + assertEq(kms.rootEpochCount(), 0); vm.startPrank(owner); vm.expectEmit(true, true, true, true); - emit RootEpochRegistered(1, hex"a0", hex"b0"); - assertEq(kms.registerRootEpoch(hex"a0", hex"b0", hex"c0"), 1); - assertEq(kms.registerRootEpoch(hex"a1", hex"b1", hex"c1"), 2); + emit RootEpochRegistered(0, hex"a0", hex"b0"); + assertEq(kms.registerRootEpoch(hex"a0", hex"b0", hex"c0"), 0); + assertEq(kms.registerRootEpoch(hex"a1", hex"b1", hex"c1"), 1); vm.stopPrank(); - assertEq(kms.latestRootEpoch(), 2); - (bytes memory ca, bytes memory k256, bytes memory attestation) = kms.rootEpoch(1); + assertEq(kms.rootEpochCount(), 2); + (bytes memory ca, bytes memory k256, bytes memory attestation) = kms.rootEpoch(0); assertEq(ca, hex"a0"); assertEq(k256, hex"b0"); assertEq(attestation, hex"c0"); - (ca, k256, attestation) = kms.rootEpoch(2); + (ca, k256, attestation) = kms.rootEpoch(1); assertEq(ca, hex"a1"); assertEq(k256, hex"b1"); assertEq(attestation, hex"c1"); vm.expectRevert("Unknown root epoch"); - kms.rootEpoch(0); - vm.expectRevert("Unknown root epoch"); - kms.rootEpoch(3); + kms.rootEpoch(2); } function test_RegisterRootEpochOnlyOwner() public { diff --git a/dstack/kms/auth-mock/README.md b/dstack/kms/auth-mock/README.md index 51aff51f6..93be1eaa8 100644 --- a/dstack/kms/auth-mock/README.md +++ b/dstack/kms/auth-mock/README.md @@ -72,7 +72,7 @@ bun run check - `MOCK_GATEWAY_APP_ID` - mock gateway app ID (default: 0xmockgateway1234567890123456789012345678) - `MOCK_CHAIN_ID` - mock blockchain chain ID (default: 1337) - `MOCK_APP_IMPLEMENTATION` - mock app implementation address (default: 0xmockapp9876543210987654321098765432109) -- `MOCK_ROOT_EPOCHS` - JSON array of `{"caPubkey", "k256Pubkey"}` for KMS root epochs 1, 2, ... (default: none registered) +- `MOCK_ROOT_EPOCHS` - JSON array of `{"caPubkey", "k256Pubkey"}`, indexed by KMS root epoch (default: none registered) ## API endpoints diff --git a/dstack/kms/auth-mock/index.test.ts b/dstack/kms/auth-mock/index.test.ts index 872705107..3cc672660 100644 --- a/dstack/kms/auth-mock/index.test.ts +++ b/dstack/kms/auth-mock/index.test.ts @@ -173,7 +173,7 @@ describe('Mock Backend Tests', () => { composeHash: '0x04', instanceId: '0x05', deviceId: '0x06', }), })); - expect((await boot.json()).rootEpoch).toBe(1); + expect((await boot.json()).rootEpoch).toBe(0); const info = await appFetch(new Request('http://localhost:3002/')); expect((await info.json()).rootEpochs).toEqual(rootEpochs); } finally { diff --git a/dstack/kms/auth-mock/index.ts b/dstack/kms/auth-mock/index.ts index 23f84944b..eff46cf85 100644 --- a/dstack/kms/auth-mock/index.ts +++ b/dstack/kms/auth-mock/index.ts @@ -34,7 +34,7 @@ const BootResponseSchema = z.object({ const RootEpochsSchema = z.array(z.object({ caPubkey: z.string(), k256Pubkey: z.string() })); -// MOCK_ROOT_EPOCHS: JSON array of {caPubkey, k256Pubkey} for root epochs 1, 2, ... +// MOCK_ROOT_EPOCHS: JSON array of {caPubkey, k256Pubkey}, indexed by root epoch function getRootEpochs(): z.infer { return RootEpochsSchema.parse(JSON.parse(process.env.MOCK_ROOT_EPOCHS || '[]')); } @@ -79,8 +79,8 @@ class MockBackend { async checkBoot(bootInfo: BootInfo, isKms: boolean): Promise { const response = this.evaluateBoot(bootInfo, isKms); - const latestEpoch = getRootEpochs().length; - return latestEpoch > 0 ? { ...response, rootEpoch: latestEpoch } : response; + const epochCount = getRootEpochs().length; + return epochCount > 0 ? { ...response, rootEpoch: epochCount - 1 } : response; } private evaluateBoot(bootInfo: BootInfo, isKms: boolean): BootResponse { diff --git a/dstack/kms/auth-mock/openapi.json b/dstack/kms/auth-mock/openapi.json index 66dd817e4..f197a499a 100644 --- a/dstack/kms/auth-mock/openapi.json +++ b/dstack/kms/auth-mock/openapi.json @@ -227,7 +227,7 @@ }, "rootEpochs": { "type": "array", - "description": "KMS root epochs 1, 2, ... from MOCK_ROOT_EPOCHS", + "description": "KMS root epochs from MOCK_ROOT_EPOCHS, indexed by epoch", "items": { "type": "object", "properties": { diff --git a/dstack/kms/auth-simple/README.md b/dstack/kms/auth-simple/README.md index a7937e8d8..017d01f2d 100644 --- a/dstack/kms/auth-simple/README.md +++ b/dstack/kms/auth-simple/README.md @@ -107,7 +107,7 @@ raw NitroTPM evidence by itself. | `apps..composeHashes` | No | Allowed compose hashes for this app | | `apps..devices` | No | Allowed device IDs for this app. Empty denies all unless `allowAnyDevice` is true. | | `apps..allowAnyDevice` | No | If true, skip device ID check for this app | -| `rootEpochs` | No | KMS root public keys of root epochs 1, 2, ... in order, `[{"caPubkey": "0x...", "k256Pubkey": "0x..."}]`. Epoch 1 is the original roots. Append only. When set, the KMS checks its keyring against it and boot responses report the latest epoch as `rootEpoch`. | +| `rootEpochs` | No | KMS root public keys of each root epoch, indexed by epoch, `[{"caPubkey": "0x...", "k256Pubkey": "0x..."}]`. Epoch 0 is the original roots. Append only. When set, the KMS checks its keyring against it and boot responses report the latest epoch as `rootEpoch`. | For experimental AMD SEV-SNP dry-run authorization, keep the default fail-closed TCB policy unless you intentionally want the auth webhook to accept non-up-to-date verifier-derived SNP `BootInfo`. To exercise the dry-run path without enabling key release, allowlist the recomputed SNP `mrAggregated`, `osImageHash`, app/compose identity, device/chip identity, and any non-default `allowedTcbStatuses`/`allowedAdvisoryIds` values explicitly. KMS still rejects SNP before returning app keys, KMS keys, or app certificates. diff --git a/dstack/kms/auth-simple/index.test.ts b/dstack/kms/auth-simple/index.test.ts index 043ca22d4..4ba189e9c 100644 --- a/dstack/kms/auth-simple/index.test.ts +++ b/dstack/kms/auth-simple/index.test.ts @@ -460,7 +460,7 @@ describe('auth-simple', () => { headers: { 'Content-Type': 'application/json' }, body: JSON.stringify(baseBootInfo) })); - expect((await boot.json()).rootEpoch).toBe(2); + expect((await boot.json()).rootEpoch).toBe(1); const info = await app.fetch(new Request('http://localhost/', { method: 'GET' })); expect((await info.json()).rootEpochs).toEqual(rootEpochs); diff --git a/dstack/kms/auth-simple/index.ts b/dstack/kms/auth-simple/index.ts index a7197eea1..84e875a3e 100644 --- a/dstack/kms/auth-simple/index.ts +++ b/dstack/kms/auth-simple/index.ts @@ -59,7 +59,7 @@ const AuthConfigSchema = z.object({ allowedAdvisoryIds: z.array(z.string()).default([]), kms: KmsConfigSchema.default({}), apps: z.record(z.string(), AppConfigSchema).default({}), - // Root public keys of KMS root epochs 1, 2, ... in order; append only. + // Root public keys of each KMS root epoch, indexed by epoch; append only. rootEpochs: z.array(RootEpochSchema).default([]) }); @@ -102,8 +102,8 @@ class ConfigBackend { async checkBoot(bootInfo: BootInfo, isKms: boolean): Promise { const config = this.loadConfig(); const response = this.evaluateBoot(bootInfo, isKms, config); - const latestEpoch = config.rootEpochs.length; - return latestEpoch > 0 ? { ...response, rootEpoch: latestEpoch } : response; + const epochCount = config.rootEpochs.length; + return epochCount > 0 ? { ...response, rootEpoch: epochCount - 1 } : response; } private evaluateBoot(bootInfo: BootInfo, isKms: boolean, config: AuthConfig): BootResponse { diff --git a/dstack/kms/kms.toml b/dstack/kms/kms.toml index 48efa121f..378372278 100644 --- a/dstack/kms/kms.toml +++ b/dstack/kms/kms.toml @@ -125,10 +125,10 @@ timeout = "60s" gateway_app_id = "any" [core.key_rotation] -# Root key epoch served as current; the original roots are epoch 1. Raising it requires a KMS onboarded with +# Root key epoch served as current; the original roots are epoch 0. Raising it requires a KMS onboarded with # `rotate` (or from a source already at this epoch); editing it alone fails. # Older epochs stay available so CVMs can migrate their data at any time. -current_epoch = 1 +current_epoch = 0 [core.onboard] enabled = true diff --git a/dstack/kms/rpc/proto/kms_rpc.proto b/dstack/kms/rpc/proto/kms_rpc.proto index ca2031735..5b5ba57aa 100644 --- a/dstack/kms/rpc/proto/kms_rpc.proto +++ b/dstack/kms/rpc/proto/kms_rpc.proto @@ -11,10 +11,9 @@ package kms; message GetAppKeyRequest { uint32 api_version = 1; string vm_config = 2; - // Root key epoch to derive from: 0 (default) = the current epoch, N = epoch N. - // Epoch 1 is the original root keys. Older epochs remain available so CVMs - // can migrate. - uint32 root_epoch = 3; + // Root key epoch to derive from; unset means the current epoch. Epoch 0 is + // the original root keys. Older epochs remain available so CVMs can migrate. + optional uint32 root_epoch = 3; } message AppId { @@ -50,8 +49,7 @@ message AppKeyResponse { string gateway_app_id = 7; // OS Image hash bytes os_image_hash = 8; - // Root key epoch all epoch-dependent fields above were derived from. 0 from a - // KMS that predates root epochs, which means epoch 1. + // Root key epoch all epoch-dependent fields above were derived from. uint32 root_epoch = 9; } @@ -77,8 +75,7 @@ message GetKmsKeyRequest { message KmsKeys { string ca_key = 1; bytes k256_key = 2; - // Root key epoch of this pair. 0 from a KMS that predates root epochs, which - // means epoch 1. + // Root key epoch of this pair. uint32 epoch = 3; } @@ -168,7 +165,7 @@ message OnboardRequest { string domain = 2; // Optional bearer token for the source KMS admin listener. string source_token = 3; - // Import epochs 1..N-1 from the source and generate epoch N locally, where N + // Import epochs 0..N-1 from the source and generate epoch N locally, where N // is this KMS's configured current epoch. bool rotate = 4; } diff --git a/dstack/kms/src/config.rs b/dstack/kms/src/config.rs index 1476c98ff..4eac5be85 100644 --- a/dstack/kms/src/config.rs +++ b/dstack/kms/src/config.rs @@ -176,7 +176,7 @@ impl KmsConfig { self.cert_dir.join(K256_KEY) } - /// Records which epoch the top-level root keys belong to. Absent means 1. + /// Records which epoch the top-level root keys belong to. Absent means 0. pub fn root_epoch(&self) -> PathBuf { self.cert_dir.join(ROOT_EPOCH) } @@ -241,34 +241,22 @@ pub(crate) struct RootKeyPaths { pub k256_key: PathBuf, } -#[derive(Debug, Clone, Deserialize)] +#[derive(Debug, Clone, Default, Deserialize)] pub(crate) struct KeyRotationConfig { - /// Root key epoch served as current; epochs count from 1, the original - /// roots. Every KMS holds all epochs up to it. - #[serde(default = "first_epoch")] + /// Root key epoch served as current; epoch 0 is the original roots. Every + /// KMS holds all epochs up to it. + #[serde(default)] pub current_epoch: u32, } -impl Default for KeyRotationConfig { - fn default() -> Self { - Self { - current_epoch: first_epoch(), - } - } -} - -fn first_epoch() -> u32 { - 1 -} - impl KeyRotationConfig { - /// Resolve a root epoch selector: 0 is the current epoch. - pub fn resolve_epoch(&self, requested: u32) -> Result { + /// Resolve a root epoch selector: unset is the current epoch. + pub fn resolve_epoch(&self, requested: Option) -> Result { let current = self.current_epoch; match requested { - 0 => Ok(current), - epoch if epoch <= current => Ok(epoch), - epoch => bail!("unknown root epoch {epoch}; the current epoch is {current}"), + None => Ok(current), + Some(epoch) if epoch <= current => Ok(epoch), + Some(epoch) => bail!("unknown root epoch {epoch}; the current epoch is {current}"), } } } @@ -305,11 +293,11 @@ mod tests { #[test] fn root_epoch_selector_resolution() { let cfg = KeyRotationConfig { current_epoch: 2 }; - assert_eq!(cfg.resolve_epoch(0).unwrap(), 2); - assert_eq!(cfg.resolve_epoch(1).unwrap(), 1); - assert_eq!(cfg.resolve_epoch(2).unwrap(), 2); - assert!(cfg.resolve_epoch(3).is_err()); - assert_eq!(KeyRotationConfig::default().resolve_epoch(0).unwrap(), 1); + assert_eq!(cfg.resolve_epoch(None).unwrap(), 2); + assert_eq!(cfg.resolve_epoch(Some(0)).unwrap(), 0); + assert_eq!(cfg.resolve_epoch(Some(2)).unwrap(), 2); + assert!(cfg.resolve_epoch(Some(3)).is_err()); + assert_eq!(KeyRotationConfig::default().resolve_epoch(None).unwrap(), 0); } #[test] diff --git a/dstack/kms/src/main_service.rs b/dstack/kms/src/main_service.rs index bd21b3236..d159599c7 100644 --- a/dstack/kms/src/main_service.rs +++ b/dstack/kms/src/main_service.rs @@ -93,7 +93,7 @@ impl RootKeys { pub struct KmsStateInner { config: KmsConfig, - /// Root keys of epochs 1, 2, ... in order; the last is current. + /// Root keys of every epoch, indexed by epoch; the last is current. roots: Vec, temp_ca_cert: String, temp_ca_key: String, @@ -153,9 +153,6 @@ impl KmsState { pub fn new(config: KmsConfig) -> Result { let current_epoch = config.key_rotation.current_epoch; - if current_epoch == 0 { - bail!("key_rotation.current_epoch counts from 1"); - } let stored_epoch = read_root_epoch(&config)?; if stored_epoch != current_epoch { bail!( @@ -163,7 +160,7 @@ impl KmsState { is {current_epoch}; a new epoch must be imported or generated by onboarding" ); } - let roots = (1..=current_epoch) + let roots = (0..=current_epoch) .map(|epoch| { RootKeys::load(&config.root_key_paths(epoch)) .with_context(|| format!("failed to load root keys of epoch {epoch}")) @@ -217,21 +214,21 @@ impl KmsState { /// The current epoch's root keys. fn root(&self) -> &RootKeys { - self.roots.last().expect("the keyring always holds epoch 1") + self.roots.last().expect("the keyring always holds epoch 0") } /// The root keys of `epoch`, which must be one this KMS serves. fn epoch_root(&self, epoch: u32) -> &RootKeys { - &self.roots[epoch as usize - 1] + &self.roots[epoch as usize] } /// Refuse to serve a keyring governance has not registered. Deployments - /// that never registered epochs keep working while they stay at epoch 1. + /// that never registered epochs keep working while they stay at epoch 0. pub async fn ensure_keyring_registered(&self) -> Result<()> { let current = self.config.key_rotation.current_epoch; let info = match self.config.auth_api.get_info().await { Ok(info) => info, - Err(err) if current == 1 => { + Err(err) if current == 0 => { warn!("skipping key epoch registration check: {err:#}"); return Ok(()); } @@ -240,16 +237,16 @@ impl KmsState { let Some(registered) = info.root_epochs else { return Ok(()); }; - if registered.is_empty() && current == 1 { + if registered.is_empty() && current == 0 { return Ok(()); } - if registered.len() != current as usize { + if registered.len() != current as usize + 1 { bail!( "this KMS serves key epoch {current}, but {} epochs are registered", registered.len() ); } - for (epoch, (root, registered)) in (1..).zip(self.roots.iter().zip(®istered)) { + for (epoch, (root, registered)) in (0..).zip(self.roots.iter().zip(®istered)) { if root.pubkeys() != *registered { bail!("root keys of epoch {epoch} do not match the registered public keys"); } @@ -264,7 +261,7 @@ impl KmsState { Some(epoch) if epoch != current => { bail!("the registered key epoch is {epoch}, but this KMS serves epoch {current}") } - None if current > 1 && !self.config.auth_api.is_dev() => { + None if current > 0 && !self.config.auth_api.is_dev() => { bail!("the auth API reports no key epoch, but this KMS serves epoch {current}") } _ => Ok(()), @@ -275,7 +272,7 @@ impl KmsState { pub(crate) fn read_root_epoch(config: &KmsConfig) -> Result { match fs::read_to_string(config.root_epoch()) { Ok(epoch) => epoch.trim().parse().context("invalid root epoch file"), - Err(err) if err.kind() == std::io::ErrorKind::NotFound => Ok(1), + Err(err) if err.kind() == std::io::ErrorKind::NotFound => Ok(0), Err(err) => Err(err).context("failed to read root epoch file"), } } @@ -301,7 +298,7 @@ impl RpcHandler { ensure_key_release_allowed(&info, (&self.state.config).into())?; Ok(KmsKeyResponse { temp_ca_key: self.state.inner.temp_ca_key.clone(), - keys: (1..) + keys: (0..) .zip(&self.state.roots) .map(|(epoch, root)| KmsKeys { ca_key: root.ca.key.serialize_pem(), diff --git a/dstack/kms/src/onboard_service.rs b/dstack/kms/src/onboard_service.rs index fe914070b..0ed12c838 100644 --- a/dstack/kms/src/onboard_service.rs +++ b/dstack/kms/src/onboard_service.rs @@ -552,7 +552,7 @@ mod tests { temp_ca_key: KeyPair::generate_for(&PKCS_ECDSA_P256_SHA256) .unwrap() .serialize_pem(), - keys: (1..=epochs) + keys: (0..epochs) .map(|epoch| KmsKeys { ca_key: KeyPair::generate_for(&PKCS_ECDSA_P256_SHA256) .unwrap() @@ -575,7 +575,7 @@ mod tests { #[rocket::async_test] async fn join_imports_and_stores_every_epoch() { let cert_dir = tempfile::tempdir().unwrap(); - let cfg = epoch_config(cert_dir.path(), 3); + let cfg = epoch_config(cert_dir.path(), 2); let res = handover(3); let sent: Vec<_> = res.keys.iter().map(|k| k.k256_key.clone()).collect(); @@ -584,30 +584,18 @@ mod tests { .unwrap(); keys.store(&cfg).unwrap(); - assert_eq!(crate::main_service::read_root_epoch(&cfg).unwrap(), 3); - for (epoch, k256) in (1..).zip(&sent) { + assert_eq!(crate::main_service::read_root_epoch(&cfg).unwrap(), 2); + for (epoch, k256) in (0..).zip(&sent) { assert_eq!(&fs::read(cfg.root_key_paths(epoch).k256_key).unwrap(), k256); } assert_eq!(&fs::read(cfg.k256_key()).unwrap(), &sent[2]); crate::main_service::KmsState::new(cfg.clone()).unwrap(); // Editing the configured epoch cannot relabel the stored roots. - let err = crate::main_service::KmsState::new(epoch_config(cert_dir.path(), 4)) + let err = crate::main_service::KmsState::new(epoch_config(cert_dir.path(), 3)) .err() .unwrap(); - assert!(format!("{err:#}").contains("belong to epoch 3"), "{err:#}"); - } - - #[rocket::async_test] - async fn join_accepts_a_source_predating_root_epochs() { - let cert_dir = tempfile::tempdir().unwrap(); - let cfg = epoch_config(cert_dir.path(), 1); - let mut res = handover(1); - res.keys[0].epoch = 0; - let keys = Keys::from_key_response(&cfg, res, "kms.example.com", false, None) - .await - .unwrap(); - assert_eq!(keys.epoch, 1); + assert!(format!("{err:#}").contains("belong to epoch 2"), "{err:#}"); } fn pubkeys_of(res: &KmsKeyResponse) -> Vec { @@ -626,7 +614,7 @@ mod tests { async fn joined_state(res: KmsKeyResponse, url: String) -> crate::main_service::KmsState { let cert_dir = tempfile::tempdir().unwrap(); - let mut cfg = epoch_config(cert_dir.path(), res.keys.len() as u32); + let mut cfg = epoch_config(cert_dir.path(), res.keys.len() as u32 - 1); Keys::from_key_response(&cfg, res, "kms.example.com", false, None) .await .unwrap() @@ -687,15 +675,15 @@ mod tests { reason: String::new(), root_epoch, }; - assert!(state.ensure_epoch_current(&response(Some(2))).is_ok()); - assert!(state.ensure_epoch_current(&response(Some(3))).is_err()); + assert!(state.ensure_epoch_current(&response(Some(1))).is_ok()); + assert!(state.ensure_epoch_current(&response(Some(2))).is_err()); assert!(state.ensure_epoch_current(&response(None)).is_err()); } #[rocket::async_test] async fn rotation_imports_history_and_generates_the_new_epoch() { let cert_dir = tempfile::tempdir().unwrap(); - let cfg = epoch_config(cert_dir.path(), 3); + let cfg = epoch_config(cert_dir.path(), 2); let res = handover(2); let sent: Vec<_> = res.keys.iter().map(|k| k.k256_key.clone()).collect(); @@ -704,7 +692,7 @@ mod tests { .unwrap(); keys.store(&cfg).unwrap(); - assert_eq!(fs::read(cfg.root_key_paths(2).k256_key).unwrap(), sent[1]); + assert_eq!(fs::read(cfg.root_key_paths(1).k256_key).unwrap(), sent[1]); let current = fs::read(cfg.k256_key()).unwrap(); assert!(!sent.contains(¤t)); } @@ -712,7 +700,7 @@ mod tests { #[rocket::async_test] async fn onboarding_rejects_incomplete_or_extra_history() { let cert_dir = tempfile::tempdir().unwrap(); - let cfg = epoch_config(cert_dir.path(), 2); + let cfg = epoch_config(cert_dir.path(), 1); for (epochs, rotate) in [(1, false), (3, false), (2, true), (0, true)] { let res = handover(epochs); assert!( @@ -722,7 +710,7 @@ mod tests { "epochs={epochs} rotate={rotate}" ); } - let cfg = epoch_config(cert_dir.path(), 1); + let cfg = epoch_config(cert_dir.path(), 0); assert!( Keys::from_key_response(&cfg, handover(0), "kms.example.com", true, None) .await @@ -733,8 +721,8 @@ mod tests { #[rocket::async_test] async fn onboarding_requires_the_registered_keyring() { let cert_dir = tempfile::tempdir().unwrap(); - let join = epoch_config(cert_dir.path(), 2); - let rotate = epoch_config(cert_dir.path(), 3); + let join = epoch_config(cert_dir.path(), 1); + let rotate = epoch_config(cert_dir.path(), 2); let res = handover(2); let registered = pubkeys_of(&res); let onboard = |cfg: &KmsConfig, rotate: bool, registered: Vec| { @@ -756,14 +744,14 @@ mod tests { other[0] = pubkeys_of(&handover(1))[0].clone(); assert!(onboard(&join, false, other.clone()).await.is_err()); assert!(onboard(&rotate, true, other).await.is_err()); - // Epochs must be registered before rotating past epoch 1. + // Epochs must be registered before rotating past epoch 0. assert!(onboard(&rotate, true, Vec::new()).await.is_err()); } #[rocket::async_test] - async fn an_unregistered_deployment_can_still_join_at_epoch_1() { + async fn an_unregistered_deployment_can_still_join_at_epoch_0() { let cert_dir = tempfile::tempdir().unwrap(); - let cfg = epoch_config(cert_dir.path(), 1); + let cfg = epoch_config(cert_dir.path(), 0); assert!( Keys::from_key_response(&cfg, handover(1), "kms.example.com", false, Some(&[])) .await @@ -774,7 +762,7 @@ mod tests { #[rocket::async_test] async fn bootstrap_refuses_a_later_epoch() { let cert_dir = tempfile::tempdir().unwrap(); - let cfg = epoch_config(cert_dir.path(), 2); + let cfg = epoch_config(cert_dir.path(), 1); let verifier = AttestationVerifier::load(&cfg.attestation).unwrap(); assert!(bootstrap_keys(&cfg, &verifier).await.is_err()); assert!(!cfg.root_keys_exist()); @@ -870,7 +858,7 @@ impl HistoricalRoot { } struct Keys { - /// Epoch of the root keys below; `history` holds epochs `1..epoch`. + /// Epoch of the root keys below; `history` holds epochs `0..epoch`. epoch: u32, history: Vec, k256_key: SigningKey, @@ -900,7 +888,7 @@ impl Keys { .await } - /// Attest the current root public keys. Epoch 1 keeps the original + /// Attest the current root public keys. Epoch 0 keeps the original /// bootstrap encoding; later epochs also bind the epoch number. async fn attest(&self) -> Result { let ca_pubkey = self.ca_key.public_key_der(); @@ -960,7 +948,7 @@ impl Keys { .build() .signed_by(&ca_cert, &ca_key)?; Ok(Keys { - epoch: 1, + epoch: 0, history: Vec::new(), k256_key, tmp_ca_key, @@ -1085,7 +1073,7 @@ impl Keys { /// `registered` is the governance registry; `None` means the auth backend /// has none (dev mode), and an empty one is accepted only while joining at - /// epoch 1, as deployments that never registered epochs did before. + /// epoch 0, as deployments that never registered epochs did before. async fn from_key_response( cfg: &KmsConfig, keys_res: KmsKeyResponse, @@ -1096,17 +1084,15 @@ impl Keys { let epoch = cfg.key_rotation.current_epoch; // A rotation imports every epoch before the new one; a join imports all. let last = if rotate { - if epoch < 2 { - bail!("rotation requires key_rotation.current_epoch > 1"); - } - epoch - 1 + epoch + .checked_sub(1) + .context("rotation requires key_rotation.current_epoch > 0")? } else { epoch }; - // A source predating root epochs sends its single pair as epoch 0. - let received: Vec = keys_res.keys.iter().map(|k| k.epoch.max(1)).collect(); - if !received.iter().copied().eq(1..=last) { - bail!("source KMS must hand over root key epochs 1..={last}, got {received:?}"); + let received: Vec = keys_res.keys.iter().map(|k| k.epoch).collect(); + if !received.iter().copied().eq(0..=last) { + bail!("source KMS must hand over root key epochs 0..={last}, got {received:?}"); } let mut roots = keys_res .keys @@ -1153,7 +1139,7 @@ impl Keys { /// The root keys go last: see [`KmsConfig::root_keys_exist`]. fn store(&self, cfg: &KmsConfig) -> Result<()> { - for (epoch, root) in (1..).zip(&self.history) { + for (epoch, root) in (0..).zip(&self.history) { let paths = cfg.root_key_paths(epoch); if let Some(dir) = paths.ca_key.parent() { fs::create_dir_all(dir)?; @@ -1285,7 +1271,7 @@ pub(crate) async fn bootstrap_keys(cfg: &KmsConfig, verifier: &AttestationVerifi Ok(()) } -/// Check that the handed-over roots are exactly the registered epochs 1..=last, +/// Check that the handed-over roots are exactly the registered epochs 0..=last, /// which also proves the source is at the expected epoch. fn ensure_handover_registered( roots: &[(KeyPair, SigningKey)], @@ -1295,16 +1281,16 @@ fn ensure_handover_registered( let Some(registered) = registered else { return Ok(()); }; - if registered.is_empty() && last == 1 { + if registered.is_empty() && last == 0 { return Ok(()); } - if registered.len() != last as usize { + if registered.len() != last as usize + 1 { bail!( - "the source hands over root epochs 1..={last}, but {} epochs are registered", + "the source hands over root epochs 0..={last}, but {} epochs are registered", registered.len() ); } - for (epoch, ((ca_key, k256_key), registered)) in (1..).zip(roots.iter().zip(registered)) { + for (epoch, ((ca_key, k256_key), registered)) in (0..).zip(roots.iter().zip(registered)) { if ca_key.public_key_der() != registered.ca_pubkey || k256_key.verifying_key().to_sec1_bytes()[..] != registered.k256_pubkey[..] { @@ -1314,10 +1300,10 @@ fn ensure_handover_registered( Ok(()) } -/// Bootstrapping creates epoch 1 only; later epochs come from rotation onboarding. +/// Bootstrapping creates epoch 0 only; later epochs come from rotation onboarding. fn ensure_bootstrap_epoch(cfg: &KmsConfig) -> Result<()> { - if cfg.key_rotation.current_epoch != 1 { - bail!("bootstrap only creates epoch 1; onboard with rotate to create a later epoch"); + if cfg.key_rotation.current_epoch != 0 { + bail!("bootstrap only creates epoch 0; onboard with rotate to create a later epoch"); } Ok(()) } @@ -1335,7 +1321,7 @@ fn self_signed_root_ca(ca_key: &KeyPair) -> Result { async fn attest_keys(epoch: u32, p256_pubkey: &[u8], k256_pubkey: &[u8]) -> Result> { let p256_hex = hex::encode(p256_pubkey); let k256_hex = hex::encode(k256_pubkey); - let content_to_quote = if epoch == 1 { + let content_to_quote = if epoch == 0 { format!("dstack-kms-genereted-keys-v1:{p256_hex};{k256_hex};") } else { format!("dstack-kms-generated-keys-v2:{epoch};{p256_hex};{k256_hex};") diff --git a/sdk/go/dstack/client_v1.go b/sdk/go/dstack/client_v1.go index eb718aa74..5b27c29da 100644 --- a/sdk/go/dstack/client_v1.go +++ b/sdk/go/dstack/client_v1.go @@ -42,8 +42,7 @@ type GetKeyV1Response struct { // key over the app root public key. See `docs/guest-api-v1.md` for the claim // encoding and the verification steps. SignatureChain [][]byte - // KMS root key epoch the key was derived from; 0 from an agent that - // predates root epochs, which means epoch 1. + // KMS root key epoch the key was derived from. RootEpoch uint32 } @@ -445,21 +444,28 @@ func (c *DstackClientV1) IssueCert(ctx context.Context, options ...IssueCertV1Op // Derivation is flat -- `a/b` is not a child of `a` -- and binds both arguments, // so these keys differ from the v0 keys of the same name. func (c *DstackClientV1) GetKey(ctx context.Context, domain string, algorithm string) (*GetKeyV1Response, error) { - return c.GetKeyAtRootEpoch(ctx, domain, algorithm, 0) + return c.getKey(ctx, domain, algorithm, nil) } -// GetKeyAtRootEpoch is GetKey under the KMS root key epoch rootEpoch, counting -// from 1; 0 is the current epoch. Name an older epoch only to reach a key from -// before a KMS root key rotation, for example to migrate off it. +// GetKeyAtRootEpoch is GetKey under the KMS root key epoch rootEpoch instead of +// the current one; epoch 0 is the original root keys. Name an older epoch only +// to reach a key from before a KMS root key rotation, for example to migrate +// off it. func (c *DstackClientV1) GetKeyAtRootEpoch(ctx context.Context, domain string, algorithm string, rootEpoch uint32) (*GetKeyV1Response, error) { + return c.getKey(ctx, domain, algorithm, &rootEpoch) +} + +func (c *DstackClientV1) getKey(ctx context.Context, domain string, algorithm string, rootEpoch *uint32) (*GetKeyV1Response, error) { if algorithm == "" { return nil, fmt.Errorf("algorithm is required, use `secp256k1` or `ed25519`") } payload := map[string]interface{}{ - "domain": domain, - "algorithm": algorithm, - "root_epoch": rootEpoch, + "domain": domain, + "algorithm": algorithm, + } + if rootEpoch != nil { + payload["root_epoch"] = *rootEpoch } data, err := c.sendRPCRequest(ctx, "/v1/GetKey", payload) diff --git a/sdk/js/src/client-v1.ts b/sdk/js/src/client-v1.ts index f743aebd5..fbf3156f5 100644 --- a/sdk/js/src/client-v1.ts +++ b/sdk/js/src/client-v1.ts @@ -68,7 +68,7 @@ export interface GetKeyResponseV1 { public_key: Uint8Array /** Two links: the app root key over the v1 key claim, then the KMS root key. */ signature_chain: Uint8Array[] - /** KMS root key epoch the key was derived from; 0 from an agent predating root epochs, meaning 1. */ + /** KMS root key epoch the key was derived from. */ root_epoch: number } @@ -335,17 +335,18 @@ export class DstackClientV1 { * @param domain Caller-chosen domain-separation string. May be empty. * @param algorithm Exactly `secp256k1` or `ed25519`. No default, no `k256` alias. */ - async getKey(domain: string, algorithm: string, rootEpoch = 0): Promise { + async getKey(domain: string, algorithm: string, rootEpoch?: number): Promise { // v0 defaulted an empty algorithm to secp256k1, which let a typo hand back a // key of the wrong type under a name the caller thought meant something else. if (!algorithm) { throw new Error('algorithm is required, use "secp256k1" or "ed25519"') } - // 0 is the current KMS root key epoch; name an older one only to migrate off it. - if (!Number.isInteger(rootEpoch) || rootEpoch < 0) { - throw new Error('rootEpoch must be 0 (current) or an epoch number') + // Unset is the current KMS root key epoch; name an older one only to migrate off it. + if (rootEpoch !== undefined && (!Number.isInteger(rootEpoch) || rootEpoch < 0)) { + throw new Error('rootEpoch must be a non-negative epoch number') } - const payload = JSON.stringify({ domain, algorithm, root_epoch: rootEpoch }) + const payload = JSON.stringify( + rootEpoch === undefined ? { domain, algorithm } : { domain, algorithm, root_epoch: rootEpoch }) const result = await send_rpc_request<{ key: string, public_key: string, signature_chain: string[], root_epoch?: number }>( this.endpoint, '/v1/GetKey', payload) throwOnRpcError(result) diff --git a/sdk/python/src/dstack_sdk/dstack_client_v1.py b/sdk/python/src/dstack_sdk/dstack_client_v1.py index a1ebe2c4f..77c1438d0 100644 --- a/sdk/python/src/dstack_sdk/dstack_client_v1.py +++ b/sdk/python/src/dstack_sdk/dstack_client_v1.py @@ -107,8 +107,7 @@ class GetKeyResponseV1(BaseModel): # Two links: the app root key's signature over the v1 key claim, then the # KMS root key's signature over the app root public key. signature_chain: List[HexBytes] - # KMS root key epoch the key was derived from. 0 from an agent that - # predates root epochs, which means epoch 1. + # KMS root key epoch the key was derived from. root_epoch: int = 0 @@ -241,7 +240,7 @@ async def issue_cert( return IssueCertResponseV1(**result) async def get_key( - self, domain: str, algorithm: str, root_epoch: int = 0 + self, domain: str, algorithm: str, root_epoch: Optional[int] = None ) -> GetKeyResponseV1: """Derive an application key from ``(domain, algorithm)``. @@ -251,19 +250,18 @@ async def get_key( ``k256`` alias, because in v0 a typo silently produced a key of the wrong type under a name the caller thought meant something else. - ``root_epoch`` selects the KMS root key epoch, counting from 1; the - default 0 is the current epoch. Name an older epoch only to reach a key - from before a KMS root key rotation, for example to migrate off it. + ``root_epoch`` selects the KMS root key epoch; the default ``None`` is + the current epoch, and epoch 0 is the original root keys. Name an older + epoch only to reach a key from before a KMS root key rotation, for + example to migrate off it. """ if not algorithm: raise ValueError("algorithm is required, use `secp256k1` or `ed25519`") - if root_epoch < 0: - raise ValueError("root_epoch must be 0 (current) or an epoch number") - data: Dict[str, Any] = { - "domain": domain, - "algorithm": algorithm, - "root_epoch": root_epoch, - } + data: Dict[str, Any] = {"domain": domain, "algorithm": algorithm} + if root_epoch is not None: + if root_epoch < 0: + raise ValueError("root_epoch must be a non-negative epoch number") + data["root_epoch"] = root_epoch result = await self._send_rpc_request("GetKey", data) return GetKeyResponseV1(**result) @@ -366,7 +364,7 @@ def issue_cert( @call_async def get_key( - self, domain: str, algorithm: str, root_epoch: int = 0 + self, domain: str, algorithm: str, root_epoch: Optional[int] = None ) -> GetKeyResponseV1: """Derive an application key from ``(domain, algorithm)``.""" raise NotImplementedError diff --git a/sdk/rust/src/dstack_client_v1.rs b/sdk/rust/src/dstack_client_v1.rs index 3ab11e83d..71f9c5f33 100644 --- a/sdk/rust/src/dstack_client_v1.rs +++ b/sdk/rust/src/dstack_client_v1.rs @@ -95,22 +95,35 @@ impl DstackClientV1 { /// default and no `k256` alias -- a typo is an error rather than a key of /// the wrong type under a name the caller misread. pub async fn get_key(&self, domain: &str, algorithm: &str) -> Result { - self.get_key_at_root_epoch(domain, algorithm, 0).await + self.request_key(domain, algorithm, None).await } - /// [`Self::get_key`] under the KMS root key epoch `root_epoch` (counting - /// from 1; 0 is the current epoch). Name an older epoch only to reach a key - /// from before a KMS root key rotation, for example to migrate off it. + /// [`Self::get_key`] under the KMS root key epoch `root_epoch` instead of + /// the current one; epoch 0 is the original root keys. Name an older epoch + /// only to reach a key from before a KMS root key rotation, for example to + /// migrate off it. pub async fn get_key_at_root_epoch( &self, domain: &str, algorithm: &str, root_epoch: u32, + ) -> Result { + self.request_key(domain, algorithm, Some(root_epoch)).await + } + + async fn request_key( + &self, + domain: &str, + algorithm: &str, + root_epoch: Option, ) -> Result { if algorithm.is_empty() { anyhow::bail!("algorithm is required, use `secp256k1` or `ed25519`") } - let data = json!({ "domain": domain, "algorithm": algorithm, "root_epoch": root_epoch }); + let mut data = json!({ "domain": domain, "algorithm": algorithm }); + if let Some(root_epoch) = root_epoch { + data["root_epoch"] = root_epoch.into(); + } let response = self.send_rpc_request("GetKey", &data).await?; Ok(serde_json::from_value::(response)?) } diff --git a/sdk/rust/types/src/dstack_v1.rs b/sdk/rust/types/src/dstack_v1.rs index 02a11a3e9..3812b2e4f 100644 --- a/sdk/rust/types/src/dstack_v1.rs +++ b/sdk/rust/types/src/dstack_v1.rs @@ -125,8 +125,7 @@ pub struct GetKeyResponse { /// steps. Verifying is the relying party's job; this SDK does not do it. #[serde(with = "hex_vec")] pub signature_chain: Vec>, - /// KMS root key epoch the key was derived from. 0 from an agent that - /// predates root epochs, which means epoch 1. + /// KMS root key epoch the key was derived from. #[serde(default)] pub root_epoch: u32, } From 385e5d3898b6e47a8ad9d93990859caf9a5d17ba Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Fri, 6 Nov 2026 02:24:27 -0800 Subject: [PATCH 14/17] kms: serve the root epoch governance activates, checked by chain hash DstackKms now keeps a chain hash per root epoch, committing to the public keys of that epoch and every earlier one, and separates registering an epoch from making it current: activateRootEpoch moves the current epoch forward one step at a time, and registering epoch 0 makes it current. Auth backends report just the current epoch and its chain hash. The KMS drops its key_rotation config. It serves the current epoch on chain, provided its keyring holds it and reproduces its chain hash, and waits rather than exits until it does. A move of the current epoch makes it restart, so it re-signs its RPC certificate under the new root CA. Onboarding imports the source's whole keyring and checks it by chain hash; a rotation must start from the current epoch and generates the next one. A whole cluster can therefore be prepared at a new epoch while still serving the old one, then switched by one on-chain activation. GetAppKey can return every older epoch's app keys in one response, so a guest makes one request per boot however many epochs exist. Signed-off-by: Kevin Wang --- docs/deployment.md | 31 +- dstack/Cargo.lock | 2 + dstack/crates/dstack-auth/Cargo.toml | 2 + dstack/crates/dstack-auth/src/main.rs | 77 +++- dstack/dstack-types/src/lib.rs | 1 + dstack/dstack-types/src/root_epoch.rs | 49 ++ dstack/dstack-util/src/main.rs | 1 + dstack/dstack-util/src/system_setup.rs | 84 ++-- dstack/kms/auth-eth-bun/index.test.ts | 29 +- dstack/kms/auth-eth-bun/index.ts | 66 +-- dstack/kms/auth-eth-bun/openapi.json | 24 +- dstack/kms/auth-eth/contracts/DstackKms.sol | 52 ++- dstack/kms/auth-eth/script/Manage.s.sol | 12 + dstack/kms/auth-eth/script/Query.s.sol | 8 +- dstack/kms/auth-eth/script/README.md | 12 +- dstack/kms/auth-eth/src/ethereum.ts | 28 +- dstack/kms/auth-eth/src/main.test.ts | 11 +- dstack/kms/auth-eth/src/server.ts | 5 +- dstack/kms/auth-eth/src/types.ts | 10 +- dstack/kms/auth-eth/test/DstackKms.t.sol | 69 ++- dstack/kms/auth-mock/README.md | 2 +- dstack/kms/auth-mock/index.test.ts | 14 +- dstack/kms/auth-mock/index.ts | 21 +- dstack/kms/auth-mock/openapi.json | 19 +- dstack/kms/auth-simple/README.md | 3 +- dstack/kms/auth-simple/bun.lock | 3 + dstack/kms/auth-simple/index.test.ts | 36 +- dstack/kms/auth-simple/index.ts | 42 +- dstack/kms/auth-simple/package.json | 1 + dstack/kms/kms.toml | 6 - dstack/kms/rpc/proto/kms_rpc.proto | 14 + dstack/kms/src/config.rs | 46 +- dstack/kms/src/main.rs | 42 +- dstack/kms/src/main_service.rs | 290 ++++++++---- .../kms/src/main_service/upgrade_authority.rs | 106 +++-- dstack/kms/src/onboard_service.rs | 434 +++++++++--------- 36 files changed, 1004 insertions(+), 648 deletions(-) create mode 100644 dstack/dstack-types/src/root_epoch.rs diff --git a/docs/deployment.md b/docs/deployment.md index 12bb1efd2..52a684768 100644 --- a/docs/deployment.md +++ b/docs/deployment.md @@ -515,8 +515,9 @@ curl http://:9203/finish #### Rotating KMS Root Keys -KMS root keys are versioned by epoch, and every KMS keeps all epochs so CVMs -can migrate from any of them. See the [rotation design](./kms-emergency-key-rotation-protocol.md) +KMS root keys are versioned by epoch; epoch 0 is the original roots. Every KMS +keeps all epochs so CVMs can migrate from any of them, and serves the epoch +governance made current. See the [rotation design](./kms-emergency-key-rotation-protocol.md) for the full procedure; the KMS steps are: 1. If epochs were never registered, register the current root public keys as @@ -524,20 +525,26 @@ for the full procedure; the KMS steps are: `bootstrap-info.json` on the bootstrapped node): `DstackKms.registerRootEpoch(caPubkey, k256Pubkey, attestation)`, via `script/Manage.s.sol:RegisterRootEpoch`, or a `rootEpochs` entry in the - auth-simple config. -2. Start a fresh KMS with `[core.key_rotation] current_epoch = N` and onboard it - with `"rotate": true` (the **Rotate** checkbox in the onboard UI) from a KMS - at epoch N-1. It imports epochs 0..N-1, generates epoch N, and returns the new - public keys with an attestation whose report data is + auth-simple config. Epoch 0 is current once registered. +2. Onboard a fresh KMS with `"rotate": true` (the **Rotate** checkbox in the + onboard UI) from a KMS at the current epoch N-1. It imports epochs 0..N-1, + generates epoch N, and returns the new public keys with an attestation whose + report data is `keccak256("dstack-kms-generated-keys-v2:{N};{ca_pubkey_hex};{k256_pubkey_hex};")`. -3. Verify the attestation and register epoch N together with it. A KMS refuses to start until its - keyring matches the registered epochs, and a KMS still at an older epoch stops - serving once a newer one is registered. -4. Onboard the other KMS nodes from the rotated one with `current_epoch = N`. -5. Boot CVMs on an OS that supports key epochs. They re-key their data disk and +3. Verify the attestation and register epoch N. The cluster keeps serving N-1. +4. Onboard the other KMS nodes from the rotated one; they import 0..N and keep + serving N-1. +5. Activate epoch N (`ActivateRootEpoch`, or `currentRootEpoch` for + auth-simple). Every KMS restarts onto N; a KMS that does not hold N waits + until it is onboarded again. +6. Boot CVMs on an OS that supports root epochs. They re-key their data disk and re-encrypt their env at boot; apply the new env with `vmm-cli.py apply-reencrypted-env `. +The auth backend reports the current epoch and the chain hash committing to the +root public keys of epochs 0..current; a KMS serves only while its own keys +produce that hash. + --- ## Deploying Apps diff --git a/dstack/Cargo.lock b/dstack/Cargo.lock index f4ff68f9a..3b10af8fc 100644 --- a/dstack/Cargo.lock +++ b/dstack/Cargo.lock @@ -1899,6 +1899,8 @@ version = "0.6.3" dependencies = [ "anyhow", "clap", + "dstack-types", + "hex", "rocket", "serde", "serde_json", diff --git a/dstack/crates/dstack-auth/Cargo.toml b/dstack/crates/dstack-auth/Cargo.toml index d9fb1879b..9a374b988 100644 --- a/dstack/crates/dstack-auth/Cargo.toml +++ b/dstack/crates/dstack-auth/Cargo.toml @@ -14,6 +14,8 @@ path = "src/main.rs" [dependencies] anyhow.workspace = true +dstack-types.workspace = true +hex.workspace = true clap.workspace = true rocket = { workspace = true, features = ["json"] } serde.workspace = true diff --git a/dstack/crates/dstack-auth/src/main.rs b/dstack/crates/dstack-auth/src/main.rs index 35b705b4a..9c6845b05 100644 --- a/dstack/crates/dstack-auth/src/main.rs +++ b/dstack/crates/dstack-auth/src/main.rs @@ -66,9 +66,17 @@ struct BootResponse { is_allowed: bool, gateway_app_id: String, reason: String, - /// latest registered KMS root epoch; absent when none is. - #[serde(skip_serializing_if = "Option::is_none")] - root_epoch: Option, + #[serde(flatten)] + root_epoch: Option, +} + +/// the KMS root epoch currently served and its chain hash, which commits to +/// the root public keys of epochs 0..=root_epoch. +#[derive(Serialize, Clone, PartialEq, Debug)] +#[serde(rename_all = "camelCase")] +struct CurrentRootEpoch { + root_epoch: usize, + root_epoch_hash: String, } #[derive(Deserialize, Default)] @@ -80,9 +88,11 @@ struct Allowlist { apps: HashMap, /// root public keys of each KMS root epoch, indexed by epoch; append only. root_epochs: Vec, + /// the root epoch KMS instances serve; raise it once every node holds it. + current_root_epoch: usize, } -#[derive(Deserialize, Serialize, Default, Clone)] +#[derive(Deserialize, Default, Clone)] #[serde(rename_all = "camelCase", default)] struct RootEpoch { ca_pubkey: String, @@ -90,8 +100,26 @@ struct RootEpoch { } impl Allowlist { - fn latest_root_epoch(&self) -> Option { - self.root_epochs.len().checked_sub(1) + /// `None` when no epoch is configured, or the configuration is invalid, + /// which a KMS past epoch 0 treats as fail-closed. + fn current_root_epoch(&self) -> Option { + let epochs = self.root_epochs.get(..=self.current_root_epoch)?; + let keys = epochs + .iter() + .map(|e| { + Some(( + hex::decode(norm(&e.ca_pubkey)).ok()?, + hex::decode(norm(&e.k256_pubkey)).ok()?, + )) + }) + .collect::>>()?; + let hashes = dstack_types::root_epoch::chain_hashes( + keys.iter().map(|(ca, k256)| (&ca[..], &k256[..])), + ); + Some(CurrentRootEpoch { + root_epoch: self.current_root_epoch, + root_epoch_hash: format!("0x{}", hex::encode(hashes.last()?)), + }) } } @@ -140,7 +168,7 @@ fn deny(al: &Allowlist, reason: &str) -> BootResponse { is_allowed: false, gateway_app_id: al.gateway_app_id.clone(), reason: reason.to_string(), - root_epoch: al.latest_root_epoch(), + root_epoch: al.current_root_epoch(), } } @@ -149,7 +177,7 @@ fn allow(al: &Allowlist) -> BootResponse { is_allowed: true, gateway_app_id: al.gateway_app_id.clone(), reason: "ok".to_string(), - root_epoch: al.latest_root_epoch(), + root_epoch: al.current_root_epoch(), } } @@ -231,15 +259,19 @@ fn info(cli: &State) -> Json { } fn info_json(al: &Allowlist) -> serde_json::Value { - json!({ + let mut info = json!({ "status": "ok", "kmsContractAddr": "", "ethRpcUrl": "", "gatewayAppId": "", "chainId": 0, "appImplementation": "", - "rootEpochs": al.root_epochs, - }) + }); + if let Some(epoch) = al.current_root_epoch() { + info["rootEpoch"] = epoch.root_epoch.into(); + info["rootEpochHash"] = epoch.root_epoch_hash.into(); + } + info } #[rocket::main] @@ -309,20 +341,31 @@ mod tests { } #[test] - fn reports_registered_root_epochs() { - let al: Allowlist = serde_json::from_str( + fn reports_the_current_root_epoch() { + let mut al: Allowlist = serde_json::from_str( r#"{"kms":{"mrAggregated":["0xMR"]}, - "rootEpochs":[{"caPubkey":"0xa1","k256Pubkey":"0xb1"}, - {"caPubkey":"0xa2","k256Pubkey":"0xb2"}]}"#, + "rootEpochs":[{"caPubkey":"0xa0","k256Pubkey":"0xb0"}, + {"caPubkey":"0xa1","k256Pubkey":"0xb1"}]}"#, ) .unwrap(); let info = BootInfo { mr_aggregated: "0xMR".into(), ..Default::default() }; + // Shared with the DstackKms contract tests. let response = serde_json::to_value(check_kms(&info, &al)).unwrap(); - assert_eq!(response["rootEpoch"], 1); - assert_eq!(info_json(&al)["rootEpochs"][1]["k256Pubkey"], "0xb2"); + assert_eq!(response["rootEpoch"], 0); + assert_eq!( + response["rootEpochHash"], + "0xf5f9a7455970becbb0d8441e53cf06df6290b34419cf783eb5869107867ec9ae" + ); + al.current_root_epoch = 1; + assert_eq!( + info_json(&al)["rootEpochHash"], + "0x2219fc74439588d8f9d514f4e897d4a1dc692519b0899d93f5ac5661b8e2850a" + ); + al.current_root_epoch = 2; + assert!(al.current_root_epoch().is_none()); let response = serde_json::to_value(check_kms(&info, &allowlist())).unwrap(); assert!(response.get("rootEpoch").is_none()); diff --git a/dstack/dstack-types/src/lib.rs b/dstack/dstack-types/src/lib.rs index 532dc9911..c76e655a5 100644 --- a/dstack/dstack-types/src/lib.rs +++ b/dstack/dstack-types/src/lib.rs @@ -2567,6 +2567,7 @@ pub struct ImageInfo { } pub mod mr_config; +pub mod root_epoch; pub mod sha256sum; pub mod shared_filenames; pub mod version; diff --git a/dstack/dstack-types/src/root_epoch.rs b/dstack/dstack-types/src/root_epoch.rs new file mode 100644 index 000000000..a5d386522 --- /dev/null +++ b/dstack/dstack-types/src/root_epoch.rs @@ -0,0 +1,49 @@ +// SPDX-FileCopyrightText: © 2026 Phala Network +// +// SPDX-License-Identifier: Apache-2.0 + +//! KMS root epoch chain hashes, computed as the DstackKms contract does. + +use sha3::{Digest, Keccak256}; + +/// Chain hashes of root epochs 0, 1, ... given each epoch's +/// `(ca_pubkey, k256_pubkey)`. Epoch `n`'s hash is +/// `keccak256(previous || uint256(n) || keccak256(ca_pubkey) || keccak256(k256_pubkey))` +/// with a zero previous hash for epoch 0, so it commits to every key up to `n`. +pub fn chain_hashes<'a>(epochs: impl IntoIterator) -> Vec<[u8; 32]> { + let mut previous = [0u8; 32]; + (0u64..) + .zip(epochs) + .map(|(epoch, (ca_pubkey, k256_pubkey))| { + let mut epoch_word = [0u8; 32]; + epoch_word[24..].copy_from_slice(&epoch.to_be_bytes()); + previous = Keccak256::new() + .chain_update(previous) + .chain_update(epoch_word) + .chain_update(Keccak256::digest(ca_pubkey)) + .chain_update(Keccak256::digest(k256_pubkey)) + .finalize() + .into(); + previous + }) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Shared with the DstackKms contract and auth-simple tests. + #[test] + fn matches_the_contract_vectors() { + let hashes = chain_hashes([(&[0xa0u8][..], &[0xb0u8][..]), (&[0xa1][..], &[0xb1][..])]); + assert_eq!( + hex::encode(hashes[0]), + "f5f9a7455970becbb0d8441e53cf06df6290b34419cf783eb5869107867ec9ae" + ); + assert_eq!( + hex::encode(hashes[1]), + "2219fc74439588d8f9d514f4e897d4a1dc692519b0899d93f5ac5661b8e2850a" + ); + } +} diff --git a/dstack/dstack-util/src/main.rs b/dstack/dstack-util/src/main.rs index 05e8fdc34..ee2f7ee7d 100644 --- a/dstack/dstack-util/src/main.rs +++ b/dstack/dstack-util/src/main.rs @@ -702,6 +702,7 @@ async fn cmd_get_keys(args: GetKeysArgs) -> Result<()> { api_version: 1, vm_config: "".to_string(), root_epoch: None, + include_history: false, }) .await .context("Failed to get app key")?; diff --git a/dstack/dstack-util/src/system_setup.rs b/dstack/dstack-util/src/system_setup.rs index f8a18419b..35a374908 100644 --- a/dstack/dstack-util/src/system_setup.rs +++ b/dstack/dstack-util/src/system_setup.rs @@ -2355,13 +2355,10 @@ struct AppInfo { } impl HostShared { - /// Request app keys of `root_epoch` (`None` = current) from one KMS. - async fn request_kms_keys( - &self, - kms_url: String, - root_epoch: Option, - ) -> Result<(AppKeys, Vec)> { - info!("Requesting app keys of root epoch {root_epoch:?} from KMS: {kms_url}"); + /// Request the current app keys from one KMS, with those of every older + /// root epoch in order. + async fn request_kms_keys(&self, kms_url: String) -> Result { + info!("Requesting app keys from KMS: {kms_url}"); let tmp_ca = { info!("Getting temp ca cert"); let client = RaClient::new(kms_url.clone(), true)?; @@ -2390,14 +2387,15 @@ impl HostShared { .get_app_key(rpc::GetAppKeyRequest { api_version: 1, vm_config: self.sys_config.vm_config.clone(), - root_epoch, + root_epoch: None, + include_history: true, }) .await .context("Failed to get app key")?; - // A KMS predating root epochs ignores the selector and serves epoch 0. + // A KMS predating root epochs serves epoch 0 with no history. let served = response.root_epoch; - if root_epoch.is_some_and(|epoch| epoch != served) { - bail!("KMS returned keys of epoch {served} instead of {root_epoch:?}"); + if !response.history.iter().map(|h| h.root_epoch).eq(0..served) { + bail!("KMS returned incomplete key history for root epoch {served}"); } let (_, ca_pem) = x509_parser::pem::parse_x509_pem(tmp_ca.ca_cert.as_bytes()) @@ -2421,24 +2419,19 @@ impl HostShared { root_epoch: served, }, }; - Ok((keys, response.os_image_hash)) + Ok(KmsKeys { + keys, + history: response.history, + os_image_hash: response.os_image_hash, + }) } +} - /// Request the keys of every root epoch before the one `keys` came from, - /// oldest first, from the same KMS. Empty for other key providers. - async fn request_key_history(&self, keys: &AppKeys) -> Result> { - let KeyProvider::Kms { - url, root_epoch, .. - } = &keys.key_provider - else { - return Ok(Vec::new()); - }; - let mut history = Vec::new(); - for epoch in 0..*root_epoch { - history.push(self.request_kms_keys(url.clone(), Some(epoch)).await?.0); - } - Ok(history) - } +struct KmsKeys { + keys: AppKeys, + /// App keys of the root epochs before `keys.root_epoch()`, indexed by epoch. + history: Vec, + os_image_hash: Vec, } struct Stage0<'a> { @@ -2453,7 +2446,7 @@ struct Stage1<'a> { shared: HostShared, keys: AppKeys, /// Keys of the root epochs before `keys.root_epoch()`, indexed by epoch. - key_history: Vec, + key_history: Vec, } fn validate_key_provider_inputs(kind: KeyProviderKind, kms_urls: &[String]) -> Result<()> { @@ -2505,20 +2498,24 @@ impl<'a> Stage0<'a> { self.shared.dir.join(APP_KEYS) } - async fn request_app_keys_from_kms(&self) -> Result { + async fn request_app_keys_from_kms(&self) -> Result<(AppKeys, Vec)> { if self.shared.sys_config.kms_urls.is_empty() { bail!("No KMS URLs are set"); } let mut errors = vec![]; for kms_url in &self.shared.sys_config.kms_urls { let kms_url = kms_rpc_url(kms_url); - match self.shared.request_kms_keys(kms_url.clone(), None).await { - Ok((keys, os_image_hash)) => { + match self.shared.request_kms_keys(kms_url.clone()).await { + Ok(KmsKeys { + keys, + history, + os_image_hash, + }) => { // Measured here, outside the per-URL request, so that a // failover can never extend os-image-hash twice. emit_runtime_event("os-image-hash", &os_image_hash) .context("failed to extend os-image-hash to the launch measurement")?; - return Ok(keys); + return Ok((keys, history)); } Err(err) => { warn!("failed to get app keys from KMS {kms_url}: {err:?}"); @@ -2597,11 +2594,12 @@ impl<'a> Stage0<'a> { .context("failed to generate TPM app keys") } - async fn request_app_keys(&self) -> Result { + /// The app keys, with those of older KMS root epochs when the KMS has any. + async fn request_app_keys(&self) -> Result<(AppKeys, Vec)> { let key_provider = self.shared.app_compose.key_provider(); validate_key_provider_inputs(key_provider, &self.shared.sys_config.kms_urls)?; - match key_provider { - KeyProviderKind::Kms => self.request_app_keys_from_kms().await, + let keys = match key_provider { + KeyProviderKind::Kms => return self.request_app_keys_from_kms().await, KeyProviderKind::Local => self.get_keys_from_local_key_provider().await, KeyProviderKind::None => { info!("No key provider is enabled, generating temporary app keys"); @@ -2613,7 +2611,8 @@ impl<'a> Stage0<'a> { info!("Generating app keys from TPM"); self.generate_tpm_app_keys() } - } + }; + Ok((keys?, Vec::new())) } async fn setup_swap(&self, swap_size: u64, opts: &DstackOptions) -> Result<()> { @@ -2758,7 +2757,7 @@ impl<'a> Stage0<'a> { async fn mount_data_disk( &self, keys: &AppKeys, - key_history: &[AppKeys], + key_history: &[rpc::HistoricalAppKey], opts: &DstackOptions, ) -> Result<()> { let disk_crypt_key = &hex::encode(&keys.disk_crypt_key); @@ -2994,7 +2993,7 @@ impl<'a> Stage0<'a> { /// readable by older OS images. fn migrate_disk_key( &self, - key_history: &[AppKeys], + key_history: &[rpc::HistoricalAppKey], current_key: &str, current: u32, ) -> Result<()> { @@ -3244,19 +3243,14 @@ impl<'a> Stage0<'a> { self.vmm .notify_q("boot.progress", "requesting app keys") .await; - let mut app_keys = self + let (mut app_keys, key_history) = self .request_app_keys() .await .context("Failed to request app keys")?; - let key_history = self - .shared - .request_key_history(&app_keys) - .await - .context("Failed to request keys of older root epochs")?; app_keys.historical_roots = key_history .iter() .map(|keys| HistoricalAppRoot { - root_epoch: keys.root_epoch(), + root_epoch: keys.root_epoch, k256_key: keys.k256_key.clone(), k256_signature: keys.k256_signature.clone(), }) diff --git a/dstack/kms/auth-eth-bun/index.test.ts b/dstack/kms/auth-eth-bun/index.test.ts index 658c06fc7..aefdaa5a8 100644 --- a/dstack/kms/auth-eth-bun/index.test.ts +++ b/dstack/kms/auth-eth-bun/index.test.ts @@ -240,14 +240,13 @@ describe('API Compatibility Tests', () => { deviceId: '0x' + '44'.repeat(32), }; - it('reports the latest registered root epoch and lists all epochs', async () => { + it('reports the current root epoch and its chain hash', async () => { mockReadContract.mockImplementation((params) => { switch (params.functionName) { case 'isAppAllowed': return [true, '']; case 'gatewayAppId': return '0xgateway'; case 'appImplementation': return '0x9876543210987654321098765432109876543210'; - case 'rootEpochCount': return 2n; - case 'rootEpoch': return [`0xa${params.args[0]}`, `0xb${params.args[0]}`, '0xc0']; + case 'currentRootEpoch': return [1n, '0x' + '11'.repeat(32)]; } }); @@ -256,13 +255,11 @@ describe('API Compatibility Tests', () => { headers: { 'Content-Type': 'application/json' }, body: JSON.stringify(validBootInfo), })); - expect((await boot.json()).rootEpoch).toBe(1); + const expected = { rootEpoch: 1, rootEpochHash: '0x' + '11'.repeat(32) }; + expect(await boot.json()).toMatchObject(expected); const info = await appFetch(new Request('http://localhost:3001/')); - expect((await info.json()).rootEpochs).toEqual([ - { caPubkey: '0xa0', k256Pubkey: '0xb0' }, - { caPubkey: '0xa1', k256Pubkey: '0xb1' }, - ]); + expect(await info.json()).toMatchObject(expected); }); it('treats a contract without the registry as having no epochs', async () => { @@ -270,7 +267,7 @@ describe('API Compatibility Tests', () => { switch (params.functionName) { case 'isAppAllowed': return [true, '']; case 'gatewayAppId': return '0xgateway'; - case 'rootEpochCount': + case 'currentRootEpoch': throw Object.assign(new Error('call failed'), { cause: { name: 'ContractFunctionZeroDataError' }, }); @@ -507,7 +504,7 @@ describe('Authorization freshness and domain binding', () => { return decisions === 1 ? [true, 'initial allow'] : [false, 'policy changed']; } if (params.functionName === 'gatewayAppId') return 'gateway-app'; - if (params.functionName === 'rootEpochCount') return 0n; + if (params.functionName === 'currentRootEpoch') return undefined; throw new Error(`unexpected function ${params.functionName}`); }); @@ -527,7 +524,7 @@ describe('Authorization freshness and domain binding', () => { return [true, 'allowed']; } if (params.functionName === 'gatewayAppId') return 'gateway-app'; - if (params.functionName === 'rootEpochCount') return 0n; + if (params.functionName === 'currentRootEpoch') return undefined; throw new Error(`unexpected function ${params.functionName}`); }); @@ -552,7 +549,7 @@ describe('Authorization freshness and domain binding', () => { mockReadContract.mockImplementation((params) => { if (params.functionName === 'isAppAllowed') return [true, 'recovered']; if (params.functionName === 'gatewayAppId') return 'gateway-app'; - if (params.functionName === 'rootEpochCount') return 0n; + if (params.functionName === 'currentRootEpoch') return undefined; throw new Error(`unexpected function ${params.functionName}`); }); const recovered = await postApp(); @@ -583,7 +580,7 @@ describe('Ethereum finalized snapshot authorization', () => { expect(params.blockNumber).toBe(98n); if (params.functionName === 'isAppAllowed') return [true, 'finalized allow']; if (params.functionName === 'gatewayAppId') return 'gateway-app'; - if (params.functionName === 'rootEpochCount') return 0n; + if (params.functionName === 'currentRootEpoch') return undefined; throw new Error(`unexpected function ${params.functionName}`); }); @@ -603,7 +600,7 @@ describe('Ethereum finalized snapshot authorization', () => { return decisions === 1 ? [true, 'old canonical allow'] : [false, 'new canonical deny']; } if (params.functionName === 'gatewayAppId') return 'gateway-app'; - if (params.functionName === 'rootEpochCount') return 0n; + if (params.functionName === 'currentRootEpoch') return undefined; throw new Error(`unexpected function ${params.functionName}`); }); @@ -632,7 +629,7 @@ describe('Ethereum finalized snapshot authorization', () => { mockReadContract.mockImplementation((params) => { if (params.functionName === 'isAppAllowed') return [true, 'recovered']; if (params.functionName === 'gatewayAppId') return 'gateway-app'; - if (params.functionName === 'rootEpochCount') return 0n; + if (params.functionName === 'currentRootEpoch') return undefined; throw new Error(`unexpected function ${params.functionName}`); }); const recovered = await authorize(); @@ -675,7 +672,7 @@ describe('dstack-kms 0.6 bootAuth payload', () => { return [true, '']; } if (params.functionName === 'gatewayAppId') return 'gateway-app'; - if (params.functionName === 'rootEpochCount') return 0n; + if (params.functionName === 'currentRootEpoch') return undefined; throw new Error(`unexpected function ${params.functionName}`); }); }); diff --git a/dstack/kms/auth-eth-bun/index.ts b/dstack/kms/auth-eth-bun/index.ts index e9be9c35a..4cc8f9a9a 100644 --- a/dstack/kms/auth-eth-bun/index.ts +++ b/dstack/kms/auth-eth-bun/index.ts @@ -36,11 +36,13 @@ const BootResponseSchema = z.object({ isAllowed: z.boolean(), reason: z.string(), gatewayAppId: z.string(), - // Latest registered KMS root epoch; absent when none is registered. - rootEpoch: z.number().optional() + // KMS root epoch currently served and its chain hash, which commits to the + // root public keys of epochs 0..=rootEpoch; absent when none is registered. + rootEpoch: z.number().optional(), + rootEpochHash: z.string().optional() }); -type RootEpoch = { caPubkey: Hex; k256Pubkey: Hex }; +type RootEpochInfo = { rootEpoch?: number; rootEpochHash?: Hex }; type BootInfo = z.infer; type BootResponse = z.infer; @@ -114,26 +116,19 @@ const DSTACK_KMS_ABI = [ outputs: [{ name: '', type: 'address' }] }, { - name: 'rootEpochCount', + name: 'currentRootEpoch', type: 'function', stateMutability: 'view', inputs: [], - outputs: [{ name: '', type: 'uint256' }] - }, - { - name: 'rootEpoch', - type: 'function', - stateMutability: 'view', - inputs: [{ name: 'epoch', type: 'uint256' }], outputs: [ - { name: 'caPubkey', type: 'bytes' }, - { name: 'k256Pubkey', type: 'bytes' }, - { name: 'attestation', type: 'bytes' } + { name: 'epoch', type: 'uint256' }, + { name: 'chainHash', type: 'bytes32' } ] } ] as const; -// Contracts deployed before root epochs lack the registry getters. +// Contracts deployed before root epochs lack the registry getters, and a +// registry with no epoch reverts. const isMissingFunction = (error: unknown): boolean => { for (let cause: any = error; cause; cause = cause.cause) { if (cause.name === 'ContractFunctionZeroDataError' || cause.name === 'ContractFunctionRevertedError') { @@ -225,49 +220,34 @@ class EthereumBackend { blockNumber }); - const epochCount = await this.getRootEpochCount(blockNumber); - return { isAllowed, reason, gatewayAppId: gatewayAppId as string, - ...(epochCount > 0 && { rootEpoch: epochCount - 1 }), + ...await this.getRootEpoch(blockNumber), }; } - private async getRootEpochCount(blockNumber?: bigint): Promise { + // The current root epoch and its chain hash; empty when none is registered + // or the contract predates root epochs. + async getRootEpoch(blockNumber?: bigint): Promise { + blockNumber ??= await this.finalizedBlockNumber(); try { - const count = await this.client.readContract({ + const current = await this.client.readContract({ address: this.kmsContractAddr, abi: DSTACK_KMS_ABI, - functionName: 'rootEpochCount', + functionName: 'currentRootEpoch', blockNumber }); - return Number(count ?? 0); + if (!current) return {}; + const [epoch, chainHash] = current; + return { rootEpoch: Number(epoch), rootEpochHash: chainHash }; } catch (error) { - if (isMissingFunction(error)) return 0; + if (isMissingFunction(error)) return {}; throw error; } } - // Registered root epochs, indexed by epoch. - async getRootEpochs(): Promise { - const blockNumber = await this.finalizedBlockNumber(); - const count = await this.getRootEpochCount(blockNumber); - return Promise.all( - Array.from({ length: count }, async (_, epoch) => { - const [caPubkey, k256Pubkey] = await this.client.readContract({ - address: this.kmsContractAddr, - abi: DSTACK_KMS_ABI, - functionName: 'rootEpoch', - args: [BigInt(epoch)], - blockNumber - }); - return { caPubkey, k256Pubkey }; - }) - ); - } - async getGatewayAppId(): Promise { const result = await this.client.readContract({ address: this.kmsContractAddr, @@ -346,7 +326,7 @@ app.get('/', async (c) => { ethereum.getGatewayAppId(), ethereum.getChainId(), ethereum.getAppImplementation(), - ethereum.getRootEpochs(), + ethereum.getRootEpoch(), ]); console.log('batch', batch); @@ -358,7 +338,7 @@ app.get('/', async (c) => { chainId: batch[1], appAuthImplementation: batch[2], // NOTE: for backward compatibility appImplementation: batch[2], - rootEpochs: batch[3], + ...batch[3], }); } catch { console.error('authorization backend health check failed'); diff --git a/dstack/kms/auth-eth-bun/openapi.json b/dstack/kms/auth-eth-bun/openapi.json index 5ca950e5d..9f838cfa4 100644 --- a/dstack/kms/auth-eth-bun/openapi.json +++ b/dstack/kms/auth-eth-bun/openapi.json @@ -176,7 +176,11 @@ }, "rootEpoch": { "type": "integer", - "description": "Latest registered KMS root epoch (the original roots are 0); absent when none is registered" + "description": "KMS root epoch currently served (the original roots are 0); absent when none is registered" + }, + "rootEpochHash": { + "type": "string", + "description": "Chain hash committing to the root public keys of epochs 0..=rootEpoch (hex)" } } }, @@ -221,17 +225,13 @@ "type": "string", "description": "App implementation address" }, - "rootEpochs": { - "type": "array", - "description": "Registered KMS root epochs, indexed by epoch", - "items": { - "type": "object", - "required": ["caPubkey", "k256Pubkey"], - "properties": { - "caPubkey": { "type": "string", "description": "Root CA public key (hex)" }, - "k256Pubkey": { "type": "string", "description": "Root secp256k1 public key (hex)" } - } - } + "rootEpoch": { + "type": "integer", + "description": "KMS root epoch currently served; absent when none is registered" + }, + "rootEpochHash": { + "type": "string", + "description": "Chain hash committing to the root public keys of epochs 0..=rootEpoch (hex)" } } }, diff --git a/dstack/kms/auth-eth/contracts/DstackKms.sol b/dstack/kms/auth-eth/contracts/DstackKms.sol index 930512d24..06dd743f4 100644 --- a/dstack/kms/auth-eth/contracts/DstackKms.sol +++ b/dstack/kms/auth-eth/contracts/DstackKms.sol @@ -51,13 +51,20 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E // and both public keys, so anyone can check the keys came from an // approved KMS in a TEE. bytes attestation; + // Commits to the public keys of this epoch and every earlier one: + // keccak256(previous chainHash || epoch || keccak256(caPubkey) || + // keccak256(k256Pubkey)), with a zero previous hash for epoch 0. + bytes32 chainHash; } - // Root public keys of each KMS root epoch, indexed by epoch; epoch 0 is the - // original roots. Append-only: a registered epoch is never changed, so the - // last entry is current. + // Root keys of each KMS root epoch, indexed by epoch; epoch 0 is the + // original roots. Append-only: a registered epoch is never changed. RootEpoch[] private _rootEpochs; + // The epoch KMS instances serve. Registering epoch 0 makes it current; + // later epochs take effect only when activated, one at a time. + uint256 private _currentRootEpoch; + // Events // appId is unindexed for backward compatibility with deployed log indexers. // slither-disable-next-line unindexed-event-address @@ -72,7 +79,8 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E event GatewayAppIdSet(string gatewayAppId); // slither-disable-next-line unindexed-event-address event AppImplementationSet(address implementation); - event RootEpochRegistered(uint256 indexed epoch, bytes caPubkey, bytes k256Pubkey); + event RootEpochRegistered(uint256 indexed epoch, bytes32 chainHash, bytes caPubkey, bytes k256Pubkey); + event RootEpochActivated(uint256 indexed epoch); event AppDeployedViaFactory(address indexed appId, address indexed deployer); /// @notice Additive audit event for reconstructing authorization policy. /// @dev `value` is the affected bytes32 value, address, or hash of dynamic public data. @@ -131,7 +139,8 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E } // Register the root public keys of the next KMS root epoch, with the - // attestation proving where they were generated + // attestation proving where they were generated. It is not served until + // activated, except epoch 0, which is current from the start. function registerRootEpoch( bytes calldata caPubkey, bytes calldata k256Pubkey, @@ -144,24 +153,43 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E require(caPubkey.length > 0 && k256Pubkey.length > 0, "Empty epoch public key"); require(attestation.length > 0, "Empty epoch attestation"); epoch = _rootEpochs.length; - _rootEpochs.push(RootEpoch(caPubkey, k256Pubkey, attestation)); - emit RootEpochRegistered(epoch, caPubkey, k256Pubkey); - _emitPolicy("root-epoch", keccak256(abi.encode(caPubkey, k256Pubkey)), true); + bytes32 previous = epoch == 0 ? bytes32(0) : _rootEpochs[epoch - 1].chainHash; + bytes32 chainHash = keccak256(abi.encodePacked(previous, epoch, keccak256(caPubkey), keccak256(k256Pubkey))); + _rootEpochs.push(RootEpoch(caPubkey, k256Pubkey, attestation, chainHash)); + emit RootEpochRegistered(epoch, chainHash, caPubkey, k256Pubkey); + _emitPolicy("root-epoch", chainHash, true); + if (epoch == 0) { + emit RootEpochActivated(0); + } + } + + // Make the next registered root epoch the one KMS instances serve + function activateRootEpoch(uint256 epoch) external onlyOwner { + require(epoch == _currentRootEpoch + 1 && epoch < _rootEpochs.length, "Invalid root epoch activation"); + _currentRootEpoch = epoch; + emit RootEpochActivated(epoch); + _emitPolicy("root-epoch-activated", bytes32(epoch), true); } - // Number of registered root epochs; the latest is this minus one function rootEpochCount() external view returns (uint256) { return _rootEpochs.length; } + // The root epoch KMS instances serve and its chain hash + function currentRootEpoch() external view returns (uint256 epoch, bytes32 chainHash) { + require(_rootEpochs.length > 0, "No root epoch registered"); + epoch = _currentRootEpoch; + chainHash = _rootEpochs[epoch].chainHash; + } + function rootEpoch(uint256 epoch) external view - returns (bytes memory caPubkey, bytes memory k256Pubkey, bytes memory attestation) + returns (bytes memory caPubkey, bytes memory k256Pubkey, bytes memory attestation, bytes32 chainHash) { require(epoch < _rootEpochs.length, "Unknown root epoch"); RootEpoch storage keys = _rootEpochs[epoch]; - return (keys.caPubkey, keys.k256Pubkey, keys.attestation); + return (keys.caPubkey, keys.k256Pubkey, keys.attestation, keys.chainHash); } // Function to set KMS quote @@ -355,7 +383,7 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E } // Add storage gap for upgradeable contracts - uint256[49] private __gap; + uint256[48] private __gap; } function isContract(address addr) view returns (bool) { diff --git a/dstack/kms/auth-eth/script/Manage.s.sol b/dstack/kms/auth-eth/script/Manage.s.sol index c0a84ed7a..17c7bbef1 100644 --- a/dstack/kms/auth-eth/script/Manage.s.sol +++ b/dstack/kms/auth-eth/script/Manage.s.sol @@ -123,6 +123,18 @@ contract RegisterRootEpoch is BaseScript { } } +contract ActivateRootEpoch is BaseScript { + function run() external { + uint256 epoch = vm.envUint("ROOT_EPOCH"); + + vm.startBroadcast(vm.envUint("PRIVATE_KEY")); + kms.activateRootEpoch(epoch); + vm.stopBroadcast(); + + console.log("Activated KMS root epoch", epoch); + } +} + // App Management Scripts contract AddComposeHash is BaseScript { function run() external { diff --git a/dstack/kms/auth-eth/script/Query.s.sol b/dstack/kms/auth-eth/script/Query.s.sol index c4bc628be..7548c21cb 100644 --- a/dstack/kms/auth-eth/script/Query.s.sol +++ b/dstack/kms/auth-eth/script/Query.s.sol @@ -213,12 +213,18 @@ contract GetKmsSettings is BaseQueryScript { uint256 epochCount = kms.rootEpochCount(); console.log("\nRegistered root epochs:", epochCount); + if (epochCount > 0) { + (uint256 current,) = kms.currentRootEpoch(); + console.log("Current root epoch:", current); + } for (uint256 epoch = 0; epoch < epochCount; epoch++) { - (bytes memory epochCa, bytes memory epochK256, bytes memory attestation) = kms.rootEpoch(epoch); + (bytes memory epochCa, bytes memory epochK256, bytes memory attestation, bytes32 chainHash) = + kms.rootEpoch(epoch); console.log(" Epoch", epoch); console.log(" CA Pubkey:", vm.toString(epochCa)); console.log(" K256 Pubkey:", vm.toString(epochK256)); console.log(" Attestation length:", attestation.length); + console.log(" Chain hash:", vm.toString(chainHash)); } // Check implementation via storage diff --git a/dstack/kms/auth-eth/script/README.md b/dstack/kms/auth-eth/script/README.md index 37ff17ea1..472f1aca8 100644 --- a/dstack/kms/auth-eth/script/README.md +++ b/dstack/kms/auth-eth/script/README.md @@ -59,7 +59,8 @@ forge script script/Manage.s.sol:SetKmsInfo --broadcast --rpc-url $RPC_URL ``` #### Register a KMS Root Epoch -Appends the next root epoch (the original roots are epoch 0). `ATTESTATION` is +Appends the next root epoch (the original roots are epoch 0). Epoch 0 takes +effect when registered; later epochs take effect when activated. `ATTESTATION` is the one the KMS returned when it generated the keys (Bootstrap, or Onboard with `rotate`); verify it before registering. Registered epochs cannot be changed. ```bash @@ -67,6 +68,15 @@ KMS_CONTRACT_ADDR=0x... CA_PUBKEY=0x... K256_PUBKEY=0x... ATTESTATION=0x... \ forge script script/Manage.s.sol:RegisterRootEpoch --broadcast --rpc-url $RPC_URL ``` +#### Activate a KMS Root Epoch +Makes the next registered root epoch the one KMS instances serve. Activate it +only once every KMS node holds it; nodes still at the previous epoch stop +serving, and nodes holding it restart to switch. Epochs only move forward. +```bash +KMS_CONTRACT_ADDR=0x... ROOT_EPOCH=1 \ +forge script script/Manage.s.sol:ActivateRootEpoch --broadcast --rpc-url $RPC_URL +``` + #### Remove KMS Aggregated MR ```bash KMS_CONTRACT_ADDR=0x... MR_AGGREGATED=0x1234... \ diff --git a/dstack/kms/auth-eth/src/ethereum.ts b/dstack/kms/auth-eth/src/ethereum.ts index e44b88d73..c9f8b8bcb 100644 --- a/dstack/kms/auth-eth/src/ethereum.ts +++ b/dstack/kms/auth-eth/src/ethereum.ts @@ -3,7 +3,7 @@ // SPDX-License-Identifier: Apache-2.0 import { ethers } from 'ethers'; -import { BootInfo, BootResponse, RootEpoch } from './types'; +import { BootInfo, BootResponse, RootEpochInfo } from './types'; // Minimal ABI for DstackKms contract const DSTACK_KMS_ABI = [ @@ -11,8 +11,7 @@ const DSTACK_KMS_ABI = [ "function isKmsAllowed((address appId,bytes32 composeHash,address instanceId,bytes32 deviceId,bytes32 mrAggregated,bytes32 mrSystem,bytes32 osImageHash,string tcbStatus,string[] advisoryIds) bootInfo) view returns (bool, string)", "function gatewayAppId() view returns (string)", "function appImplementation() view returns (address)", - "function rootEpochCount() view returns (uint256)", - "function rootEpoch(uint256 epoch) view returns (bytes caPubkey, bytes k256Pubkey, bytes attestation)" + "function currentRootEpoch() view returns (uint256 epoch, bytes32 chainHash)" ]; // Contracts deployed before key epochs existed lack the registry getters. @@ -67,35 +66,26 @@ export class EthereumBackend { } const [isAllowed, reason] = response; const gatewayAppId = await this.kmsContract.gatewayAppId({ blockTag }); - const epochCount = await this.getRootEpochCount(blockTag); return { isAllowed, reason, gatewayAppId, - ...(epochCount > 0 && { rootEpoch: epochCount - 1 }), + ...await this.getRootEpoch(blockTag), } } - private async getRootEpochCount(blockTag?: number): Promise { + // The current root epoch and its chain hash; empty when none is registered + // or the contract predates root epochs. + async getRootEpoch(blockTag?: number): Promise { try { - return Number(await this.kmsContract.rootEpochCount({ blockTag })); + const [epoch, chainHash] = await this.kmsContract.currentRootEpoch({ blockTag }); + return { rootEpoch: Number(epoch), rootEpochHash: chainHash }; } catch (error) { - if (isMissingFunction(error)) return 0; + if (isMissingFunction(error)) return {}; throw error; } } - async getRootEpochs(): Promise { - const blockTag = await this.provider.getBlockNumber(); - const count = await this.getRootEpochCount(blockTag); - return Promise.all( - Array.from({ length: count }, async (_, epoch) => { - const [caPubkey, k256Pubkey] = await this.kmsContract.rootEpoch(epoch, { blockTag }); - return { caPubkey, k256Pubkey }; - }) - ); - } - async getGatewayAppId(): Promise { return await this.kmsContract.gatewayAppId(); } diff --git a/dstack/kms/auth-eth/src/main.test.ts b/dstack/kms/auth-eth/src/main.test.ts index 1a6beff9d..2d999a635 100644 --- a/dstack/kms/auth-eth/src/main.test.ts +++ b/dstack/kms/auth-eth/src/main.test.ts @@ -94,7 +94,8 @@ describe('Server', () => { isAllowed: true, reason: '', gatewayAppId: '', - rootEpoch: 2, + rootEpoch: 1, + rootEpochHash: '0x' + '11'.repeat(32), }); const response = await app.inject({ @@ -103,22 +104,22 @@ describe('Server', () => { payload: mockBootInfo }); - expect(JSON.parse(response.payload).rootEpoch).toBe(2); + expect(JSON.parse(response.payload)).toMatchObject({ rootEpoch: 1, rootEpochHash: '0x' + '11'.repeat(32) }); }); it('should list registered key epochs in the info response', async () => { - const rootEpochs = [{ caPubkey: '0xa0', k256Pubkey: '0xb0' }]; + const rootEpoch = { rootEpoch: 1, rootEpochHash: '0x' + '11'.repeat(32) }; Object.assign(app.ethereum, { getGatewayAppId: jest.fn().mockResolvedValue('gw'), getChainId: jest.fn().mockResolvedValue(1), getAppImplementation: jest.fn().mockResolvedValue('0x00'), - getRootEpochs: jest.fn().mockResolvedValue(rootEpochs), + getRootEpoch: jest.fn().mockResolvedValue(rootEpoch), }); const response = await app.inject({ method: 'GET', url: '/' }); expect(response.statusCode).toBe(200); - expect(JSON.parse(response.payload).rootEpochs).toEqual(rootEpochs); + expect(JSON.parse(response.payload)).toMatchObject(rootEpoch); }); }); diff --git a/dstack/kms/auth-eth/src/server.ts b/dstack/kms/auth-eth/src/server.ts index 747c4f3d5..cc52e90f8 100644 --- a/dstack/kms/auth-eth/src/server.ts +++ b/dstack/kms/auth-eth/src/server.ts @@ -51,6 +51,7 @@ export async function build(): Promise { reason: { type: 'string' }, gatewayAppId: { type: 'string' }, rootEpoch: { type: 'integer' }, + rootEpochHash: { type: 'string' }, } }); @@ -85,7 +86,7 @@ export async function build(): Promise { server.ethereum.getGatewayAppId(), server.ethereum.getChainId(), server.ethereum.getAppImplementation(), - server.ethereum.getRootEpochs(), + server.ethereum.getRootEpoch(), ]); return { status: 'ok', @@ -95,7 +96,7 @@ export async function build(): Promise { chainId: batch[1], appAuthImplementation: batch[2], // NOTE: for backward compatibility appImplementation: batch[2], - rootEpochs: batch[3], + ...batch[3], }; } catch { _request.log.error('authorization backend health check failed'); diff --git a/dstack/kms/auth-eth/src/types.ts b/dstack/kms/auth-eth/src/types.ts index 1ed4ad1a2..57ca87720 100644 --- a/dstack/kms/auth-eth/src/types.ts +++ b/dstack/kms/auth-eth/src/types.ts @@ -18,13 +18,15 @@ export interface BootResponse { isAllowed: boolean; gatewayAppId: string; reason: string; - // Latest registered KMS root epoch; absent when none is registered. rootEpoch?: number; + rootEpochHash?: string; } -export interface RootEpoch { - caPubkey: string; - k256Pubkey: string; +// The KMS root epoch currently served and its chain hash, which commits to the +// root public keys of epochs 0..=rootEpoch. Both absent when none is registered. +export interface RootEpochInfo { + rootEpoch?: number; + rootEpochHash?: string; } // Removed KMS_CONTRACT_ABI and APP_CONTRACT_ABI since we're using typechain types now diff --git a/dstack/kms/auth-eth/test/DstackKms.t.sol b/dstack/kms/auth-eth/test/DstackKms.t.sol index 6278c663e..6c9910f92 100644 --- a/dstack/kms/auth-eth/test/DstackKms.t.sol +++ b/dstack/kms/auth-eth/test/DstackKms.t.sol @@ -19,7 +19,8 @@ contract DstackKmsTest is Test { address public user; event KmsInfoSet(bytes k256Pubkey); - event RootEpochRegistered(uint256 indexed epoch, bytes caPubkey, bytes k256Pubkey); + event RootEpochRegistered(uint256 indexed epoch, bytes32 chainHash, bytes caPubkey, bytes k256Pubkey); + event RootEpochActivated(uint256 indexed epoch); event AppRegistered(address appId); @@ -67,27 +68,77 @@ contract DstackKmsTest is Test { assertEq(storedEventlog, eventlog); } - function test_RegisterRootEpochAppends() public { - assertEq(kms.rootEpochCount(), 0); + function chainHash( + bytes32 previous, + uint256 epoch, + bytes memory ca, + bytes memory k256 + ) + internal + pure + returns (bytes32) + { + return keccak256(abi.encodePacked(previous, epoch, keccak256(ca), keccak256(k256))); + } + function test_RegisterRootEpochAppendsAndChains() public { + assertEq(kms.rootEpochCount(), 0); + vm.expectRevert("No root epoch registered"); + kms.currentRootEpoch(); + + bytes32 h0 = chainHash(bytes32(0), 0, hex"a0", hex"b0"); + bytes32 h1 = chainHash(h0, 1, hex"a1", hex"b1"); + // Shared with the off-chain implementations' tests. + assertEq(h0, 0xf5f9a7455970becbb0d8441e53cf06df6290b34419cf783eb5869107867ec9ae); + assertEq(h1, 0x2219fc74439588d8f9d514f4e897d4a1dc692519b0899d93f5ac5661b8e2850a); vm.startPrank(owner); vm.expectEmit(true, true, true, true); - emit RootEpochRegistered(0, hex"a0", hex"b0"); + emit RootEpochRegistered(0, h0, hex"a0", hex"b0"); + vm.expectEmit(true, true, true, true); + emit RootEpochActivated(0); assertEq(kms.registerRootEpoch(hex"a0", hex"b0", hex"c0"), 0); assertEq(kms.registerRootEpoch(hex"a1", hex"b1", hex"c1"), 1); vm.stopPrank(); assertEq(kms.rootEpochCount(), 2); - (bytes memory ca, bytes memory k256, bytes memory attestation) = kms.rootEpoch(0); - assertEq(ca, hex"a0"); - assertEq(k256, hex"b0"); - assertEq(attestation, hex"c0"); - (ca, k256, attestation) = kms.rootEpoch(1); + (bytes memory ca, bytes memory k256, bytes memory attestation, bytes32 h) = kms.rootEpoch(1); assertEq(ca, hex"a1"); assertEq(k256, hex"b1"); assertEq(attestation, hex"c1"); + assertEq(h, h1); vm.expectRevert("Unknown root epoch"); kms.rootEpoch(2); + + // Epoch 1 is registered but epoch 0 is still served. + (uint256 current, bytes32 currentHash) = kms.currentRootEpoch(); + assertEq(current, 0); + assertEq(currentHash, h0); + } + + function test_ActivateRootEpochMovesForwardOneAtATime() public { + vm.startPrank(owner); + kms.registerRootEpoch(hex"a0", hex"b0", hex"c0"); + vm.expectRevert("Invalid root epoch activation"); + kms.activateRootEpoch(1); + kms.registerRootEpoch(hex"a1", hex"b1", hex"c1"); + kms.registerRootEpoch(hex"a2", hex"b2", hex"c2"); + vm.expectRevert("Invalid root epoch activation"); + kms.activateRootEpoch(2); + vm.expectEmit(true, true, true, true); + emit RootEpochActivated(1); + kms.activateRootEpoch(1); + vm.expectRevert("Invalid root epoch activation"); + kms.activateRootEpoch(1); + vm.expectRevert("Invalid root epoch activation"); + kms.activateRootEpoch(0); + vm.stopPrank(); + + (uint256 current,) = kms.currentRootEpoch(); + assertEq(current, 1); + + vm.prank(user); + vm.expectRevert(abi.encodeWithSelector(OwnableUpgradeable.OwnableUnauthorizedAccount.selector, user)); + kms.activateRootEpoch(2); } function test_RegisterRootEpochOnlyOwner() public { diff --git a/dstack/kms/auth-mock/README.md b/dstack/kms/auth-mock/README.md index 93be1eaa8..912d8f88a 100644 --- a/dstack/kms/auth-mock/README.md +++ b/dstack/kms/auth-mock/README.md @@ -72,7 +72,7 @@ bun run check - `MOCK_GATEWAY_APP_ID` - mock gateway app ID (default: 0xmockgateway1234567890123456789012345678) - `MOCK_CHAIN_ID` - mock blockchain chain ID (default: 1337) - `MOCK_APP_IMPLEMENTATION` - mock app implementation address (default: 0xmockapp9876543210987654321098765432109) -- `MOCK_ROOT_EPOCHS` - JSON array of `{"caPubkey", "k256Pubkey"}`, indexed by KMS root epoch (default: none registered) +- `MOCK_ROOT_EPOCH`, `MOCK_ROOT_EPOCH_HASH` - current KMS root epoch and its chain hash reported as `rootEpoch`/`rootEpochHash` (default: none registered) ## API endpoints diff --git a/dstack/kms/auth-mock/index.test.ts b/dstack/kms/auth-mock/index.test.ts index 3cc672660..b56c8c132 100644 --- a/dstack/kms/auth-mock/index.test.ts +++ b/dstack/kms/auth-mock/index.test.ts @@ -161,9 +161,10 @@ describe('Mock Backend Tests', () => { }); describe('root epochs', () => { - it('reports MOCK_ROOT_EPOCHS', async () => { - const rootEpochs = [{ caPubkey: '0xa1', k256Pubkey: '0xb1' }]; - process.env.MOCK_ROOT_EPOCHS = JSON.stringify(rootEpochs); + it('reports MOCK_ROOT_EPOCH and MOCK_ROOT_EPOCH_HASH', async () => { + const expected = { rootEpoch: 1, rootEpochHash: '0x' + '11'.repeat(32) }; + process.env.MOCK_ROOT_EPOCH = '1'; + process.env.MOCK_ROOT_EPOCH_HASH = expected.rootEpochHash; try { const boot = await appFetch(new Request('http://localhost:3002/bootAuth/kms', { method: 'POST', @@ -173,11 +174,12 @@ describe('Mock Backend Tests', () => { composeHash: '0x04', instanceId: '0x05', deviceId: '0x06', }), })); - expect((await boot.json()).rootEpoch).toBe(0); + expect(await boot.json()).toMatchObject(expected); const info = await appFetch(new Request('http://localhost:3002/')); - expect((await info.json()).rootEpochs).toEqual(rootEpochs); + expect(await info.json()).toMatchObject(expected); } finally { - delete process.env.MOCK_ROOT_EPOCHS; + delete process.env.MOCK_ROOT_EPOCH; + delete process.env.MOCK_ROOT_EPOCH_HASH; } }); }); diff --git a/dstack/kms/auth-mock/index.ts b/dstack/kms/auth-mock/index.ts index eff46cf85..b49f56994 100644 --- a/dstack/kms/auth-mock/index.ts +++ b/dstack/kms/auth-mock/index.ts @@ -29,14 +29,18 @@ const BootResponseSchema = z.object({ isAllowed: z.boolean(), reason: z.string(), gatewayAppId: z.string(), - rootEpoch: z.number().optional() + rootEpoch: z.number().optional(), + rootEpochHash: z.string().optional() }); -const RootEpochsSchema = z.array(z.object({ caPubkey: z.string(), k256Pubkey: z.string() })); - -// MOCK_ROOT_EPOCHS: JSON array of {caPubkey, k256Pubkey}, indexed by root epoch -function getRootEpochs(): z.infer { - return RootEpochsSchema.parse(JSON.parse(process.env.MOCK_ROOT_EPOCHS || '[]')); +// MOCK_ROOT_EPOCH and MOCK_ROOT_EPOCH_HASH: the current KMS root epoch and its +// chain hash; both unset means none is registered +function getRootEpoch(): { rootEpoch?: number; rootEpochHash?: string } { + const epoch = process.env.MOCK_ROOT_EPOCH; + const hash = process.env.MOCK_ROOT_EPOCH_HASH; + if (!epoch && !hash) return {}; + if (!epoch || !hash) throw new Error('set both MOCK_ROOT_EPOCH and MOCK_ROOT_EPOCH_HASH'); + return { rootEpoch: z.coerce.number().int().min(0).parse(epoch), rootEpochHash: hash }; } type BootInfo = z.infer; @@ -79,8 +83,7 @@ class MockBackend { async checkBoot(bootInfo: BootInfo, isKms: boolean): Promise { const response = this.evaluateBoot(bootInfo, isKms); - const epochCount = getRootEpochs().length; - return epochCount > 0 ? { ...response, rootEpoch: epochCount - 1 } : response; + return { ...response, ...getRootEpoch() }; } private evaluateBoot(bootInfo: BootInfo, isKms: boolean): BootResponse { @@ -161,7 +164,7 @@ app.get('/', async (c) => { chainId: batch[1], appAuthImplementation: batch[2], // NOTE: for backward compatibility appImplementation: batch[2], - rootEpochs: getRootEpochs(), + ...getRootEpoch(), note: 'this is a mock backend - all authentications will succeed' }); } catch (error) { diff --git a/dstack/kms/auth-mock/openapi.json b/dstack/kms/auth-mock/openapi.json index f197a499a..e3e767409 100644 --- a/dstack/kms/auth-mock/openapi.json +++ b/dstack/kms/auth-mock/openapi.json @@ -179,7 +179,11 @@ }, "rootEpoch": { "type": "integer", - "description": "Latest KMS root epoch from MOCK_ROOT_EPOCHS; absent when none" + "description": "Current KMS root epoch from MOCK_ROOT_EPOCH; absent when none" + }, + "rootEpochHash": { + "type": "string", + "description": "Its chain hash from MOCK_ROOT_EPOCH_HASH" } } }, @@ -225,17 +229,8 @@ "description": "Mock app implementation address", "example": "0xmockapp9876543210987654321098765432109" }, - "rootEpochs": { - "type": "array", - "description": "KMS root epochs from MOCK_ROOT_EPOCHS, indexed by epoch", - "items": { - "type": "object", - "properties": { - "caPubkey": { "type": "string" }, - "k256Pubkey": { "type": "string" } - } - } - }, + "rootEpoch": { "type": "integer", "description": "MOCK_ROOT_EPOCH" }, + "rootEpochHash": { "type": "string", "description": "MOCK_ROOT_EPOCH_HASH" }, "note": { "type": "string", "description": "Note about mock backend behavior", diff --git a/dstack/kms/auth-simple/README.md b/dstack/kms/auth-simple/README.md index 017d01f2d..be387fcc0 100644 --- a/dstack/kms/auth-simple/README.md +++ b/dstack/kms/auth-simple/README.md @@ -107,7 +107,8 @@ raw NitroTPM evidence by itself. | `apps..composeHashes` | No | Allowed compose hashes for this app | | `apps..devices` | No | Allowed device IDs for this app. Empty denies all unless `allowAnyDevice` is true. | | `apps..allowAnyDevice` | No | If true, skip device ID check for this app | -| `rootEpochs` | No | KMS root public keys of each root epoch, indexed by epoch, `[{"caPubkey": "0x...", "k256Pubkey": "0x..."}]`. Epoch 0 is the original roots. Append only. When set, the KMS checks its keyring against it and boot responses report the latest epoch as `rootEpoch`. | +| `rootEpochs` | No | KMS root public keys of each root epoch, indexed by epoch, `[{"caPubkey": "0x...", "k256Pubkey": "0x..."}]`. Epoch 0 is the original roots. Append only. Boot and info responses report the current epoch as `rootEpoch` with its chain hash as `rootEpochHash`. | +| `currentRootEpoch` | No | Root epoch KMS instances serve (default 0). Raise it only once every KMS node holds that epoch. | For experimental AMD SEV-SNP dry-run authorization, keep the default fail-closed TCB policy unless you intentionally want the auth webhook to accept non-up-to-date verifier-derived SNP `BootInfo`. To exercise the dry-run path without enabling key release, allowlist the recomputed SNP `mrAggregated`, `osImageHash`, app/compose identity, device/chip identity, and any non-default `allowedTcbStatuses`/`allowedAdvisoryIds` values explicitly. KMS still rejects SNP before returning app keys, KMS keys, or app certificates. diff --git a/dstack/kms/auth-simple/bun.lock b/dstack/kms/auth-simple/bun.lock index daaf50f92..6350fe8d5 100644 --- a/dstack/kms/auth-simple/bun.lock +++ b/dstack/kms/auth-simple/bun.lock @@ -6,6 +6,7 @@ "name": "auth-simple", "dependencies": { "@hono/zod-validator": "0.9.1", + "@noble/hashes": "2.4.0", "hono": "4.13.14", "zod": "4.6.7", }, @@ -29,6 +30,8 @@ "@jridgewell/trace-mapping": ["@jridgewell/trace-mapping@0.3.31", "", { "dependencies": { "@jridgewell/resolve-uri": "^3.1.0", "@jridgewell/sourcemap-codec": "^1.4.14" } }, "sha512-zzNR+SdQSDJzc8joaeP8QQoCQr8NuYx2dIIytl1QeBEZHJ9uW6hebsrYgbz8hJwUQao3TWCMtmfV8Nu1twOLAw=="], + "@noble/hashes": ["@noble/hashes@2.4.0", "", {}, "sha512-X5XaVWZIBCT7HHZGm5I7ZQXDwLG+bGXuSrMQAW+7Zvl87h1kmc1ZB1VSRJcpUfoUrGQp4Fkoxm5kZ+Ms+aW+eA=="], + "@oxc-project/types": ["@oxc-project/types@0.155.0", "", {}, "sha512-SeE8JPwW1xs1tHeNycjUtUBdhorRa8o4jlvjwk49R7/Z291i+4DQa7JJIgOm/TaYQZx+QrPQbVJ/9c8dMIZN8Q=="], "@oxlint/darwin-arm64": ["@oxlint/darwin-arm64@0.9.10", "", { "os": "darwin", "cpu": "arm64" }, "sha512-eOXKZYq5bnCSgDefgM5bzAg+4Fc//Rc4yjgKN8iDWUARweCaChiQXb6TXX8MfEfs6qayEMy6yVj0pqoFz0B1aw=="], diff --git a/dstack/kms/auth-simple/index.test.ts b/dstack/kms/auth-simple/index.test.ts index 4ba189e9c..dc3572e68 100644 --- a/dstack/kms/auth-simple/index.test.ts +++ b/dstack/kms/auth-simple/index.test.ts @@ -446,35 +446,41 @@ describe('auth-simple', () => { }); }); - describe('key epochs', () => { + describe('root epochs', () => { const rootEpochs = [ { caPubkey: '0xa0', k256Pubkey: '0xb0' }, { caPubkey: '0xa1', k256Pubkey: '0xb1' } ]; + // Shared with the DstackKms contract tests. + const h0 = '0xf5f9a7455970becbb0d8441e53cf06df6290b34419cf783eb5869107867ec9ae'; + const h1 = '0x2219fc74439588d8f9d514f4e897d4a1dc692519b0899d93f5ac5661b8e2850a'; - it('reports the latest key epoch in boot responses and lists all epochs', async () => { - writeTestConfig({ gatewayAppId: '0xgateway', rootEpochs }); - + const bootRootEpoch = async () => { const boot = await app.fetch(new Request('http://localhost/bootAuth/app', { method: 'POST', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify(baseBootInfo) })); - expect((await boot.json()).rootEpoch).toBe(1); + const { rootEpoch, rootEpochHash } = await boot.json(); + return { rootEpoch, rootEpochHash }; + }; + const info = async () => + (await app.fetch(new Request('http://localhost/', { method: 'GET' }))).json(); - const info = await app.fetch(new Request('http://localhost/', { method: 'GET' })); - expect((await info.json()).rootEpochs).toEqual(rootEpochs); + it('reports the current root epoch and its chain hash', async () => { + writeTestConfig({ gatewayAppId: '0xgateway', rootEpochs }); + expect(await bootRootEpoch()).toEqual({ rootEpoch: 0, rootEpochHash: h0 }); + expect(await info()).toMatchObject({ rootEpoch: 0, rootEpochHash: h0 }); + + writeTestConfig({ gatewayAppId: '0xgateway', rootEpochs, currentRootEpoch: 1 }); + expect(await bootRootEpoch()).toEqual({ rootEpoch: 1, rootEpochHash: h1 }); + expect(await info()).toMatchObject({ rootEpoch: 1, rootEpochHash: h1 }); }); - it('omits the key epoch when none is configured', async () => { + it('omits the root epoch when none is configured', async () => { writeTestConfig({ gatewayAppId: '0xgateway' }); - - const boot = await app.fetch(new Request('http://localhost/bootAuth/app', { - method: 'POST', - headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify(baseBootInfo) - })); - expect((await boot.json()).rootEpoch).toBeUndefined(); + expect(await bootRootEpoch()).toEqual({ rootEpoch: undefined, rootEpochHash: undefined }); + expect((await info()).rootEpoch).toBeUndefined(); }); }); diff --git a/dstack/kms/auth-simple/index.ts b/dstack/kms/auth-simple/index.ts index 84e875a3e..eaea5da0b 100644 --- a/dstack/kms/auth-simple/index.ts +++ b/dstack/kms/auth-simple/index.ts @@ -6,6 +6,8 @@ import { Hono } from 'hono'; import { zValidator } from '@hono/zod-validator'; import { z } from 'zod'; import { readFileSync, existsSync } from 'fs'; +import { keccak_256 } from '@noble/hashes/sha3.js'; +import { bytesToHex, concatBytes, hexToBytes } from '@noble/hashes/utils.js'; // zod schemas for validation - compatible with auth-eth implementation const BootInfoSchema = z.object({ @@ -25,7 +27,8 @@ const BootResponseSchema = z.object({ isAllowed: z.boolean(), reason: z.string(), gatewayAppId: z.string(), - rootEpoch: z.number().optional() + rootEpoch: z.number().optional(), + rootEpochHash: z.string().optional() }); const RootEpochSchema = z.object({ @@ -60,13 +63,37 @@ const AuthConfigSchema = z.object({ kms: KmsConfigSchema.default({}), apps: z.record(z.string(), AppConfigSchema).default({}), // Root public keys of each KMS root epoch, indexed by epoch; append only. - rootEpochs: z.array(RootEpochSchema).default([]) + rootEpochs: z.array(RootEpochSchema).default([]), + // The root epoch KMS instances serve; raise it once every node holds it. + currentRootEpoch: z.number().int().min(0).default(0) }); type BootInfo = z.infer; type BootResponse = z.infer; type AuthConfig = z.infer; +const bytesOf = (hex: string) => hexToBytes(hex.replace(/^0x/i, '')); + +// The current root epoch and its chain hash, as DstackKms computes it: +// keccak256(previous || epoch || keccak256(caPubkey) || keccak256(k256Pubkey)), +// with a zero previous hash for epoch 0. Empty when no epoch is configured. +function currentRootEpoch(config: AuthConfig): { rootEpoch?: number; rootEpochHash?: string } { + const current = config.currentRootEpoch; + if (config.rootEpochs.length === 0) return {}; + if (current >= config.rootEpochs.length) { + throw new Error('currentRootEpoch is not a configured root epoch'); + } + let hash = new Uint8Array(32); + config.rootEpochs.slice(0, current + 1).forEach((keys, epoch) => { + const epochBytes = new Uint8Array(32); + new DataView(epochBytes.buffer).setUint32(28, epoch); + hash = keccak_256(concatBytes( + hash, epochBytes, keccak_256(bytesOf(keys.caPubkey)), keccak_256(bytesOf(keys.k256Pubkey)) + )); + }); + return { rootEpoch: current, rootEpochHash: '0x' + bytesToHex(hash) }; +} + // normalize hex string to lowercase with 0x prefix function normalizeHex(hex: string): string { hex = hex.toLowerCase(); @@ -102,8 +129,7 @@ class ConfigBackend { async checkBoot(bootInfo: BootInfo, isKms: boolean): Promise { const config = this.loadConfig(); const response = this.evaluateBoot(bootInfo, isKms, config); - const epochCount = config.rootEpochs.length; - return epochCount > 0 ? { ...response, rootEpoch: epochCount - 1 } : response; + return { ...response, ...currentRootEpoch(config) }; } private evaluateBoot(bootInfo: BootInfo, isKms: boolean, config: AuthConfig): BootResponse { @@ -234,7 +260,8 @@ class ConfigBackend { kmsContractAddr: string; chainId: number; appImplementation: string; - rootEpochs: AuthConfig['rootEpochs']; + rootEpoch?: number; + rootEpochHash?: string; }> { const config = this.loadConfig(); return { @@ -242,7 +269,7 @@ class ConfigBackend { kmsContractAddr: config.kmsContractAddr, chainId: config.chainId, appImplementation: config.appImplementation, - rootEpochs: config.rootEpochs, + ...currentRootEpoch(config), }; } } @@ -264,7 +291,8 @@ app.get('/', async (c) => { chainId: info.chainId, appAuthImplementation: info.appImplementation, // backward compat appImplementation: info.appImplementation, - rootEpochs: info.rootEpochs, + rootEpoch: info.rootEpoch, + rootEpochHash: info.rootEpochHash, }); } catch (error) { console.error('error in health check:', error); diff --git a/dstack/kms/auth-simple/package.json b/dstack/kms/auth-simple/package.json index bb0bbb7e6..eba28e617 100644 --- a/dstack/kms/auth-simple/package.json +++ b/dstack/kms/auth-simple/package.json @@ -16,6 +16,7 @@ }, "dependencies": { "@hono/zod-validator": "0.9.1", + "@noble/hashes": "2.4.0", "hono": "4.13.14", "zod": "4.6.7" }, diff --git a/dstack/kms/kms.toml b/dstack/kms/kms.toml index 378372278..22ee6d7b5 100644 --- a/dstack/kms/kms.toml +++ b/dstack/kms/kms.toml @@ -124,12 +124,6 @@ timeout = "60s" [core.auth_api.dev] gateway_app_id = "any" -[core.key_rotation] -# Root key epoch served as current; the original roots are epoch 0. Raising it requires a KMS onboarded with -# `rotate` (or from a source already at this epoch); editing it alone fails. -# Older epochs stay available so CVMs can migrate their data at any time. -current_epoch = 0 - [core.onboard] enabled = true auto_bootstrap_domain = "" diff --git a/dstack/kms/rpc/proto/kms_rpc.proto b/dstack/kms/rpc/proto/kms_rpc.proto index 5b5ba57aa..fac7aad15 100644 --- a/dstack/kms/rpc/proto/kms_rpc.proto +++ b/dstack/kms/rpc/proto/kms_rpc.proto @@ -14,6 +14,18 @@ message GetAppKeyRequest { // Root key epoch to derive from; unset means the current epoch. Epoch 0 is // the original root keys. Older epochs remain available so CVMs can migrate. optional uint32 root_epoch = 3; + // Also return the app keys of every epoch before the selected one, so a CVM + // migrating off older epochs needs a single request. + bool include_history = 4; +} + +// App keys derived from one older root epoch. +message HistoricalAppKey { + uint32 root_epoch = 1; + bytes disk_crypt_key = 2; + bytes env_crypt_key = 3; + bytes k256_key = 4; + bytes k256_signature = 5; } message AppId { @@ -51,6 +63,8 @@ message AppKeyResponse { bytes os_image_hash = 8; // Root key epoch all epoch-dependent fields above were derived from. uint32 root_epoch = 9; + // With `include_history`, the app keys of epochs 0..root_epoch in order. + repeated HistoricalAppKey history = 10; } message GetMetaResponse { diff --git a/dstack/kms/src/config.rs b/dstack/kms/src/config.rs index 4eac5be85..f90bdf9ce 100644 --- a/dstack/kms/src/config.rs +++ b/dstack/kms/src/config.rs @@ -2,7 +2,6 @@ // // SPDX-License-Identifier: Apache-2.0 -use anyhow::{bail, Result}; use dstack_attest::attestation::AttestationVerifierConfig; use load_config::load_config; use rocket::figment::Figment; @@ -41,8 +40,6 @@ pub(crate) struct KmsConfig { pub attestation: AttestationVerifierConfig, pub auth_api: AuthApi, pub onboard: OnboardConfig, - #[serde(default)] - pub key_rotation: KeyRotationConfig, pub image: ImageConfig, /// Whether to enable the additional local release gate for AMD SEV-SNP /// key/cert material. This is separate from the auth API so production @@ -176,16 +173,17 @@ impl KmsConfig { self.cert_dir.join(K256_KEY) } - /// Records which epoch the top-level root keys belong to. Absent means 0. + /// Records the newest epoch this KMS holds, whose roots are the top-level + /// root files. Absent means 0. pub fn root_epoch(&self) -> PathBuf { self.cert_dir.join(ROOT_EPOCH) } - /// The current epoch's roots are the top-level root files, so every path - /// that only needs the current roots is unchanged. Older epochs live in - /// `epochs//`. - pub fn root_key_paths(&self, epoch: u32) -> RootKeyPaths { - let dir = if epoch == self.key_rotation.current_epoch { + /// Root key paths of `epoch` in a keyring whose newest epoch is `latest`: + /// the newest roots are the top-level root files, so paths that predate + /// epochs are unchanged. Older epochs live in `epochs//`. + pub fn root_key_paths(&self, epoch: u32, latest: u32) -> RootKeyPaths { + let dir = if epoch == latest { self.cert_dir.clone() } else { self.cert_dir.join(EPOCHS_DIR).join(epoch.to_string()) @@ -241,26 +239,6 @@ pub(crate) struct RootKeyPaths { pub k256_key: PathBuf, } -#[derive(Debug, Clone, Default, Deserialize)] -pub(crate) struct KeyRotationConfig { - /// Root key epoch served as current; epoch 0 is the original roots. Every - /// KMS holds all epochs up to it. - #[serde(default)] - pub current_epoch: u32, -} - -impl KeyRotationConfig { - /// Resolve a root epoch selector: unset is the current epoch. - pub fn resolve_epoch(&self, requested: Option) -> Result { - let current = self.current_epoch; - match requested { - None => Ok(current), - Some(epoch) if epoch <= current => Ok(epoch), - Some(epoch) => bail!("unknown root epoch {epoch}; the current epoch is {current}"), - } - } -} - #[derive(Debug, Clone, Deserialize)] pub(crate) struct OnboardConfig { pub enabled: bool, @@ -290,16 +268,6 @@ mod tests { assert!(!config.admin.insecure_no_auth); } - #[test] - fn root_epoch_selector_resolution() { - let cfg = KeyRotationConfig { current_epoch: 2 }; - assert_eq!(cfg.resolve_epoch(None).unwrap(), 2); - assert_eq!(cfg.resolve_epoch(Some(0)).unwrap(), 0); - assert_eq!(cfg.resolve_epoch(Some(2)).unwrap(), 2); - assert!(cfg.resolve_epoch(Some(3)).is_err()); - assert_eq!(KeyRotationConfig::default().resolve_epoch(None).unwrap(), 0); - } - #[test] fn rpc_cert_is_attested_by_default() { let figment = load_config_figment(None); diff --git a/dstack/kms/src/main.rs b/dstack/kms/src/main.rs index 526fd7f9b..2da3f26ef 100644 --- a/dstack/kms/src/main.rs +++ b/dstack/kms/src/main.rs @@ -143,8 +143,13 @@ async fn main() -> Result<()> { } } + let keyring = main_service::Keyring::load(&config).context("Failed to load the root keys")?; + let serving_epoch = main_service::wait_for_servable_epoch(&config, &keyring).await; + info!("Serving root epoch {serving_epoch}"); + info!("Updating certs"); - if let Err(err) = onboard_service::update_certs(&config).await { + if let Err(err) = onboard_service::update_certs(&config, serving_epoch, keyring.latest()).await + { if config.attest_rpc_cert { return Err(err).context( "Failed to reissue the attested KMS RPC certificate; refusing to start with a \ @@ -176,11 +181,8 @@ async fn main() -> Result<()> { } else { None }; - let state = main_service::KmsState::new(config).context("Failed to initialize KMS state")?; - state - .ensure_keyring_registered() - .await - .context("KMS keyring does not match the registered key epochs")?; + let state = main_service::KmsState::new(config, keyring, serving_epoch) + .context("Failed to initialize KMS state")?; let quote_verifier = QuoteVerifier::new(state.attestation_verifier()); let admin_quote_verifier = QuoteVerifier::new(state.attestation_verifier()); let figment = figment @@ -217,6 +219,12 @@ async fn main() -> Result<()> { rocket = rocket.manage(quote_verifier); let main_srv = rocket.launch(); + // A move of the current root epoch on chain ends the process so it can be + // restarted onto the new epoch. + let restart = { + let state = state.clone(); + async move { state.restart_requested().await } + }; match admin_setup { Some((admin_figment, admin_fairing)) => { if admin_config.insecure_no_auth { @@ -237,14 +245,22 @@ async fn main() -> Result<()> { .manage(state) .manage(admin_quote_verifier) .launch(); - tokio::try_join!( - async { main_srv.await.map_err(|err| anyhow!(err.to_string())) }, - async { admin_srv.await.map_err(|err| anyhow!(err.to_string())) }, - )?; - } - None => { - main_srv.await.map_err(|err| anyhow!(err.to_string()))?; + let servers = async { + tokio::try_join!( + async { main_srv.await.map_err(|err| anyhow!(err.to_string())) }, + async { admin_srv.await.map_err(|err| anyhow!(err.to_string())) }, + ) + .map(drop) + }; + tokio::select! { + result = servers => result?, + () = restart => bail!("the current root epoch changed; restarting to follow it"), + } } + None => tokio::select! { + result = main_srv => result.map(drop).map_err(|err| anyhow!(err.to_string()))?, + () = restart => bail!("the current root epoch changed; restarting to follow it"), + }, } Ok(()) } diff --git a/dstack/kms/src/main_service.rs b/dstack/kms/src/main_service.rs index d159599c7..8030f7b1e 100644 --- a/dstack/kms/src/main_service.rs +++ b/dstack/kms/src/main_service.rs @@ -14,8 +14,8 @@ use anyhow::{anyhow, bail, Context, Result}; use dstack_kms_rpc::{ kms_server::{KmsRpc, KmsServer}, AppId, AppKeyResponse, GetAppKeyRequest, GetKmsKeyRequest, GetMetaResponse, - GetTempCaCertResponse, KmsKeyResponse, KmsKeys, PublicKeyResponse, SignCertRequest, - SignCertResponse, + GetTempCaCertResponse, HistoricalAppKey, KmsKeyResponse, KmsKeys, PublicKeyResponse, + SignCertRequest, SignCertResponse, }; use dstack_verifier::{CvmVerifier, VerificationDetails}; use fs_err as fs; @@ -31,7 +31,7 @@ use scale::Decode; use tracing::warn; use upgrade_authority::{ build_boot_info, ensure_app_id_len, local_kms_boot_info, BootInfo, BootResponse, - GetInfoResponse, RootEpochPubkeys, + GetInfoResponse, RootEpochRegistry, }; use crate::{ @@ -82,19 +82,82 @@ impl RootKeys { let k256 = SigningKey::from_slice(&key_bytes).context("failed to load ECDSA root key")?; Ok(Self { ca, k256 }) } +} + +/// The root keys of epochs 0..=latest this KMS holds, and their chain hashes. +pub(crate) struct Keyring { + roots: Vec, + chain_hashes: Vec<[u8; 32]>, +} - fn pubkeys(&self) -> RootEpochPubkeys { - RootEpochPubkeys { - ca_pubkey: self.ca.key.public_key_der(), - k256_pubkey: self.k256.verifying_key().to_sec1_bytes().to_vec(), +impl Keyring { + pub(crate) fn load(config: &KmsConfig) -> Result { + let latest = read_root_epoch(config)?; + let roots = (0..=latest) + .map(|epoch| { + RootKeys::load(&config.root_key_paths(epoch, latest)) + .with_context(|| format!("failed to load root keys of epoch {epoch}")) + }) + .collect::>>()?; + let pubkeys: Vec<_> = roots + .iter() + .map(|root| { + ( + root.ca.key.public_key_der(), + root.k256.verifying_key().to_sec1_bytes().to_vec(), + ) + }) + .collect(); + let chain_hashes = dstack_types::root_epoch::chain_hashes( + pubkeys.iter().map(|(ca, k256)| (&ca[..], &k256[..])), + ); + Ok(Self { + roots, + chain_hashes, + }) + } + + pub(crate) fn latest(&self) -> u32 { + self.roots.len() as u32 - 1 + } + + /// The epoch to serve under `registry`: its current epoch, which this + /// keyring must hold with a matching chain hash. A keyring nobody + /// registered is servable only at epoch 0, as before root epochs existed. + pub(crate) fn servable_epoch(&self, registry: RootEpochRegistry) -> Result { + let latest = self.latest(); + match registry { + RootEpochRegistry::Unchecked => Ok(latest), + RootEpochRegistry::Empty if latest == 0 => Ok(0), + RootEpochRegistry::Empty => { + bail!("this KMS holds root epoch {latest}, but no root epoch is registered") + } + RootEpochRegistry::Current(current) => { + let Some(local) = self.chain_hashes.get(current.epoch as usize) else { + bail!( + "root epoch {} is current, but this KMS holds only up to {latest}", + current.epoch + ); + }; + if *local != current.chain_hash { + bail!( + "the root keys up to epoch {} do not match the registered ones", + current.epoch + ); + } + Ok(current.epoch) + } } } } pub struct KmsStateInner { config: KmsConfig, - /// Root keys of every epoch, indexed by epoch; the last is current. - roots: Vec, + keyring: Keyring, + /// The epoch served for this process's lifetime; a newer current epoch on + /// chain makes the process restart. + serving_epoch: u32, + restart: tokio::sync::Notify, temp_ca_cert: String, temp_ca_key: String, verifier: CvmVerifier, @@ -151,21 +214,10 @@ impl KmsState { self.verifier.clear_cache(image_hash, config_hash) } - pub fn new(config: KmsConfig) -> Result { - let current_epoch = config.key_rotation.current_epoch; - let stored_epoch = read_root_epoch(&config)?; - if stored_epoch != current_epoch { - bail!( - "stored root keys belong to epoch {stored_epoch}, but key_rotation.current_epoch \ - is {current_epoch}; a new epoch must be imported or generated by onboarding" - ); + pub fn new(config: KmsConfig, keyring: Keyring, serving_epoch: u32) -> Result { + if serving_epoch > keyring.latest() { + bail!("cannot serve root epoch {serving_epoch}, which this KMS does not hold"); } - let roots = (0..=current_epoch) - .map(|epoch| { - RootKeys::load(&config.root_key_paths(epoch)) - .with_context(|| format!("failed to load root keys of epoch {epoch}")) - }) - .collect::>>()?; let temp_ca_key = fs::read_to_string(config.tmp_ca_key()).context("Faeild to read temp ca key")?; let temp_ca_cert = @@ -188,7 +240,9 @@ impl KmsState { Ok(Self { inner: Arc::new(KmsStateInner { config, - roots, + keyring, + serving_epoch, + restart: tokio::sync::Notify::new(), temp_ca_cert, temp_ca_key, verifier, @@ -212,63 +266,105 @@ impl KmsState { self.inner.attestation_verifier.clone() } - /// The current epoch's root keys. + /// The served epoch's root keys. fn root(&self) -> &RootKeys { - self.roots.last().expect("the keyring always holds epoch 0") + self.epoch_root(self.serving_epoch) } - /// The root keys of `epoch`, which must be one this KMS serves. fn epoch_root(&self, epoch: u32) -> &RootKeys { - &self.roots[epoch as usize] - } - - /// Refuse to serve a keyring governance has not registered. Deployments - /// that never registered epochs keep working while they stay at epoch 0. - pub async fn ensure_keyring_registered(&self) -> Result<()> { - let current = self.config.key_rotation.current_epoch; - let info = match self.config.auth_api.get_info().await { - Ok(info) => info, - Err(err) if current == 0 => { - warn!("skipping key epoch registration check: {err:#}"); - return Ok(()); - } - Err(err) => return Err(err).context("failed to read registered key epochs"), - }; - let Some(registered) = info.root_epochs else { - return Ok(()); - }; - if registered.is_empty() && current == 0 { - return Ok(()); + &self.keyring.roots[epoch as usize] + } + + /// Resolve a root epoch selector: unset is the served epoch. Epochs held + /// but not yet current are not served. + fn resolve_epoch(&self, requested: Option) -> Result { + let serving = self.serving_epoch; + match requested { + None => Ok(serving), + Some(epoch) if epoch <= serving => Ok(epoch), + Some(epoch) => bail!("root epoch {epoch} is not served; the current one is {serving}"), } - if registered.len() != current as usize + 1 { - bail!( - "this KMS serves key epoch {current}, but {} epochs are registered", - registered.len() - ); + } + + /// Resolves once governance moves the current root epoch away from the + /// served one, so the process can restart to follow it. + pub async fn restart_requested(&self) { + self.restart.notified().await + } + + /// Serve only while governance keeps the served epoch current. A move to + /// another epoch requests a restart, which re-checks the keyring and + /// re-signs the RPC certificate. + pub(crate) fn ensure_epoch_current(&self, response: &BootResponse) -> Result<()> { + if self.config.auth_api.is_dev() { + return Ok(()); } - for (epoch, (root, registered)) in (0..).zip(self.roots.iter().zip(®istered)) { - if root.pubkeys() != *registered { - bail!("root keys of epoch {epoch} do not match the registered public keys"); + let serving = self.serving_epoch; + match response.root_epoch.registry()? { + RootEpochRegistry::Current(current) + if current.epoch == serving + && current.chain_hash == self.keyring.chain_hashes[serving as usize] => + { + Ok(()) + } + RootEpochRegistry::Empty if serving == 0 => Ok(()), + registry => { + warn!("governance moved off root epoch {serving} ({registry:?}); restarting"); + self.restart.notify_one(); + bail!("root epoch {serving} is no longer current; the KMS is restarting") } } - Ok(()) } +} - /// Stop serving once governance registers a newer epoch than this KMS has. - pub(crate) fn ensure_epoch_current(&self, response: &BootResponse) -> Result<()> { - let current = self.config.key_rotation.current_epoch; - match response.root_epoch { - Some(epoch) if epoch != current => { - bail!("the registered key epoch is {epoch}, but this KMS serves epoch {current}") +/// Wait until governance lets `keyring` serve, and return the epoch to serve. +/// A keyring that predates root epochs keeps serving epoch 0 when the auth API +/// is unreachable, as before. +pub(crate) async fn wait_for_servable_epoch(config: &KmsConfig, keyring: &Keyring) -> u32 { + const RETRY: Duration = Duration::from_secs(10); + loop { + let result = match config.auth_api.get_info().await { + Ok(info) => keyring.servable_epoch(info.root_epochs), + Err(err) if keyring.latest() == 0 => { + warn!("skipping the root epoch check: {err:#}"); + return 0; } - None if current > 0 && !self.config.auth_api.is_dev() => { - bail!("the auth API reports no key epoch, but this KMS serves epoch {current}") + Err(err) => Err(err.context("failed to read the current root epoch")), + }; + match result { + Ok(epoch) => return epoch, + Err(err) => { + warn!("not serving yet: {err:#}; retrying in {}s", RETRY.as_secs()); + tokio::time::sleep(RETRY).await; } - _ => Ok(()), } } } +struct AppKeys { + disk_crypt_key: Vec, + env_crypt_key: Vec, + k256_key: Vec, + k256_signature: Vec, +} + +fn derive_app_keys(root: &RootKeys, app_id: &[u8], instance_id: &[u8]) -> Result { + let disk_crypt_key = + kdf::derive_dh_secret(&root.ca.key, &[app_id, instance_id, b"app-disk-crypt-key"]) + .context("Failed to derive app disk key")?; + let env_crypt_key = kdf::derive_dh_secret(&root.ca.key, &[app_id, b"env-encrypt-key"]) + .context("Failed to derive env encrypt key")?; + let env_crypt_key = x25519_dalek::StaticSecret::from(env_crypt_key).to_bytes(); + let (k256_key, k256_signature) = + derive_k256_key(&root.k256, app_id).context("Failed to derive app ecdsa key")?; + Ok(AppKeys { + disk_crypt_key: disk_crypt_key.to_vec(), + env_crypt_key: env_crypt_key.to_vec(), + k256_key: k256_key.to_bytes().to_vec(), + k256_signature, + }) +} + pub(crate) fn read_root_epoch(config: &KmsConfig) -> Result { match fs::read_to_string(config.root_epoch()) { Ok(epoch) => epoch.trim().parse().context("invalid root epoch file"), @@ -299,7 +395,7 @@ impl RpcHandler { Ok(KmsKeyResponse { temp_ca_key: self.state.inner.temp_ca_key.clone(), keys: (0..) - .zip(&self.state.roots) + .zip(&self.state.keyring.roots) .map(|(epoch, root)| KmsKeys { ca_key: root.ca.key.serialize_pem(), k256_key: root.k256.to_bytes().to_vec(), @@ -506,42 +602,43 @@ impl KmsRpc for RpcHandler { .await .context("App not allowed")?; ensure_key_release_allowed(&boot_info, (&self.state.config).into())?; - let root_epoch = self - .state - .config - .key_rotation - .resolve_epoch(request.root_epoch)?; - let root = self.state.epoch_root(root_epoch); - let app_id = boot_info.app_id; - let instance_id = boot_info.instance_id; - let os_image_hash = boot_info.os_image_hash; - - let context_data = vec![&app_id[..], &instance_id[..], b"app-disk-crypt-key"]; - let app_disk_key = kdf::derive_dh_secret(&root.ca.key, &context_data) - .context("Failed to derive app disk key")?; - let env_crypt_key = { - let secret = kdf::derive_dh_secret(&root.ca.key, &[&app_id[..], b"env-encrypt-key"]) - .context("Failed to derive env encrypt key")?; - let secret = x25519_dalek::StaticSecret::from(secret); - secret.to_bytes() + let root_epoch = self.state.resolve_epoch(request.root_epoch)?; + let derive = |epoch| { + derive_app_keys( + self.state.epoch_root(epoch), + &boot_info.app_id, + &boot_info.instance_id, + ) }; - - let (k256_key, k256_signature) = { - let (k256_app_key, signature) = - derive_k256_key(&root.k256, &app_id).context("Failed to derive app ecdsa key")?; - (k256_app_key.to_bytes().to_vec(), signature) + let keys = derive(root_epoch)?; + let history = if request.include_history { + (0..root_epoch) + .map(|epoch| { + let keys = derive(epoch)?; + Ok(HistoricalAppKey { + root_epoch: epoch, + disk_crypt_key: keys.disk_crypt_key, + env_crypt_key: keys.env_crypt_key, + k256_key: keys.k256_key, + k256_signature: keys.k256_signature, + }) + }) + .collect::>()? + } else { + Vec::new() }; Ok(AppKeyResponse { - ca_cert: root.ca.pem_cert.clone(), - disk_crypt_key: app_disk_key.to_vec(), - env_crypt_key: env_crypt_key.to_vec(), - k256_key, - k256_signature, + ca_cert: self.state.epoch_root(root_epoch).ca.pem_cert.clone(), + disk_crypt_key: keys.disk_crypt_key, + env_crypt_key: keys.env_crypt_key, + k256_key: keys.k256_key, + k256_signature: keys.k256_signature, tproxy_app_id: gateway_app_id.clone(), gateway_app_id, - os_image_hash, + os_image_hash: boot_info.os_image_hash, root_epoch, + history, }) } @@ -1441,7 +1538,8 @@ mod tests { .focus("core") .extract() .unwrap(); - KmsState::new(config).unwrap() + let keyring = Keyring::load(&config).unwrap(); + KmsState::new(config, keyring, 0).unwrap() } const AUTH_API_INFO: &str = r#"{"status":"ok","kmsContractAddr":"0xkms","ethRpcUrl":"https://rpc.example","gatewayAppId":"0xgateway","chainId":1,"appImplementation":"0ximpl"}"#; diff --git a/dstack/kms/src/main_service/upgrade_authority.rs b/dstack/kms/src/main_service/upgrade_authority.rs index 93df207b3..6d9379151 100644 --- a/dstack/kms/src/main_service/upgrade_authority.rs +++ b/dstack/kms/src/main_service/upgrade_authority.rs @@ -84,24 +84,55 @@ pub(crate) struct BootResponse { pub is_allowed: bool, pub gateway_app_id: String, pub reason: String, - /// Latest KMS key epoch registered by governance, if any. - #[serde(default)] - pub root_epoch: Option, + #[serde(flatten)] + pub root_epoch: RootEpochFields, } -/// Root public keys governance registered for one KMS key epoch. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +/// The root epoch fields auth backends add to boot and info responses. +#[derive(Debug, Clone, Default, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] -pub(crate) struct RootEpochPubkeys { - #[serde(deserialize_with = "deserialize_hex")] - pub ca_pubkey: Vec, - #[serde(deserialize_with = "deserialize_hex")] - pub k256_pubkey: Vec, +pub(crate) struct RootEpochFields { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub root_epoch: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub root_epoch_hash: Option, +} + +impl RootEpochFields { + pub fn registry(&self) -> Result { + match (self.root_epoch, &self.root_epoch_hash) { + (None, None) => Ok(RootEpochRegistry::Empty), + (Some(epoch), Some(hash)) => { + let chain_hash = hex::decode(hash.trim_start_matches("0x")) + .ok() + .and_then(|hash| hash.try_into().ok()) + .context("invalid root epoch chain hash")?; + Ok(RootEpochRegistry::Current(CurrentRootEpoch { + epoch, + chain_hash, + })) + } + _ => bail!("the auth API reports a root epoch without its chain hash"), + } + } +} + +/// The root epoch governance made current, with the chain hash committing to +/// the root public keys of epochs 0..=epoch. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub(crate) struct CurrentRootEpoch { + pub epoch: u32, + pub chain_hash: [u8; 32], } -fn deserialize_hex<'de, D: serde::Deserializer<'de>>(deserializer: D) -> Result, D::Error> { - let value = String::deserialize(deserializer)?; - hex::decode(value.trim_start_matches("0x")).map_err(serde::de::Error::custom) +/// What governance says about root epochs. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub(crate) enum RootEpochRegistry { + /// The dev backend: there is nothing to check against. + Unchecked, + /// No epoch is registered, or the backend predates root epochs. + Empty, + Current(CurrentRootEpoch), } #[derive(Debug, Serialize, Deserialize)] @@ -114,8 +145,8 @@ pub(crate) struct AuthApiInfoResponse { pub gateway_app_id: String, pub chain_id: u64, pub app_implementation: String, - #[serde(default)] - pub root_epochs: Vec, + #[serde(flatten)] + pub root_epoch: RootEpochFields, } #[derive(Debug, Clone, Serialize, Deserialize)] @@ -127,8 +158,7 @@ pub(crate) struct GetInfoResponse { pub eth_rpc_url: Option, pub chain_id: Option, pub app_implementation: Option, - /// Registered key epochs, indexed by epoch. `None` means no registry. - pub root_epochs: Option>, + pub root_epochs: RootEpochRegistry, } const AUTH_API_CONNECT_TIMEOUT: Duration = Duration::from_secs(5); @@ -196,7 +226,7 @@ impl AuthApi { is_allowed: true, reason: "".to_string(), gateway_app_id: dev.gateway_app_id.clone(), - root_epoch: None, + root_epoch: RootEpochFields::default(), }), AuthApi::Webhook { webhook } => { let path = if is_kms { @@ -219,7 +249,7 @@ impl AuthApi { gateway_app_id: Some(dev.gateway_app_id.clone()), chain_id: None, app_implementation: None, - root_epochs: None, + root_epochs: RootEpochRegistry::Unchecked, }), AuthApi::Webhook { webhook } => { let info: AuthApiInfoResponse = http_get(&webhook.url, webhook.timeout).await?; @@ -235,7 +265,7 @@ impl AuthApi { chain_id: Some(info.chain_id), gateway_app_id: Some(info.gateway_app_id.clone()), app_implementation: Some(info.app_implementation.clone()), - root_epochs: Some(info.root_epochs), + root_epochs: info.root_epoch.registry()?, }) } } @@ -547,29 +577,29 @@ pub(crate) mod tests { async fn root_epochs_are_parsed_from_info_and_boot_responses() { let (url, server) = serve(vec![ r#"{"status":"ok","kmsContractAddr":"0x00","gatewayAppId":"gw","chainId":1, - "appImplementation":"0x00","rootEpochs":[{"caPubkey":"0xa0","k256Pubkey":"b0"}]}"#, - r#"{"isAllowed":true,"gatewayAppId":"gw","reason":"","rootEpoch":3}"#, + "appImplementation":"0x00","rootEpoch":3,"rootEpochHash":"0x1111111111111111111111111111111111111111111111111111111111111111"}"#, + r#"{"isAllowed":true,"gatewayAppId":"gw","reason":"","rootEpoch":3,"rootEpochHash":"1111111111111111111111111111111111111111111111111111111111111111"}"#, r#"{"isAllowed":true,"gatewayAppId":"gw","reason":""}"#, + r#"{"isAllowed":true,"gatewayAppId":"gw","reason":"","rootEpoch":3}"#, ]); let auth = webhook(url); + let current = RootEpochRegistry::Current(CurrentRootEpoch { + epoch: 3, + chain_hash: [0x11; 32], + }); - let info = auth.get_info().await.unwrap(); - assert_eq!( - info.root_epochs, - Some(vec![RootEpochPubkeys { - ca_pubkey: vec![0xa0], - k256_pubkey: vec![0xb0], - }]) - ); + assert_eq!(auth.get_info().await.unwrap().root_epochs, current); let info = boot_info(50); - assert_eq!( - auth.is_app_allowed(&info, false).await.unwrap().root_epoch, - Some(3) - ); - assert_eq!( - auth.is_app_allowed(&info, false).await.unwrap().root_epoch, - None - ); + let registry = || async { + auth.is_app_allowed(&info, false) + .await + .unwrap() + .root_epoch + .registry() + }; + assert_eq!(registry().await.unwrap(), current); + assert_eq!(registry().await.unwrap(), RootEpochRegistry::Empty); + assert!(registry().await.is_err()); server.join().unwrap(); } diff --git a/dstack/kms/src/onboard_service.rs b/dstack/kms/src/onboard_service.rs index 0ed12c838..db54129d9 100644 --- a/dstack/kms/src/onboard_service.rs +++ b/dstack/kms/src/onboard_service.rs @@ -39,7 +39,7 @@ use crate::{ build_boot_info_for_attestation, upgrade_authority::{ app_attest, dstack_client, ensure_kms_allowed, ensure_self_kms_allowed, pad64, - RootEpochPubkeys, + RootEpochRegistry, }, }, }; @@ -141,7 +141,6 @@ impl OnboardRpc for OnboardHandler { if cfg.root_keys_exist() { bail!("KMS has already been bootstrapped"); } - ensure_bootstrap_epoch(cfg)?; ensure_self_kms_allowed(cfg, &self.state.attestation_verifier) .await .context("KMS is not allowed to bootstrap")?; @@ -296,7 +295,7 @@ mod tests { use crate::main_service::amd_attest::{ compute_expected_measurement, MeasurementInput, OvmfSectionParam, }; - use crate::main_service::upgrade_authority::BootResponse; + use crate::main_service::upgrade_authority::{BootResponse, CurrentRootEpoch, RootEpochFields}; use dstack_kms_rpc::KmsKeys; use sha2::Digest; @@ -566,206 +565,214 @@ mod tests { } } - fn epoch_config(cert_dir: &std::path::Path, current_epoch: u32) -> KmsConfig { - let mut cfg = auto_bootstrap_config(cert_dir); - cfg.key_rotation.current_epoch = current_epoch; - cfg + /// The registry a correct handover of `res` would be checked against. + fn current_of(res: &KmsKeyResponse, epoch: u32) -> RootEpochRegistry { + let roots: Vec<_> = res + .keys + .iter() + .map(|k| { + ( + KeyPair::from_pem(&k.ca_key).unwrap().public_key_der(), + SigningKey::from_slice(&k.k256_key) + .unwrap() + .verifying_key() + .to_sec1_bytes() + .to_vec(), + ) + }) + .collect(); + let hashes = dstack_types::root_epoch::chain_hashes( + roots.iter().map(|(ca, k256)| (&ca[..], &k256[..])), + ); + RootEpochRegistry::Current(CurrentRootEpoch { + epoch, + chain_hash: hashes[epoch as usize], + }) + } + + async fn onboard( + cfg: &KmsConfig, + res: KmsKeyResponse, + rotate: bool, + registry: RootEpochRegistry, + ) -> Result { + Keys::from_key_response(cfg, res, "kms.example.com", rotate, registry).await } #[rocket::async_test] async fn join_imports_and_stores_every_epoch() { let cert_dir = tempfile::tempdir().unwrap(); - let cfg = epoch_config(cert_dir.path(), 2); + let cfg = auto_bootstrap_config(cert_dir.path()); let res = handover(3); let sent: Vec<_> = res.keys.iter().map(|k| k.k256_key.clone()).collect(); - let keys = Keys::from_key_response(&cfg, res, "kms.example.com", false, None) + let keys = onboard(&cfg, res, false, RootEpochRegistry::Unchecked) .await .unwrap(); keys.store(&cfg).unwrap(); assert_eq!(crate::main_service::read_root_epoch(&cfg).unwrap(), 2); for (epoch, k256) in (0..).zip(&sent) { - assert_eq!(&fs::read(cfg.root_key_paths(epoch).k256_key).unwrap(), k256); - } - assert_eq!(&fs::read(cfg.k256_key()).unwrap(), &sent[2]); - - crate::main_service::KmsState::new(cfg.clone()).unwrap(); - // Editing the configured epoch cannot relabel the stored roots. - let err = crate::main_service::KmsState::new(epoch_config(cert_dir.path(), 3)) - .err() - .unwrap(); - assert!(format!("{err:#}").contains("belong to epoch 2"), "{err:#}"); - } - - fn pubkeys_of(res: &KmsKeyResponse) -> Vec { - res.keys - .iter() - .map(|k| RootEpochPubkeys { - ca_pubkey: KeyPair::from_pem(&k.ca_key).unwrap().public_key_der(), - k256_pubkey: SigningKey::from_slice(&k.k256_key) - .unwrap() - .verifying_key() - .to_sec1_bytes() - .to_vec(), - }) - .collect() - } - - async fn joined_state(res: KmsKeyResponse, url: String) -> crate::main_service::KmsState { - let cert_dir = tempfile::tempdir().unwrap(); - let mut cfg = epoch_config(cert_dir.path(), res.keys.len() as u32 - 1); - Keys::from_key_response(&cfg, res, "kms.example.com", false, None) - .await - .unwrap() - .store(&cfg) - .unwrap(); - cfg.auth_api = crate::main_service::upgrade_authority::tests::webhook(url); - crate::main_service::KmsState::new(cfg).unwrap() - } - - fn info_with_epochs(epochs: &[RootEpochPubkeys]) -> String { - let epochs: Vec<_> = epochs - .iter() - .map(|e| { - serde_json::json!({ - "caPubkey": hex::encode(&e.ca_pubkey), - "k256Pubkey": hex::encode(&e.k256_pubkey), - }) - }) - .collect(); - serde_json::json!({ - "status": "ok", "kmsContractAddr": "0x00", "gatewayAppId": "", "chainId": 1, - "appImplementation": "0x00", "rootEpochs": epochs, - }) - .to_string() - } - - #[rocket::async_test] - async fn keyring_must_match_the_registered_epochs() { - use crate::main_service::upgrade_authority::tests::serve; - - type Registered = fn(Vec) -> Vec; - let cases: [(Registered, bool); 4] = [ - (|keys| keys, true), - (|keys| keys[..1].to_vec(), false), - (|keys| vec![keys[0].clone(), keys[0].clone()], false), - (|_| Vec::new(), false), - ]; - for (registered, ok) in cases { - let res = handover(2); - let epochs = registered(pubkeys_of(&res)); - let (url, server) = serve(vec![info_with_epochs(&epochs)]); - let state = joined_state(res, url).await; assert_eq!( - state.ensure_keyring_registered().await.is_ok(), - ok, - "{epochs:?}" + &fs::read(cfg.root_key_paths(epoch, 2).k256_key).unwrap(), + k256 ); - server.join().unwrap(); } + assert_eq!(&fs::read(cfg.k256_key()).unwrap(), &sent[2]); } #[rocket::async_test] - async fn a_newer_registered_epoch_stops_a_stale_kms() { - let state = joined_state(handover(2), "http://127.0.0.1:1".into()).await; - let response = |root_epoch| BootResponse { - is_allowed: true, - gateway_app_id: String::new(), - reason: String::new(), - root_epoch, - }; - assert!(state.ensure_epoch_current(&response(Some(1))).is_ok()); - assert!(state.ensure_epoch_current(&response(Some(2))).is_err()); - assert!(state.ensure_epoch_current(&response(None)).is_err()); - } - - #[rocket::async_test] - async fn rotation_imports_history_and_generates_the_new_epoch() { + async fn rotation_imports_the_keyring_and_generates_the_next_epoch() { let cert_dir = tempfile::tempdir().unwrap(); - let cfg = epoch_config(cert_dir.path(), 2); + let cfg = auto_bootstrap_config(cert_dir.path()); let res = handover(2); + let registry = current_of(&res, 1); let sent: Vec<_> = res.keys.iter().map(|k| k.k256_key.clone()).collect(); - let keys = Keys::from_key_response(&cfg, res, "kms.example.com", true, None) - .await - .unwrap(); + let keys = onboard(&cfg, res, true, registry).await.unwrap(); + assert_eq!(keys.epoch, 2); keys.store(&cfg).unwrap(); - assert_eq!(fs::read(cfg.root_key_paths(1).k256_key).unwrap(), sent[1]); + assert_eq!( + fs::read(cfg.root_key_paths(1, 2).k256_key).unwrap(), + sent[1] + ); let current = fs::read(cfg.k256_key()).unwrap(); assert!(!sent.contains(¤t)); } #[rocket::async_test] - async fn onboarding_rejects_incomplete_or_extra_history() { + async fn onboarding_checks_the_handover_against_the_current_epoch() { let cert_dir = tempfile::tempdir().unwrap(); - let cfg = epoch_config(cert_dir.path(), 1); - for (epochs, rotate) in [(1, false), (3, false), (2, true), (0, true)] { - let res = handover(epochs); - assert!( - Keys::from_key_response(&cfg, res, "kms.example.com", rotate, None) - .await - .is_err(), - "epochs={epochs} rotate={rotate}" - ); - } - let cfg = epoch_config(cert_dir.path(), 0); + let cfg = auto_bootstrap_config(cert_dir.path()); + let res = handover(2); + let other = handover(2); + + // A join may hold an epoch pending activation beyond the current one. + assert!(onboard(&cfg, res.clone(), false, current_of(&res, 1)) + .await + .is_ok()); + assert!(onboard(&cfg, res.clone(), false, current_of(&res, 0)) + .await + .is_ok()); + // A rotation must start at the current epoch. + assert!(onboard(&cfg, res.clone(), true, current_of(&res, 0)) + .await + .is_err()); + // A keyring other than the registered one. + assert!(onboard(&cfg, res.clone(), false, current_of(&other, 1)) + .await + .is_err()); + assert!(onboard(&cfg, res.clone(), true, current_of(&other, 1)) + .await + .is_err()); + // A source behind the current epoch. + assert!(onboard(&cfg, handover(1), false, current_of(&res, 1)) + .await + .is_err()); + // Out-of-order or missing epochs. + let mut gap = res.clone(); + gap.keys.remove(0); + assert!(onboard(&cfg, gap, false, RootEpochRegistry::Unchecked) + .await + .is_err()); assert!( - Keys::from_key_response(&cfg, handover(0), "kms.example.com", true, None) + onboard(&cfg, handover(0), false, RootEpochRegistry::Unchecked) .await .is_err() ); } #[rocket::async_test] - async fn onboarding_requires_the_registered_keyring() { + async fn an_unregistered_deployment_can_only_join_at_epoch_0() { let cert_dir = tempfile::tempdir().unwrap(); - let join = epoch_config(cert_dir.path(), 1); - let rotate = epoch_config(cert_dir.path(), 2); - let res = handover(2); - let registered = pubkeys_of(&res); - let onboard = |cfg: &KmsConfig, rotate: bool, registered: Vec| { - let (cfg, res) = (cfg.clone(), res.clone()); - async move { - Keys::from_key_response(&cfg, res, "kms.example.com", rotate, Some(®istered)) - .await - } - }; + let cfg = auto_bootstrap_config(cert_dir.path()); + assert!(onboard(&cfg, handover(1), false, RootEpochRegistry::Empty) + .await + .is_ok()); + assert!(onboard(&cfg, handover(2), false, RootEpochRegistry::Empty) + .await + .is_err()); + assert!(onboard(&cfg, handover(1), true, RootEpochRegistry::Empty) + .await + .is_err()); + } - assert!(onboard(&join, false, registered.clone()).await.is_ok()); - assert!(onboard(&rotate, true, registered.clone()).await.is_ok()); - // A source at an older epoch than the registry. - let mut newer = registered.clone(); - newer.push(registered[0].clone()); - assert!(onboard(&join, false, newer).await.is_err()); - // A source holding a keyring other than the registered one. - let mut other = registered.clone(); - other[0] = pubkeys_of(&handover(1))[0].clone(); - assert!(onboard(&join, false, other.clone()).await.is_err()); - assert!(onboard(&rotate, true, other).await.is_err()); - // Epochs must be registered before rotating past epoch 0. - assert!(onboard(&rotate, true, Vec::new()).await.is_err()); + /// A stored keyring at `latest`, loaded back. + async fn stored_keyring(epochs: u32) -> (tempfile::TempDir, KmsConfig, KmsKeyResponse) { + let cert_dir = tempfile::tempdir().unwrap(); + let cfg = auto_bootstrap_config(cert_dir.path()); + let res = handover(epochs); + onboard(&cfg, res.clone(), false, RootEpochRegistry::Unchecked) + .await + .unwrap() + .store(&cfg) + .unwrap(); + (cert_dir, cfg, res) } #[rocket::async_test] - async fn an_unregistered_deployment_can_still_join_at_epoch_0() { - let cert_dir = tempfile::tempdir().unwrap(); - let cfg = epoch_config(cert_dir.path(), 0); - assert!( - Keys::from_key_response(&cfg, handover(1), "kms.example.com", false, Some(&[])) - .await - .is_ok() + async fn a_keyring_serves_the_current_epoch_it_holds() { + use crate::main_service::Keyring; + + let (_dir, cfg, res) = stored_keyring(3).await; + let keyring = Keyring::load(&cfg).unwrap(); + assert_eq!(keyring.latest(), 2); + assert_eq!(keyring.servable_epoch(current_of(&res, 1)).unwrap(), 1); + assert_eq!(keyring.servable_epoch(current_of(&res, 2)).unwrap(), 2); + assert_eq!( + keyring + .servable_epoch(RootEpochRegistry::Unchecked) + .unwrap(), + 2 ); + assert!(keyring.servable_epoch(RootEpochRegistry::Empty).is_err()); + assert!(keyring.servable_epoch(current_of(&handover(3), 1)).is_err()); + let mut ahead = handover(4); + ahead.keys[..3].clone_from_slice(&res.keys); + assert!(keyring.servable_epoch(current_of(&ahead, 3)).is_err()); + + let (_dir, cfg, _) = stored_keyring(1).await; + let keyring = Keyring::load(&cfg).unwrap(); + assert_eq!(keyring.servable_epoch(RootEpochRegistry::Empty).unwrap(), 0); } #[rocket::async_test] - async fn bootstrap_refuses_a_later_epoch() { - let cert_dir = tempfile::tempdir().unwrap(); - let cfg = epoch_config(cert_dir.path(), 1); - let verifier = AttestationVerifier::load(&cfg.attestation).unwrap(); - assert!(bootstrap_keys(&cfg, &verifier).await.is_err()); - assert!(!cfg.root_keys_exist()); + async fn a_moved_current_epoch_stops_serving_and_requests_a_restart() { + use crate::main_service::{Keyring, KmsState}; + + let (_dir, mut cfg, res) = stored_keyring(3).await; + let keyring = Keyring::load(&cfg).unwrap(); + cfg.auth_api = + crate::main_service::upgrade_authority::tests::webhook("http://127.0.0.1:1".into()); + let state = KmsState::new(cfg, keyring, 1).unwrap(); + let response = |registry| BootResponse { + is_allowed: true, + gateway_app_id: String::new(), + reason: String::new(), + root_epoch: match registry { + RootEpochRegistry::Current(current) => RootEpochFields { + root_epoch: Some(current.epoch), + root_epoch_hash: Some(hex::encode(current.chain_hash)), + }, + _ => RootEpochFields::default(), + }, + }; + + assert!(state + .ensure_epoch_current(&response(current_of(&res, 1))) + .is_ok()); + assert!(state + .ensure_epoch_current(&response(RootEpochRegistry::Empty)) + .is_err()); + assert!(state + .ensure_epoch_current(&response(current_of(&res, 2))) + .is_err()); + let restart = tokio::time::timeout(Duration::from_secs(1), state.restart_requested()); + assert!( + restart.await.is_ok(), + "a moved epoch must request a restart" + ); } #[rocket::async_test] @@ -1062,37 +1069,33 @@ impl Keys { domain: &str, rotate: bool, ) -> Result { - let registered = cfg + let registry = cfg .auth_api .get_info() .await - .context("failed to read the registered root epochs")? + .context("failed to read the current root epoch")? .root_epochs; - Self::from_key_response(cfg, keys_res, domain, rotate, registered.as_deref()).await + Self::from_key_response(cfg, keys_res, domain, rotate, registry).await } - /// `registered` is the governance registry; `None` means the auth backend - /// has none (dev mode), and an empty one is accepted only while joining at - /// epoch 0, as deployments that never registered epochs did before. + /// Import the source's whole keyring, epochs 0..=latest; a rotation then + /// generates epoch latest+1. async fn from_key_response( cfg: &KmsConfig, keys_res: KmsKeyResponse, domain: &str, rotate: bool, - registered: Option<&[RootEpochPubkeys]>, + registry: RootEpochRegistry, ) -> Result { - let epoch = cfg.key_rotation.current_epoch; - // A rotation imports every epoch before the new one; a join imports all. - let last = if rotate { - epoch - .checked_sub(1) - .context("rotation requires key_rotation.current_epoch > 0")? - } else { - epoch - }; let received: Vec = keys_res.keys.iter().map(|k| k.epoch).collect(); - if !received.iter().copied().eq(0..=last) { - bail!("source KMS must hand over root key epochs 0..={last}, got {received:?}"); + let latest = received + .len() + .checked_sub(1) + .context("source KMS handed over no keys")? as u32; + if !received.iter().copied().eq(0..=latest) { + bail!( + "source KMS must hand over root key epochs 0..={latest} in order, got {received:?}" + ); } let mut roots = keys_res .keys @@ -1104,7 +1107,8 @@ impl Keys { Ok((ca_key, k256_key)) }) .collect::>>()?; - ensure_handover_registered(&roots, registered, last)?; + ensure_handover_registered(&roots, registry, rotate)?; + let epoch = if rotate { latest + 1 } else { latest }; let (ca_key, k256_key) = if rotate { ( KeyPair::generate_for(&PKCS_ECDSA_P256_SHA256)?, @@ -1140,7 +1144,7 @@ impl Keys { /// The root keys go last: see [`KmsConfig::root_keys_exist`]. fn store(&self, cfg: &KmsConfig) -> Result<()> { for (epoch, root) in (0..).zip(&self.history) { - let paths = cfg.root_key_paths(epoch); + let paths = cfg.root_key_paths(epoch, self.epoch); if let Some(dir) = paths.ca_key.parent() { fs::create_dir_all(dir)?; } @@ -1171,14 +1175,17 @@ impl Keys { } } -pub(crate) async fn update_certs(cfg: &KmsConfig) -> Result<()> { +/// Reissue the RPC certificate under the root CA of the served `epoch`, so +/// guests reach this KMS through the CA of the keys it serves. +pub(crate) async fn update_certs(cfg: &KmsConfig, epoch: u32, latest: u32) -> Result<()> { + let roots = cfg.root_key_paths(epoch, latest); // 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 ca_key = KeyPair::from_pem(&fs::read_to_string(&roots.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_bytes = fs::read(&roots.k256_key)?; let k256_key = SigningKey::from_slice(&k256_key_bytes)?; let domain = if cfg.onboard.auto_bootstrap_domain.is_empty() { @@ -1201,7 +1208,7 @@ pub(crate) async fn update_certs(cfg: &KmsConfig) -> Result<()> { .context("Failed to regenerate certificates")?; renew_ca_cert_if_expiring( - cfg.root_ca_cert(), + &roots.ca_cert, keys.ca_cert.pem(), "KMS root CA certificate", )?; @@ -1260,7 +1267,6 @@ pub(crate) async fn bootstrap_keys(cfg: &KmsConfig, verifier: &AttestationVerifi if cfg.root_keys_exist() { bail!("KMS has already been bootstrapped"); } - ensure_bootstrap_epoch(cfg)?; ensure_self_kms_allowed(cfg, verifier) .await .context("KMS is not allowed to auto-bootstrap")?; @@ -1271,41 +1277,55 @@ pub(crate) async fn bootstrap_keys(cfg: &KmsConfig, verifier: &AttestationVerifi Ok(()) } -/// Check that the handed-over roots are exactly the registered epochs 0..=last, -/// which also proves the source is at the expected epoch. +/// Check the handed-over roots against governance: they must include the +/// current epoch with a matching chain hash. A rotation must start from the +/// current epoch with nothing pending activation, or it would collide with the +/// next registered epoch. Deployments that never registered epochs can still +/// join at epoch 0. fn ensure_handover_registered( roots: &[(KeyPair, SigningKey)], - registered: Option<&[RootEpochPubkeys]>, - last: u32, + registry: RootEpochRegistry, + rotate: bool, ) -> Result<()> { - let Some(registered) = registered else { - return Ok(()); + let latest = roots.len() as u32 - 1; + let current = match registry { + RootEpochRegistry::Unchecked => return Ok(()), + RootEpochRegistry::Empty if latest == 0 && !rotate => return Ok(()), + RootEpochRegistry::Empty => { + bail!("register root epoch 0 before rotating or joining past it") + } + RootEpochRegistry::Current(current) => current, }; - if registered.is_empty() && last == 0 { - return Ok(()); - } - if registered.len() != last as usize + 1 { + if rotate && current.epoch != latest { bail!( - "the source hands over root epochs 0..={last}, but {} epochs are registered", - registered.len() + "the source holds root epochs up to {latest}, but {} is current; rotate from a \ + source at the current epoch with no epoch pending activation", + current.epoch ); } - for (epoch, ((ca_key, k256_key), registered)) in (0..).zip(roots.iter().zip(registered)) { - if ca_key.public_key_der() != registered.ca_pubkey - || k256_key.verifying_key().to_sec1_bytes()[..] != registered.k256_pubkey[..] - { - bail!("the source's root keys of epoch {epoch} are not the registered ones"); - } - } - Ok(()) -} - -/// Bootstrapping creates epoch 0 only; later epochs come from rotation onboarding. -fn ensure_bootstrap_epoch(cfg: &KmsConfig) -> Result<()> { - if cfg.key_rotation.current_epoch != 0 { - bail!("bootstrap only creates epoch 0; onboard with rotate to create a later epoch"); + let pubkeys: Vec<_> = roots + .iter() + .map(|(ca_key, k256_key)| { + ( + ca_key.public_key_der(), + k256_key.verifying_key().to_sec1_bytes().to_vec(), + ) + }) + .collect(); + let hashes = dstack_types::root_epoch::chain_hashes( + pubkeys.iter().map(|(ca, k256)| (&ca[..], &k256[..])), + ); + match hashes.get(current.epoch as usize) { + Some(hash) if *hash == current.chain_hash => Ok(()), + Some(_) => bail!( + "the source's root keys up to epoch {} are not the registered ones", + current.epoch + ), + None => bail!( + "the source holds root epochs up to {latest}, but {} is current", + current.epoch + ), } - Ok(()) } fn self_signed_root_ca(ca_key: &KeyPair) -> Result { From ce4755e43a79ff636f96ec8523c97c982aa667be Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Fri, 6 Nov 2026 02:57:15 -0800 Subject: [PATCH 15/17] Fetch older root epochs' keys by number when they are needed Drop GetAppKey's include_history and the historical app roots in the keys file. The guest now asks the KMS for one older epoch only when it needs it: the epoch recorded on the data disk, or older epochs one by one until the env decrypts. The guest agent fetches an older app root when GetKey names its epoch, through its existing KMS client, and keeps it in memory. Old roots are no longer resident in every CVM, and a boot with nothing to migrate makes a single request as before. Signed-off-by: Kevin Wang --- dstack/cert-client/src/lib.rs | 24 ++++- dstack/dstack-types/src/lib.rs | 14 --- dstack/dstack-util/src/main.rs | 3 - dstack/dstack-util/src/system_setup.rs | 110 +++++++++-------------- dstack/guest-agent/src/rpc_service.rs | 110 +++++++++++++++++------ dstack/guest-agent/src/rpc_service_v1.rs | 43 ++++----- dstack/kms/rpc/proto/kms_rpc.proto | 14 --- dstack/kms/src/main_service.rs | 33 ++----- 8 files changed, 175 insertions(+), 176 deletions(-) diff --git a/dstack/cert-client/src/lib.rs b/dstack/cert-client/src/lib.rs index da48efa05..74c29ecbe 100644 --- a/dstack/cert-client/src/lib.rs +++ b/dstack/cert-client/src/lib.rs @@ -5,7 +5,7 @@ use std::sync::Arc; use anyhow::{bail, Context, Result}; -use dstack_kms_rpc::{kms_client::KmsClient, SignCertRequest}; +use dstack_kms_rpc::{kms_client::KmsClient, AppKeyResponse, GetAppKeyRequest, SignCertRequest}; use dstack_types::{AppKeys, KeyProvider}; use ra_rpc::client::{CertInfo, RaClient, RaClientConfig}; use ra_tls::{ @@ -101,6 +101,28 @@ impl CertRequestClient { } } + /// The app keys of an older KMS root epoch, from the KMS that served the + /// current ones. + pub async fn get_app_key(&self, root_epoch: u32) -> Result { + let CertRequestClient::Kms { client, vm_config } = self else { + bail!("only KMS-provided keys have older root epochs"); + }; + let response = client + .get_app_key(GetAppKeyRequest { + api_version: 1, + vm_config: vm_config.clone(), + root_epoch: Some(root_epoch), + }) + .await?; + if response.root_epoch != root_epoch { + bail!( + "KMS returned keys of root epoch {} instead of {root_epoch}", + response.root_epoch + ); + } + Ok(response) + } + pub async fn get_root_ca(&self) -> Result { match self { CertRequestClient::Local { ca } => Ok(ca.pem_cert.clone()), diff --git a/dstack/dstack-types/src/lib.rs b/dstack/dstack-types/src/lib.rs index c76e655a5..eb49e9cee 100644 --- a/dstack/dstack-types/src/lib.rs +++ b/dstack/dstack-types/src/lib.rs @@ -2405,16 +2405,6 @@ pub struct ReencryptedEnv { pub root_epoch: u32, } -/// The app root key of an older KMS root epoch. -#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq)] -pub struct HistoricalAppRoot { - pub root_epoch: u32, - #[serde(with = "hex_bytes")] - pub k256_key: Vec, - #[serde(with = "hex_bytes")] - pub k256_signature: Vec, -} - #[derive(Serialize, Deserialize, Debug, Clone)] pub struct AppKeys { #[serde(with = "hex_bytes")] @@ -2428,10 +2418,6 @@ pub struct AppKeys { pub gateway_app_id: String, pub ca_cert: String, pub key_provider: KeyProvider, - /// App root keys of the KMS root epochs before the current one, indexed by - /// epoch, for key requests that select one explicitly. - #[serde(default)] - pub historical_roots: Vec, } impl AppKeys { diff --git a/dstack/dstack-util/src/main.rs b/dstack/dstack-util/src/main.rs index ee2f7ee7d..4d76237fd 100644 --- a/dstack/dstack-util/src/main.rs +++ b/dstack/dstack-util/src/main.rs @@ -702,7 +702,6 @@ async fn cmd_get_keys(args: GetKeysArgs) -> Result<()> { api_version: 1, vm_config: "".to_string(), root_epoch: None, - include_history: false, }) .await .context("Failed to get app key")?; @@ -720,7 +719,6 @@ async fn cmd_get_keys(args: GetKeysArgs) -> Result<()> { k256_key: response.k256_key, k256_signature: response.k256_signature, gateway_app_id: response.gateway_app_id, - historical_roots: Vec::new(), key_provider: KeyProvider::Kms { url: kms_url, pubkey: root_pubkey, @@ -1238,7 +1236,6 @@ fn make_app_keys( k256_signature: vec![], gateway_app_id: "".to_string(), ca_cert: cert.pem(), - historical_roots: Vec::new(), key_provider, }) } diff --git a/dstack/dstack-util/src/system_setup.rs b/dstack/dstack-util/src/system_setup.rs index 35a374908..770283379 100644 --- a/dstack/dstack-util/src/system_setup.rs +++ b/dstack/dstack-util/src/system_setup.rs @@ -24,8 +24,7 @@ use dstack_types::{ APP_COMPOSE, APP_KEYS, DECRYPTED_ENV, DECRYPTED_ENV_JSON, ENCRYPTED_ENV, HOST_SHARED_DIR_NAME, INSTANCE_INFO, REENCRYPTED_ENV, SYS_CONFIG, USER_CONFIG, }, - GpuPolicy, HistoricalAppRoot, KeyProvider, KeyProviderInfo, ReencryptedEnv, - GPU_ATTESTATION_OUTPUT, + GpuPolicy, KeyProvider, KeyProviderInfo, ReencryptedEnv, GPU_ATTESTATION_OUTPUT, }; use fs_err as fs; use luks::{ @@ -2355,10 +2354,13 @@ struct AppInfo { } impl HostShared { - /// Request the current app keys from one KMS, with those of every older - /// root epoch in order. - async fn request_kms_keys(&self, kms_url: String) -> Result { - info!("Requesting app keys from KMS: {kms_url}"); + /// Request app keys of `root_epoch` (`None` = current) from one KMS. + async fn request_kms_keys( + &self, + kms_url: String, + root_epoch: Option, + ) -> Result<(AppKeys, Vec)> { + info!("Requesting app keys of root epoch {root_epoch:?} from KMS: {kms_url}"); let tmp_ca = { info!("Getting temp ca cert"); let client = RaClient::new(kms_url.clone(), true)?; @@ -2387,15 +2389,14 @@ impl HostShared { .get_app_key(rpc::GetAppKeyRequest { api_version: 1, vm_config: self.sys_config.vm_config.clone(), - root_epoch: None, - include_history: true, + root_epoch, }) .await .context("Failed to get app key")?; - // A KMS predating root epochs serves epoch 0 with no history. + // A KMS predating root epochs ignores the selector and serves epoch 0. let served = response.root_epoch; - if !response.history.iter().map(|h| h.root_epoch).eq(0..served) { - bail!("KMS returned incomplete key history for root epoch {served}"); + if root_epoch.is_some_and(|epoch| epoch != served) { + bail!("KMS returned keys of epoch {served} instead of {root_epoch:?}"); } let (_, ca_pem) = x509_parser::pem::parse_x509_pem(tmp_ca.ca_cert.as_bytes()) @@ -2410,7 +2411,6 @@ impl HostShared { k256_key: response.k256_key, k256_signature: response.k256_signature, gateway_app_id: response.gateway_app_id, - historical_roots: Vec::new(), key_provider: KeyProvider::Kms { url: kms_url, pubkey: root_pubkey, @@ -2419,19 +2419,17 @@ impl HostShared { root_epoch: served, }, }; - Ok(KmsKeys { - keys, - history: response.history, - os_image_hash: response.os_image_hash, - }) + Ok((keys, response.os_image_hash)) } -} -struct KmsKeys { - keys: AppKeys, - /// App keys of the root epochs before `keys.root_epoch()`, indexed by epoch. - history: Vec, - os_image_hash: Vec, + /// Request the app keys of an older root epoch from the KMS that served + /// `keys`. + async fn request_historical_keys(&self, keys: &AppKeys, epoch: u32) -> Result { + let KeyProvider::Kms { url, .. } = &keys.key_provider else { + bail!("only KMS-provided keys have older root epochs"); + }; + Ok(self.request_kms_keys(url.clone(), Some(epoch)).await?.0) + } } struct Stage0<'a> { @@ -2445,8 +2443,6 @@ struct Stage1<'a> { vmm: HostApi, shared: HostShared, keys: AppKeys, - /// Keys of the root epochs before `keys.root_epoch()`, indexed by epoch. - key_history: Vec, } fn validate_key_provider_inputs(kind: KeyProviderKind, kms_urls: &[String]) -> Result<()> { @@ -2498,24 +2494,20 @@ impl<'a> Stage0<'a> { self.shared.dir.join(APP_KEYS) } - async fn request_app_keys_from_kms(&self) -> Result<(AppKeys, Vec)> { + async fn request_app_keys_from_kms(&self) -> Result { if self.shared.sys_config.kms_urls.is_empty() { bail!("No KMS URLs are set"); } let mut errors = vec![]; for kms_url in &self.shared.sys_config.kms_urls { let kms_url = kms_rpc_url(kms_url); - match self.shared.request_kms_keys(kms_url.clone()).await { - Ok(KmsKeys { - keys, - history, - os_image_hash, - }) => { + match self.shared.request_kms_keys(kms_url.clone(), None).await { + Ok((keys, os_image_hash)) => { // Measured here, outside the per-URL request, so that a // failover can never extend os-image-hash twice. emit_runtime_event("os-image-hash", &os_image_hash) .context("failed to extend os-image-hash to the launch measurement")?; - return Ok((keys, history)); + return Ok(keys); } Err(err) => { warn!("failed to get app keys from KMS {kms_url}: {err:?}"); @@ -2594,12 +2586,11 @@ impl<'a> Stage0<'a> { .context("failed to generate TPM app keys") } - /// The app keys, with those of older KMS root epochs when the KMS has any. - async fn request_app_keys(&self) -> Result<(AppKeys, Vec)> { + async fn request_app_keys(&self) -> Result { let key_provider = self.shared.app_compose.key_provider(); validate_key_provider_inputs(key_provider, &self.shared.sys_config.kms_urls)?; - let keys = match key_provider { - KeyProviderKind::Kms => return self.request_app_keys_from_kms().await, + match key_provider { + KeyProviderKind::Kms => self.request_app_keys_from_kms().await, KeyProviderKind::Local => self.get_keys_from_local_key_provider().await, KeyProviderKind::None => { info!("No key provider is enabled, generating temporary app keys"); @@ -2611,8 +2602,7 @@ impl<'a> Stage0<'a> { info!("Generating app keys from TPM"); self.generate_tpm_app_keys() } - }; - Ok((keys?, Vec::new())) + } } async fn setup_swap(&self, swap_size: u64, opts: &DstackOptions) -> Result<()> { @@ -2754,12 +2744,7 @@ impl<'a> Stage0<'a> { initialized } - async fn mount_data_disk( - &self, - keys: &AppKeys, - key_history: &[rpc::HistoricalAppKey], - opts: &DstackOptions, - ) -> Result<()> { + async fn mount_data_disk(&self, keys: &AppKeys, opts: &DstackOptions) -> Result<()> { let disk_crypt_key = &hex::encode(&keys.disk_crypt_key); let root_epoch = kms_root_epoch(keys); let name = "dstack_data_disk"; @@ -2783,7 +2768,8 @@ impl<'a> Stage0<'a> { let has_filesystem = if opts.storage_encrypted && disk_initialized { info!("Mounting encrypted data disk"); if let Some(epoch) = root_epoch { - self.migrate_disk_key(key_history, disk_crypt_key, epoch) + self.migrate_disk_key(keys, disk_crypt_key, epoch) + .await .context("failed to migrate the data disk key")?; } let initializing = @@ -2991,9 +2977,9 @@ impl<'a> Stage0<'a> { /// is recorded in the header so the right old key is used after the next /// rotation. An unrotated epoch-0 disk is left untouched, so it stays /// readable by older OS images. - fn migrate_disk_key( + async fn migrate_disk_key( &self, - key_history: &[rpc::HistoricalAppKey], + keys: &AppKeys, current_key: &str, current: u32, ) -> Result<()> { @@ -3005,9 +2991,7 @@ impl<'a> Stage0<'a> { bail!("the data disk does not unlock with the key of epoch {current}"); } info!("re-keying the data disk from key epoch {epoch} to {current}"); - let old = key_history - .get(epoch as usize) - .with_context(|| format!("no keys for root epoch {epoch}"))?; + let old = self.shared.request_historical_keys(keys, epoch).await?; let old_key = hex::encode(&old.disk_crypt_key); luks_change_key(device, &old_key, current_key, Path::new(REKEY_DIR))?; if !luks_unlocks(device, current_key)? { @@ -3243,18 +3227,10 @@ impl<'a> Stage0<'a> { self.vmm .notify_q("boot.progress", "requesting app keys") .await; - let (mut app_keys, key_history) = self + let app_keys = self .request_app_keys() .await .context("Failed to request app keys")?; - app_keys.historical_roots = key_history - .iter() - .map(|keys| HistoricalAppRoot { - root_epoch: keys.root_epoch, - k256_key: keys.k256_key.clone(), - k256_signature: keys.k256_signature.clone(), - }) - .collect(); if app_keys.disk_crypt_key.is_empty() { bail!("Failed to get valid key phrase from KMS"); } @@ -3289,7 +3265,7 @@ impl<'a> Stage0<'a> { opts.storage_encrypted, opts.storage_fs ); - self.mount_data_disk(&app_keys, &key_history, &opts).await?; + self.mount_data_disk(&app_keys, &opts).await?; self.setup_swap(self.shared.app_compose.swap_size, &opts) .await?; self.vmm @@ -3309,7 +3285,6 @@ impl<'a> Stage0<'a> { shared: self.shared, vmm: self.vmm, keys: app_keys, - key_history, }) } } @@ -3317,7 +3292,7 @@ impl<'a> Stage0<'a> { impl Stage1<'_> { /// Decrypt the host-supplied env. An env still encrypted to an older KMS /// key epoch is re-encrypted to the current one for the host to swap in. - fn decrypt_env(&self) -> Result> { + async fn decrypt_env(&self) -> Result> { let ciphertext = &self.shared.encrypted_env; let current_key = env_crypt_key(&self.keys.env_crypt_key)?; let error = match dh_decrypt(current_key, ciphertext) { @@ -3325,8 +3300,11 @@ impl Stage1<'_> { Err(error) => error, }; let current = self.keys.root_epoch(); - for (index, old) in self.key_history.iter().enumerate().rev() { - let epoch = index as u32; + for epoch in (0..current).rev() { + let old = self + .shared + .request_historical_keys(&self.keys, epoch) + .await?; let Ok(plaintext) = dh_decrypt(env_crypt_key(&old.env_crypt_key)?, ciphertext) else { continue; }; @@ -3380,7 +3358,7 @@ impl Stage1<'_> { Default::default() } else { info!("Processing encrypted env"); - let plaintext = self.decrypt_env()?; + let plaintext = self.decrypt_env().await?; crate::parse_env_file::parse_env(&plaintext, &allowed_envs)? }; self.write_env_file(&decrypted_env)?; diff --git a/dstack/guest-agent/src/rpc_service.rs b/dstack/guest-agent/src/rpc_service.rs index 3a3fcbb5c..838939034 100644 --- a/dstack/guest-agent/src/rpc_service.rs +++ b/dstack/guest-agent/src/rpc_service.rs @@ -2,6 +2,8 @@ // // SPDX-License-Identifier: Apache-2.0 +use std::borrow::Cow; +use std::collections::HashMap; use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; @@ -95,9 +97,40 @@ pub struct AppState { /// The app root key of one KMS root epoch. pub(crate) struct AppRoot<'a> { pub epoch: u32, - pub k256_key: &'a [u8], - pub signing_key: std::borrow::Cow<'a, SigningKey>, - pub kms_signature: &'a [u8], + pub k256_key: Cow<'a, [u8]>, + pub signing_key: Cow<'a, SigningKey>, + pub kms_signature: Cow<'a, [u8]>, +} + +/// The app root key of an older KMS root epoch, fetched from the KMS on first +/// use. +#[derive(Clone)] +pub(crate) struct HistoricalRoot { + k256_key: Vec, + signing_key: SigningKey, + kms_signature: Vec, +} + +impl HistoricalRoot { + #[cfg(test)] + pub(crate) fn k256_key(&self) -> &[u8] { + &self.k256_key + } + + #[cfg(test)] + pub(crate) fn kms_signature(&self) -> &[u8] { + &self.kms_signature + } + + fn new(k256_key: Vec, kms_signature: Vec) -> Result { + let signing_key = SigningKey::from_slice(&k256_key) + .context("the app root k256 key is not a valid secp256k1 scalar")?; + Ok(Self { + k256_key, + signing_key, + kms_signature, + }) + } } struct AppStateInner { @@ -105,6 +138,8 @@ struct AppStateInner { keys: AppKeys, vm_config: String, cert_client: CertRequestClient, + /// Held across a fetch so concurrent requests for one epoch fetch it once. + historical_roots: tokio::sync::Mutex>, demo_cert: tokio::sync::OnceCell, /// When the last demo-cert request failed; held while one is running. demo_cert_last_failure: tokio::sync::Mutex>, @@ -308,6 +343,7 @@ impl AppState { config, keys, cert_client, + historical_roots: Default::default(), demo_cert: Default::default(), demo_cert_last_failure: Default::default(), vm_config, @@ -391,36 +427,58 @@ impl AppState { } /// The app root key and its KMS signature of `root_epoch`; the current - /// ones when unset. The key is parsed for older epochs only, which are - /// requested rarely. - pub(crate) fn app_root(&self, root_epoch: Option) -> Result> { - let keys = &self.inner.keys; - let current = keys.root_epoch(); + /// ones when unset. An older epoch is requested from the KMS once and kept + /// in memory. + pub(crate) async fn app_root(&self, root_epoch: Option) -> Result> { + let current = self.root_epoch(); let epoch = root_epoch.unwrap_or(current); if epoch == current { return Ok(AppRoot { epoch, - k256_key: self.app_root_k256_key(), - signing_key: std::borrow::Cow::Borrowed(self.app_root_signing_key()?), - kms_signature: self.kms_k256_signature(), + k256_key: Cow::Borrowed(self.app_root_k256_key()), + signing_key: Cow::Borrowed(self.app_root_signing_key()?), + kms_signature: Cow::Borrowed(self.kms_k256_signature()), }); } - let root = keys - .historical_roots - .iter() - .find(|root| root.root_epoch == epoch) - .with_context(|| format!("root key epoch {epoch} is not available"))?; + if epoch > current { + anyhow::bail!("root epoch {epoch} is newer than the current one, {current}"); + } + let root = self.historical_root(epoch).await?; Ok(AppRoot { epoch, - k256_key: &root.k256_key, - signing_key: std::borrow::Cow::Owned( - SigningKey::from_slice(&root.k256_key) - .context("the app root k256 key is not a valid secp256k1 scalar")?, - ), - kms_signature: &root.k256_signature, + k256_key: Cow::Owned(root.k256_key), + signing_key: Cow::Owned(root.signing_key), + kms_signature: Cow::Owned(root.kms_signature), }) } + pub(crate) async fn historical_root(&self, epoch: u32) -> Result { + let mut roots = self.inner.historical_roots.lock().await; + if let Some(root) = roots.get(&epoch) { + return Ok(root.clone()); + } + let keys = self + .inner + .cert_client + .get_app_key(epoch) + .await + .with_context(|| format!("failed to fetch the app keys of root epoch {epoch}"))?; + let root = HistoricalRoot::new(keys.k256_key, keys.k256_signature)?; + roots.insert(epoch, root.clone()); + Ok(root) + } + + #[cfg(test)] + pub(crate) async fn insert_historical_root( + &self, + epoch: u32, + k256_key: Vec, + kms_signature: Vec, + ) { + let root = HistoricalRoot::new(k256_key, kms_signature).unwrap(); + self.inner.historical_roots.lock().await.insert(epoch, root); + } + /// The KMS root key's signature over the app root public key: the second /// link of every signature chain, produced outside this agent and passed /// through byte-for-byte on both API surfaces. @@ -1401,13 +1459,6 @@ pNs85uhOZE8z2jr8Pg== k256_signature: Vec::new(), gateway_app_id: String::new(), ca_cert: DUMMY_PEM_CERT.to_string(), - // Only selected by tests that name root epoch 1 explicitly; the - // fixture's current root is epoch 0. - historical_roots: vec![dstack_types::HistoricalAppRoot { - root_epoch: 1, - k256_key: HISTORICAL_K256_KEY.to_vec(), - k256_signature: vec![0xab; 65], - }], key_provider: KeyProvider::None { key: DUMMY_PEM_KEY.to_string(), }, @@ -1495,6 +1546,7 @@ pNs85uhOZE8z2jr8Pg== keys: dummy_keys, vm_config: String::new(), cert_client: dummy_cert_client, + historical_roots: Default::default(), demo_cert: Default::default(), demo_cert_last_failure: Default::default(), platform: Arc::new(TestSimulatorPlatform { diff --git a/dstack/guest-agent/src/rpc_service_v1.rs b/dstack/guest-agent/src/rpc_service_v1.rs index da6ee9765..fa7edd7f0 100644 --- a/dstack/guest-agent/src/rpc_service_v1.rs +++ b/dstack/guest-agent/src/rpc_service_v1.rs @@ -161,15 +161,15 @@ impl V1RpcHandler { /// Derive the key named by `(domain, algorithm)` under the app root of /// `root_epoch`, together with its signature chain and that epoch. - fn derive( + async fn derive( &self, domain: &str, algorithm: &str, root_epoch: Option, ) -> Result<(AppKey, Vec>, u32)> { let algorithm = Algorithm::parse(algorithm)?; - let root = self.state.app_root(root_epoch)?; - let key = AppKey::derive(root.k256_key, domain, algorithm)?; + let root = self.state.app_root(root_epoch).await?; + let key = AppKey::derive(&root.k256_key, domain, algorithm)?; let chain = vec![ key.claim_signature(&root.signing_key)?, root.kms_signature.to_vec(), @@ -202,8 +202,9 @@ impl DstackGuestRpc for V1RpcHandler { } async fn get_key(self, request: GetKeyRequest) -> Result { - let (key, signature_chain, root_epoch) = - self.derive(&request.domain, &request.algorithm, request.root_epoch)?; + let (key, signature_chain, root_epoch) = self + .derive(&request.domain, &request.algorithm, request.root_epoch) + .await?; Ok(GetKeyResponse { key: key.secret(), public_key: key.public_key(), @@ -369,23 +370,23 @@ mod tests { handler().get_key(request(Some(0))).await.unwrap().key, current.key ); + assert!(handler().get_key(request(Some(1))).await.is_err()); + } - let old = handler().get_key(request(Some(1))).await.unwrap(); - assert_eq!(old.root_epoch, 1); - assert_ne!(old.key, current.key); - assert_eq!( - old.key, - AppKey::derive( - &crate::rpc_service::tests::HISTORICAL_K256_KEY, - "storage-encryption", - Algorithm::Secp256k1 - ) - .unwrap() - .secret() - ); - assert_eq!(old.signature_chain[1], vec![0xab; 65]); - - assert!(handler().get_key(request(Some(2))).await.is_err()); + /// Older epochs come from the KMS once and are then served from memory; + /// the fixture has no KMS, so the cache is filled directly. + #[tokio::test] + async fn an_older_root_is_fetched_once_and_cached() { + let (state, _guard) = state().await; + let key = crate::rpc_service::tests::HISTORICAL_K256_KEY.to_vec(); + state + .insert_historical_root(1, key.clone(), vec![0xab; 65]) + .await; + let root = state.historical_root(1).await.unwrap(); + assert_eq!(root.k256_key(), &key[..]); + assert_eq!(root.kms_signature(), &[0xab; 65][..]); + // An uncached epoch needs the KMS, which a local key provider lacks. + assert!(state.historical_root(2).await.is_err()); } /// The fixture's app root key is what `keys::tests` committed its vectors diff --git a/dstack/kms/rpc/proto/kms_rpc.proto b/dstack/kms/rpc/proto/kms_rpc.proto index fac7aad15..5b5ba57aa 100644 --- a/dstack/kms/rpc/proto/kms_rpc.proto +++ b/dstack/kms/rpc/proto/kms_rpc.proto @@ -14,18 +14,6 @@ message GetAppKeyRequest { // Root key epoch to derive from; unset means the current epoch. Epoch 0 is // the original root keys. Older epochs remain available so CVMs can migrate. optional uint32 root_epoch = 3; - // Also return the app keys of every epoch before the selected one, so a CVM - // migrating off older epochs needs a single request. - bool include_history = 4; -} - -// App keys derived from one older root epoch. -message HistoricalAppKey { - uint32 root_epoch = 1; - bytes disk_crypt_key = 2; - bytes env_crypt_key = 3; - bytes k256_key = 4; - bytes k256_signature = 5; } message AppId { @@ -63,8 +51,6 @@ message AppKeyResponse { bytes os_image_hash = 8; // Root key epoch all epoch-dependent fields above were derived from. uint32 root_epoch = 9; - // With `include_history`, the app keys of epochs 0..root_epoch in order. - repeated HistoricalAppKey history = 10; } message GetMetaResponse { diff --git a/dstack/kms/src/main_service.rs b/dstack/kms/src/main_service.rs index 8030f7b1e..942290f18 100644 --- a/dstack/kms/src/main_service.rs +++ b/dstack/kms/src/main_service.rs @@ -14,8 +14,8 @@ use anyhow::{anyhow, bail, Context, Result}; use dstack_kms_rpc::{ kms_server::{KmsRpc, KmsServer}, AppId, AppKeyResponse, GetAppKeyRequest, GetKmsKeyRequest, GetMetaResponse, - GetTempCaCertResponse, HistoricalAppKey, KmsKeyResponse, KmsKeys, PublicKeyResponse, - SignCertRequest, SignCertResponse, + GetTempCaCertResponse, KmsKeyResponse, KmsKeys, PublicKeyResponse, SignCertRequest, + SignCertResponse, }; use dstack_verifier::{CvmVerifier, VerificationDetails}; use fs_err as fs; @@ -603,33 +603,11 @@ impl KmsRpc for RpcHandler { .context("App not allowed")?; ensure_key_release_allowed(&boot_info, (&self.state.config).into())?; let root_epoch = self.state.resolve_epoch(request.root_epoch)?; - let derive = |epoch| { - derive_app_keys( - self.state.epoch_root(epoch), - &boot_info.app_id, - &boot_info.instance_id, - ) - }; - let keys = derive(root_epoch)?; - let history = if request.include_history { - (0..root_epoch) - .map(|epoch| { - let keys = derive(epoch)?; - Ok(HistoricalAppKey { - root_epoch: epoch, - disk_crypt_key: keys.disk_crypt_key, - env_crypt_key: keys.env_crypt_key, - k256_key: keys.k256_key, - k256_signature: keys.k256_signature, - }) - }) - .collect::>()? - } else { - Vec::new() - }; + let root = self.state.epoch_root(root_epoch); + let keys = derive_app_keys(root, &boot_info.app_id, &boot_info.instance_id)?; Ok(AppKeyResponse { - ca_cert: self.state.epoch_root(root_epoch).ca.pem_cert.clone(), + ca_cert: root.ca.pem_cert.clone(), disk_crypt_key: keys.disk_crypt_key, env_crypt_key: keys.env_crypt_key, k256_key: keys.k256_key, @@ -638,7 +616,6 @@ impl KmsRpc for RpcHandler { gateway_app_id, os_image_hash: boot_info.os_image_hash, root_epoch, - history, }) } From 04d51fed667767557edef269351c3132ef9d7c03 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Fri, 6 Nov 2026 03:19:40 -0800 Subject: [PATCH 16/17] Push the re-encrypted env to the VMM once per boot instead of pulling it The guest now hands the env it re-encrypted to the current root epoch to the VMM through a new host API call, UpdateEncryptedEnv, made on every boot before any application code runs (empty when nothing changed). The VMM records the env each boot starts with and accepts one update per boot, only while the boot has not reached init scripts or containers, and only if the stored env is still the one the VM booted with. The request size is bounded by Rocket's per-method limit, 1 MiB by default under [host_api.limits]. This drops the guest-side re-encrypted env file, GuestApi.GetEncryptedEnv and its VMM proxy, Vmm.ApplyReencryptedEnv, and the vmm-cli command, so no operator step is needed. Signed-off-by: Kevin Wang --- docs/deployment.md | 6 +- docs/vmm-cli-user-guide.md | 11 -- dstack/dstack-types/src/lib.rs | 13 --- dstack/dstack-types/src/shared_filenames.rs | 1 - dstack/dstack-util/src/host_api.rs | 18 ++++ dstack/dstack-util/src/system_setup.rs | 43 ++++---- dstack/guest-agent/src/guest_api_service.rs | 30 +----- dstack/guest-api/proto/guest_api.proto | 15 --- dstack/host-api/proto/host_api.proto | 9 ++ dstack/vmm/rpc/proto/vmm_rpc.proto | 11 -- dstack/vmm/src/app.rs | 111 +++++++++++++++++++- dstack/vmm/src/config.rs | 13 +++ dstack/vmm/src/guest_api_service.rs | 9 +- dstack/vmm/src/host_api_service.rs | 7 +- dstack/vmm/src/main_service.rs | 34 +----- dstack/vmm/src/vmm-cli.py | 16 --- dstack/vmm/vmm.toml | 4 + 17 files changed, 193 insertions(+), 158 deletions(-) diff --git a/docs/deployment.md b/docs/deployment.md index 52a684768..f3766ad37 100644 --- a/docs/deployment.md +++ b/docs/deployment.md @@ -538,8 +538,10 @@ for the full procedure; the KMS steps are: auth-simple). Every KMS restarts onto N; a KMS that does not hold N waits until it is onboarded again. 6. Boot CVMs on an OS that supports root epochs. They re-key their data disk and - re-encrypt their env at boot; apply the new env with - `vmm-cli.py apply-reencrypted-env `. + re-encrypt their env at boot; the VMM stores the re-encrypted env, provided + it was not changed while the VM booted. The VMM accepts one such update per + boot, before application code runs, of at most + `[host_api.limits] UpdateEncryptedEnv` (1 MiB by default). The auth backend reports the current epoch and the chain hash committing to the root public keys of epochs 0..current; a KMS serves only while its own keys diff --git a/docs/vmm-cli-user-guide.md b/docs/vmm-cli-user-guide.md index 94e335c1e..c2953d42f 100644 --- a/docs/vmm-cli-user-guide.md +++ b/docs/vmm-cli-user-guide.md @@ -534,17 +534,6 @@ The whitelist is stored in `~/.dstack-vmm/kms-whitelist.json`. ./vmm-cli.py update-env --env-file ./new-secrets.env ``` -After a KMS key rotation, a VM whose env is still encrypted to an older key -epoch re-encrypts it to the current epoch at boot. Replace the stored env with -that copy so later boots no longer need the old key: - -```bash -./vmm-cli.py apply-reencrypted-env -``` - -The env is only replaced if it has not changed since the VM booted. Rotate -any secrets the old key could expose separately. - #### Update Application Compose ```bash diff --git a/dstack/dstack-types/src/lib.rs b/dstack/dstack-types/src/lib.rs index eb49e9cee..edb667c81 100644 --- a/dstack/dstack-types/src/lib.rs +++ b/dstack/dstack-types/src/lib.rs @@ -2392,19 +2392,6 @@ impl OsImageMeasurementDocument { } } -/// Encrypted env that boot re-encrypted from an older KMS key epoch to the -/// current one, for the host to swap in place of the env it supplied. -#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq)] -pub struct ReencryptedEnv { - /// SHA-256 of the encrypted env the host supplied for this boot. - #[serde(with = "hex_bytes")] - pub source_sha256: Vec, - /// The same env encrypted to the current env public key. - #[serde(with = "hex_bytes")] - pub encrypted_env: Vec, - pub root_epoch: u32, -} - #[derive(Serialize, Deserialize, Debug, Clone)] pub struct AppKeys { #[serde(with = "hex_bytes")] diff --git a/dstack/dstack-types/src/shared_filenames.rs b/dstack/dstack-types/src/shared_filenames.rs index a653490a1..79944cff6 100644 --- a/dstack/dstack-types/src/shared_filenames.rs +++ b/dstack/dstack-types/src/shared_filenames.rs @@ -10,7 +10,6 @@ pub const USER_CONFIG: &str = ".user-config"; pub const ENCRYPTED_ENV: &str = ".encrypted-env"; pub const DECRYPTED_ENV: &str = ".decrypted-env"; pub const DECRYPTED_ENV_JSON: &str = ".decrypted-env.json"; -pub const REENCRYPTED_ENV: &str = ".reencrypted-env.json"; pub const INSTANCE_INFO: &str = ".instance_info"; pub const HOST_SHARED_DIR: &str = "/dstack/.host-shared"; pub const HOST_SHARED_DIR_NAME: &str = ".host-shared"; diff --git a/dstack/dstack-util/src/host_api.rs b/dstack/dstack-util/src/host_api.rs index e038ea3b2..abc3e43fd 100644 --- a/dstack/dstack-util/src/host_api.rs +++ b/dstack/dstack-util/src/host_api.rs @@ -80,6 +80,24 @@ impl HostApi { Ok(()) } + /// Hand the host the env re-encrypted to the current root epoch, or an + /// empty one when nothing changed. + pub async fn update_encrypted_env(&self, encrypted_env: Vec) -> Result<()> { + if !matches!(Platform::detect_or_dstack(), Platform::Dstack) { + return Ok(()); + } + let Some(client) = &self.client else { + return Ok(()); + }; + tokio::time::timeout( + HOST_API_TIMEOUT, + client.update_encrypted_env(host_api::EncryptedEnv { encrypted_env }), + ) + .await + .context("Timed out updating the encrypted env on the host")??; + Ok(()) + } + pub async fn notify_q(&self, event: &str, payload: &str) { if let Err(err) = self.notify(event, payload).await { warn!("Failed to notify event {event} to host: {:?}", err); diff --git a/dstack/dstack-util/src/system_setup.rs b/dstack/dstack-util/src/system_setup.rs index 770283379..fefc66036 100644 --- a/dstack/dstack-util/src/system_setup.rs +++ b/dstack/dstack-util/src/system_setup.rs @@ -22,9 +22,9 @@ use dstack_types::{ gpu_policy_hash, shared_filenames::{ APP_COMPOSE, APP_KEYS, DECRYPTED_ENV, DECRYPTED_ENV_JSON, ENCRYPTED_ENV, - HOST_SHARED_DIR_NAME, INSTANCE_INFO, REENCRYPTED_ENV, SYS_CONFIG, USER_CONFIG, + HOST_SHARED_DIR_NAME, INSTANCE_INFO, SYS_CONFIG, USER_CONFIG, }, - GpuPolicy, KeyProvider, KeyProviderInfo, ReencryptedEnv, GPU_ATTESTATION_OUTPUT, + GpuPolicy, KeyProvider, KeyProviderInfo, GPU_ATTESTATION_OUTPUT, }; use fs_err as fs; use luks::{ @@ -3291,12 +3291,12 @@ impl<'a> Stage0<'a> { impl Stage1<'_> { /// Decrypt the host-supplied env. An env still encrypted to an older KMS - /// key epoch is re-encrypted to the current one for the host to swap in. - async fn decrypt_env(&self) -> Result> { + /// root epoch is also returned re-encrypted to the current one. + async fn decrypt_env(&self) -> Result<(Vec, Option>)> { let ciphertext = &self.shared.encrypted_env; let current_key = env_crypt_key(&self.keys.env_crypt_key)?; let error = match dh_decrypt(current_key, ciphertext) { - Ok(plaintext) => return Ok(plaintext), + Ok(plaintext) => return Ok((plaintext, None)), Err(error) => error, }; let current = self.keys.root_epoch(); @@ -3315,18 +3315,7 @@ impl Stage1<'_> { if dh_decrypt(current_key, &encrypted_env)? != plaintext { bail!("re-encrypted env does not round-trip"); } - let migrated = ReencryptedEnv { - source_sha256: sha256(ciphertext).to_vec(), - encrypted_env, - root_epoch: current, - }; - safe_write_with_mode( - self.shared.dir.join(REENCRYPTED_ENV), - serde_json::to_vec(&migrated)?, - 0o600, - ) - .context("failed to write the re-encrypted env")?; - return Ok(plaintext); + return Ok((plaintext, Some(encrypted_env))); } Err(error).context("failed to decrypt env file") } @@ -3353,15 +3342,27 @@ impl Stage1<'_> { .iter() .cloned() .collect(); - let decrypted_env = if self.shared.encrypted_env.is_empty() { + let (decrypted_env, reencrypted) = if self.shared.encrypted_env.is_empty() { info!("No encrypted env, using default"); - Default::default() + (Default::default(), None) } else { info!("Processing encrypted env"); - let plaintext = self.decrypt_env().await?; - crate::parse_env_file::parse_env(&plaintext, &allowed_envs)? + let (plaintext, reencrypted) = self.decrypt_env().await?; + ( + crate::parse_env_file::parse_env(&plaintext, &allowed_envs)?, + reencrypted, + ) }; self.write_env_file(&decrypted_env)?; + // Sent on every boot, empty when nothing changed: the host accepts one + // update per boot, and this one comes before any application code. + if let Err(err) = self + .vmm + .update_encrypted_env(reencrypted.unwrap_or_default()) + .await + { + warn!("failed to store the re-encrypted env on the host: {err:#}"); + } Ok(decrypted_env) } diff --git a/dstack/guest-agent/src/guest_api_service.rs b/dstack/guest-agent/src/guest_api_service.rs index f26ca4640..cf3e96163 100644 --- a/dstack/guest-agent/src/guest_api_service.rs +++ b/dstack/guest-agent/src/guest_api_service.rs @@ -7,17 +7,16 @@ use std::{collections::HashSet, fmt::Debug, path::PathBuf, process::ExitStatus, use anyhow::{Context, Result}; use bollard::{container::ListContainersOptions, Docker}; use dstack_guest_agent_rpc::v0::worker_server::WorkerRpc as _; -use dstack_types::shared_filenames::{ENCRYPTED_ENV, HOST_SHARED_DIR, REENCRYPTED_ENV, SYS_CONFIG}; +use dstack_types::shared_filenames::{HOST_SHARED_DIR, SYS_CONFIG}; use dstack_types::SysConfig; use fs_err as fs; use guest_api::{ guest_api_server::{GuestApiRpc, GuestApiServer}, - Container, DiskInfo, EncryptedEnv, Gateway, GpuInfoResponse, GuestInfo, Interface, IpAddress, + Container, DiskInfo, Gateway, GpuInfoResponse, GuestInfo, Interface, IpAddress, ListContainersResponse, NetworkInformation, SystemInfo, }; use host_api::Notification; use ra_rpc::{CallContext, RpcCall}; -use sha2::Digest as _; use tokio::{process::Command, task::spawn_blocking, time::timeout}; use tracing::error; @@ -45,31 +44,6 @@ impl RpcCall for GuestApiHandler { } impl GuestApiRpc for GuestApiHandler { - async fn get_encrypted_env(self) -> Result { - let read = |name: &str| match fs::read(PathBuf::from(HOST_SHARED_DIR).join(name)) { - Ok(data) => Ok(Some(data)), - Err(err) if err.kind() == std::io::ErrorKind::NotFound => Ok(None), - Err(err) => Err(err), - }; - if let Some(json) = read(REENCRYPTED_ENV).context("failed to read the re-encrypted env")? { - let migrated: dstack_types::ReencryptedEnv = - serde_json::from_slice(&json).context("invalid re-encrypted env")?; - return Ok(EncryptedEnv { - source_sha256: migrated.source_sha256, - encrypted_env: migrated.encrypted_env, - root_epoch: migrated.root_epoch, - }); - } - let Some(env) = read(ENCRYPTED_ENV).context("failed to read the encrypted env")? else { - return Ok(Default::default()); - }; - Ok(EncryptedEnv { - source_sha256: sha2::Sha256::digest(&env).to_vec(), - encrypted_env: env, - root_epoch: self.state.root_epoch(), - }) - } - async fn info(self) -> Result { let ext_rpc = ExternalRpcHandler::new(self.state); let info = timeout(GUEST_API_HANDLER_TIMEOUT, ext_rpc.info()) diff --git a/dstack/guest-api/proto/guest_api.proto b/dstack/guest-api/proto/guest_api.proto index 60f214452..9e2a9f990 100644 --- a/dstack/guest-api/proto/guest_api.proto +++ b/dstack/guest-api/proto/guest_api.proto @@ -185,18 +185,6 @@ message GpuInfoResponse { optional uint64 sample_age_ms = 5; } -// The encrypted env the host should store for this VM: the env this boot -// started from, or, if it was encrypted to an older KMS root epoch, the copy -// boot re-encrypted to the current one. Empty when the VM has no env. -message EncryptedEnv { - // SHA-256 of the encrypted env the host supplied for this boot. - bytes source_sha256 = 1; - // The env encrypted to the current env public key. - bytes encrypted_env = 2; - // KMS root epoch `encrypted_env` is encrypted for. - uint32 root_epoch = 3; -} - // Direct gRPC surface exposed by the in-guest agent. service GuestApi { // Returns attestation material and identifiers for the calling guest. @@ -211,8 +199,6 @@ service GuestApi { rpc ListContainers(google.protobuf.Empty) returns (ListContainersResponse); // Initiates a graceful shutdown inside the guest VM. rpc Shutdown(google.protobuf.Empty) returns (google.protobuf.Empty); - // Returns the encrypted env the host should store for this VM. - rpc GetEncryptedEnv(google.protobuf.Empty) returns (EncryptedEnv); } // Same API surface as GuestApi but multiplexed by VM ID through VMM. @@ -223,5 +209,4 @@ service ProxiedGuestApi { rpc NetworkInfo(Id) returns (NetworkInformation); rpc ListContainers(Id) returns (ListContainersResponse); rpc Shutdown(Id) returns (google.protobuf.Empty); - rpc GetEncryptedEnv(Id) returns (EncryptedEnv); } diff --git a/dstack/host-api/proto/host_api.proto b/dstack/host-api/proto/host_api.proto index f8864f516..ba7941d0c 100644 --- a/dstack/host-api/proto/host_api.proto +++ b/dstack/host-api/proto/host_api.proto @@ -27,8 +27,17 @@ message GetSealingKeyResponse { bytes provider_quote = 2; } +// The env the guest re-encrypted to the current KMS root epoch at boot; empty +// when it needed no change. +message EncryptedEnv { + bytes encrypted_env = 1; +} + service HostApi { rpc Info(google.protobuf.Empty) returns (HostInfo); rpc Notify(Notification) returns (google.protobuf.Empty); rpc GetSealingKey(GetSealingKeyRequest) returns (GetSealingKeyResponse); + // Replace this VM's stored encrypted env. Accepted once per boot, before + // application code runs, and only if the stored env is the one it booted with. + rpc UpdateEncryptedEnv(EncryptedEnv) returns (google.protobuf.Empty); } diff --git a/dstack/vmm/rpc/proto/vmm_rpc.proto b/dstack/vmm/rpc/proto/vmm_rpc.proto index e71743b87..84b66b2d6 100644 --- a/dstack/vmm/rpc/proto/vmm_rpc.proto +++ b/dstack/vmm/rpc/proto/vmm_rpc.proto @@ -448,17 +448,6 @@ service Vmm { // Delete a local guest image by name. rpc DeleteImage(Id) returns (google.protobuf.Empty); - - // Replace the VM's encrypted env with the copy the guest re-encrypted to - // the current KMS key epoch at boot, if the env has not changed since. - rpc ApplyReencryptedEnv(Id) returns (ApplyReencryptedEnvResponse); -} - -message ApplyReencryptedEnvResponse { - // Whether the stored encrypted env was replaced. - bool replaced = 1; - // KMS key epoch the stored env is now encrypted for, when replaced. - uint32 root_epoch = 2; } // Response containing a list of supervisor processes. diff --git a/dstack/vmm/src/app.rs b/dstack/vmm/src/app.rs index 379bea90a..afb34581f 100644 --- a/dstack/vmm/src/app.rs +++ b/dstack/vmm/src/app.rs @@ -530,9 +530,16 @@ impl App { } self.sync_dynamic_config(id)?; self.set_started(id, true)?; + // A guest the VMM did not launch may already run application code. + let env_update = if is_running { + None + } else { + Some(sha256_of_env(&self.encrypted_env_path(id)?)?) + }; let vm_config = { let mut state = self.lock(); let vm_state = state.get_mut(id).context("VM not found")?; + vm_state.state.env_update = env_update; // Older images does not support for progress reporting if vm_state.config.image.info.shared_ro { vm_state.state.start(is_running); @@ -1561,7 +1568,22 @@ impl App { vm.state.events.pop_front(); } match event { - "boot.progress" => vm.state.boot_progress = body, + "boot.progress" => { + // Application code (init and pre-launch scripts, containers) + // runs from here on, so the guest OS can no longer vouch for + // an env update. + if matches!( + body.as_str(), + "init-script" + | "pre-launch" + | "starting containers" + | "running main script" + | "done" + ) { + vm.state.env_update = None; + } + vm.state.boot_progress = body; + } "boot.error" => vm.state.boot_error = body, "shutdown.progress" => { let powering_off = body == "powering off"; @@ -1582,6 +1604,35 @@ impl App { Ok(()) } + /// Store the encrypted env a guest re-encrypted at boot. Accepted once per + /// boot, before application code runs, and only over the env the VM booted + /// with, so an update made since then is never overwritten. + pub(crate) fn update_encrypted_env(&self, cid: u32, encrypted_env: Vec) -> Result<()> { + let (id, boot_env) = { + let mut state = self.lock(); + let vm = state + .vms + .values_mut() + .find(|vm| vm.config.cid == cid) + .context("VM not found")?; + let boot_env = vm.state.env_update.take(); + (vm.config.manifest.id.clone(), boot_env) + }; + let Some(boot_env) = boot_env else { + bail!("the encrypted env can only be updated once per boot, before applications start"); + }; + if encrypted_env.is_empty() { + return Ok(()); + } + let path = self.encrypted_env_path(&id)?; + if sha256_of_env(&path)? != boot_env { + bail!("the encrypted env changed since the VM booted; not replacing it"); + } + safe_write::safe_write(&path, &encrypted_env).context("Failed to write encrypted env")?; + info!(id, "stored the encrypted env re-encrypted by the guest"); + Ok(()) + } + pub(crate) fn compose_file_path(&self, id: &str) -> Result { Ok(self.shared_dir(id)?.join(APP_COMPOSE)) } @@ -2103,6 +2154,15 @@ fn sha256_file(path: impl AsRef) -> Result<[u8; 32]> { Ok(out) } +/// SHA-256 of a VM's encrypted env; a VM without one hashes as empty. +fn sha256_of_env(path: &Path) -> Result<[u8; 32]> { + match fs::read(path) { + Ok(env) => Ok(Sha256::digest(env).into()), + Err(err) if err.kind() == std::io::ErrorKind::NotFound => Ok(Sha256::digest([]).into()), + Err(err) => Err(err).context("Failed to read encrypted env"), + } +} + fn image_supports_tdx_lite(image: &Image) -> bool { image .digest @@ -2508,6 +2568,51 @@ mod tests { (vm, networks) } + #[test] + fn a_guest_stores_its_env_once_per_boot_before_apps_start() { + let dir = tempfile::tempdir().unwrap(); + let app = App::new( + test_config(&dir.path().join("absent-netd.sock"), dir.path()), + SupervisorClient::new("http://127.0.0.1:0"), + ); + let (vm, _) = bridge_vm(&app, "vm-1"); + let cid = vm.cid; + app.lock().vms.insert("vm-1".into(), VmState::new(vm)); + let env_path = app.encrypted_env_path("vm-1").unwrap(); + std::fs::create_dir_all(env_path.parent().unwrap()).unwrap(); + std::fs::write(&env_path, b"old").unwrap(); + let boot = || { + let hash = sha256_of_env(&env_path).unwrap(); + app.lock().vms.get_mut("vm-1").unwrap().state.env_update = Some(hash); + }; + let stored = || std::fs::read(&env_path).unwrap(); + + // Nothing is accepted outside a boot the VMM launched. + assert!(app.update_encrypted_env(cid, b"new".to_vec()).is_err()); + + boot(); + app.update_encrypted_env(cid, b"new".to_vec()).unwrap(); + assert_eq!(stored(), b"new"); + assert!(app.update_encrypted_env(cid, b"newer".to_vec()).is_err()); + + // An empty update consumes the boot's one chance. + boot(); + app.update_encrypted_env(cid, Vec::new()).unwrap(); + assert!(app.update_encrypted_env(cid, b"newer".to_vec()).is_err()); + + // An env updated since the boot is never overwritten. + boot(); + std::fs::write(&env_path, b"updated by the owner").unwrap(); + assert!(app.update_encrypted_env(cid, b"newer".to_vec()).is_err()); + assert_eq!(stored(), b"updated by the owner"); + + // Once application code runs, the guest OS no longer vouches for it. + boot(); + app.vm_event_report(cid, "boot.progress", "init-script".into()) + .unwrap(); + assert!(app.update_encrypted_env(cid, b"newer".to_vec()).is_err()); + } + /// A launch that never reached a netd built nothing, so it must not leave /// a marker that makes every later stop and removal wait for one. What /// an earlier boot may still hold keeps its marker. @@ -3950,6 +4055,10 @@ struct VmStateMut { auto_restart: AutoRestartState, /// True when the VM is being removed (cleanup in progress). removing: bool, + /// SHA-256 of the encrypted env this boot started with, while the guest + /// may still replace it; taken by the first update and cleared once the + /// boot reaches application code. + env_update: Option<[u8; 32]>, } impl VmStateMut { diff --git a/dstack/vmm/src/config.rs b/dstack/vmm/src/config.rs index c923db519..4f1d7e9d8 100644 --- a/dstack/vmm/src/config.rs +++ b/dstack/vmm/src/config.rs @@ -1432,6 +1432,19 @@ impl Config { mod tests { use super::*; + #[test] + fn host_api_limits_encrypted_env_updates() { + use rocket::data::ToByteUnit; + use rocket::figment::providers::Serialized; + // Built the way `run_host_api` builds the host API server's config. + let figment = load_config_figment(None); + let figment = figment.clone().merge(Serialized::defaults( + figment.find_value("host_api").unwrap(), + )); + let config: rocket::Config = figment.extract().unwrap(); + assert_eq!(config.limits.get("UpdateEncryptedEnv"), Some(1.mebibytes())); + } + #[test] fn image_dir_stays_a_direct_child_of_the_image_root() { let cfg = ImageConfig { diff --git a/dstack/vmm/src/guest_api_service.rs b/dstack/vmm/src/guest_api_service.rs index 763df5005..439df00d1 100644 --- a/dstack/vmm/src/guest_api_service.rs +++ b/dstack/vmm/src/guest_api_service.rs @@ -6,8 +6,7 @@ use crate::App as AppState; use anyhow::Result; use guest_api::{ proxied_guest_api_server::{ProxiedGuestApiRpc, ProxiedGuestApiServer}, - EncryptedEnv, GpuInfoResponse, GuestInfo, Id, ListContainersResponse, NetworkInformation, - SystemInfo, + GpuInfoResponse, GuestInfo, Id, ListContainersResponse, NetworkInformation, SystemInfo, }; use ra_rpc::{CallContext, RpcCall}; use std::ops::Deref; @@ -60,10 +59,4 @@ impl ProxiedGuestApiRpc for GuestApiHandler { async fn shutdown(self, request: Id) -> Result<()> { self.guest_agent_client(&request.id)?.shutdown().await } - - async fn get_encrypted_env(self, request: Id) -> Result { - self.guest_agent_client(&request.id)? - .get_encrypted_env() - .await - } } diff --git a/dstack/vmm/src/host_api_service.rs b/dstack/vmm/src/host_api_service.rs index fc4dcc704..25ec97751 100644 --- a/dstack/vmm/src/host_api_service.rs +++ b/dstack/vmm/src/host_api_service.rs @@ -5,7 +5,7 @@ use anyhow::{bail, Context, Result}; use host_api::{ host_api_server::{HostApiRpc, HostApiServer}, - GetSealingKeyRequest, GetSealingKeyResponse, HostInfo, Notification, + EncryptedEnv, GetSealingKeyRequest, GetSealingKeyResponse, HostInfo, Notification, }; use ra_rpc::{CallContext, RemoteEndpoint, RpcCall}; use rocket_vsock_listener::VsockEndpoint; @@ -46,6 +46,11 @@ impl HostApiRpc for HostApiHandler { .vm_event_report(self.endpoint.cid, &request.event, request.payload) } + async fn update_encrypted_env(self, request: EncryptedEnv) -> Result<()> { + self.app + .update_encrypted_env(self.endpoint.cid, request.encrypted_env) + } + async fn get_sealing_key(self, request: GetSealingKeyRequest) -> Result { let key_provider = &self.app.config.key_provider; if !key_provider.enabled { diff --git a/dstack/vmm/src/main_service.rs b/dstack/vmm/src/main_service.rs index cca11541a..ffd06ad61 100644 --- a/dstack/vmm/src/main_service.rs +++ b/dstack/vmm/src/main_service.rs @@ -12,11 +12,10 @@ use dstack_types::AppCompose; use dstack_vmm_rpc as rpc; use dstack_vmm_rpc::vmm_server::{VmmRpc, VmmServer}; use dstack_vmm_rpc::{ - AppId, ApplyReencryptedEnvResponse, ComposeHash as RpcComposeHash, GatewaySettings, - GetInfoResponse, GetMetaResponse, Id, ImageInfo as RpcImageInfo, ImageListResponse, - KmsSettings, ListGpusResponse, PublicKeyResponse, ReloadVmsResponse, ResizeVmRequest, - ResourcesSettings, StatusRequest, StatusResponse, SvListResponse, SvProcessInfo, - UpdateVmRequest, VersionResponse, VmConfiguration, + AppId, ComposeHash as RpcComposeHash, GatewaySettings, GetInfoResponse, GetMetaResponse, Id, + ImageInfo as RpcImageInfo, ImageListResponse, KmsSettings, ListGpusResponse, PublicKeyResponse, + ReloadVmsResponse, ResizeVmRequest, ResourcesSettings, StatusRequest, StatusResponse, + SvListResponse, SvProcessInfo, UpdateVmRequest, VersionResponse, VmConfiguration, }; use fs_err as fs; use path_absolutize::Absolutize; @@ -1298,31 +1297,6 @@ impl VmmRpc for RpcHandler { Ok(()) } - async fn apply_reencrypted_env(self, request: Id) -> Result { - use sha2::Digest; - let migrated = self - .guest_agent_client(&request.id)? - .get_encrypted_env() - .await?; - if migrated.encrypted_env.is_empty() - || sha2::Sha256::digest(&migrated.encrypted_env)[..] == migrated.source_sha256[..] - { - return Ok(ApplyReencryptedEnvResponse::default()); - } - let path = self.encrypted_env_path(&request.id)?; - let stored = fs::read(&path).context("Failed to read encrypted env")?; - if sha2::Sha256::digest(&stored)[..] != migrated.source_sha256[..] { - bail!("the encrypted env changed since the VM booted; not replacing it"); - } - safe_write::safe_write(&path, &migrated.encrypted_env) - .context("Failed to write encrypted env")?; - info!(vm_id = %request.id, root_epoch = migrated.root_epoch, "applied re-encrypted env"); - Ok(ApplyReencryptedEnvResponse { - replaced: true, - root_epoch: migrated.root_epoch, - }) - } - async fn shutdown_vm(self, request: Id) -> Result<()> { info!(vm_id = %request.id, "shutdown_vm RPC called"); self.guest_agent_client(&request.id)?.shutdown().await?; diff --git a/dstack/vmm/src/vmm-cli.py b/dstack/vmm/src/vmm-cli.py index f43c58114..ad469f981 100755 --- a/dstack/vmm/src/vmm-cli.py +++ b/dstack/vmm/src/vmm-cli.py @@ -692,14 +692,6 @@ def stop_vm(self, vm_id: str, force: bool = False) -> None: self.rpc_call("ShutdownVm", {"id": vm_id}) print(f"Gracefully shutting down VM {vm_id}") - def apply_reencrypted_env(self, vm_id: str) -> None: - """Swap in the env the guest re-encrypted to the current KMS key epoch.""" - result = self.rpc_call("ApplyReencryptedEnv", {"id": vm_id}) - if result.get("replaced"): - print(f"Replaced env of VM {vm_id} (key epoch {result.get('root_epoch', 0)})") - else: - print(f"VM {vm_id} has no re-encrypted env to apply") - def remove_vm(self, vm_id: str) -> None: """Remove a VM.""" self.rpc_call("RemoveVm", {"id": vm_id}) @@ -2167,12 +2159,6 @@ def _patched_format_help(): "--kms-url", action="append", type=str, help="KMS URL" ) - apply_env_parser = subparsers.add_parser( - "apply-reencrypted-env", - help="Replace a VM's env with the copy it re-encrypted to the current KMS key epoch", - ) - apply_env_parser.add_argument("vm_id", help="VM ID") - # Whitelist command kms_parser = subparsers.add_parser("kms", help="Manage trusted KMS whitelist") kms_subparsers = kms_parser.add_subparsers(dest="kms_action", help="KMS actions") @@ -2409,8 +2395,6 @@ def _patched_format_help(): cli.update_vm_env( args.vm_id, parse_env_file(args.env_file), kms_urls=args.kms_url ) - elif args.command == "apply-reencrypted-env": - cli.apply_reencrypted_env(args.vm_id) elif args.command == "update-user-config": cli.update_vm_user_config(args.vm_id, open(args.user_config, "r").read()) elif args.command == "update-app-compose": diff --git a/dstack/vmm/vmm.toml b/dstack/vmm/vmm.toml index 6db4441f6..177ef0ad3 100644 --- a/dstack/vmm/vmm.toml +++ b/dstack/vmm/vmm.toml @@ -346,6 +346,10 @@ ident = "dstack VMM" address = "vsock:2" port = 10000 +[host_api.limits] +# Largest encrypted env a guest may store after re-encrypting it at boot. +UpdateEncryptedEnv = "1MiB" + [key_provider] enabled = true address = "127.0.0.1" From 2e22df8f78a7afd4f7a424425b24c74e3bec9efa Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Fri, 6 Nov 2026 04:19:24 -0800 Subject: [PATCH 17/17] Simplify root epoch checks, app roots, and disk re-keying Share one chain-hash check between serving, onboarding, and the per-request epoch check. Collapse the guest agent's borrowed and historical app roots into one owned type. Re-key the data disk on the validated in-memory header copy and restore it to the disk in one step, so the boot never acts on an unvalidated header and the epoch token always matches slot 0. This drops the luksDump and test-passphrase probes. Signed-off-by: Kevin Wang --- dstack/dstack-util/src/system_setup.rs | 224 +++++++++++------------ dstack/guest-agent/src/rpc_service.rs | 60 ++---- dstack/guest-agent/src/rpc_service_v1.rs | 9 +- dstack/kms/src/main_service.rs | 103 ++++++----- dstack/kms/src/onboard_service.rs | 49 ++--- 5 files changed, 187 insertions(+), 258 deletions(-) diff --git a/dstack/dstack-util/src/system_setup.rs b/dstack/dstack-util/src/system_setup.rs index fefc66036..7177fcf02 100644 --- a/dstack/dstack-util/src/system_setup.rs +++ b/dstack/dstack-util/src/system_setup.rs @@ -2745,8 +2745,6 @@ impl<'a> Stage0<'a> { } async fn mount_data_disk(&self, keys: &AppKeys, opts: &DstackOptions) -> Result<()> { - let disk_crypt_key = &hex::encode(&keys.disk_crypt_key); - let root_epoch = kms_root_epoch(keys); let name = "dstack_data_disk"; let mount_point = &self.args.mount_point; @@ -2767,13 +2765,9 @@ impl<'a> Stage0<'a> { let has_filesystem = if opts.storage_encrypted && disk_initialized { info!("Mounting encrypted data disk"); - if let Some(epoch) = root_epoch { - self.migrate_disk_key(keys, disk_crypt_key, epoch) - .await - .context("failed to migrate the data disk key")?; - } - let initializing = - self.open_encrypted_volume(disk_crypt_key, name, opts.storage_discard)?; + let initializing = self + .open_encrypted_volume(keys, name, opts.storage_discard) + .await?; if initializing { warn!("data disk initialization was interrupted, initializing it again"); } @@ -2791,7 +2785,7 @@ impl<'a> Stage0<'a> { info!("Skipping disk encryption as requested by kernel cmdline"); } else if !disk_initialized { info!("Setting up disk encryption"); - self.luks_setup(disk_crypt_key, root_epoch, name, opts.storage_discard)?; + self.luks_setup(keys, name, opts.storage_discard).await?; } match opts.storage_fs { @@ -2920,13 +2914,8 @@ impl<'a> Stage0<'a> { Ok(()) } - fn luks_setup( - &self, - disk_crypt_key: &str, - root_epoch: Option, - name: &str, - discard: bool, - ) -> Result<()> { + async fn luks_setup(&self, keys: &AppKeys, name: &str, discard: bool) -> Result<()> { + let disk_crypt_key = hex::encode(&keys.disk_crypt_key); let root_hd = &self.args.device; let sector_offset = PAYLOAD_OFFSET / 512; info!("Formatting encrypted disk"); @@ -2964,14 +2953,15 @@ impl<'a> Stage0<'a> { { bail!("Failed to setup luks volume"); } - if let Some(epoch) = root_epoch.filter(|epoch| *epoch > 0) { - write_luks_root_epoch(root_hd, epoch, false)?; + if let Some(epoch) = kms_root_epoch(keys).filter(|epoch| *epoch > 0) { + write_luks_root_epoch(root_hd, epoch)?; } - self.open_encrypted_volume(disk_crypt_key, name, discard)?; + self.open_encrypted_volume(keys, name, discard).await?; Ok(()) } - /// Re-key a data disk still protected by an older KMS epoch. + /// Re-key the in-memory `header` of a data disk still protected by an + /// older KMS epoch, and write it back to the disk. /// /// Only the keyslot passphrase changes; the volume key stays. The epoch /// is recorded in the header so the right old key is used after the next @@ -2980,39 +2970,32 @@ impl<'a> Stage0<'a> { async fn migrate_disk_key( &self, keys: &AppKeys, - current_key: &str, - current: u32, + header: &Path, + from: u32, + to: u32, ) -> Result<()> { - let device = &self.args.device; - let stored = read_luks_root_epoch(device)?; - if !luks_unlocks(device, current_key)? { - let epoch = stored.unwrap_or(0); - if epoch >= current { - bail!("the data disk does not unlock with the key of epoch {current}"); - } - info!("re-keying the data disk from key epoch {epoch} to {current}"); - let old = self.shared.request_historical_keys(keys, epoch).await?; - let old_key = hex::encode(&old.disk_crypt_key); - luks_change_key(device, &old_key, current_key, Path::new(REKEY_DIR))?; - if !luks_unlocks(device, current_key)? { - bail!("the data disk does not unlock after re-keying"); - } - } - if stored.unwrap_or(0) != current { - write_luks_root_epoch(device, current, stored.is_some())?; - } + info!("re-keying the data disk from root epoch {from} to {to}"); + let old = self.shared.request_historical_keys(keys, from).await?; + rekey_luks_header( + header, + &hex::encode(&old.disk_crypt_key), + &hex::encode(&keys.disk_crypt_key), + to, + )?; + let root_hd = &self.args.device; + cmd!(cryptsetup luksHeaderRestore --batch-mode --header-backup-file=$header $root_hd) + .context("failed to write the re-keyed LUKS header")?; Ok(()) } - /// Returns whether the header still carries the initializing label. - fn open_encrypted_volume( + async fn open_encrypted_volume( &self, - disk_crypt_key: &str, + keys: &AppKeys, name: &str, discard: bool, ) -> Result { let root_hd = &self.args.device; - let disk_crypt_key = disk_crypt_key.trim(); + let disk_crypt_key = hex::encode(&keys.disk_crypt_key); // Create a private tmpfs mount to ensure the header stays in-memory. let tmp_hdr_dir = "/tmp/dstack-luks-header"; let in_mem_hdr = format!("{tmp_hdr_dir}/luks-header"); @@ -3034,9 +3017,27 @@ impl<'a> Stage0<'a> { } .context("Failed to load LUKS2 header")?; - let hdr_file = fs::File::open(&in_mem_hdr).context("Failed to open LUKS2 header")?; - let LuksHeaderState { initializing, .. } = - validate_luks2_headers(hdr_file).context("Failed to validate LUKS2 header")?; + let validate = || { + let hdr_file = fs::File::open(&in_mem_hdr).context("Failed to open LUKS2 header")?; + validate_luks2_headers(hdr_file).context("Failed to validate LUKS2 header") + }; + let mut header = validate()?; + if let Some(current) = kms_root_epoch(keys) { + let stored = header.root_epoch.unwrap_or(0); + if stored > current { + bail!("the data disk is at root epoch {stored}, newer than the current {current}"); + } + if stored < current { + self.migrate_disk_key(keys, Path::new(&in_mem_hdr), stored, current) + .await + .context("failed to migrate the data disk key")?; + header = validate()?; + if header.root_epoch != Some(current) { + bail!("the re-keyed LUKS header does not record root epoch {current}"); + } + } + } + let initializing = header.initializing; info!("Opening the device"); let mut command = Command::new("cryptsetup"); @@ -3518,73 +3519,32 @@ fn run_with_input(mut command: Command, input: &[u8]) -> Result> { .code()) } -/// Whether `key` unlocks keyslot 0. cryptsetup exits with 2 on a wrong key. -fn luks_unlocks(device: &Path, key: &str) -> Result { - let command = cryptsetup( - &["open", "--test-passphrase", "--key-slot", "0", "-d-"], - device, - ); - match run_with_input(command, key.as_bytes())? { - Some(0) => Ok(true), - Some(2) => Ok(false), - code => bail!("cryptsetup failed to test the data disk key: {code:?}"), - } -} - -fn read_luks_root_epoch(device: &Path) -> Result> { - #[derive(Deserialize)] - struct Metadata { - tokens: BTreeMap, - } - let output = cryptsetup(&["luksDump", "--dump-json-metadata"], device) - .output() - .context("failed to run cryptsetup luksDump")?; - if !output.status.success() { - bail!("cryptsetup luksDump failed"); - } - let metadata: Metadata = - serde_json::from_slice(&output.stdout).context("failed to parse LUKS metadata")?; - Ok(metadata - .tokens - .values() - .find(|token| token.kind == ROOT_EPOCH_TOKEN_TYPE) - .and_then(|token| token.root_epoch)) -} - -fn write_luks_root_epoch(device: &Path, epoch: u32, replace: bool) -> Result<()> { +fn write_luks_root_epoch(device: &Path, epoch: u32) -> Result<()> { let token = serde_json::json!({ "type": ROOT_EPOCH_TOKEN_TYPE, "keyslots": ["0"], "root_epoch": epoch, }); - let args: &[&str] = if replace { - &["token", "import", "--token-id", "0", "--token-replace"] - } else { - &["token", "import", "--token-id", "0"] - }; - let code = run_with_input(cryptsetup(args, device), token.to_string().as_bytes())?; + let args = ["token", "import", "--token-id", "0", "--token-replace"]; + let code = run_with_input(cryptsetup(&args, device), token.to_string().as_bytes())?; if code != Some(0) { - bail!("failed to record key epoch {epoch} in the LUKS header: {code:?}"); + bail!("failed to record root epoch {epoch} in the LUKS header: {code:?}"); } Ok(()) } -/// Directory on the /run tmpfs for the old key handed to cryptsetup as a file. -const REKEY_DIR: &str = "/run/dstack-rekey"; +/// Replace the keyslot 0 passphrase of a detached LUKS `header` and record +/// `epoch` in it. The new key goes over stdin; the old one is written next to +/// the header, which must be in a private in-memory directory. +fn rekey_luks_header(header: &Path, old_key: &str, new_key: &str, epoch: u32) -> Result<()> { + use std::os::unix::fs::PermissionsExt; -/// Replace the keyslot 0 passphrase in place. The new key goes over stdin; -/// the old one is written to a 0700 directory in `secret_dir`, which is -/// removed before returning. -fn luks_change_key(device: &Path, old_key: &str, new_key: &str, secret_dir: &Path) -> Result<()> { - fs::create_dir_all(secret_dir)?; + // luksHeaderBackup leaves the copy read-only. + fs::set_permissions(header, std::fs::Permissions::from_mode(0o600))?; + let old_file = header.with_file_name("old-key"); defer! { - fs::remove_dir_all(secret_dir).ok(); + fs::remove_file(&old_file).ok(); } - fs::set_permissions( - secret_dir, - std::os::unix::fs::PermissionsExt::from_mode(0o700), - )?; - let old_file = secret_dir.join("old-key"); safe_write_with_mode(&old_file, old_key, 0o600)?; let mut command = cryptsetup( &[ @@ -3595,13 +3555,13 @@ fn luks_change_key(device: &Path, old_key: &str, new_key: &str, secret_dir: &Pat "--key-slot", "0", ], - device, + header, ); command.arg("--key-file").arg(&old_file); if run_with_input(command, new_key.as_bytes())? != Some(0) { bail!("cryptsetup luksChangeKey failed"); } - Ok(()) + write_luks_root_epoch(header, epoch) } /// Set by luksFormat and cleared once the filesystem exists, so a disk still @@ -3929,29 +3889,49 @@ fn test_disk_key_migration_with_cryptsetup() { &image, ); assert_eq!(run_with_input(format, b"old").unwrap(), Some(0)); + let unlocks = |key: &str| { + let test = cryptsetup( + &["open", "--test-passphrase", "--key-slot", "0", "-d-"], + &image, + ); + run_with_input(test, key.as_bytes()).unwrap() == Some(0) + }; - assert!(luks_unlocks(&image, "old").unwrap()); - assert!(!luks_unlocks(&image, "new").unwrap()); - assert_eq!(read_luks_root_epoch(&image).unwrap(), None); - - let secret_dir = dir.path().join("rekey"); - luks_change_key(&image, "old", "new", &secret_dir).unwrap(); - assert!(!secret_dir.exists()); - assert!(luks_unlocks(&image, "new").unwrap()); - assert!(!luks_unlocks(&image, "old").unwrap()); - - write_luks_root_epoch(&image, 2, false).unwrap(); - assert_eq!(read_luks_root_epoch(&image).unwrap(), Some(2)); - write_luks_root_epoch(&image, 3, true).unwrap(); - assert_eq!(read_luks_root_epoch(&image).unwrap(), Some(3)); - - // The re-keyed, tagged header still passes the boot-time validation. + // Re-key a detached copy of the header, as the boot does, and restore it. let header = dir.path().join("header"); - let status = cryptsetup(&["luksHeaderBackup", "--header-backup-file"], &header) + let backup = cryptsetup(&["luksHeaderBackup", "--header-backup-file"], &header) + .arg(&image) + .status() + .unwrap(); + assert!(backup.success()); + rekey_luks_header(&header, "old", "new", 2).unwrap(); + assert!(!dir.path().join("old-key").exists()); + assert!( + unlocks("old"), + "the disk is unchanged until the header is restored" + ); + let state = validate_luks2_headers(fs::File::open(&header).unwrap()).unwrap(); + assert_eq!(state.root_epoch, Some(2)); + + let restore = cryptsetup( + &["luksHeaderRestore", "--batch-mode", "--header-backup-file"], + &header, + ) + .arg(&image) + .status() + .unwrap(); + assert!(restore.success()); + assert!(unlocks("new")); + assert!(!unlocks("old")); + + // A later rotation replaces the recorded epoch. + write_luks_root_epoch(&image, 3).unwrap(); + fs::remove_file(&header).unwrap(); + let backup = cryptsetup(&["luksHeaderBackup", "--header-backup-file"], &header) .arg(&image) .status() .unwrap(); - assert!(status.success()); + assert!(backup.success()); let state = validate_luks2_headers(fs::File::open(&header).unwrap()).unwrap(); assert_eq!(state.root_epoch, Some(3)); } diff --git a/dstack/guest-agent/src/rpc_service.rs b/dstack/guest-agent/src/rpc_service.rs index 838939034..068d8b53d 100644 --- a/dstack/guest-agent/src/rpc_service.rs +++ b/dstack/guest-agent/src/rpc_service.rs @@ -2,7 +2,6 @@ // // SPDX-License-Identifier: Apache-2.0 -use std::borrow::Cow; use std::collections::HashMap; use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; @@ -94,38 +93,21 @@ pub struct AppState { inner: Arc, } -/// The app root key of one KMS root epoch. -pub(crate) struct AppRoot<'a> { - pub epoch: u32, - pub k256_key: Cow<'a, [u8]>, - pub signing_key: Cow<'a, SigningKey>, - pub kms_signature: Cow<'a, [u8]>, -} - -/// The app root key of an older KMS root epoch, fetched from the KMS on first -/// use. +/// The app root key of one KMS root epoch, with its KMS signature. #[derive(Clone)] -pub(crate) struct HistoricalRoot { - k256_key: Vec, - signing_key: SigningKey, - kms_signature: Vec, +pub(crate) struct AppRoot { + pub epoch: u32, + pub k256_key: Vec, + pub signing_key: SigningKey, + pub kms_signature: Vec, } -impl HistoricalRoot { - #[cfg(test)] - pub(crate) fn k256_key(&self) -> &[u8] { - &self.k256_key - } - - #[cfg(test)] - pub(crate) fn kms_signature(&self) -> &[u8] { - &self.kms_signature - } - - fn new(k256_key: Vec, kms_signature: Vec) -> Result { +impl AppRoot { + fn new(epoch: u32, k256_key: Vec, kms_signature: Vec) -> Result { let signing_key = SigningKey::from_slice(&k256_key) .context("the app root k256 key is not a valid secp256k1 scalar")?; Ok(Self { + epoch, k256_key, signing_key, kms_signature, @@ -139,7 +121,7 @@ struct AppStateInner { vm_config: String, cert_client: CertRequestClient, /// Held across a fetch so concurrent requests for one epoch fetch it once. - historical_roots: tokio::sync::Mutex>, + historical_roots: tokio::sync::Mutex>, demo_cert: tokio::sync::OnceCell, /// When the last demo-cert request failed; held while one is running. demo_cert_last_failure: tokio::sync::Mutex>, @@ -429,30 +411,24 @@ impl AppState { /// The app root key and its KMS signature of `root_epoch`; the current /// ones when unset. An older epoch is requested from the KMS once and kept /// in memory. - pub(crate) async fn app_root(&self, root_epoch: Option) -> Result> { + pub(crate) async fn app_root(&self, root_epoch: Option) -> Result { let current = self.root_epoch(); let epoch = root_epoch.unwrap_or(current); if epoch == current { return Ok(AppRoot { epoch, - k256_key: Cow::Borrowed(self.app_root_k256_key()), - signing_key: Cow::Borrowed(self.app_root_signing_key()?), - kms_signature: Cow::Borrowed(self.kms_k256_signature()), + k256_key: self.app_root_k256_key().to_vec(), + signing_key: self.app_root_signing_key()?.clone(), + kms_signature: self.kms_k256_signature().to_vec(), }); } if epoch > current { anyhow::bail!("root epoch {epoch} is newer than the current one, {current}"); } - let root = self.historical_root(epoch).await?; - Ok(AppRoot { - epoch, - k256_key: Cow::Owned(root.k256_key), - signing_key: Cow::Owned(root.signing_key), - kms_signature: Cow::Owned(root.kms_signature), - }) + self.historical_root(epoch).await } - pub(crate) async fn historical_root(&self, epoch: u32) -> Result { + pub(crate) async fn historical_root(&self, epoch: u32) -> Result { let mut roots = self.inner.historical_roots.lock().await; if let Some(root) = roots.get(&epoch) { return Ok(root.clone()); @@ -463,7 +439,7 @@ impl AppState { .get_app_key(epoch) .await .with_context(|| format!("failed to fetch the app keys of root epoch {epoch}"))?; - let root = HistoricalRoot::new(keys.k256_key, keys.k256_signature)?; + let root = AppRoot::new(epoch, keys.k256_key, keys.k256_signature)?; roots.insert(epoch, root.clone()); Ok(root) } @@ -475,7 +451,7 @@ impl AppState { k256_key: Vec, kms_signature: Vec, ) { - let root = HistoricalRoot::new(k256_key, kms_signature).unwrap(); + let root = AppRoot::new(epoch, k256_key, kms_signature).unwrap(); self.inner.historical_roots.lock().await.insert(epoch, root); } diff --git a/dstack/guest-agent/src/rpc_service_v1.rs b/dstack/guest-agent/src/rpc_service_v1.rs index fa7edd7f0..8f178ccbf 100644 --- a/dstack/guest-agent/src/rpc_service_v1.rs +++ b/dstack/guest-agent/src/rpc_service_v1.rs @@ -170,10 +170,7 @@ impl V1RpcHandler { let algorithm = Algorithm::parse(algorithm)?; let root = self.state.app_root(root_epoch).await?; let key = AppKey::derive(&root.k256_key, domain, algorithm)?; - let chain = vec![ - key.claim_signature(&root.signing_key)?, - root.kms_signature.to_vec(), - ]; + let chain = vec![key.claim_signature(&root.signing_key)?, root.kms_signature]; Ok((key, chain, root.epoch)) } } @@ -383,8 +380,8 @@ mod tests { .insert_historical_root(1, key.clone(), vec![0xab; 65]) .await; let root = state.historical_root(1).await.unwrap(); - assert_eq!(root.k256_key(), &key[..]); - assert_eq!(root.kms_signature(), &[0xab; 65][..]); + assert_eq!(root.k256_key, key); + assert_eq!(root.kms_signature, [0xab; 65]); // An uncached epoch needs the KMS, which a local key provider lacks. assert!(state.historical_root(2).await.is_err()); } diff --git a/dstack/kms/src/main_service.rs b/dstack/kms/src/main_service.rs index 942290f18..feeb04e19 100644 --- a/dstack/kms/src/main_service.rs +++ b/dstack/kms/src/main_service.rs @@ -26,6 +26,7 @@ use ra_tls::{ attestation::{resolve_vm_config, AttestationVerifier, TeeVariant, VerifiedAttestation}, cert::{CaCert, CertRequest, CertSigningRequestV1, CertSigningRequestV2, Csr}, kdf, + rcgen::KeyPair, }; use scale::Decode; use tracing::warn; @@ -99,18 +100,7 @@ impl Keyring { .with_context(|| format!("failed to load root keys of epoch {epoch}")) }) .collect::>>()?; - let pubkeys: Vec<_> = roots - .iter() - .map(|root| { - ( - root.ca.key.public_key_der(), - root.k256.verifying_key().to_sec1_bytes().to_vec(), - ) - }) - .collect(); - let chain_hashes = dstack_types::root_epoch::chain_hashes( - pubkeys.iter().map(|(ca, k256)| (&ca[..], &k256[..])), - ); + let chain_hashes = root_chain_hashes(roots.iter().map(|root| (&root.ca.key, &root.k256))); Ok(Self { roots, chain_hashes, @@ -121,33 +111,53 @@ impl Keyring { self.roots.len() as u32 - 1 } - /// The epoch to serve under `registry`: its current epoch, which this - /// keyring must hold with a matching chain hash. A keyring nobody - /// registered is servable only at epoch 0, as before root epochs existed. pub(crate) fn servable_epoch(&self, registry: RootEpochRegistry) -> Result { - let latest = self.latest(); - match registry { - RootEpochRegistry::Unchecked => Ok(latest), - RootEpochRegistry::Empty if latest == 0 => Ok(0), - RootEpochRegistry::Empty => { - bail!("this KMS holds root epoch {latest}, but no root epoch is registered") - } - RootEpochRegistry::Current(current) => { - let Some(local) = self.chain_hashes.get(current.epoch as usize) else { - bail!( - "root epoch {} is current, but this KMS holds only up to {latest}", - current.epoch - ); - }; - if *local != current.chain_hash { - bail!( - "the root keys up to epoch {} do not match the registered ones", - current.epoch - ); - } - Ok(current.epoch) - } + servable_epoch(&self.chain_hashes, registry) + } +} + +/// Chain hashes of the root keys of epochs 0, 1, ... +pub(crate) fn root_chain_hashes<'a>( + roots: impl Iterator, +) -> Vec<[u8; 32]> { + let pubkeys: Vec<_> = roots + .map(|(ca, k256)| { + ( + ca.public_key_der(), + k256.verifying_key().to_sec1_bytes().to_vec(), + ) + }) + .collect(); + dstack_types::root_epoch::chain_hashes(pubkeys.iter().map(|(ca, k256)| (&ca[..], &k256[..]))) +} + +/// The epoch a keyring with `chain_hashes` serves under `registry`: the +/// current epoch, which the keyring must hold with a matching chain hash. A +/// keyring nobody registered is servable only at epoch 0, as before root +/// epochs existed. +pub(crate) fn servable_epoch( + chain_hashes: &[[u8; 32]], + registry: RootEpochRegistry, +) -> Result { + let latest = chain_hashes.len() as u32 - 1; + let current = match registry { + RootEpochRegistry::Unchecked => return Ok(latest), + RootEpochRegistry::Empty if latest == 0 => return Ok(0), + RootEpochRegistry::Empty => { + bail!("the keyring holds root epoch {latest}, but no root epoch is registered") } + RootEpochRegistry::Current(current) => current, + }; + match chain_hashes.get(current.epoch as usize) { + Some(hash) if *hash == current.chain_hash => Ok(current.epoch), + Some(_) => bail!( + "the root keys up to epoch {} are not the registered ones", + current.epoch + ), + None => bail!( + "root epoch {} is current, but the keyring holds only up to {latest}", + current.epoch + ), } } @@ -300,20 +310,13 @@ impl KmsState { return Ok(()); } let serving = self.serving_epoch; - match response.root_epoch.registry()? { - RootEpochRegistry::Current(current) - if current.epoch == serving - && current.chain_hash == self.keyring.chain_hashes[serving as usize] => - { - Ok(()) - } - RootEpochRegistry::Empty if serving == 0 => Ok(()), - registry => { - warn!("governance moved off root epoch {serving} ({registry:?}); restarting"); - self.restart.notify_one(); - bail!("root epoch {serving} is no longer current; the KMS is restarting") - } + let registry = response.root_epoch.registry()?; + if self.keyring.servable_epoch(registry).ok() == Some(serving) { + return Ok(()); } + warn!("governance moved off root epoch {serving} ({registry:?}); restarting"); + self.restart.notify_one(); + bail!("root epoch {serving} is no longer current; the KMS is restarting") } } diff --git a/dstack/kms/src/onboard_service.rs b/dstack/kms/src/onboard_service.rs index db54129d9..311057e6d 100644 --- a/dstack/kms/src/onboard_service.rs +++ b/dstack/kms/src/onboard_service.rs @@ -36,7 +36,7 @@ use tracing::info; use crate::{ config::KmsConfig, main_service::{ - build_boot_info_for_attestation, + build_boot_info_for_attestation, root_chain_hashes, servable_epoch, upgrade_authority::{ app_attest, dstack_client, ensure_kms_allowed, ensure_self_kms_allowed, pad64, RootEpochRegistry, @@ -1280,52 +1280,25 @@ pub(crate) async fn bootstrap_keys(cfg: &KmsConfig, verifier: &AttestationVerifi /// Check the handed-over roots against governance: they must include the /// current epoch with a matching chain hash. A rotation must start from the /// current epoch with nothing pending activation, or it would collide with the -/// next registered epoch. Deployments that never registered epochs can still -/// join at epoch 0. +/// next registered epoch. fn ensure_handover_registered( roots: &[(KeyPair, SigningKey)], registry: RootEpochRegistry, rotate: bool, ) -> Result<()> { + if rotate && registry == RootEpochRegistry::Empty { + bail!("register root epoch 0 before rotating"); + } + let hashes = root_chain_hashes(roots.iter().map(|(ca, k256)| (ca, k256))); + let current = servable_epoch(&hashes, registry)?; let latest = roots.len() as u32 - 1; - let current = match registry { - RootEpochRegistry::Unchecked => return Ok(()), - RootEpochRegistry::Empty if latest == 0 && !rotate => return Ok(()), - RootEpochRegistry::Empty => { - bail!("register root epoch 0 before rotating or joining past it") - } - RootEpochRegistry::Current(current) => current, - }; - if rotate && current.epoch != latest { + if rotate && current != latest { bail!( - "the source holds root epochs up to {latest}, but {} is current; rotate from a \ - source at the current epoch with no epoch pending activation", - current.epoch + "the source holds root epochs up to {latest}, but {current} is current; rotate from \ + a source at the current epoch with no epoch pending activation" ); } - let pubkeys: Vec<_> = roots - .iter() - .map(|(ca_key, k256_key)| { - ( - ca_key.public_key_der(), - k256_key.verifying_key().to_sec1_bytes().to_vec(), - ) - }) - .collect(); - let hashes = dstack_types::root_epoch::chain_hashes( - pubkeys.iter().map(|(ca, k256)| (&ca[..], &k256[..])), - ); - match hashes.get(current.epoch as usize) { - Some(hash) if *hash == current.chain_hash => Ok(()), - Some(_) => bail!( - "the source's root keys up to epoch {} are not the registered ones", - current.epoch - ), - None => bail!( - "the source holds root epochs up to {latest}, but {} is current", - current.epoch - ), - } + Ok(()) } fn self_signed_root_ca(ca_key: &KeyPair) -> Result {