diff --git a/docs/deployment.md b/docs/deployment.md index 4f6656668..f3766ad37 100644 --- a/docs/deployment.md +++ b/docs/deployment.md @@ -513,6 +513,40 @@ 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; 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 + 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. 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. 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; 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 +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/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/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 2a752e4b0..9c6845b05 100644 --- a/dstack/crates/dstack-auth/src/main.rs +++ b/dstack/crates/dstack-auth/src/main.rs @@ -66,6 +66,17 @@ struct BootResponse { is_allowed: bool, gateway_app_id: String, reason: String, + #[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)] @@ -75,6 +86,41 @@ struct Allowlist { gateway_app_id: String, kms: KmsRules, 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, Default, Clone)] +#[serde(rename_all = "camelCase", default)] +struct RootEpoch { + ca_pubkey: String, + k256_pubkey: String, +} + +impl Allowlist { + /// `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()?)), + }) + } } #[derive(Deserialize, Default)] @@ -122,6 +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.current_root_epoch(), } } @@ -130,6 +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.current_root_epoch(), } } @@ -206,15 +254,24 @@ 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 { + let mut info = json!({ "status": "ok", "kmsContractAddr": "", "ethRpcUrl": "", "gatewayAppId": "", "chainId": 0, - "appImplementation": "" - })) + "appImplementation": "", + }); + if let Some(epoch) = al.current_root_epoch() { + info["rootEpoch"] = epoch.root_epoch.into(); + info["rootEpochHash"] = epoch.root_epoch_hash.into(); + } + info } #[rocket::main] @@ -283,6 +340,37 @@ mod tests { assert!(!check_kms(&info, &Allowlist::default()).is_allowed); } + #[test] + fn reports_the_current_root_epoch() { + let mut al: Allowlist = serde_json::from_str( + r#"{"kms":{"mrAggregated":["0xMR"]}, + "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"], 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()); + } + // 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/dstack-types/src/lib.rs b/dstack/dstack-types/src/lib.rs index 68369be18..edb667c81 100644 --- a/dstack/dstack-types/src/lib.rs +++ b/dstack/dstack-types/src/lib.rs @@ -2407,6 +2407,17 @@ pub struct AppKeys { pub key_provider: KeyProvider, } +impl AppKeys { + /// 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, + _ => 0, + } + } +} + /// Maximum length of a framed message exchanged with the local key provider. pub const KEY_PROVIDER_MAX_FRAME_SIZE: usize = 8 * 1024 * 1024; @@ -2431,6 +2442,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)] + root_epoch: u32, }, } @@ -2494,6 +2508,7 @@ mod key_provider_tests { pubkey: vec![0xab; 32], tmp_ca_key: String::new(), tmp_ca_cert: String::new(), + root_epoch: 0, }; assert_eq!(kms.id(), &[0xab; 32]); } @@ -2525,6 +2540,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/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/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/main.rs b/dstack/dstack-util/src/main.rs index 87867085b..4d76237fd 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(), + root_epoch: None, }) .await .context("Failed to get app key")?; @@ -723,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, + root_epoch: response.root_epoch, }, }; diff --git a/dstack/dstack-util/src/system_setup.rs b/dstack/dstack-util/src/system_setup.rs index 358a8a003..7177fcf02 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, ROOT_EPOCH_TOKEN_TYPE, }; use ra_rpc::{ client::{CertInfo, RaClient, RaClientConfig}, @@ -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}, @@ -2353,6 +2353,85 @@ struct AppInfo { gpu_attestation_event: Option>, } +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}"); + 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(), + root_epoch, + }) + .await + .context("Failed to get app key")?; + // 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()) + .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, + root_epoch: served, + }, + }; + Ok((keys, response.os_image_hash)) + } + + /// 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> { args: &'a SetupArgs, shared: HostShared, @@ -2415,62 +2494,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(), - }) - .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"); @@ -2478,7 +2501,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(), 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. @@ -2721,7 +2744,7 @@ 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 name = "dstack_data_disk"; let mount_point = &self.args.mount_point; @@ -2742,8 +2765,9 @@ impl<'a> Stage0<'a> { let has_filesystem = if opts.storage_encrypted && disk_initialized { info!("Mounting encrypted data disk"); - 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"); } @@ -2761,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, name, opts.storage_discard)?; + self.luks_setup(keys, name, opts.storage_discard).await?; } match opts.storage_fs { @@ -2890,7 +2914,8 @@ impl<'a> Stage0<'a> { Ok(()) } - fn luks_setup(&self, disk_crypt_key: &str, 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"); @@ -2928,19 +2953,49 @@ impl<'a> Stage0<'a> { { bail!("Failed to setup luks volume"); } - self.open_encrypted_volume(disk_crypt_key, name, discard)?; + if let Some(epoch) = kms_root_epoch(keys).filter(|epoch| *epoch > 0) { + write_luks_root_epoch(root_hd, epoch)?; + } + self.open_encrypted_volume(keys, name, discard).await?; Ok(()) } - /// Returns whether the header still carries the initializing label. - fn open_encrypted_volume( + /// 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 + /// rotation. An unrotated epoch-0 disk is left untouched, so it stays + /// readable by older OS images. + async fn migrate_disk_key( &self, - disk_crypt_key: &str, + keys: &AppKeys, + header: &Path, + from: u32, + to: u32, + ) -> Result<()> { + 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(()) + } + + async fn open_encrypted_volume( + &self, + 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"); @@ -2962,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 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"); @@ -3193,8 +3266,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 @@ -3219,25 +3291,34 @@ 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 + /// 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, None)), + Err(error) => error, }; - Ok(vars) + let current = self.keys.root_epoch(); + 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"); + } + return Ok((plaintext, Some(encrypted_env))); + } + Err(error).context("failed to decrypt env file") } fn write_env_file(&self, env_vars: &BTreeMap) -> Result<()> { @@ -3254,7 +3335,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 @@ -3262,18 +3343,32 @@ 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, reencrypted) = if self.shared.encrypted_env.is_empty() { + info!("No encrypted env, using default"); + (Default::default(), None) + } else { + info!("Processing encrypted env"); + 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) } 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()?; @@ -3387,18 +3482,113 @@ 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_root_epoch(keys: &AppKeys) -> Option { + match keys.key_provider { + KeyProvider::Kms { root_epoch, .. } => Some(root_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()) +} + +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 = ["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 root epoch {epoch} in the LUKS header: {code:?}"); + } + Ok(()) +} + +/// 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; + + // 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_file(&old_file).ok(); + } + safe_write_with_mode(&old_file, old_key, 0o600)?; + let mut command = cryptsetup( + &[ + "luksChangeKey", + "--batch-mode", + "--pbkdf", + "pbkdf2", + "--key-slot", + "0", + ], + header, + ); + command.arg("--key-file").arg(&old_file); + if run_with_input(command, new_key.as_bytes())? != Some(0) { + bail!("cryptsetup luksChangeKey failed"); + } + write_luks_root_epoch(header, epoch) +} + /// 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, + /// Root epoch recorded in the dstack token; `None` means epoch 0. + 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.root_epoch != secondary.root_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 { +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) @@ -3469,9 +3659,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 root_epoch = match tokens.len() { + 0 => None, + 1 => { + let token = tokens.get(&0).context("Invalid LUKS token id")?; + if token.kind != ROOT_EPOCH_TOKEN_TYPE || token.keyslots != ["0"] { + bail!("Invalid LUKS token"); + } + Some(token.root_epoch.context("LUKS token has no root epoch")?) + } + _ => bail!("Invalid LUKS tokens"), + }; if segments.len() != 1 { bail!("Invalid LUKS segments"); } @@ -3587,18 +3785,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, + root_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, + root_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 @@ -3642,6 +3853,109 @@ 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)); + 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) + }; + + // Re-key a detached copy of the header, as the boot does, and restore it. + let header = dir.path().join("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!(backup.success()); + let state = validate_luks2_headers(fs::File::open(&header).unwrap()).unwrap(); + assert_eq!(state.root_epoch, Some(3)); +} + +#[test] +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}}}"#); + 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"],"root_epoch":2}"#).unwrap(); + assert_eq!(state.root_epoch, Some(2)); + for token in [ + 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}"); + } +} + #[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..aaa184a1c 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 ROOT_EPOCH_TOKEN_TYPE: &str = "dstack-kms"; + +#[derive(Deserialize)] +pub(super) struct LuksToken { + #[serde(rename = "type")] + pub kind: String, + pub keyslots: Vec, + #[serde(default)] + pub root_epoch: Option, +} + #[derive(Deserialize)] #[serde(tag = "type", rename_all = "lowercase")] pub(super) enum LuksKeyslot { diff --git a/dstack/guest-agent/rpc/proto/agent_rpc_v1.proto b/dstack/guest-agent/rpc/proto/agent_rpc_v1.proto index 3df408ad8..bebc85f99 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. 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 { @@ -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/rpc_service.rs b/dstack/guest-agent/src/rpc_service.rs index a6dcd8e5f..068d8b53d 100644 --- a/dstack/guest-agent/src/rpc_service.rs +++ b/dstack/guest-agent/src/rpc_service.rs @@ -2,6 +2,7 @@ // // SPDX-License-Identifier: Apache-2.0 +use std::collections::HashMap; use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; @@ -92,11 +93,35 @@ pub struct AppState { inner: Arc, } +/// The app root key of one KMS root epoch, with its KMS signature. +#[derive(Clone)] +pub(crate) struct AppRoot { + pub epoch: u32, + pub k256_key: Vec, + pub signing_key: SigningKey, + pub kms_signature: Vec, +} + +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, + }) + } +} + struct AppStateInner { config: Config, 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>, @@ -300,6 +325,7 @@ impl AppState { config, keys, cert_client, + historical_roots: Default::default(), demo_cert: Default::default(), demo_cert_last_failure: Default::default(), vm_config, @@ -377,6 +403,58 @@ 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 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().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}"); + } + self.historical_root(epoch).await + } + + 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 = AppRoot::new(epoch, 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 = AppRoot::new(epoch, 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. @@ -1152,6 +1230,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 } @@ -1442,6 +1522,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 354e94618..8f178ccbf 100644 --- a/dstack/guest-agent/src/rpc_service_v1.rs +++ b/dstack/guest-agent/src/rpc_service_v1.rs @@ -159,16 +159,19 @@ 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. + async fn derive( + &self, + domain: &str, + algorithm: &str, + root_epoch: Option, + ) -> Result<(AppKey, Vec>, u32)> { let algorithm = Algorithm::parse(algorithm)?; - let key = AppKey::derive(self.state.app_root_k256_key(), domain, algorithm)?; - let chain = vec![ - key.claim_signature(self.state.app_root_signing_key()?)?, - self.state.kms_k256_signature().to_vec(), - ]; - Ok((key, chain)) + 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]; + Ok((key, chain, root.epoch)) } } @@ -196,11 +199,14 @@ 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) + .await?; Ok(GetKeyResponse { key: key.secret(), public_key: key.public_key(), signature_chain, + root_epoch, }) } @@ -342,9 +348,44 @@ mod tests { GetKeyRequest { domain: domain.to_string(), algorithm: algorithm.to_string(), + root_epoch: None, } } + #[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(None)).await.unwrap(); + assert_eq!(current.root_epoch, 0); + assert_eq!( + handler().get_key(request(Some(0))).await.unwrap().key, + current.key + ); + assert!(handler().get_key(request(Some(1))).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 /// 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/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/kms/auth-eth-bun/index.test.ts b/dstack/kms/auth-eth-bun/index.test.ts index 4d8a85f45..aefdaa5a8 100644 --- a/dstack/kms/auth-eth-bun/index.test.ts +++ b/dstack/kms/auth-eth-bun/index.test.ts @@ -230,6 +230,61 @@ 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 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 'currentRootEpoch': return [1n, '0x' + '11'.repeat(32)]; + } + }); + + const boot = await appFetch(new Request('http://localhost:3001/bootAuth/app', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify(validBootInfo), + })); + 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()).toMatchObject(expected); + }); + + 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 'currentRootEpoch': + 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 +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 === 'currentRootEpoch') return undefined; throw new Error(`unexpected function ${params.functionName}`); }); @@ -468,6 +524,7 @@ describe('Authorization freshness and domain binding', () => { return [true, 'allowed']; } if (params.functionName === 'gatewayAppId') return 'gateway-app'; + if (params.functionName === 'currentRootEpoch') return undefined; throw new Error(`unexpected function ${params.functionName}`); }); @@ -492,6 +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 === 'currentRootEpoch') return undefined; throw new Error(`unexpected function ${params.functionName}`); }); const recovered = await postApp(); @@ -522,6 +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 === 'currentRootEpoch') return undefined; throw new Error(`unexpected function ${params.functionName}`); }); @@ -541,6 +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 === 'currentRootEpoch') return undefined; throw new Error(`unexpected function ${params.functionName}`); }); @@ -569,6 +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 === 'currentRootEpoch') return undefined; throw new Error(`unexpected function ${params.functionName}`); }); const recovered = await authorize(); @@ -611,6 +672,7 @@ describe('dstack-kms 0.6 bootAuth payload', () => { return [true, '']; } if (params.functionName === 'gatewayAppId') return 'gateway-app'; + 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 b192d9f0f..4cc8f9a9a 100644 --- a/dstack/kms/auth-eth-bun/index.ts +++ b/dstack/kms/auth-eth-bun/index.ts @@ -35,9 +35,15 @@ const BootInfoSchema = z.object({ const BootResponseSchema = z.object({ isAllowed: z.boolean(), reason: z.string(), - gatewayAppId: z.string() + gatewayAppId: z.string(), + // 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 RootEpochInfo = { rootEpoch?: number; rootEpochHash?: Hex }; + type BootInfo = z.infer; type BootResponse = z.infer; @@ -108,9 +114,30 @@ const DSTACK_KMS_ABI = [ stateMutability: 'view', inputs: [], outputs: [{ name: '', type: 'address' }] + }, + { + name: 'currentRootEpoch', + type: 'function', + stateMutability: 'view', + inputs: [], + outputs: [ + { name: 'epoch', type: 'uint256' }, + { name: 'chainHash', type: 'bytes32' } + ] } ] as const; +// 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') { + return true; + } + } + return false; +}; + // ethereum backend class class EthereumBackend { private client: ReturnType; @@ -197,9 +224,30 @@ class EthereumBackend { isAllowed, reason, gatewayAppId: gatewayAppId as string, + ...await this.getRootEpoch(blockNumber), }; } + // 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 current = await this.client.readContract({ + address: this.kmsContractAddr, + abi: DSTACK_KMS_ABI, + functionName: 'currentRootEpoch', + blockNumber + }); + if (!current) return {}; + const [epoch, chainHash] = current; + return { rootEpoch: Number(epoch), rootEpochHash: chainHash }; + } catch (error) { + if (isMissingFunction(error)) return {}; + throw error; + } + } + async getGatewayAppId(): Promise { const result = await this.client.readContract({ address: this.kmsContractAddr, @@ -278,6 +326,7 @@ app.get('/', async (c) => { ethereum.getGatewayAppId(), ethereum.getChainId(), ethereum.getAppImplementation(), + ethereum.getRootEpoch(), ]); console.log('batch', batch); @@ -289,6 +338,7 @@ app.get('/', async (c) => { chainId: batch[1], appAuthImplementation: batch[2], // NOTE: for backward compatibility appImplementation: batch[2], + ...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..9f838cfa4 100644 --- a/dstack/kms/auth-eth-bun/openapi.json +++ b/dstack/kms/auth-eth-bun/openapi.json @@ -173,6 +173,14 @@ "gatewayAppId": { "type": "string", "description": "Gateway application ID" + }, + "rootEpoch": { + "type": "integer", + "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)" } } }, @@ -216,6 +224,14 @@ "appImplementation": { "type": "string", "description": "App implementation address" + }, + "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 f79325b89..06dd743f4 100644 --- a/dstack/kms/auth-eth/contracts/DstackKms.sol +++ b/dstack/kms/auth-eth/contracts/DstackKms.sol @@ -44,6 +44,27 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E // DstackApp implementation contract address for factory deployment address public appImplementation; + 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; + // 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 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 @@ -58,6 +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, 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. @@ -115,6 +138,60 @@ 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, with the + // 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, + bytes calldata attestation + ) + external + onlyOwner + returns (uint256 epoch) + { + require(caPubkey.length > 0 && k256Pubkey.length > 0, "Empty epoch public key"); + require(attestation.length > 0, "Empty epoch attestation"); + epoch = _rootEpochs.length; + 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); + } + + 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, bytes32 chainHash) + { + require(epoch < _rootEpochs.length, "Unknown root epoch"); + RootEpoch storage keys = _rootEpochs[epoch]; + return (keys.caPubkey, keys.k256Pubkey, keys.attestation, keys.chainHash); + } + // Function to set KMS quote function setKmsQuote(bytes memory quote) external onlyOwner { kmsInfo.quote = quote; @@ -306,7 +383,7 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E } // Add storage gap for upgradeable contracts - uint256[50] 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 110612b7e..17c7bbef1 100644 --- a/dstack/kms/auth-eth/script/Manage.s.sol +++ b/dstack/kms/auth-eth/script/Manage.s.sol @@ -109,6 +109,32 @@ 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); + } +} + +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 4a9a01496..7548c21cb 100644 --- a/dstack/kms/auth-eth/script/Query.s.sol +++ b/dstack/kms/auth-eth/script/Query.s.sol @@ -211,6 +211,22 @@ contract GetKmsSettings is BaseQueryScript { console.log(" Quote length:", quote.length); console.log(" Eventlog length:", eventlog.length); + 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, 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 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..472f1aca8 100644 --- a/dstack/kms/auth-eth/script/README.md +++ b/dstack/kms/auth-eth/script/README.md @@ -58,6 +58,25 @@ 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 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 +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 ebd655140..c9f8b8bcb 100644 --- a/dstack/kms/auth-eth/src/ethereum.ts +++ b/dstack/kms/auth-eth/src/ethereum.ts @@ -3,16 +3,22 @@ // SPDX-License-Identifier: Apache-2.0 import { ethers } from 'ethers'; -import { BootInfo, BootResponse } from './types'; +import { BootInfo, BootResponse, RootEpochInfo } 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 currentRootEpoch() view returns (uint256 epoch, bytes32 chainHash)" ]; +// 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; @@ -64,6 +70,19 @@ export class EthereumBackend { isAllowed, reason, gatewayAppId, + ...await this.getRootEpoch(blockTag), + } + } + + // 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 { + const [epoch, chainHash] = await this.kmsContract.currentRootEpoch({ blockTag }); + return { rootEpoch: Number(epoch), rootEpochHash: chainHash }; + } catch (error) { + if (isMissingFunction(error)) return {}; + throw error; } } diff --git a/dstack/kms/auth-eth/src/main.test.ts b/dstack/kms/auth-eth/src/main.test.ts index 842f1e36b..2d999a635 100644 --- a/dstack/kms/auth-eth/src/main.test.ts +++ b/dstack/kms/auth-eth/src/main.test.ts @@ -88,6 +88,41 @@ describe('Server', () => { }); }); + describe('key epochs', () => { + it('should report the latest registered key epoch', async () => { + app.ethereum.checkBoot = jest.fn().mockResolvedValue({ + isAllowed: true, + reason: '', + gatewayAppId: '', + rootEpoch: 1, + rootEpochHash: '0x' + '11'.repeat(32), + }); + + const response = await app.inject({ + method: 'POST', + url: '/bootAuth/app', + payload: mockBootInfo + }); + + 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 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'), + getRootEpoch: jest.fn().mockResolvedValue(rootEpoch), + }); + + const response = await app.inject({ method: 'GET', url: '/' }); + + expect(response.statusCode).toBe(200); + expect(JSON.parse(response.payload)).toMatchObject(rootEpoch); + }); + }); + 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..cc52e90f8 100644 --- a/dstack/kms/auth-eth/src/server.ts +++ b/dstack/kms/auth-eth/src/server.ts @@ -50,6 +50,8 @@ export async function build(): Promise { isAllowed: { type: 'boolean' }, reason: { type: 'string' }, gatewayAppId: { type: 'string' }, + rootEpoch: { type: 'integer' }, + rootEpochHash: { type: 'string' }, } }); @@ -84,6 +86,7 @@ export async function build(): Promise { server.ethereum.getGatewayAppId(), server.ethereum.getChainId(), server.ethereum.getAppImplementation(), + server.ethereum.getRootEpoch(), ]); return { status: 'ok', @@ -93,6 +96,7 @@ export async function build(): Promise { chainId: batch[1], appAuthImplementation: batch[2], // NOTE: for backward compatibility appImplementation: batch[2], + ...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..57ca87720 100644 --- a/dstack/kms/auth-eth/src/types.ts +++ b/dstack/kms/auth-eth/src/types.ts @@ -18,6 +18,15 @@ export interface BootResponse { isAllowed: boolean; gatewayAppId: string; reason: string; + rootEpoch?: number; + rootEpochHash?: 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 899b2b7b9..6c9910f92 100644 --- a/dstack/kms/auth-eth/test/DstackKms.t.sol +++ b/dstack/kms/auth-eth/test/DstackKms.t.sol @@ -19,6 +19,8 @@ contract DstackKmsTest is Test { address public user; event KmsInfoSet(bytes k256Pubkey); + event RootEpochRegistered(uint256 indexed epoch, bytes32 chainHash, bytes caPubkey, bytes k256Pubkey); + event RootEpochActivated(uint256 indexed epoch); event AppRegistered(address appId); @@ -66,6 +68,96 @@ contract DstackKmsTest is Test { assertEq(storedEventlog, eventlog); } + 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, 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, 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 { + vm.prank(user); + vm.expectRevert(abi.encodeWithSelector(OwnableUpgradeable.OwnableUnauthorizedAccount.selector, user)); + 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", hex"c0"); + vm.expectRevert("Empty epoch public key"); + kms.registerRootEpoch(hex"a0", "", hex"c0"); + vm.expectRevert("Empty epoch attestation"); + kms.registerRootEpoch(hex"a0", hex"b0", ""); + vm.stopPrank(); + } + function test_SetKmsInfoOnlyOwner() public { bytes memory k256Pubkey = hex"123456"; bytes memory caPubkey = hex"789abc"; diff --git a/dstack/kms/auth-mock/README.md b/dstack/kms/auth-mock/README.md index c26ed0e19..912d8f88a 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_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 493745613..b56c8c132 100644 --- a/dstack/kms/auth-mock/index.test.ts +++ b/dstack/kms/auth-mock/index.test.ts @@ -160,6 +160,30 @@ describe('Mock Backend Tests', () => { }); }); + describe('root epochs', () => { + 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', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + mrAggregated: '0x01', osImageHash: '0x02', appId: '0x03', + composeHash: '0x04', instanceId: '0x05', deviceId: '0x06', + }), + })); + expect(await boot.json()).toMatchObject(expected); + const info = await appFetch(new Request('http://localhost:3002/')); + expect(await info.json()).toMatchObject(expected); + } finally { + delete process.env.MOCK_ROOT_EPOCH; + delete process.env.MOCK_ROOT_EPOCH_HASH; + } + }); + }); + 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..b49f56994 100644 --- a/dstack/kms/auth-mock/index.ts +++ b/dstack/kms/auth-mock/index.ts @@ -28,9 +28,21 @@ const BootInfoSchema = z.object({ const BootResponseSchema = z.object({ isAllowed: z.boolean(), reason: z.string(), - gatewayAppId: z.string() + gatewayAppId: z.string(), + rootEpoch: z.number().optional(), + rootEpochHash: z.string().optional() }); +// 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; type BootResponse = z.infer; @@ -70,6 +82,11 @@ class MockBackend { } async checkBoot(bootInfo: BootInfo, isKms: boolean): Promise { + const response = this.evaluateBoot(bootInfo, isKms); + return { ...response, ...getRootEpoch() }; + } + + private evaluateBoot(bootInfo: BootInfo, isKms: boolean): BootResponse { const policy = getPolicy(); const deny = (reason: string): BootResponse => ({ isAllowed: false, @@ -147,6 +164,7 @@ app.get('/', async (c) => { chainId: batch[1], appAuthImplementation: batch[2], // NOTE: for backward compatibility appImplementation: batch[2], + ...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 1689308d0..e3e767409 100644 --- a/dstack/kms/auth-mock/openapi.json +++ b/dstack/kms/auth-mock/openapi.json @@ -176,6 +176,14 @@ "type": "string", "description": "Mock gateway application ID", "example": "0xmockgateway1234567890123456789012345678" + }, + "rootEpoch": { + "type": "integer", + "description": "Current KMS root epoch from MOCK_ROOT_EPOCH; absent when none" + }, + "rootEpochHash": { + "type": "string", + "description": "Its chain hash from MOCK_ROOT_EPOCH_HASH" } } }, @@ -221,6 +229,8 @@ "description": "Mock app implementation address", "example": "0xmockapp9876543210987654321098765432109" }, + "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 efd2efcb5..be387fcc0 100644 --- a/dstack/kms/auth-simple/README.md +++ b/dstack/kms/auth-simple/README.md @@ -107,6 +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. 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 fee441712..dc3572e68 100644 --- a/dstack/kms/auth-simple/index.test.ts +++ b/dstack/kms/auth-simple/index.test.ts @@ -446,6 +446,44 @@ describe('auth-simple', () => { }); }); + 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'; + + 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) + })); + const { rootEpoch, rootEpochHash } = await boot.json(); + return { rootEpoch, rootEpochHash }; + }; + const info = async () => + (await app.fetch(new Request('http://localhost/', { method: 'GET' }))).json(); + + 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 root epoch when none is configured', async () => { + writeTestConfig({ gatewayAppId: '0xgateway' }); + expect(await bootRootEpoch()).toEqual({ rootEpoch: undefined, rootEpochHash: undefined }); + expect((await info()).rootEpoch).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..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({ @@ -24,7 +26,14 @@ const BootInfoSchema = z.object({ const BootResponseSchema = z.object({ isAllowed: z.boolean(), reason: z.string(), - gatewayAppId: z.string() + gatewayAppId: z.string(), + rootEpoch: z.number().optional(), + rootEpochHash: z.string().optional() +}); + +const RootEpochSchema = z.object({ + caPubkey: z.string(), + k256Pubkey: z.string() }); // config schema @@ -52,13 +61,39 @@ 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 root epoch, indexed by epoch; append only. + 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(); @@ -93,6 +128,11 @@ class ConfigBackend { async checkBoot(bootInfo: BootInfo, isKms: boolean): Promise { const config = this.loadConfig(); + const response = this.evaluateBoot(bootInfo, isKms, config); + return { ...response, ...currentRootEpoch(config) }; + } + + private evaluateBoot(bootInfo: BootInfo, isKms: boolean, config: AuthConfig): BootResponse { const osImageHash = normalizeHex(bootInfo.osImageHash); const deviceId = normalizeHex(bootInfo.deviceId); @@ -220,6 +260,8 @@ class ConfigBackend { kmsContractAddr: string; chainId: number; appImplementation: string; + rootEpoch?: number; + rootEpochHash?: string; }> { const config = this.loadConfig(); return { @@ -227,6 +269,7 @@ class ConfigBackend { kmsContractAddr: config.kmsContractAddr, chainId: config.chainId, appImplementation: config.appImplementation, + ...currentRootEpoch(config), }; } } @@ -248,6 +291,8 @@ app.get('/', async (c) => { chainId: info.chainId, appAuthImplementation: info.appImplementation, // backward compat appImplementation: info.appImplementation, + 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/rpc/proto/kms_rpc.proto b/dstack/kms/rpc/proto/kms_rpc.proto index 6853df2be..5b5ba57aa 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; 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 { @@ -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 root_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..f90bdf9ce 100644 --- a/dstack/kms/src/config.rs +++ b/dstack/kms/src/config.rs @@ -21,6 +21,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)] @@ -171,6 +173,28 @@ impl KmsConfig { self.cert_dir.join(K256_KEY) } + /// 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) + } + + /// 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()) + }; + 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 +233,12 @@ 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, Deserialize)] pub(crate) struct OnboardConfig { pub enabled: bool, diff --git a/dstack/kms/src/main.rs b/dstack/kms/src/main.rs index f4bf428a8..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,7 +181,8 @@ async fn main() -> Result<()> { } else { None }; - let state = main_service::KmsState::new(config).context("Failed to initialize KMS state")?; + 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 @@ -213,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 { @@ -233,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 baac0147d..feeb04e19 100644 --- a/dstack/kms/src/main_service.rs +++ b/dstack/kms/src/main_service.rs @@ -26,15 +26,17 @@ 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; 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, RootEpochRegistry, }; use crate::{ - config::KmsConfig, + config::{KmsConfig, RootKeyPaths}, crypto::{derive_k256_key, sign_message, sign_message_with_timestamp}, }; @@ -67,10 +69,105 @@ 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 }) + } +} + +/// 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]>, +} + +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 chain_hashes = root_chain_hashes(roots.iter().map(|root| (&root.ca.key, &root.k256))); + Ok(Self { + roots, + chain_hashes, + }) + } + + pub(crate) fn latest(&self) -> u32 { + self.roots.len() as u32 - 1 + } + + pub(crate) fn servable_epoch(&self, registry: RootEpochRegistry) -> Result { + 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 + ), + } +} + pub struct KmsStateInner { config: KmsConfig, - root_ca: CaCert, - k256_key: SigningKey, + 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, @@ -127,12 +224,10 @@ impl KmsState { self.verifier.clear_cache(image_hash, config_hash) } - 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")?; + 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 temp_ca_key = fs::read_to_string(config.tmp_ca_key()).context("Faeild to read temp ca key")?; let temp_ca_cert = @@ -155,8 +250,9 @@ impl KmsState { Ok(Self { inner: Arc::new(KmsStateInner { config, - root_ca, - k256_key, + keyring, + serving_epoch, + restart: tokio::sync::Notify::new(), temp_ca_cert, temp_ca_key, verifier, @@ -179,6 +275,105 @@ impl KmsState { pub(crate) fn attestation_verifier(&self) -> Arc { self.inner.attestation_verifier.clone() } + + /// The served epoch's root keys. + fn root(&self) -> &RootKeys { + self.epoch_root(self.serving_epoch) + } + + fn epoch_root(&self, epoch: u32) -> &RootKeys { + &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}"), + } + } + + /// 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(()); + } + let serving = self.serving_epoch; + 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") + } +} + +/// 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; + } + 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; + } + } + } +} + +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"), + 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 +397,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.keyring.roots) + .map(|(epoch, root)| KmsKeys { + ca_key: root.ca.key.serialize_pem(), + k256_key: root.k256.to_bytes().to_vec(), + epoch, + }) + .collect(), }) } } @@ -303,7 +502,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:#}"))?; @@ -347,6 +546,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. @@ -369,7 +569,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 +581,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,36 +605,20 @@ impl KmsRpc for RpcHandler { .await .context("App not allowed")?; ensure_key_release_allowed(&boot_info, (&self.state.config).into())?; - 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) - .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 = 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")?; - (k256_app_key.to_bytes().to_vec(), signature) - }; + let root_epoch = self.state.resolve_epoch(request.root_epoch)?; + 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.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: root.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, }) } @@ -443,7 +628,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 +643,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 +652,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 +679,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 +728,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 +777,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(), ], }) } @@ -1333,7 +1518,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 79f4a7c23..6d9379151 100644 --- a/dstack/kms/src/main_service/upgrade_authority.rs +++ b/dstack/kms/src/main_service/upgrade_authority.rs @@ -84,6 +84,55 @@ pub(crate) struct BootResponse { pub is_allowed: bool, pub gateway_app_id: String, pub reason: String, + #[serde(flatten)] + pub root_epoch: RootEpochFields, +} + +/// 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 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], +} + +/// 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)] @@ -96,6 +145,8 @@ pub(crate) struct AuthApiInfoResponse { pub gateway_app_id: String, pub chain_id: u64, pub app_implementation: String, + #[serde(flatten)] + pub root_epoch: RootEpochFields, } #[derive(Debug, Clone, Serialize, Deserialize)] @@ -107,6 +158,7 @@ pub(crate) struct GetInfoResponse { pub eth_rpc_url: Option, pub chain_id: Option, pub app_implementation: Option, + pub root_epochs: RootEpochRegistry, } const AUTH_API_CONNECT_TIMEOUT: Duration = Duration::from_secs(5); @@ -174,6 +226,7 @@ impl AuthApi { is_allowed: true, reason: "".to_string(), gateway_app_id: dev.gateway_app_id.clone(), + root_epoch: RootEpochFields::default(), }), AuthApi::Webhook { webhook } => { let path = if is_kms { @@ -196,6 +249,7 @@ impl AuthApi { gateway_app_id: Some(dev.gateway_app_id.clone()), chain_id: None, app_implementation: None, + root_epochs: RootEpochRegistry::Unchecked, }), AuthApi::Webhook { webhook } => { let info: AuthApiInfoResponse = http_get(&webhook.url, webhook.timeout).await?; @@ -211,6 +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: info.root_epoch.registry()?, }) } } @@ -264,10 +319,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")?; @@ -277,6 +335,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 @@ -295,7 +356,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 +381,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 +425,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 +445,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 +573,36 @@ mod tests { assert_ne!(requests[0].1["mrAggregated"], requests[1].1["mrAggregated"]); } + #[rocket::async_test] + 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","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], + }); + + assert_eq!(auth.get_info().await.unwrap().root_epochs, current); + let info = boot_info(50); + 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(); + } + #[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 c16b67d4d..311057e6d 100644 --- a/dstack/kms/src/onboard_service.rs +++ b/dstack/kms/src/onboard_service.rs @@ -36,9 +36,10 @@ 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, }, }, }; @@ -146,16 +147,7 @@ impl OnboardRpc for OnboardHandler { 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,8 @@ mod tests { use crate::main_service::amd_attest::{ compute_expected_measurement, MeasurementInput, OvmfSectionParam, }; + use crate::main_service::upgrade_authority::{BootResponse, CurrentRootEpoch, RootEpochFields}; + use dstack_kms_rpc::KmsKeys; use sha2::Digest; fn hex_of(byte: u8, len: usize) -> String { @@ -536,6 +546,235 @@ 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(), + } + } + + /// 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 = 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 = 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, 2).k256_key).unwrap(), + k256 + ); + } + assert_eq!(&fs::read(cfg.k256_key()).unwrap(), &sent[2]); + } + + #[rocket::async_test] + async fn rotation_imports_the_keyring_and_generates_the_next_epoch() { + let cert_dir = tempfile::tempdir().unwrap(); + 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 = 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, 2).k256_key).unwrap(), + sent[1] + ); + let current = fs::read(cfg.k256_key()).unwrap(); + assert!(!sent.contains(¤t)); + } + + #[rocket::async_test] + async fn onboarding_checks_the_handover_against_the_current_epoch() { + let cert_dir = tempfile::tempdir().unwrap(); + 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!( + onboard(&cfg, handover(0), false, RootEpochRegistry::Unchecked) + .await + .is_err() + ); + } + + #[rocket::async_test] + async fn an_unregistered_deployment_can_only_join_at_epoch_0() { + let cert_dir = tempfile::tempdir().unwrap(); + 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()); + } + + /// 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 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 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] async fn setup_status_includes_bootstrap_info_when_present() { let cert_dir = tempfile::tempdir().unwrap(); @@ -607,7 +846,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 +895,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 +924,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 +955,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 +973,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::)); @@ -753,10 +1022,10 @@ 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_key_response(cfg, keys_res, domain).await; + return Self::from_registered_key_response(cfg, keys_res, domain, rotate).await; } let client = RaClientConfig::builder() @@ -779,7 +1048,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")?; @@ -789,43 +1058,103 @@ impl Keys { vm_config: info.vm_config, }) .await?; - Self::from_key_response(cfg, keys_res, domain).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 registry = cfg + .auth_api + .get_info() + .await + .context("failed to read the current root epoch")? + .root_epochs; + Self::from_key_response(cfg, keys_res, domain, rotate, registry).await + } + + /// 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, + registry: RootEpochRegistry, ) -> Result { - if keys_res.keys.len() != 1 { - return Err(anyhow::anyhow!("Invalid keys")); + let received: Vec = keys_res.keys.iter().map(|k| k.epoch).collect(); + 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 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::>>()?; + 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)?, + 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, self.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(()) } @@ -846,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() { @@ -876,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", )?; @@ -945,10 +1277,48 @@ pub(crate) async fn bootstrap_keys(cfg: &KmsConfig, verifier: &AttestationVerifi Ok(()) } -async fn attest_keys(p256_pubkey: &[u8], k256_pubkey: &[u8]) -> Result> { +/// 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. +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; + if rotate && current != latest { + bail!( + "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" + ); + } + 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); 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/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/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" diff --git a/sdk/go/dstack/client_v1.go b/sdk/go/dstack/client_v1.go index fd0f98ace..5b27c29da 100644 --- a/sdk/go/dstack/client_v1.go +++ b/sdk/go/dstack/client_v1.go @@ -42,6 +42,8 @@ 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. + RootEpoch uint32 } // Response from a v1 attestation request. @@ -442,6 +444,18 @@ 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.getKey(ctx, domain, algorithm, nil) +} + +// 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`") } @@ -450,6 +464,9 @@ func (c *DstackClientV1) GetKey(ctx context.Context, domain string, algorithm st "domain": domain, "algorithm": algorithm, } + if rootEpoch != nil { + payload["root_epoch"] = *rootEpoch + } data, err := c.sendRPCRequest(ctx, "/v1/GetKey", payload) if err != nil { @@ -464,6 +481,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 +507,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..fbf3156f5 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. */ + root_epoch: number } export interface AttestResponseV1 { @@ -333,14 +335,19 @@ 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?: 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"') } - const payload = JSON.stringify({ domain, algorithm }) - const result = await send_rpc_request<{ key: string, public_key: string, signature_chain: string[] }>( + // 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( + 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) return Object.freeze({ @@ -348,6 +355,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..77c1438d0 100644 --- a/sdk/python/src/dstack_sdk/dstack_client_v1.py +++ b/sdk/python/src/dstack_sdk/dstack_client_v1.py @@ -107,6 +107,8 @@ 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. + root_epoch: int = 0 class GpuEvidenceBundleV1(BaseModel): @@ -237,7 +239,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: Optional[int] = None + ) -> GetKeyResponseV1: """Derive an application key from ``(domain, algorithm)``. Both arguments are required. Derivation is flat: two domains yield @@ -245,10 +249,19 @@ 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; 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`") 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) @@ -350,7 +363,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: 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 efcc78236..71f9c5f33 100644 --- a/sdk/rust/src/dstack_client_v1.rs +++ b/sdk/rust/src/dstack_client_v1.rs @@ -95,10 +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.request_key(domain, algorithm, None).await + } + + /// [`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 }); + 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 917db972e..3812b2e4f 100644 --- a/sdk/rust/types/src/dstack_v1.rs +++ b/sdk/rust/types/src/dstack_v1.rs @@ -125,6 +125,9 @@ 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. + #[serde(default)] + pub root_epoch: u32, } /// The value an attestation binds to: bytes, and only bytes.