backport(v0.5.x): migrate the data disk and env across KMS key epochs

Cargo.toml

@@ -128,6 +128,7 @@ scale = { version = "3.7.4", package = "parity-scale-codec", features = [
 ] }
 serde = { version = "1.0.228", features = ["derive"], default-features = false }
 serde-human-bytes = "0.1.2"
+serde_with = "3.14.0"
 rmp-serde = "1.3.1"
 serde_json = { version = "1.0.140", default-features = false }
 serde_ini = "0.2.0"
@@ -195,6 +196,7 @@ sha3 = "0.10.8"
 blake2 = "0.10.6"
 tokio-rustls = { version = "0.26.2", features = ["ring"] }
 x25519-dalek = { version = "2.0.1", features = ["static_secrets"] }
+hpke = { version = "0.15.1", default-features = false, features = ["alloc", "getrandom", "x25519", "aes"] }
 k256 = "0.13.4"
 ed25519-dalek = { version = "2.2.0", features = ["rand_core"] }
 # Additional RustCrypto dependencies for sealed box
@@ -244,7 +246,6 @@ smallvec = "1.14.0"
 cmd_lib = "1.9.5"
 yaml-rust2 = "0.10.4"
 
-luks2 = "0.5.0"
 scopeguard = "1.2.0"
 tar = "0.4"
 proxy-protocol = "0.5.0"

dstack-types/src/lib.rs

@@ -234,6 +234,8 @@ pub struct AppKeys {
     pub disk_crypt_key: Vec<u8>,
     #[serde(with = "hex_bytes", default)]
     pub env_crypt_key: Vec<u8>,
+    #[serde(with = "hex_bytes", default)]
+    pub sealed_env_key: Vec<u8>,
     #[serde(with = "hex_bytes")]
     pub k256_key: Vec<u8>,
     #[serde(with = "hex_bytes")]
@@ -264,9 +266,30 @@ 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,
+        #[serde(default)]
+        app_epoch: u32,
     },
 }
 
