From 2947dcb78fca70b11065d1bab82542a2ffb5e435 Mon Sep 17 00:00:00 2001
From: Kevin Wang <wy721@qq.com>
Date: Thu, 13 Aug 2026 19:25:18 -0700
Subject: [PATCH 1/9] feat(dstack-util): add decrypt command

---
 dstack/dstack-util/src/main.rs | 70 ++++++++++++++++++++++++++++++++++
 1 file changed, 70 insertions(+)

diff --git a/dstack/dstack-util/src/main.rs b/dstack/dstack-util/src/main.rs
index b094d27c9..768f5f517 100644
--- a/dstack/dstack-util/src/main.rs
+++ b/dstack/dstack-util/src/main.rs
@@ -97,6 +97,8 @@ enum Commands {
     AttestStrip(AttestStripArgs),
     /// Get app keys from a KMS server
     GetKeys(GetKeysArgs),
+    /// Decrypt data encrypted with the app's environment encryption public key
+    Decrypt(DecryptArgs),
 }
 
 #[derive(Parser)]
@@ -366,6 +368,26 @@ struct GetKeysArgs {
     root_ca: Option<PathBuf>,
 }
 
+#[derive(Parser)]
+/// Decrypt data encrypted with the app's environment encryption public key
+struct DecryptArgs {
+    /// Input file (default: stdin)
+    #[arg(short, long)]
+    input: Option<PathBuf>,
+
+    /// Output file (default: stdout)
+    #[arg(short, long)]
+    output: Option<PathBuf>,
+
+    /// App keys file containing env_crypt_key
+    #[arg(long)]
+    key_file: Option<PathBuf>,
+
+    /// Decode the input as hexadecimal text before decrypting
+    #[arg(long)]
+    hex: bool,
+}
+
 fn pad64(data: &[u8]) -> Result<[u8; 64]> {
     if data.len() > 64 {
         anyhow::bail!("report_data must be at most 64 bytes");
@@ -662,6 +684,51 @@ async fn cmd_get_keys(args: GetKeysArgs) -> Result<()> {
     Ok(())
 }
 
+fn cmd_decrypt(args: DecryptArgs) -> Result<()> {
+    use dstack_types::shared_filenames::{host_shared_dir, APP_KEYS};
+
+    let key_file = args
+        .key_file
+        .unwrap_or_else(|| host_shared_dir().join(APP_KEYS));
+    let keys: AppKeys = utils::deserialize_json_file(&key_file)
+        .with_context(|| format!("Failed to load app keys from {}", key_file.display()))?;
+    let env_crypt_key: [u8; 32] = keys
+        .env_crypt_key
+        .try_into()
+        .map_err(|key: Vec<u8>| anyhow::anyhow!("Invalid env crypt key length: {}", key.len()))?;
+
+    let mut input = match args.input {
+        Some(path) => fs::read(&path)
+            .with_context(|| format!("Failed to read ciphertext from {}", path.display()))?,
+        None => {
+            let mut input = Vec::new();
+            io::stdin()
+                .read_to_end(&mut input)
+                .context("Failed to read ciphertext from stdin")?;
+            input
+        }
+    };
+    if args.hex {
+        input = hex_decode(
+            std::str::from_utf8(&input)
+                .context("Hex ciphertext is not valid UTF-8")?
+                .trim(),
+        )
+        .context("Failed to decode hex ciphertext")?;
+    }
+
+    let plaintext = crypto::dh_decrypt(env_crypt_key, &input).context("Failed to decrypt input")?;
+    if let Some(output) = args.output {
+        safe_write_with_mode(&output, &plaintext, 0o600)
+            .with_context(|| format!("Failed to write plaintext to {}", output.display()))?;
+    } else {
+        io::stdout()
+            .write_all(&plaintext)
+            .context("Failed to write plaintext to stdout")?;
+    }
+    Ok(())
+}
+
 fn cmd_quote() -> Result<()> {
     let mut input = Vec::with_capacity(65);
     io::stdin()
@@ -1374,6 +1441,9 @@ async fn main() -> Result<()> {
         Commands::GetKeys(args) => {
             cmd_get_keys(args).await?;
         }
+        Commands::Decrypt(args) => {
+            cmd_decrypt(args)?;
+        }
     }
 
     Ok(())

From ba017ad254a6b9ef086bfd32e4e834d06fb4c62f Mon Sep 17 00:00:00 2001
From: Kevin Wang <wy721@qq.com>
Date: Thu, 13 Aug 2026 19:41:04 -0700
Subject: [PATCH 2/9] feat(dstack-util): stream encrypted data in chunks

---
 docs/stream-encryption.md        |  65 ++++++++
 dstack/dstack-util/src/crypto.rs | 245 ++++++++++++++++++++++++++++++-
 dstack/dstack-util/src/main.rs   | 171 ++++++++++++++++++---
 3 files changed, 461 insertions(+), 20 deletions(-)
 create mode 100644 docs/stream-encryption.md

diff --git a/docs/stream-encryption.md b/docs/stream-encryption.md
new file mode 100644
index 000000000..c683969e7
--- /dev/null
+++ b/docs/stream-encryption.md
@@ -0,0 +1,65 @@
+# dstack Chunked Encryption Format
+
+`dstack-util encrypt` and `dstack-util decrypt` use a chunked format for
+bounded-memory encryption of arbitrary data. It uses the same app-scoped X25519
+key pair as encrypted environment variables, but it is a separate wire format.
+
+All integers are unsigned and big-endian.
+
+## Header
+
+| Field | Size | Description |
+|---|---:|---|
+| Magic | 16 bytes | ASCII `dstack-stream-v1` |
+| Ephemeral public key | 32 bytes | X25519 public key generated by the sender |
+| Nonce prefix | 8 bytes | Random prefix shared by all chunks |
+| Chunk size | 4 bytes | Maximum plaintext bytes in each chunk |
+
+The X25519 shared secret is used directly as the AES-256-GCM key, matching the
+encrypted environment variable protocol.
+
+## Frames
+
+Each frame contains:
+
+| Field | Size | Description |
+|---|---:|---|
+| Flags | 1 byte | Bit 0 marks the final chunk; all other bits must be zero |
+| Plaintext length | 4 bytes | Number of plaintext bytes in this chunk |
+| Ciphertext and tag | `plaintext length + 16` bytes | AES-256-GCM output |
+
+The 12-byte nonce is `nonce_prefix || chunk_index`, where `chunk_index` is a
+4-byte integer starting at zero. The authenticated additional data is:
+
+```text
+header || chunk_index || flags || plaintext_length
+```
+
+Every non-final frame must contain exactly `chunk_size` plaintext bytes. The
+final frame may be shorter or empty. A final frame is always emitted, including
+for empty input and for input whose length is an exact multiple of the chunk
+size. Missing final frames, trailing data, unknown flags, and authentication
+failures are rejected.
+
+## CLI
+
+Encrypt data after retrieving the app public key over verified TLS:
+
+```bash
+dstack-util encrypt \
+  --kms-url https://kms.example.com \
+  --app-id "$APP_ID" \
+  --input plaintext.bin \
+  --output ciphertext.bin
+```
+
+For a KMS using a private CA, pass `--root-ca ca.pem`. Decrypt inside the CVM:
+
+```bash
+dstack-util decrypt --input ciphertext.bin --output plaintext.bin
+```
+
+`decrypt` detects the magic string automatically. Inputs without the magic are
+handled as the legacy encrypted-environment format. Hex input remains available
+through `--hex`, but it is decoded in memory and should not be used for large
+files.
diff --git a/dstack/dstack-util/src/crypto.rs b/dstack/dstack-util/src/crypto.rs
index 047c9c1a1..dfe1593da 100644
--- a/dstack/dstack-util/src/crypto.rs
+++ b/dstack/dstack-util/src/crypto.rs
@@ -3,12 +3,19 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use aes_gcm::{
-    aead::{Aead, Nonce},
+    aead::{Aead, Nonce, Payload},
     Aes256Gcm, KeyInit,
 };
-use anyhow::{anyhow, Result};
+use anyhow::{anyhow, ensure, Context, Result};
+use std::io::{Read, Write};
 use x25519_dalek::{PublicKey, StaticSecret};
 
+pub const STREAM_MAGIC: &[u8; 16] = b"dstack-stream-v1";
+pub const DEFAULT_CHUNK_SIZE: usize = 1024 * 1024;
+pub const MAX_CHUNK_SIZE: usize = 16 * 1024 * 1024;
+const FINAL_CHUNK: u8 = 1;
+const HEADER_LEN: usize = STREAM_MAGIC.len() + 32 + 8 + 4;
+
 pub fn dh_agree(secret: [u8; 32], their_pubkey: [u8; 32]) -> [u8; 32] {
     let secret = StaticSecret::from(secret);
     let their_public = PublicKey::from(their_pubkey);
@@ -44,6 +51,196 @@ pub fn dh_decrypt(secret: [u8; 32], ciphertext: &[u8]) -> Result<Vec<u8>> {
         .map_err(|e| anyhow!("Decryption failed: {}", e))
 }
 
+fn stream_nonce(prefix: &[u8; 8], index: u32) -> [u8; 12] {
+    let mut nonce = [0u8; 12];
+    nonce[..8].copy_from_slice(prefix);
+    nonce[8..].copy_from_slice(&index.to_be_bytes());
+    nonce
+}
+
+fn stream_aad(header: &[u8; HEADER_LEN], index: u32, frame_header: &[u8; 5]) -> Vec<u8> {
+    let mut aad = Vec::with_capacity(HEADER_LEN + 4 + frame_header.len());
+    aad.extend_from_slice(header);
+    aad.extend_from_slice(&index.to_be_bytes());
+    aad.extend_from_slice(frame_header);
+    aad
+}
+
+/// Encrypts a reader as independently authenticated chunks.
+pub fn dh_encrypt_stream(
+    remote_public_key: [u8; 32],
+    mut input: impl Read,
+    mut output: impl Write,
+    chunk_size: usize,
+) -> Result<()> {
+    ensure!(
+        (1..=MAX_CHUNK_SIZE).contains(&chunk_size),
+        "Chunk size must be between 1 and {MAX_CHUNK_SIZE} bytes"
+    );
+
+    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 ephemeral_public_key = PublicKey::from(&ephemeral_secret).to_bytes();
+    let remote_public_key = PublicKey::from(remote_public_key);
+    let shared_secret = ephemeral_secret
+        .diffie_hellman(&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 nonce_prefix = [0u8; 8];
+    getrandom::fill(&mut nonce_prefix).context("Failed to generate nonce prefix")?;
+    let mut header = [0u8; HEADER_LEN];
+    header[..STREAM_MAGIC.len()].copy_from_slice(STREAM_MAGIC);
+    header[STREAM_MAGIC.len()..STREAM_MAGIC.len() + 32].copy_from_slice(&ephemeral_public_key);
+    header[STREAM_MAGIC.len() + 32..STREAM_MAGIC.len() + 40].copy_from_slice(&nonce_prefix);
+    header[STREAM_MAGIC.len() + 40..].copy_from_slice(&(chunk_size as u32).to_be_bytes());
+    output
+        .write_all(&header)
+        .context("Failed to write header")?;
+
+    let mut current = vec![0u8; chunk_size];
+    let mut next = vec![0u8; chunk_size];
+    let mut current_len = read_chunk(&mut input, &mut current)?;
+    let mut index = 0u32;
+    loop {
+        let next_len = read_chunk(&mut input, &mut next)?;
+        let final_chunk = next_len == 0;
+        let flags = if final_chunk { FINAL_CHUNK } else { 0 };
+        let mut frame_header = [0u8; 5];
+        frame_header[0] = flags;
+        frame_header[1..].copy_from_slice(&(current_len as u32).to_be_bytes());
+        let aad = stream_aad(&header, index, &frame_header);
+        let nonce = stream_nonce(&nonce_prefix, index);
+        let encrypted = cipher
+            .encrypt(
+                (&nonce).into(),
+                Payload {
+                    msg: &current[..current_len],
+                    aad: &aad,
+                },
+            )
+            .map_err(|e| anyhow!("Failed to encrypt chunk {index}: {e}"))?;
+        output
+            .write_all(&frame_header)
+            .and_then(|_| output.write_all(&encrypted))
+            .with_context(|| format!("Failed to write chunk {index}"))?;
+        if final_chunk {
+            break;
+        }
+        index = index.checked_add(1).context("Too many chunks")?;
+        std::mem::swap(&mut current, &mut next);
+        current_len = next_len;
+    }
+    output.flush().context("Failed to flush encrypted output")?;
+    Ok(())
+}
+
+fn read_chunk(input: &mut impl Read, buffer: &mut [u8]) -> Result<usize> {
+    let mut read = 0;
+    while read < buffer.len() {
+        match input
+            .read(&mut buffer[read..])
+            .context("Failed to read input")?
+        {
+            0 => break,
+            n => read += n,
+        }
+    }
+    Ok(read)
+}
+
+/// Decrypts a chunked stream after the caller has consumed [`STREAM_MAGIC`].
+pub fn dh_decrypt_stream(
+    secret: [u8; 32],
+    mut input: impl Read,
+    mut output: impl Write,
+) -> Result<()> {
+    let mut header = [0u8; HEADER_LEN];
+    header[..STREAM_MAGIC.len()].copy_from_slice(STREAM_MAGIC);
+    input
+        .read_exact(&mut header[STREAM_MAGIC.len()..])
+        .context("Truncated stream header")?;
+    let ephemeral_public_key: [u8; 32] = header[STREAM_MAGIC.len()..STREAM_MAGIC.len() + 32]
+        .try_into()
+        .expect("fixed-size header slice");
+    let nonce_prefix: [u8; 8] = header[STREAM_MAGIC.len() + 32..STREAM_MAGIC.len() + 40]
+        .try_into()
+        .expect("fixed-size header slice");
+    let chunk_size = u32::from_be_bytes(
+        header[STREAM_MAGIC.len() + 40..]
+            .try_into()
+            .expect("fixed-size header slice"),
+    ) as usize;
+    ensure!(
+        (1..=MAX_CHUNK_SIZE).contains(&chunk_size),
+        "Invalid chunk size: {chunk_size}"
+    );
+
+    let shared_secret = dh_agree(secret, ephemeral_public_key);
+    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 index = 0u32;
+    loop {
+        let mut frame_header = [0u8; 5];
+        input
+            .read_exact(&mut frame_header)
+            .with_context(|| format!("Missing final chunk at chunk {index}"))?;
+        ensure!(frame_header[0] & !FINAL_CHUNK == 0, "Invalid chunk flags");
+        let final_chunk = frame_header[0] == FINAL_CHUNK;
+        let plaintext_len = u32::from_be_bytes(
+            frame_header[1..]
+                .try_into()
+                .expect("fixed-size frame header"),
+        ) as usize;
+        ensure!(plaintext_len <= chunk_size, "Chunk {index} is too large");
+        ensure!(
+            final_chunk || plaintext_len == chunk_size,
+            "Non-final chunk {index} has an invalid length"
+        );
+
+        let mut encrypted = vec![0u8; plaintext_len + 16];
+        input
+            .read_exact(&mut encrypted)
+            .with_context(|| format!("Truncated chunk {index}"))?;
+        let nonce = stream_nonce(&nonce_prefix, index);
+        let aad = stream_aad(&header, index, &frame_header);
+        let plaintext = cipher
+            .decrypt(
+                (&nonce).into(),
+                Payload {
+                    msg: &encrypted,
+                    aad: &aad,
+                },
+            )
+            .map_err(|e| anyhow!("Failed to decrypt chunk {index}: {e}"))?;
+        output
+            .write_all(&plaintext)
+            .with_context(|| format!("Failed to write chunk {index}"))?;
+
+        if final_chunk {
+            let mut trailing = [0u8; 1];
+            ensure!(
+                input.read(&mut trailing).context("Failed to read input")? == 0,
+                "Trailing data after final chunk"
+            );
+            output.flush().context("Failed to flush plaintext output")?;
+            return Ok(());
+        }
+        index = index.checked_add(1).context("Too many chunks")?;
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -89,4 +286,48 @@ mod tests {
         let decrypted_str = String::from_utf8(decrypted).unwrap();
         assert_eq!(decrypted_str, "[{\"key\":\"\",\"value\":\"\"}]");
     }
+
+    #[test]
+    fn test_stream_roundtrip() {
+        let secret = StaticSecret::random_from_rng(rand::thread_rng());
+        let public_key = PublicKey::from(&secret).to_bytes();
+        let plaintext = vec![0x5a; 2500];
+        let mut encrypted = Vec::new();
+        dh_encrypt_stream(public_key, plaintext.as_slice(), &mut encrypted, 1024).unwrap();
+        assert_eq!(&encrypted[..STREAM_MAGIC.len()], STREAM_MAGIC);
+
+        let mut decrypted = Vec::new();
+        dh_decrypt_stream(
+            secret.to_bytes(),
+            &encrypted[STREAM_MAGIC.len()..],
+            &mut decrypted,
+        )
+        .unwrap();
+        assert_eq!(decrypted, plaintext);
+    }
+
+    #[test]
+    fn test_stream_rejects_tampering_and_truncation() {
+        let secret = StaticSecret::random_from_rng(rand::thread_rng());
+        let public_key = PublicKey::from(&secret).to_bytes();
+        let mut encrypted = Vec::new();
+        dh_encrypt_stream(public_key, b"hello".as_slice(), &mut encrypted, 4).unwrap();
+
+        let mut tampered = encrypted.clone();
+        *tampered.last_mut().unwrap() ^= 1;
+        assert!(dh_decrypt_stream(
+            secret.to_bytes(),
+            &tampered[STREAM_MAGIC.len()..],
+            Vec::new(),
+        )
+        .is_err());
+
+        encrypted.truncate(encrypted.len() - 1);
+        assert!(dh_decrypt_stream(
+            secret.to_bytes(),
+            &encrypted[STREAM_MAGIC.len()..],
+            Vec::new(),
+        )
+        .is_err());
+    }
 }
diff --git a/dstack/dstack-util/src/main.rs b/dstack/dstack-util/src/main.rs
index 768f5f517..fc950f57d 100644
--- a/dstack/dstack-util/src/main.rs
+++ b/dstack/dstack-util/src/main.rs
@@ -99,6 +99,8 @@ enum Commands {
     GetKeys(GetKeysArgs),
     /// Decrypt data encrypted with the app's environment encryption public key
     Decrypt(DecryptArgs),
+    /// Encrypt data for an app using its KMS-provided environment encryption key
+    Encrypt(EncryptArgs),
 }
 
 #[derive(Parser)]
@@ -388,6 +390,34 @@ struct DecryptArgs {
     hex: bool,
 }
 
+#[derive(Parser)]
+/// Encrypt data for an app using its KMS-provided environment encryption key
+struct EncryptArgs {
+    /// KMS server URL
+    #[arg(short, long)]
+    kms_url: String,
+
+    /// Application ID (20 bytes in hex)
+    #[arg(long)]
+    app_id: String,
+
+    /// Input file (default: stdin)
+    #[arg(short, long)]
+    input: Option<PathBuf>,
+
+    /// Output file (default: stdout)
+    #[arg(short, long)]
+    output: Option<PathBuf>,
+
+    /// Plaintext bytes per independently authenticated chunk
+    #[arg(long, default_value_t = crypto::DEFAULT_CHUNK_SIZE)]
+    chunk_size: usize,
+
+    /// Root CA certificate (PEM format) used to verify the KMS TLS certificate
+    #[arg(long)]
+    root_ca: Option<PathBuf>,
+}
+
 fn pad64(data: &[u8]) -> Result<[u8; 64]> {
     if data.len() > 64 {
         anyhow::bail!("report_data must be at most 64 bytes");
@@ -697,38 +727,140 @@ fn cmd_decrypt(args: DecryptArgs) -> Result<()> {
         .try_into()
         .map_err(|key: Vec<u8>| anyhow::anyhow!("Invalid env crypt key length: {}", key.len()))?;
 
-    let mut input = match args.input {
-        Some(path) => fs::read(&path)
-            .with_context(|| format!("Failed to read ciphertext from {}", path.display()))?,
-        None => {
-            let mut input = Vec::new();
-            io::stdin()
-                .read_to_end(&mut input)
-                .context("Failed to read ciphertext from stdin")?;
-            input
-        }
-    };
     if args.hex {
-        input = hex_decode(
+        let input = read_all_input(args.input.as_deref())?;
+        let input = hex_decode(
             std::str::from_utf8(&input)
                 .context("Hex ciphertext is not valid UTF-8")?
                 .trim(),
         )
         .context("Failed to decode hex ciphertext")?;
+        return decrypt_auto(
+            env_crypt_key,
+            input.as_slice(),
+            open_output(args.output.as_deref())?,
+        );
     }
 
-    let plaintext = crypto::dh_decrypt(env_crypt_key, &input).context("Failed to decrypt input")?;
-    if let Some(output) = args.output {
-        safe_write_with_mode(&output, &plaintext, 0o600)
-            .with_context(|| format!("Failed to write plaintext to {}", output.display()))?;
+    let input = open_input(args.input.as_deref())?;
+    decrypt_auto(env_crypt_key, input, open_output(args.output.as_deref())?)
+}
+
+fn decrypt_auto(
+    env_crypt_key: [u8; 32],
+    mut input: impl Read,
+    mut output: impl Write,
+) -> Result<()> {
+    let mut prefix = Vec::with_capacity(crypto::STREAM_MAGIC.len());
+    input
+        .by_ref()
+        .take(crypto::STREAM_MAGIC.len() as u64)
+        .read_to_end(&mut prefix)
+        .context("Failed to read ciphertext")?;
+    if prefix == crypto::STREAM_MAGIC {
+        crypto::dh_decrypt_stream(env_crypt_key, input, output)
+            .context("Failed to decrypt stream")?;
     } else {
-        io::stdout()
+        let mut ciphertext = prefix;
+        input
+            .read_to_end(&mut ciphertext)
+            .context("Failed to read ciphertext")?;
+        let plaintext = crypto::dh_decrypt(env_crypt_key, &ciphertext)
+            .context("Failed to decrypt legacy input")?;
+        output
             .write_all(&plaintext)
-            .context("Failed to write plaintext to stdout")?;
+            .context("Failed to write plaintext")?;
     }
     Ok(())
 }
 
+async fn cmd_encrypt(args: EncryptArgs) -> Result<()> {
+    use dstack_kms_rpc::kms_client::KmsClient;
+    use ra_rpc::client::RaClientConfig;
+
+    let app_id = decode_app_id(Some(&args.app_id))?.expect("app_id is required");
+    let kms_url = if args.kms_url.ends_with("/prpc") {
+        args.kms_url
+    } else {
+        format!("{}/prpc", args.kms_url.trim_end_matches('/'))
+    };
+    let root_ca_pem = args
+        .root_ca
+        .as_ref()
+        .map(|path| {
+            fs::read_to_string(path)
+                .with_context(|| format!("Failed to read root CA from {}", path.display()))
+        })
+        .transpose()?;
+    let client = RaClientConfig::builder()
+        .remote_uri(kms_url)
+        .tls_no_check(false)
+        .tls_built_in_root_certs(root_ca_pem.is_none())
+        .maybe_tls_ca_cert(root_ca_pem)
+        .build()
+        .into_client()
+        .context("Failed to create KMS client")?;
+    let response = KmsClient::new(client)
+        .get_app_env_encrypt_pub_key(dstack_kms_rpc::AppId {
+            app_id: app_id.to_vec(),
+        })
+        .await
+        .context("Failed to get app environment encryption public key")?;
+    let public_key: [u8; 32] = response
+        .public_key
+        .try_into()
+        .map_err(|key: Vec<u8>| anyhow::anyhow!("Invalid public key length: {}", key.len()))?;
+
+    crypto::dh_encrypt_stream(
+        public_key,
+        open_input(args.input.as_deref())?,
+        open_output(args.output.as_deref())?,
+        args.chunk_size,
+    )
+    .context("Failed to encrypt stream")
+}
+
+fn read_all_input(path: Option<&Path>) -> Result<Vec<u8>> {
+    let mut input = open_input(path)?;
+    let mut data = Vec::new();
+    input
+        .read_to_end(&mut data)
+        .context("Failed to read input")?;
+    Ok(data)
+}
+
+fn open_input(path: Option<&Path>) -> Result<Box<dyn Read>> {
+    match path {
+        Some(path) => {
+            Ok(Box::new(fs::File::open(path).with_context(|| {
+                format!("Failed to open input {}", path.display())
+            })?))
+        }
+        None => Ok(Box::new(io::stdin())),
+    }
+}
+
+fn open_output(path: Option<&Path>) -> Result<Box<dyn Write>> {
+    use fs_err::os::unix::fs::OpenOptionsExt;
+    use std::os::unix::fs::PermissionsExt;
+
+    match path {
+        Some(path) => {
+            let file = fs::OpenOptions::new()
+                .write(true)
+                .create(true)
+                .truncate(true)
+                .mode(0o600)
+                .open(path)
+                .with_context(|| format!("Failed to open output {}", path.display()))?;
+            file.set_permissions(std::fs::Permissions::from_mode(0o600))
+                .with_context(|| format!("Failed to set permissions on {}", path.display()))?;
+            Ok(Box::new(file))
+        }
+        None => Ok(Box::new(io::stdout())),
+    }
+}
+
 fn cmd_quote() -> Result<()> {
     let mut input = Vec::with_capacity(65);
     io::stdin()
@@ -1444,6 +1576,9 @@ async fn main() -> Result<()> {
         Commands::Decrypt(args) => {
             cmd_decrypt(args)?;
         }
+        Commands::Encrypt(args) => {
+            cmd_encrypt(args).await?;
+        }
     }
 
     Ok(())

From 4cb2d4f858a64e179dbd07a58f0040684493fef1 Mon Sep 17 00:00:00 2001
From: Kevin Wang <wy721@qq.com>
Date: Thu, 13 Aug 2026 20:12:18 -0700
Subject: [PATCH 3/9] fix(dstack-util): satisfy strict clippy checks

---
 dstack/dstack-util/src/crypto.rs | 26 ++++++++++----------------
 dstack/dstack-util/src/main.rs   |  2 +-
 2 files changed, 11 insertions(+), 17 deletions(-)

diff --git a/dstack/dstack-util/src/crypto.rs b/dstack/dstack-util/src/crypto.rs
index dfe1593da..9f2b22848 100644
--- a/dstack/dstack-util/src/crypto.rs
+++ b/dstack/dstack-util/src/crypto.rs
@@ -166,17 +166,13 @@ pub fn dh_decrypt_stream(
     input
         .read_exact(&mut header[STREAM_MAGIC.len()..])
         .context("Truncated stream header")?;
-    let ephemeral_public_key: [u8; 32] = header[STREAM_MAGIC.len()..STREAM_MAGIC.len() + 32]
-        .try_into()
-        .expect("fixed-size header slice");
-    let nonce_prefix: [u8; 8] = header[STREAM_MAGIC.len() + 32..STREAM_MAGIC.len() + 40]
-        .try_into()
-        .expect("fixed-size header slice");
-    let chunk_size = u32::from_be_bytes(
-        header[STREAM_MAGIC.len() + 40..]
-            .try_into()
-            .expect("fixed-size header slice"),
-    ) as usize;
+    let mut ephemeral_public_key = [0u8; 32];
+    ephemeral_public_key.copy_from_slice(&header[STREAM_MAGIC.len()..STREAM_MAGIC.len() + 32]);
+    let mut nonce_prefix = [0u8; 8];
+    nonce_prefix.copy_from_slice(&header[STREAM_MAGIC.len() + 32..STREAM_MAGIC.len() + 40]);
+    let mut chunk_size_bytes = [0u8; 4];
+    chunk_size_bytes.copy_from_slice(&header[STREAM_MAGIC.len() + 40..]);
+    let chunk_size = u32::from_be_bytes(chunk_size_bytes) as usize;
     ensure!(
         (1..=MAX_CHUNK_SIZE).contains(&chunk_size),
         "Invalid chunk size: {chunk_size}"
@@ -198,11 +194,9 @@ pub fn dh_decrypt_stream(
             .with_context(|| format!("Missing final chunk at chunk {index}"))?;
         ensure!(frame_header[0] & !FINAL_CHUNK == 0, "Invalid chunk flags");
         let final_chunk = frame_header[0] == FINAL_CHUNK;
-        let plaintext_len = u32::from_be_bytes(
-            frame_header[1..]
-                .try_into()
-                .expect("fixed-size frame header"),
-        ) as usize;
+        let mut plaintext_len_bytes = [0u8; 4];
+        plaintext_len_bytes.copy_from_slice(&frame_header[1..]);
+        let plaintext_len = u32::from_be_bytes(plaintext_len_bytes) as usize;
         ensure!(plaintext_len <= chunk_size, "Chunk {index} is too large");
         ensure!(
             final_chunk || plaintext_len == chunk_size,
diff --git a/dstack/dstack-util/src/main.rs b/dstack/dstack-util/src/main.rs
index fc950f57d..33ac7006d 100644
--- a/dstack/dstack-util/src/main.rs
+++ b/dstack/dstack-util/src/main.rs
@@ -778,7 +778,7 @@ async fn cmd_encrypt(args: EncryptArgs) -> Result<()> {
     use dstack_kms_rpc::kms_client::KmsClient;
     use ra_rpc::client::RaClientConfig;
 
-    let app_id = decode_app_id(Some(&args.app_id))?.expect("app_id is required");
+    let app_id = decode_app_id(Some(&args.app_id))?.context("app_id is required")?;
     let kms_url = if args.kms_url.ends_with("/prpc") {
         args.kms_url
     } else {

From 5adf1822f0f9325fc9f7bd605406d92441b8ac9a Mon Sep 17 00:00:00 2001
From: Kevin Wang <wy721@qq.com>
Date: Fri, 14 Aug 2026 11:41:50 +0800
Subject: [PATCH 4/9] Potential fix for pull request finding

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
---
 dstack/dstack-util/src/main.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/dstack/dstack-util/src/main.rs b/dstack/dstack-util/src/main.rs
index 33ac7006d..8344140c7 100644
--- a/dstack/dstack-util/src/main.rs
+++ b/dstack/dstack-util/src/main.rs
@@ -721,7 +721,7 @@ fn cmd_decrypt(args: DecryptArgs) -> Result<()> {
         .key_file
         .unwrap_or_else(|| host_shared_dir().join(APP_KEYS));
     let keys: AppKeys = utils::deserialize_json_file(&key_file)
-        .with_context(|| format!("Failed to load app keys from {}", key_file.display()))?;
+        .with_context(|| format!("failed to load app keys from {}", key_file.display()))?;
     let env_crypt_key: [u8; 32] = keys
         .env_crypt_key
         .try_into()

From 88bb69d481d7575895d1bf7eefbc55c2787a247e Mon Sep 17 00:00:00 2001
From: Kevin Wang <wy721@qq.com>
Date: Thu, 13 Aug 2026 21:41:11 -0700
Subject: [PATCH 5/9] fix(dstack-util): verify KMS encryption key signatures

---
 docs/stream-encryption.md        |  10 +-
 dstack/dstack-util/src/crypto.rs |  58 ++++----
 dstack/dstack-util/src/main.rs   | 224 ++++++++++++++++++++++++++-----
 3 files changed, 232 insertions(+), 60 deletions(-)

diff --git a/docs/stream-encryption.md b/docs/stream-encryption.md
index c683969e7..381c7d39c 100644
--- a/docs/stream-encryption.md
+++ b/docs/stream-encryption.md
@@ -49,11 +49,14 @@ Encrypt data after retrieving the app public key over verified TLS:
 dstack-util encrypt \
   --kms-url https://kms.example.com \
   --app-id "$APP_ID" \
+  --kms-pubkey "$TRUSTED_KMS_SIGNER_PUBKEY" \
   --input plaintext.bin \
   --output ciphertext.bin
 ```
 
-For a KMS using a private CA, pass `--root-ca ca.pem`. Decrypt inside the CVM:
+`--kms-pubkey` is the trusted compressed secp256k1 public key used to verify the
+KMS response's timestamped signature. For a KMS using a private CA, also pass
+`--root-ca ca.pem`. Decrypt inside the CVM:
 
 ```bash
 dstack-util decrypt --input ciphertext.bin --output plaintext.bin
@@ -63,3 +66,8 @@ dstack-util decrypt --input ciphertext.bin --output plaintext.bin
 handled as the legacy encrypted-environment format. Hex input remains available
 through `--hex`, but it is decoded in memory and should not be used for large
 files.
+
+Successfully authenticated chunks are written as they are processed. If a
+later chunk is corrupt or the final frame is missing, stdout or a file may
+therefore contain an authenticated but incomplete plaintext prefix. Callers
+must check the command's exit status and discard all output on failure.
diff --git a/dstack/dstack-util/src/crypto.rs b/dstack/dstack-util/src/crypto.rs
index 9f2b22848..5c6989be3 100644
--- a/dstack/dstack-util/src/crypto.rs
+++ b/dstack/dstack-util/src/crypto.rs
@@ -27,13 +27,13 @@ pub fn dh_decrypt(secret: [u8; 32], ciphertext: &[u8]) -> Result<Vec<u8>> {
     // Extract components (matching JS implementation)
     let ephemeral_pubkey = ciphertext
         .get(..32)
-        .ok_or(anyhow!("Invalid ephemeral public key length"))?
+        .ok_or(anyhow!("invalid ephemeral public key length"))?
         .try_into()
-        .map_err(|_| anyhow!("Invalid ephemeral public key length"))?;
-    let iv = &ciphertext.get(32..44).ok_or(anyhow!("Invalid IV length"))?;
+        .map_err(|_| anyhow!("invalid ephemeral public key length"))?;
+    let iv = &ciphertext.get(32..44).ok_or(anyhow!("invalid IV length"))?;
     let ciphertext = &ciphertext
         .get(44..)
-        .ok_or(anyhow!("Invalid ciphertext length"))?;
+        .ok_or(anyhow!("invalid ciphertext length"))?;
 
     // Derive shared secret using X25519
     let shared_secret = dh_agree(secret, ephemeral_pubkey);
@@ -43,7 +43,7 @@ pub fn dh_decrypt(secret: [u8; 32], ciphertext: &[u8]) -> Result<Vec<u8>> {
 
     // Create AES-GCM cipher
     let cipher = Aes256Gcm::new_from_slice(&shared_secret)
-        .map_err(|e| anyhow!("Failed to create cipher: {}", e))?;
+        .map_err(|e| anyhow!("failed to create cipher: {}", e))?;
 
     // Decrypt using AES-GCM
     cipher
@@ -75,11 +75,11 @@ pub fn dh_encrypt_stream(
 ) -> Result<()> {
     ensure!(
         (1..=MAX_CHUNK_SIZE).contains(&chunk_size),
-        "Chunk size must be between 1 and {MAX_CHUNK_SIZE} bytes"
+        "chunk size must be between 1 and {MAX_CHUNK_SIZE} bytes"
     );
 
     let mut ephemeral_secret = [0u8; 32];
-    getrandom::fill(&mut ephemeral_secret).context("Failed to generate ephemeral secret")?;
+    getrandom::fill(&mut ephemeral_secret).context("failed to generate ephemeral secret")?;
     let ephemeral_secret = StaticSecret::from(ephemeral_secret);
     let ephemeral_public_key = PublicKey::from(&ephemeral_secret).to_bytes();
     let remote_public_key = PublicKey::from(remote_public_key);
@@ -91,10 +91,10 @@ pub fn dh_encrypt_stream(
         "invalid X25519 shared secret"
     );
     let cipher = Aes256Gcm::new_from_slice(&shared_secret)
-        .map_err(|e| anyhow!("Failed to create cipher: {e}"))?;
+        .map_err(|e| anyhow!("failed to create cipher: {e}"))?;
 
     let mut nonce_prefix = [0u8; 8];
-    getrandom::fill(&mut nonce_prefix).context("Failed to generate nonce prefix")?;
+    getrandom::fill(&mut nonce_prefix).context("failed to generate nonce prefix")?;
     let mut header = [0u8; HEADER_LEN];
     header[..STREAM_MAGIC.len()].copy_from_slice(STREAM_MAGIC);
     header[STREAM_MAGIC.len()..STREAM_MAGIC.len() + 32].copy_from_slice(&ephemeral_public_key);
@@ -102,7 +102,7 @@ pub fn dh_encrypt_stream(
     header[STREAM_MAGIC.len() + 40..].copy_from_slice(&(chunk_size as u32).to_be_bytes());
     output
         .write_all(&header)
-        .context("Failed to write header")?;
+        .context("failed to write header")?;
 
     let mut current = vec![0u8; chunk_size];
     let mut next = vec![0u8; chunk_size];
@@ -125,19 +125,19 @@ pub fn dh_encrypt_stream(
                     aad: &aad,
                 },
             )
-            .map_err(|e| anyhow!("Failed to encrypt chunk {index}: {e}"))?;
+            .map_err(|e| anyhow!("failed to encrypt chunk {index}: {e}"))?;
         output
             .write_all(&frame_header)
             .and_then(|_| output.write_all(&encrypted))
-            .with_context(|| format!("Failed to write chunk {index}"))?;
+            .with_context(|| format!("failed to write chunk {index}"))?;
         if final_chunk {
             break;
         }
-        index = index.checked_add(1).context("Too many chunks")?;
+        index = index.checked_add(1).context("too many chunks")?;
         std::mem::swap(&mut current, &mut next);
         current_len = next_len;
     }
-    output.flush().context("Failed to flush encrypted output")?;
+    output.flush().context("failed to flush encrypted output")?;
     Ok(())
 }
 
@@ -146,7 +146,7 @@ fn read_chunk(input: &mut impl Read, buffer: &mut [u8]) -> Result<usize> {
     while read < buffer.len() {
         match input
             .read(&mut buffer[read..])
-            .context("Failed to read input")?
+            .context("failed to read input")?
         {
             0 => break,
             n => read += n,
@@ -165,7 +165,7 @@ pub fn dh_decrypt_stream(
     header[..STREAM_MAGIC.len()].copy_from_slice(STREAM_MAGIC);
     input
         .read_exact(&mut header[STREAM_MAGIC.len()..])
-        .context("Truncated stream header")?;
+        .context("truncated stream header")?;
     let mut ephemeral_public_key = [0u8; 32];
     ephemeral_public_key.copy_from_slice(&header[STREAM_MAGIC.len()..STREAM_MAGIC.len() + 32]);
     let mut nonce_prefix = [0u8; 8];
@@ -175,7 +175,7 @@ pub fn dh_decrypt_stream(
     let chunk_size = u32::from_be_bytes(chunk_size_bytes) as usize;
     ensure!(
         (1..=MAX_CHUNK_SIZE).contains(&chunk_size),
-        "Invalid chunk size: {chunk_size}"
+        "invalid chunk size: {chunk_size}"
     );
 
     let shared_secret = dh_agree(secret, ephemeral_public_key);
@@ -184,29 +184,29 @@ pub fn dh_decrypt_stream(
         "invalid X25519 shared secret"
     );
     let cipher = Aes256Gcm::new_from_slice(&shared_secret)
-        .map_err(|e| anyhow!("Failed to create cipher: {e}"))?;
+        .map_err(|e| anyhow!("failed to create cipher: {e}"))?;
 
     let mut index = 0u32;
     loop {
         let mut frame_header = [0u8; 5];
         input
             .read_exact(&mut frame_header)
-            .with_context(|| format!("Missing final chunk at chunk {index}"))?;
-        ensure!(frame_header[0] & !FINAL_CHUNK == 0, "Invalid chunk flags");
+            .with_context(|| format!("missing final chunk at chunk {index}"))?;
+        ensure!(frame_header[0] & !FINAL_CHUNK == 0, "invalid chunk flags");
         let final_chunk = frame_header[0] == FINAL_CHUNK;
         let mut plaintext_len_bytes = [0u8; 4];
         plaintext_len_bytes.copy_from_slice(&frame_header[1..]);
         let plaintext_len = u32::from_be_bytes(plaintext_len_bytes) as usize;
-        ensure!(plaintext_len <= chunk_size, "Chunk {index} is too large");
+        ensure!(plaintext_len <= chunk_size, "chunk {index} is too large");
         ensure!(
             final_chunk || plaintext_len == chunk_size,
-            "Non-final chunk {index} has an invalid length"
+            "non-final chunk {index} has an invalid length"
         );
 
         let mut encrypted = vec![0u8; plaintext_len + 16];
         input
             .read_exact(&mut encrypted)
-            .with_context(|| format!("Truncated chunk {index}"))?;
+            .with_context(|| format!("truncated chunk {index}"))?;
         let nonce = stream_nonce(&nonce_prefix, index);
         let aad = stream_aad(&header, index, &frame_header);
         let plaintext = cipher
@@ -217,21 +217,21 @@ pub fn dh_decrypt_stream(
                     aad: &aad,
                 },
             )
-            .map_err(|e| anyhow!("Failed to decrypt chunk {index}: {e}"))?;
+            .map_err(|e| anyhow!("failed to decrypt chunk {index}: {e}"))?;
         output
             .write_all(&plaintext)
-            .with_context(|| format!("Failed to write chunk {index}"))?;
+            .with_context(|| format!("failed to write chunk {index}"))?;
 
         if final_chunk {
             let mut trailing = [0u8; 1];
             ensure!(
-                input.read(&mut trailing).context("Failed to read input")? == 0,
-                "Trailing data after final chunk"
+                input.read(&mut trailing).context("failed to read input")? == 0,
+                "trailing data after final chunk"
             );
-            output.flush().context("Failed to flush plaintext output")?;
+            output.flush().context("failed to flush plaintext output")?;
             return Ok(());
         }
-        index = index.checked_add(1).context("Too many chunks")?;
+        index = index.checked_add(1).context("too many chunks")?;
     }
 }
 
diff --git a/dstack/dstack-util/src/main.rs b/dstack/dstack-util/src/main.rs
index 8344140c7..d4c56a93d 100644
--- a/dstack/dstack-util/src/main.rs
+++ b/dstack/dstack-util/src/main.rs
@@ -416,6 +416,14 @@ struct EncryptArgs {
     /// Root CA certificate (PEM format) used to verify the KMS TLS certificate
     #[arg(long)]
     root_ca: Option<PathBuf>,
+
+    /// Trusted compressed secp256k1 KMS signer public key (hex)
+    #[arg(long)]
+    kms_pubkey: String,
+
+    /// Maximum accepted age of the KMS public-key signature in seconds
+    #[arg(long, default_value_t = 300)]
+    max_signature_age: u64,
 }
 
 fn pad64(data: &[u8]) -> Result<[u8; 64]> {
@@ -614,11 +622,7 @@ async fn cmd_get_keys(args: GetKeysArgs) -> Result<()> {
     use dstack_kms_rpc::kms_client::KmsClient;
     use ra_rpc::client::RaClientConfig;
 
-    let kms_url = if args.kms_url.ends_with("/prpc") {
-        args.kms_url.clone()
-    } else {
-        format!("{}/prpc", args.kms_url.trim_end_matches('/'))
-    };
+    let kms_url = normalize_prpc_url(&args.kms_url);
 
     // Load root CA if provided for TLS pinning
     let root_ca_pem = if let Some(root_ca_path) = &args.root_ca {
@@ -725,16 +729,11 @@ fn cmd_decrypt(args: DecryptArgs) -> Result<()> {
     let env_crypt_key: [u8; 32] = keys
         .env_crypt_key
         .try_into()
-        .map_err(|key: Vec<u8>| anyhow::anyhow!("Invalid env crypt key length: {}", key.len()))?;
+        .map_err(|key: Vec<u8>| anyhow::anyhow!("invalid env crypt key length: {}", key.len()))?;
 
     if args.hex {
         let input = read_all_input(args.input.as_deref())?;
-        let input = hex_decode(
-            std::str::from_utf8(&input)
-                .context("Hex ciphertext is not valid UTF-8")?
-                .trim(),
-        )
-        .context("Failed to decode hex ciphertext")?;
+        let input = decode_hex_ciphertext(&input)?;
         return decrypt_auto(
             env_crypt_key,
             input.as_slice(),
@@ -756,20 +755,20 @@ fn decrypt_auto(
         .by_ref()
         .take(crypto::STREAM_MAGIC.len() as u64)
         .read_to_end(&mut prefix)
-        .context("Failed to read ciphertext")?;
+        .context("failed to read ciphertext")?;
     if prefix == crypto::STREAM_MAGIC {
         crypto::dh_decrypt_stream(env_crypt_key, input, output)
-            .context("Failed to decrypt stream")?;
+            .context("failed to decrypt stream")?;
     } else {
         let mut ciphertext = prefix;
         input
             .read_to_end(&mut ciphertext)
-            .context("Failed to read ciphertext")?;
+            .context("failed to read ciphertext")?;
         let plaintext = crypto::dh_decrypt(env_crypt_key, &ciphertext)
-            .context("Failed to decrypt legacy input")?;
+            .context("failed to decrypt legacy input")?;
         output
             .write_all(&plaintext)
-            .context("Failed to write plaintext")?;
+            .context("failed to write plaintext")?;
     }
     Ok(())
 }
@@ -779,17 +778,13 @@ async fn cmd_encrypt(args: EncryptArgs) -> Result<()> {
     use ra_rpc::client::RaClientConfig;
 
     let app_id = decode_app_id(Some(&args.app_id))?.context("app_id is required")?;
-    let kms_url = if args.kms_url.ends_with("/prpc") {
-        args.kms_url
-    } else {
-        format!("{}/prpc", args.kms_url.trim_end_matches('/'))
-    };
+    let kms_url = normalize_prpc_url(&args.kms_url);
     let root_ca_pem = args
         .root_ca
         .as_ref()
         .map(|path| {
             fs::read_to_string(path)
-                .with_context(|| format!("Failed to read root CA from {}", path.display()))
+                .with_context(|| format!("failed to read root CA from {}", path.display()))
         })
         .transpose()?;
     let client = RaClientConfig::builder()
@@ -799,17 +794,25 @@ async fn cmd_encrypt(args: EncryptArgs) -> Result<()> {
         .maybe_tls_ca_cert(root_ca_pem)
         .build()
         .into_client()
-        .context("Failed to create KMS client")?;
+        .context("failed to create KMS client")?;
     let response = KmsClient::new(client)
         .get_app_env_encrypt_pub_key(dstack_kms_rpc::AppId {
             app_id: app_id.to_vec(),
         })
         .await
-        .context("Failed to get app environment encryption public key")?;
+        .context("failed to get app environment encryption public key")?;
     let public_key: [u8; 32] = response
         .public_key
         .try_into()
-        .map_err(|key: Vec<u8>| anyhow::anyhow!("Invalid public key length: {}", key.len()))?;
+        .map_err(|key: Vec<u8>| anyhow::anyhow!("invalid public key length: {}", key.len()))?;
+    verify_env_encrypt_public_key(
+        &public_key,
+        &response.signature_v1,
+        &app_id,
+        response.timestamp,
+        &args.kms_pubkey,
+        args.max_signature_age,
+    )?;
 
     crypto::dh_encrypt_stream(
         public_key,
@@ -817,7 +820,80 @@ async fn cmd_encrypt(args: EncryptArgs) -> Result<()> {
         open_output(args.output.as_deref())?,
         args.chunk_size,
     )
-    .context("Failed to encrypt stream")
+    .context("failed to encrypt stream")
+}
+
+fn normalize_prpc_url(url: &str) -> String {
+    let url = url.trim_end_matches('/');
+    if url.ends_with("/prpc") {
+        url.to_string()
+    } else {
+        format!("{url}/prpc")
+    }
+}
+
+fn decode_hex_ciphertext(input: &[u8]) -> Result<Vec<u8>> {
+    hex_decode(
+        std::str::from_utf8(input)
+            .context("hex ciphertext is not valid UTF-8")?
+            .trim(),
+    )
+    .context("failed to decode hex ciphertext")
+}
+
+fn verify_env_encrypt_public_key(
+    public_key: &[u8; 32],
+    signature: &[u8],
+    app_id: &[u8; 20],
+    timestamp: u64,
+    trusted_pubkey: &str,
+    max_age: u64,
+) -> Result<()> {
+    use k256::ecdsa::{RecoveryId, Signature, VerifyingKey};
+    use sha3::{Digest, Keccak256};
+    use std::time::{SystemTime, UNIX_EPOCH};
+
+    const FUTURE_SKEW: u64 = 60;
+    let now = SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .context("system time is before the Unix epoch")?
+        .as_secs();
+    anyhow::ensure!(
+        timestamp <= now.saturating_add(FUTURE_SKEW),
+        "kms public-key signature timestamp is too far in the future"
+    );
+    anyhow::ensure!(
+        now.saturating_sub(timestamp) <= max_age,
+        "kms public-key signature is too old"
+    );
+    anyhow::ensure!(signature.len() == 65, "invalid KMS signature length");
+
+    let signature_value =
+        Signature::from_slice(&signature[..64]).context("invalid KMS signature")?;
+    let recovery_id = RecoveryId::from_byte(signature[64]).context("invalid KMS recovery ID")?;
+    let digest = Keccak256::new_with_prefix(
+        [
+            b"dstack-env-encrypt-pubkey".as_slice(),
+            b":".as_slice(),
+            app_id.as_slice(),
+            &timestamp.to_be_bytes(),
+            public_key.as_slice(),
+        ]
+        .concat(),
+    );
+    let recovered = VerifyingKey::recover_from_digest(digest, &signature_value, recovery_id)
+        .context("failed to recover KMS signer public key")?;
+
+    let trusted_pubkey = trusted_pubkey.strip_prefix("0x").unwrap_or(trusted_pubkey);
+    let trusted_pubkey =
+        hex_decode(trusted_pubkey).context("invalid trusted KMS public key hex")?;
+    let trusted =
+        VerifyingKey::from_sec1_bytes(&trusted_pubkey).context("invalid trusted KMS public key")?;
+    anyhow::ensure!(
+        recovered == trusted,
+        "kms public-key signature was made by an untrusted signer"
+    );
+    Ok(())
 }
 
 fn read_all_input(path: Option<&Path>) -> Result<Vec<u8>> {
@@ -825,7 +901,7 @@ fn read_all_input(path: Option<&Path>) -> Result<Vec<u8>> {
     let mut data = Vec::new();
     input
         .read_to_end(&mut data)
-        .context("Failed to read input")?;
+        .context("failed to read input")?;
     Ok(data)
 }
 
@@ -833,7 +909,7 @@ fn open_input(path: Option<&Path>) -> Result<Box<dyn Read>> {
     match path {
         Some(path) => {
             Ok(Box::new(fs::File::open(path).with_context(|| {
-                format!("Failed to open input {}", path.display())
+                format!("failed to open input {}", path.display())
             })?))
         }
         None => Ok(Box::new(io::stdin())),
@@ -852,7 +928,7 @@ fn open_output(path: Option<&Path>) -> Result<Box<dyn Write>> {
                 .truncate(true)
                 .mode(0o600)
                 .open(path)
-                .with_context(|| format!("Failed to open output {}", path.display()))?;
+                .with_context(|| format!("failed to open output {}", path.display()))?;
             file.set_permissions(std::fs::Permissions::from_mode(0o600))
                 .with_context(|| format!("Failed to set permissions on {}", path.display()))?;
             Ok(Box::new(file))
@@ -1651,4 +1727,92 @@ mod tests {
             0o600
         );
     }
+
+    #[test]
+    fn prpc_url_normalization_handles_trailing_slashes() {
+        assert_eq!(
+            normalize_prpc_url("https://kms.example.com/prpc/"),
+            "https://kms.example.com/prpc"
+        );
+        assert_eq!(
+            normalize_prpc_url("https://kms.example.com/"),
+            "https://kms.example.com/prpc"
+        );
+    }
+
+    #[test]
+    fn decrypt_auto_detects_stream_and_falls_back_to_legacy() {
+        use x25519_dalek::{PublicKey, StaticSecret};
+
+        let secret = StaticSecret::random_from_rng(rand::thread_rng());
+        let mut encrypted = Vec::new();
+        crypto::dh_encrypt_stream(
+            PublicKey::from(&secret).to_bytes(),
+            b"stream plaintext".as_slice(),
+            &mut encrypted,
+            4,
+        )
+        .unwrap();
+        let mut decrypted = Vec::new();
+        decrypt_auto(secret.to_bytes(), encrypted.as_slice(), &mut decrypted).unwrap();
+        assert_eq!(decrypted, b"stream plaintext");
+
+        let legacy_secret: [u8; 32] =
+            hex_decode("7c282bf94b35dc47801dc953bfa0896fc2bd313381d3e8eca4e42f6536d2a96f")
+                .unwrap()
+                .try_into()
+                .unwrap();
+        let legacy_ciphertext = hex_decode("0bd18749612f4c8b9dd583c7d6a646b90abd34e3c731a7708d0caf9039095641e1f0948e775f0b7351788db7f246d51806954626dcccb6a60d64665ca3715c6bef75616cab476d27bba04080361200d6a58cec").unwrap();
+        let mut legacy_plaintext = Vec::new();
+        decrypt_auto(
+            legacy_secret,
+            legacy_ciphertext.as_slice(),
+            &mut legacy_plaintext,
+        )
+        .unwrap();
+        assert_eq!(legacy_plaintext, b"[{\"key\":\"\",\"value\":\"\"}]");
+        assert_eq!(decode_hex_ciphertext(b" 00ff\n").unwrap(), [0, 255]);
+    }
+
+    #[test]
+    fn env_encrypt_public_key_requires_the_trusted_signer() {
+        use k256::ecdsa::SigningKey as EcdsaSigningKey;
+        use sha3::{Digest, Keccak256};
+        use std::time::{SystemTime, UNIX_EPOCH};
+
+        let signer = EcdsaSigningKey::random(&mut rand::thread_rng());
+        let app_id = [0x11; 20];
+        let public_key = [0x22; 32];
+        let timestamp = SystemTime::now()
+            .duration_since(UNIX_EPOCH)
+            .unwrap()
+            .as_secs();
+        let digest = Keccak256::new_with_prefix(
+            [
+                b"dstack-env-encrypt-pubkey".as_slice(),
+                b":".as_slice(),
+                app_id.as_slice(),
+                &timestamp.to_be_bytes(),
+                public_key.as_slice(),
+            ]
+            .concat(),
+        );
+        let (signature, recovery_id) = signer.sign_digest_recoverable(digest).unwrap();
+        let mut signature = signature.to_vec();
+        signature.push(recovery_id.to_byte());
+        let trusted = hex::encode(signer.verifying_key().to_sec1_bytes());
+
+        verify_env_encrypt_public_key(&public_key, &signature, &app_id, timestamp, &trusted, 300)
+            .unwrap();
+        let untrusted = EcdsaSigningKey::random(&mut rand::thread_rng());
+        assert!(verify_env_encrypt_public_key(
+            &public_key,
+            &signature,
+            &app_id,
+            timestamp,
+            &hex::encode(untrusted.verifying_key().to_sec1_bytes()),
+            300,
+        )
+        .is_err());
+    }
 }

From d312a2992423ca42852e91e9dbe46de9767d3876 Mon Sep 17 00:00:00 2001
From: Kevin Wang <wy721@qq.com>
Date: Thu, 13 Aug 2026 21:53:36 -0700
Subject: [PATCH 6/9] refactor(dstack-util): shorten encrypted stream magic

---
 docs/stream-encryption.md        | 2 +-
 dstack/dstack-util/src/crypto.rs | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/docs/stream-encryption.md b/docs/stream-encryption.md
index 381c7d39c..c51aedca0 100644
--- a/docs/stream-encryption.md
+++ b/docs/stream-encryption.md
@@ -10,7 +10,7 @@ All integers are unsigned and big-endian.
 
 | Field | Size | Description |
 |---|---:|---|
-| Magic | 16 bytes | ASCII `dstack-stream-v1` |
+| Magic | 9 bytes | ASCII `dstkscrt0` |
 | Ephemeral public key | 32 bytes | X25519 public key generated by the sender |
 | Nonce prefix | 8 bytes | Random prefix shared by all chunks |
 | Chunk size | 4 bytes | Maximum plaintext bytes in each chunk |
diff --git a/dstack/dstack-util/src/crypto.rs b/dstack/dstack-util/src/crypto.rs
index 5c6989be3..f9af919ff 100644
--- a/dstack/dstack-util/src/crypto.rs
+++ b/dstack/dstack-util/src/crypto.rs
@@ -10,7 +10,7 @@ use anyhow::{anyhow, ensure, Context, Result};
 use std::io::{Read, Write};
 use x25519_dalek::{PublicKey, StaticSecret};
 
-pub const STREAM_MAGIC: &[u8; 16] = b"dstack-stream-v1";
+pub const STREAM_MAGIC: &[u8; 9] = b"dstkscrt0";
 pub const DEFAULT_CHUNK_SIZE: usize = 1024 * 1024;
 pub const MAX_CHUNK_SIZE: usize = 16 * 1024 * 1024;
 const FINAL_CHUNK: u8 = 1;

From 610f70a0ac06949ab7bbd59e0aaeb3697e597680 Mon Sep 17 00:00:00 2001
From: Kevin Wang <wy721@qq.com>
Date: Thu, 13 Aug 2026 21:59:21 -0700
Subject: [PATCH 7/9] refactor(dstack-util): encode stream metadata with SCALE

---
 docs/stream-encryption.md        |  5 +-
 dstack/dstack-util/src/crypto.rs | 81 +++++++++++++++++---------------
 2 files changed, 47 insertions(+), 39 deletions(-)

diff --git a/docs/stream-encryption.md b/docs/stream-encryption.md
index c51aedca0..772c01451 100644
--- a/docs/stream-encryption.md
+++ b/docs/stream-encryption.md
@@ -4,7 +4,10 @@
 bounded-memory encryption of arbitrary data. It uses the same app-scoped X25519
 key pair as encrypted environment variables, but it is a separate wire format.
 
-All integers are unsigned and big-endian.
+The header and frame metadata are encoded with SCALE, the binary codec already
+used by dstack. Fixed-width integers therefore use SCALE's little-endian
+encoding. The nonce construction below deliberately uses a big-endian chunk
+index so its byte representation follows counter order.
 
 ## Header
 
diff --git a/dstack/dstack-util/src/crypto.rs b/dstack/dstack-util/src/crypto.rs
index f9af919ff..e2bdb15d7 100644
--- a/dstack/dstack-util/src/crypto.rs
+++ b/dstack/dstack-util/src/crypto.rs
@@ -7,6 +7,7 @@ use aes_gcm::{
     Aes256Gcm, KeyInit,
 };
 use anyhow::{anyhow, ensure, Context, Result};
+use scale::{Decode, Encode, IoReader};
 use std::io::{Read, Write};
 use x25519_dalek::{PublicKey, StaticSecret};
 
@@ -14,7 +15,19 @@ pub const STREAM_MAGIC: &[u8; 9] = b"dstkscrt0";
 pub const DEFAULT_CHUNK_SIZE: usize = 1024 * 1024;
 pub const MAX_CHUNK_SIZE: usize = 16 * 1024 * 1024;
 const FINAL_CHUNK: u8 = 1;
-const HEADER_LEN: usize = STREAM_MAGIC.len() + 32 + 8 + 4;
+
+#[derive(Encode, Decode)]
+struct StreamHeader {
+    ephemeral_public_key: [u8; 32],
+    nonce_prefix: [u8; 8],
+    chunk_size: u32,
+}
+
+#[derive(Encode, Decode)]
+struct FrameHeader {
+    flags: u8,
+    plaintext_len: u32,
+}
 
 pub fn dh_agree(secret: [u8; 32], their_pubkey: [u8; 32]) -> [u8; 32] {
     let secret = StaticSecret::from(secret);
@@ -58,11 +71,11 @@ fn stream_nonce(prefix: &[u8; 8], index: u32) -> [u8; 12] {
     nonce
 }
 
-fn stream_aad(header: &[u8; HEADER_LEN], index: u32, frame_header: &[u8; 5]) -> Vec<u8> {
-    let mut aad = Vec::with_capacity(HEADER_LEN + 4 + frame_header.len());
-    aad.extend_from_slice(header);
-    aad.extend_from_slice(&index.to_be_bytes());
-    aad.extend_from_slice(frame_header);
+fn stream_aad(header: &StreamHeader, index: u32, frame_header: &FrameHeader) -> Vec<u8> {
+    let mut aad = STREAM_MAGIC.to_vec();
+    header.encode_to(&mut aad);
+    index.encode_to(&mut aad);
+    frame_header.encode_to(&mut aad);
     aad
 }
 
@@ -95,13 +108,14 @@ pub fn dh_encrypt_stream(
 
     let mut nonce_prefix = [0u8; 8];
     getrandom::fill(&mut nonce_prefix).context("failed to generate nonce prefix")?;
-    let mut header = [0u8; HEADER_LEN];
-    header[..STREAM_MAGIC.len()].copy_from_slice(STREAM_MAGIC);
-    header[STREAM_MAGIC.len()..STREAM_MAGIC.len() + 32].copy_from_slice(&ephemeral_public_key);
-    header[STREAM_MAGIC.len() + 32..STREAM_MAGIC.len() + 40].copy_from_slice(&nonce_prefix);
-    header[STREAM_MAGIC.len() + 40..].copy_from_slice(&(chunk_size as u32).to_be_bytes());
+    let header = StreamHeader {
+        ephemeral_public_key,
+        nonce_prefix,
+        chunk_size: chunk_size as u32,
+    };
     output
-        .write_all(&header)
+        .write_all(STREAM_MAGIC)
+        .and_then(|_| output.write_all(&header.encode()))
         .context("failed to write header")?;
 
     let mut current = vec![0u8; chunk_size];
@@ -112,9 +126,10 @@ pub fn dh_encrypt_stream(
         let next_len = read_chunk(&mut input, &mut next)?;
         let final_chunk = next_len == 0;
         let flags = if final_chunk { FINAL_CHUNK } else { 0 };
-        let mut frame_header = [0u8; 5];
-        frame_header[0] = flags;
-        frame_header[1..].copy_from_slice(&(current_len as u32).to_be_bytes());
+        let frame_header = FrameHeader {
+            flags,
+            plaintext_len: current_len as u32,
+        };
         let aad = stream_aad(&header, index, &frame_header);
         let nonce = stream_nonce(&nonce_prefix, index);
         let encrypted = cipher
@@ -127,7 +142,7 @@ pub fn dh_encrypt_stream(
             )
             .map_err(|e| anyhow!("failed to encrypt chunk {index}: {e}"))?;
         output
-            .write_all(&frame_header)
+            .write_all(&frame_header.encode())
             .and_then(|_| output.write_all(&encrypted))
             .with_context(|| format!("failed to write chunk {index}"))?;
         if final_chunk {
@@ -161,24 +176,15 @@ pub fn dh_decrypt_stream(
     mut input: impl Read,
     mut output: impl Write,
 ) -> Result<()> {
-    let mut header = [0u8; HEADER_LEN];
-    header[..STREAM_MAGIC.len()].copy_from_slice(STREAM_MAGIC);
-    input
-        .read_exact(&mut header[STREAM_MAGIC.len()..])
-        .context("truncated stream header")?;
-    let mut ephemeral_public_key = [0u8; 32];
-    ephemeral_public_key.copy_from_slice(&header[STREAM_MAGIC.len()..STREAM_MAGIC.len() + 32]);
-    let mut nonce_prefix = [0u8; 8];
-    nonce_prefix.copy_from_slice(&header[STREAM_MAGIC.len() + 32..STREAM_MAGIC.len() + 40]);
-    let mut chunk_size_bytes = [0u8; 4];
-    chunk_size_bytes.copy_from_slice(&header[STREAM_MAGIC.len() + 40..]);
-    let chunk_size = u32::from_be_bytes(chunk_size_bytes) as usize;
+    let header =
+        StreamHeader::decode(&mut IoReader(&mut input)).context("invalid stream header")?;
+    let chunk_size = header.chunk_size as usize;
     ensure!(
         (1..=MAX_CHUNK_SIZE).contains(&chunk_size),
         "invalid chunk size: {chunk_size}"
     );
 
-    let shared_secret = dh_agree(secret, ephemeral_public_key);
+    let shared_secret = dh_agree(secret, header.ephemeral_public_key);
     ensure!(
         !shared_secret.iter().all(|byte| *byte == 0),
         "invalid X25519 shared secret"
@@ -188,15 +194,14 @@ pub fn dh_decrypt_stream(
 
     let mut index = 0u32;
     loop {
-        let mut frame_header = [0u8; 5];
-        input
-            .read_exact(&mut frame_header)
+        let frame_header = FrameHeader::decode(&mut IoReader(&mut input))
             .with_context(|| format!("missing final chunk at chunk {index}"))?;
-        ensure!(frame_header[0] & !FINAL_CHUNK == 0, "invalid chunk flags");
-        let final_chunk = frame_header[0] == FINAL_CHUNK;
-        let mut plaintext_len_bytes = [0u8; 4];
-        plaintext_len_bytes.copy_from_slice(&frame_header[1..]);
-        let plaintext_len = u32::from_be_bytes(plaintext_len_bytes) as usize;
+        ensure!(
+            frame_header.flags & !FINAL_CHUNK == 0,
+            "invalid chunk flags"
+        );
+        let final_chunk = frame_header.flags == FINAL_CHUNK;
+        let plaintext_len = frame_header.plaintext_len as usize;
         ensure!(plaintext_len <= chunk_size, "chunk {index} is too large");
         ensure!(
             final_chunk || plaintext_len == chunk_size,
@@ -207,7 +212,7 @@ pub fn dh_decrypt_stream(
         input
             .read_exact(&mut encrypted)
             .with_context(|| format!("truncated chunk {index}"))?;
-        let nonce = stream_nonce(&nonce_prefix, index);
+        let nonce = stream_nonce(&header.nonce_prefix, index);
         let aad = stream_aad(&header, index, &frame_header);
         let plaintext = cipher
             .decrypt(

From 6c266354f31736d5f9fbaff98adba7cf6635862b Mon Sep 17 00:00:00 2001
From: Kevin Wang <wy721@qq.com>
Date: Thu, 13 Aug 2026 22:20:25 -0700
Subject: [PATCH 8/9] refactor(dstack-util): define stream format with binrw

---
 docs/stream-encryption.md        |  8 ++---
 dstack/Cargo.lock                |  1 +
 dstack/dstack-util/Cargo.toml    |  1 +
 dstack/dstack-util/src/crypto.rs | 50 ++++++++++++++++++++------------
 4 files changed, 38 insertions(+), 22 deletions(-)

diff --git a/docs/stream-encryption.md b/docs/stream-encryption.md
index 772c01451..e6c35f9d0 100644
--- a/docs/stream-encryption.md
+++ b/docs/stream-encryption.md
@@ -4,10 +4,10 @@
 bounded-memory encryption of arbitrary data. It uses the same app-scoped X25519
 key pair as encrypted environment variables, but it is a separate wire format.
 
-The header and frame metadata are encoded with SCALE, the binary codec already
-used by dstack. Fixed-width integers therefore use SCALE's little-endian
-encoding. The nonce construction below deliberately uses a big-endian chunk
-index so its byte representation follows counter order.
+The header and frame metadata are defined with `binrw`, the fixed-layout binary
+codec already used by dstack. Fixed-width integers use little-endian encoding.
+The nonce construction below deliberately uses a big-endian chunk index so its
+byte representation follows counter order.
 
 ## Header
 
diff --git a/dstack/Cargo.lock b/dstack/Cargo.lock
index 9e06d5db7..1c1537999 100644
--- a/dstack/Cargo.lock
+++ b/dstack/Cargo.lock
@@ -2219,6 +2219,7 @@ version = "0.6.0"
 dependencies = [
  "aes-gcm",
  "anyhow",
+ "binrw",
  "bollard",
  "cc-eventlog",
  "cert-client",
diff --git a/dstack/dstack-util/Cargo.toml b/dstack/dstack-util/Cargo.toml
index 41c3e01cd..fe16dec66 100644
--- a/dstack/dstack-util/Cargo.toml
+++ b/dstack/dstack-util/Cargo.toml
@@ -55,6 +55,7 @@ cert-client.workspace = true
 x509-parser.workspace = true
 yaml-rust2.workspace = true
 bollard.workspace = true
+binrw.workspace = true
 sodiumbox.workspace = true
 libc.workspace = true
 luks2.workspace = true
diff --git a/dstack/dstack-util/src/crypto.rs b/dstack/dstack-util/src/crypto.rs
index e2bdb15d7..0b747b0dc 100644
--- a/dstack/dstack-util/src/crypto.rs
+++ b/dstack/dstack-util/src/crypto.rs
@@ -7,8 +7,8 @@ use aes_gcm::{
     Aes256Gcm, KeyInit,
 };
 use anyhow::{anyhow, ensure, Context, Result};
-use scale::{Decode, Encode, IoReader};
-use std::io::{Read, Write};
+use binrw::{binrw, io::NoSeek, BinRead, BinWrite};
+use std::io::{Cursor, Read, Write};
 use x25519_dalek::{PublicKey, StaticSecret};
 
 pub const STREAM_MAGIC: &[u8; 9] = b"dstkscrt0";
@@ -16,14 +16,16 @@ pub const DEFAULT_CHUNK_SIZE: usize = 1024 * 1024;
 pub const MAX_CHUNK_SIZE: usize = 16 * 1024 * 1024;
 const FINAL_CHUNK: u8 = 1;
 
-#[derive(Encode, Decode)]
+#[binrw]
+#[brw(little)]
 struct StreamHeader {
     ephemeral_public_key: [u8; 32],
     nonce_prefix: [u8; 8],
     chunk_size: u32,
 }
 
-#[derive(Encode, Decode)]
+#[binrw]
+#[brw(little)]
 struct FrameHeader {
     flags: u8,
     plaintext_len: u32,
@@ -71,12 +73,20 @@ fn stream_nonce(prefix: &[u8; 8], index: u32) -> [u8; 12] {
     nonce
 }
 
-fn stream_aad(header: &StreamHeader, index: u32, frame_header: &FrameHeader) -> Vec<u8> {
-    let mut aad = STREAM_MAGIC.to_vec();
-    header.encode_to(&mut aad);
-    index.encode_to(&mut aad);
-    frame_header.encode_to(&mut aad);
-    aad
+fn stream_aad(header: &StreamHeader, index: u32, frame_header: &FrameHeader) -> Result<Vec<u8>> {
+    let mut aad = Cursor::new(Vec::new());
+    aad.write_all(STREAM_MAGIC)
+        .context("failed to encode stream magic as AAD")?;
+    header
+        .write(&mut aad)
+        .context("failed to encode stream header as AAD")?;
+    index
+        .write_le(&mut aad)
+        .context("failed to encode chunk index as AAD")?;
+    frame_header
+        .write(&mut aad)
+        .context("failed to encode frame header as AAD")?;
+    Ok(aad.into_inner())
 }
 
 /// Encrypts a reader as independently authenticated chunks.
@@ -115,8 +125,10 @@ pub fn dh_encrypt_stream(
     };
     output
         .write_all(STREAM_MAGIC)
-        .and_then(|_| output.write_all(&header.encode()))
-        .context("failed to write header")?;
+        .context("failed to write stream magic")?;
+    header
+        .write(&mut NoSeek::new(&mut output))
+        .context("failed to write stream header")?;
 
     let mut current = vec![0u8; chunk_size];
     let mut next = vec![0u8; chunk_size];
@@ -130,7 +142,7 @@ pub fn dh_encrypt_stream(
             flags,
             plaintext_len: current_len as u32,
         };
-        let aad = stream_aad(&header, index, &frame_header);
+        let aad = stream_aad(&header, index, &frame_header)?;
         let nonce = stream_nonce(&nonce_prefix, index);
         let encrypted = cipher
             .encrypt(
@@ -141,9 +153,11 @@ pub fn dh_encrypt_stream(
                 },
             )
             .map_err(|e| anyhow!("failed to encrypt chunk {index}: {e}"))?;
+        frame_header
+            .write(&mut NoSeek::new(&mut output))
+            .with_context(|| format!("failed to write header for chunk {index}"))?;
         output
-            .write_all(&frame_header.encode())
-            .and_then(|_| output.write_all(&encrypted))
+            .write_all(&encrypted)
             .with_context(|| format!("failed to write chunk {index}"))?;
         if final_chunk {
             break;
@@ -177,7 +191,7 @@ pub fn dh_decrypt_stream(
     mut output: impl Write,
 ) -> Result<()> {
     let header =
-        StreamHeader::decode(&mut IoReader(&mut input)).context("invalid stream header")?;
+        StreamHeader::read(&mut NoSeek::new(&mut input)).context("invalid stream header")?;
     let chunk_size = header.chunk_size as usize;
     ensure!(
         (1..=MAX_CHUNK_SIZE).contains(&chunk_size),
@@ -194,7 +208,7 @@ pub fn dh_decrypt_stream(
 
     let mut index = 0u32;
     loop {
-        let frame_header = FrameHeader::decode(&mut IoReader(&mut input))
+        let frame_header = FrameHeader::read(&mut NoSeek::new(&mut input))
             .with_context(|| format!("missing final chunk at chunk {index}"))?;
         ensure!(
             frame_header.flags & !FINAL_CHUNK == 0,
@@ -213,7 +227,7 @@ pub fn dh_decrypt_stream(
             .read_exact(&mut encrypted)
             .with_context(|| format!("truncated chunk {index}"))?;
         let nonce = stream_nonce(&header.nonce_prefix, index);
-        let aad = stream_aad(&header, index, &frame_header);
+        let aad = stream_aad(&header, index, &frame_header)?;
         let plaintext = cipher
             .decrypt(
                 (&nonce).into(),

From 4c6153f89f44597fa43acfa8fbcf39e7cec67a05 Mon Sep 17 00:00:00 2001
From: Kevin Wang <wy721@qq.com>
Date: Thu, 13 Aug 2026 22:24:08 -0700
Subject: [PATCH 9/9] refactor(dstack-util): move stream version into header

---
 docs/stream-encryption.md        |  3 ++-
 dstack/dstack-util/src/crypto.rs | 19 ++++++++++++++++++-
 2 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/docs/stream-encryption.md b/docs/stream-encryption.md
index e6c35f9d0..3f1a690a3 100644
--- a/docs/stream-encryption.md
+++ b/docs/stream-encryption.md
@@ -13,7 +13,8 @@ byte representation follows counter order.
 
 | Field | Size | Description |
 |---|---:|---|
-| Magic | 9 bytes | ASCII `dstkscrt0` |
+| Magic | 8 bytes | ASCII `dstkscrt` |
+| Version | 1 byte | Format version, currently `0` |
 | Ephemeral public key | 32 bytes | X25519 public key generated by the sender |
 | Nonce prefix | 8 bytes | Random prefix shared by all chunks |
 | Chunk size | 4 bytes | Maximum plaintext bytes in each chunk |
diff --git a/dstack/dstack-util/src/crypto.rs b/dstack/dstack-util/src/crypto.rs
index 0b747b0dc..1f8227f5b 100644
--- a/dstack/dstack-util/src/crypto.rs
+++ b/dstack/dstack-util/src/crypto.rs
@@ -11,14 +11,16 @@ use binrw::{binrw, io::NoSeek, BinRead, BinWrite};
 use std::io::{Cursor, Read, Write};
 use x25519_dalek::{PublicKey, StaticSecret};
 
-pub const STREAM_MAGIC: &[u8; 9] = b"dstkscrt0";
+pub const STREAM_MAGIC: &[u8; 8] = b"dstkscrt";
 pub const DEFAULT_CHUNK_SIZE: usize = 1024 * 1024;
 pub const MAX_CHUNK_SIZE: usize = 16 * 1024 * 1024;
+const STREAM_VERSION: u8 = 0;
 const FINAL_CHUNK: u8 = 1;
 
 #[binrw]
 #[brw(little)]
 struct StreamHeader {
+    version: u8,
     ephemeral_public_key: [u8; 32],
     nonce_prefix: [u8; 8],
     chunk_size: u32,
@@ -119,6 +121,7 @@ pub fn dh_encrypt_stream(
     let mut nonce_prefix = [0u8; 8];
     getrandom::fill(&mut nonce_prefix).context("failed to generate nonce prefix")?;
     let header = StreamHeader {
+        version: STREAM_VERSION,
         ephemeral_public_key,
         nonce_prefix,
         chunk_size: chunk_size as u32,
@@ -192,6 +195,11 @@ pub fn dh_decrypt_stream(
 ) -> Result<()> {
     let header =
         StreamHeader::read(&mut NoSeek::new(&mut input)).context("invalid stream header")?;
+    ensure!(
+        header.version == STREAM_VERSION,
+        "unsupported stream version: {}",
+        header.version
+    );
     let chunk_size = header.chunk_size as usize;
     ensure!(
         (1..=MAX_CHUNK_SIZE).contains(&chunk_size),
@@ -326,6 +334,15 @@ mod tests {
         let mut encrypted = Vec::new();
         dh_encrypt_stream(public_key, b"hello".as_slice(), &mut encrypted, 4).unwrap();
 
+        let mut unknown_version = encrypted.clone();
+        unknown_version[STREAM_MAGIC.len()] = STREAM_VERSION + 1;
+        assert!(dh_decrypt_stream(
+            secret.to_bytes(),
+            &unknown_version[STREAM_MAGIC.len()..],
+            Vec::new(),
+        )
+        .is_err());
+
         let mut tampered = encrypted.clone();
         *tampered.last_mut().unwrap() ^= 1;
         assert!(dh_decrypt_stream(