feat(guest): migrate encrypted environments across KMS root epochs

docs/deployment.md

@@ -513,6 +513,40 @@ curl http://<new-kms>: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

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<AppKeyResponse> {
+        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<String> {
         match self {
             CertRequestClient::Local { ca } => Ok(ca.pem_cert.clone()),

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

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<CurrentRootEpoch>,
+}
+
+/// 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<String, AppRules>,
+    /// root public keys of each KMS root epoch, indexed by epoch; append only.
+    root_epochs: Vec<RootEpoch>,
+    /// 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<CurrentRootEpoch> {
+        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::<Option<Vec<_>>>()?;
+        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<BootInfo>, cli: &State<Cli>) -> Json<BootResponse> {
 
 /// info endpoint the KMS GETs to populate its metadata. Single-node: no chain.
 #[get("/")]
-fn info() -> Json<serde_json::Value> {
-    Json(json!({
+fn info(cli: &State<Cli>) -> Json<serde_json::Value> {
+    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 "" —

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<u8>,
         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;

dstack/dstack-types/src/root_epoch.rs

@@ -0,0 +1,49 @@
+// SPDX-FileCopyrightText: © 2026 Phala Network <dstack@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<Item = (&'a [u8], &'a [u8])>) -> 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"
+        );
+    }
+}

dstack/dstack-util/src/crypto.rs

@@ -68,6 +68,34 @@ pub fn dh_decrypt(secret: [u8; 32], ciphertext: &[u8]) -> Result<Vec<u8>> {
         .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<Vec<u8>> {
+    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::<Aes256Gcm>::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;