+impl AppKeys {
+    pub fn app_epoch(&self) -> u32 {
+        match self.key_provider {
+            KeyProvider::Kms { app_epoch, .. } => app_epoch,
+            _ => 0,
+        }
+    }
+
+    pub fn root_epoch(&self) -> u32 {
+        match self.key_provider {
+            KeyProvider::Kms { root_epoch, .. } => root_epoch,
+            _ => 0,
+        }
+    }
+}
+
 impl KeyProvider {
     pub fn kind(&self) -> KeyProviderKind {
         match self {

dstack-types/src/shared_filenames.rs

@@ -7,6 +7,7 @@ pub const APP_KEYS: &str = ".appkeys.json";
 pub const SYS_CONFIG: &str = ".sys-config.json";
 pub const USER_CONFIG: &str = ".user-config";
 pub const ENCRYPTED_ENV: &str = ".encrypted-env";
+pub const SEALED_ENV: &str = ".sealed-env";
 pub const DECRYPTED_ENV: &str = ".decrypted-env";
 pub const DECRYPTED_ENV_JSON: &str = ".decrypted-env.json";
 pub const INSTANCE_INFO: &str = ".instance_info";

dstack-util/Cargo.toml

@@ -23,6 +23,8 @@ scale = { workspace = true, features = ["derive"] }
 schnorrkel.workspace = true
 serde.workspace = true
 serde-human-bytes.workspace = true
+serde_with.workspace = true
+hpke.workspace = true
 serde_json.workspace = true
 sha2.workspace = true
 tokio = { workspace = true, features = ["full"] }
@@ -50,7 +52,6 @@ yaml-rust2.workspace = true
 bollard.workspace = true
 sodiumbox.workspace = true
 libc.workspace = true
-luks2.workspace = true
 scopeguard.workspace = true
 tempfile.workspace = true
 ez-hash.workspace = true

dstack-util/src/crypto.rs

@@ -41,6 +41,71 @@ 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.
+#[cfg(test)]
+pub fn dh_encrypt(remote_public_key: [u8; 32], plaintext: &[u8]) -> Result<Vec<u8>> {
+    let ephemeral_secret = StaticSecret::random_from_rng(rand::thread_rng());
+    let shared_secret = ephemeral_secret
+        .diffie_hellman(&PublicKey::from(remote_public_key))
+        .to_bytes();
+    let cipher = Aes256Gcm::new_from_slice(&shared_secret)
+        .map_err(|e| anyhow!("failed to create cipher: {e}"))?;
+    let iv: [u8; 12] = rand::random();
+    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())
+}
+
+/// Encrypt a sealed environment with HPKE DHKEM(X25519, HKDF-SHA256) and AES-256-GCM.
+pub fn sealed_env_encrypt(public_key: &[u8], info: &[u8], plaintext: &[u8]) -> Result<Vec<u8>> {
+    use hpke::{
+        aead::AesGcm256, kdf::HkdfSha256, kem::X25519HkdfSha256, Deserializable, OpModeS,
+        Serializable,
+    };
+    let public_key = <X25519HkdfSha256 as hpke::Kem>::PublicKey::from_bytes(public_key)
+        .map_err(|_| anyhow!("invalid sealed env public key"))?;
+    let (enc, ciphertext) = hpke::single_shot_seal::<AesGcm256, HkdfSha256, X25519HkdfSha256>(
+        &OpModeS::Base,
+        &public_key,
+        info,
+        plaintext,
+        &[],
+    )
+    .map_err(|e| anyhow!("sealed env encryption failed: {e}"))?;
+    Ok([enc.to_bytes().as_ref(), ciphertext.as_slice()].concat())
+}
+
+/// Decrypt a sealed environment HPKE payload.
+pub fn sealed_env_decrypt(secret: &[u8], info: &[u8], payload: &[u8]) -> Result<Vec<u8>> {
+    use hpke::{aead::AesGcm256, kdf::HkdfSha256, kem::X25519HkdfSha256, Deserializable, OpModeR};
+    let enc = payload
+        .get(..32)
+        .ok_or_else(|| anyhow!("sealed env payload is missing encapsulated key"))?;
+    let ciphertext = payload
+        .get(32..)
+        .ok_or_else(|| anyhow!("sealed env payload is missing ciphertext"))?;
+    let secret = <X25519HkdfSha256 as hpke::Kem>::PrivateKey::from_bytes(secret)
+        .map_err(|_| anyhow!("invalid sealed env private key"))?;
+    let enc = <X25519HkdfSha256 as hpke::Kem>::EncappedKey::from_bytes(enc)
+        .map_err(|_| anyhow!("invalid sealed env encapsulated key"))?;
+    hpke::single_shot_open::<AesGcm256, HkdfSha256, X25519HkdfSha256>(
+        &OpModeR::Base,
+        &secret,
+        &enc,
+        info,
+        ciphertext,
+        &[],
+    )
+    .map_err(|e| anyhow!("sealed env decryption failed: {e}"))
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -86,4 +151,57 @@ mod tests {
         let decrypted_str = String::from_utf8(decrypted).unwrap();
         assert_eq!(decrypted_str, "[{\"key\":\"\",\"value\":\"\"}]");
     }
+
+    #[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 opens_sdk_sealed_envs() {
+        use base64::Engine as _;
+        for source in [
+            include_str!("../tests/fixtures/sealed-env-python.json"),
+            include_str!("../tests/fixtures/sealed-env-javascript.json"),
+        ] {
+            let fixture: serde_json::Value = serde_json::from_str(source).unwrap();
+            let envelope = &fixture["sealed_env"];
+            let decode = |name: &str| {
+                base64::engine::general_purpose::STANDARD
+                    .decode(envelope[name].as_str().unwrap())
+                    .unwrap()
+            };
+            let payload = [decode("enc"), decode("ciphertext")].concat();
+            let plaintext = sealed_env_decrypt(
+                &hex::decode(fixture["private_key"].as_str().unwrap()).unwrap(),
+                b"dstack-sealed-env-v1:3:0",
+                &payload,
+            )
+            .unwrap();
+            let decoded: serde_json::Value = serde_json::from_slice(&plaintext).unwrap();
+            assert_eq!(decoded["env"], fixture["env"]);
+        }
+    }
+
+    #[test]
+    fn sealed_env_round_trips_and_binds_info() {
+        let secret = StaticSecret::random_from_rng(rand::thread_rng());
+        let public = PublicKey::from(&secret).to_bytes();
+        let info = b"dstack-sealed-env-v1:3:0";
+        let payload = sealed_env_encrypt(&public, info, b"secret env").unwrap();
+        assert_eq!(
+            sealed_env_decrypt(&secret.to_bytes(), info, &payload).unwrap(),
+            b"secret env"
+        );
+        assert!(
+            sealed_env_decrypt(&secret.to_bytes(), b"dstack-sealed-env-v1:4:0", &payload).is_err()
+        );
+        let mut tampered = payload;
+        *tampered.last_mut().unwrap() ^= 1;
+        assert!(sealed_env_decrypt(&secret.to_bytes(), info, &tampered).is_err());
+    }
 }

dstack-util/src/host_api.rs

@@ -9,13 +9,15 @@ use dstack_types::{
     Platform,
 };
 use host_api::{
-    client::{new_client, DefaultClient},
+    client::{new_client, DefaultClient, StatusError},
     Notification,
 };
 use ra_tls::attestation::validate_tcb;
 use sodiumbox::{generate_keypair, open_sealed_box, PUBLICKEYBYTES};
 use tracing::warn;
 
+const HOST_API_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30);
+
 pub(crate) struct KeyProvision {
     pub sk: [u8; 32],
     pub mr: [u8; 32],
@@ -78,6 +80,29 @@ impl HostApi {
         }
     }
 
+    /// Close the host's env update window for this boot, handing it the env
+    /// re-encrypted to the current key epochs, or an empty one when nothing
+    /// changed. An error means the window may still be open.
+    pub async fn update_sealed_env(&self, sealed_env: Vec<u8>) -> Result<()> {
+        if !matches!(Platform::detect_or_dstack(), Platform::Dstack) {
+            return Ok(());
+        }
+        let Some(client) = &self.client else {
+            return Ok(());
+        };
+        let response = tokio::time::timeout(
+            HOST_API_TIMEOUT,
+            client.update_sealed_env(host_api::SealedEnv { sealed_env }),
+        )
+        .await
+        .context("Timed out updating the sealed env on the host")?;
+        match response {
+            // A VMM predating UpdateSealedEnv has no window to close.
+            Err(err) if is_unknown_method(&err) => Ok(()),
+            response => response.map(drop),
+        }
+    }
+
     pub async fn get_sealing_key(&self) -> Result<KeyProvision> {
         let (pk, sk) = generate_keypair();
         let mut report_data = [0u8; 64];
@@ -122,3 +147,80 @@ impl HostApi {
         Ok(KeyProvision { sk, mr })
     }
 }
+
+/// Whether the host does not implement the method at all: a VMM since 0.6
+/// answers 404, and older ones 400 with the pRPC dispatcher's fixed
+/// "Service not found" message.
+fn is_unknown_method(err: &anyhow::Error) -> bool {
+    err.chain()
+        .find_map(|cause| cause.downcast_ref::<StatusError>())
+        .is_some_and(|err| {
+            err.status == 404
+                || (err.status == 400
+                    && err
+                        .message
+                        .as_deref()
+                        .is_some_and(|message| message.starts_with("Service not found: ")))
+        })
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::io::{Read, Write};
+
+    /// A host API answering one request with `status` and a JSON `body`.
+    fn host(status: u16, body: String) -> HostApi {
+        let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
+        let url = format!("http://{}", listener.local_addr().unwrap());
+        std::thread::spawn(move || {
+            let (mut stream, _) = listener.accept().unwrap();
+            let mut request = Vec::new();
+            let mut chunk = [0; 4096];
+            while !request.windows(4).any(|w| w == b"\r\n\r\n") {
+                let read = stream.read(&mut chunk).unwrap();
+                request.extend_from_slice(&chunk[..read]);
+            }
+            write!(
+                stream,
+                "HTTP/1.1 {status} X\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}",
+                body.len()
+            )
+            .unwrap();
+        });
+        HostApi::new(Some(url), None)
+    }
+
+    #[tokio::test]
+    async fn only_an_answer_or_a_host_without_a_window_closes_the_env_window() {
+        assert!(matches!(Platform::detect_or_dstack(), Platform::Dstack));
+        let unknown = r#"{"error":"Service not found: UpdateSealedEnv"}"#;
+        for (status, body, closed) in [
+            (200, "", true),
+            // A VMM since 0.6, and one before, that does not know the method.
+            (404, unknown, true),
+            (400, unknown, true),
+            // Anything else may leave the window open.
+            (400, r#"{"error":"VM not found"}"#, false),
+            (500, "{}", false),
+        ] {
+            let result = host(status, body.into())
+                .update_sealed_env(b"env".to_vec())
+                .await;
+            assert_eq!(result.is_ok(), closed, "{status} {body}");
+        }
+        // No host listening at all.
+        let gone = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
+        let url = format!("http://{}", gone.local_addr().unwrap());
+        drop(gone);
+        assert!(HostApi::new(Some(url), None)
+            .update_sealed_env(Vec::new())
+            .await
+            .is_err());
+        // No host API configured means no window.
+        HostApi::new(None, None)
+            .update_sealed_env(Vec::new())
+            .await
+            .unwrap();
+    }
+}

dstack-util/src/main.rs

@@ -431,6 +431,7 @@ fn make_app_keys(
     Ok(AppKeys {
         disk_crypt_key: sha256(&disk_key.serialize_der()).to_vec(),
         env_crypt_key: vec![],
+        sealed_env_key: vec![],
         k256_key: k256_key.to_bytes().to_vec(),
         k256_signature: vec![],
         gateway_app_id: "".to_string(),