fix(volume): address CLI and verity review findings

Author: kvinwang

dstack/Cargo.lock

@@ -12,6 +12,7 @@ license.workspace = true
 anyhow.workspace = true
 http-client = { workspace = true, features = ["prpc"] }
 dstack-vmm-rpc.workspace = true
+dstack-types.workspace = true
 serde_json.workspace = true
 # advisory file locking (flock) for the allowlist/state read-modify-write;
 # already in the dependency tree transitively, so no extra compile cost.

dstack/crates/dstack-cli-core/Cargo.toml

@@ -13,14 +13,13 @@ use serde_json::json;
 /// `kms_enabled` selects KMS mode (deterministic, upgradeable per-app keys);
 /// gateway and local-key-provider are off for the direct-port single-node flow.
 ///
-/// `verity_volumes` is a list of `(source, verity_root, target)` tuples. Each becomes a
-/// measured `verity_volumes` entry, so the CVM only seeds content matching the
-/// attested root. Empty for a normal deploy.
+/// Each `verity_volumes` entry is measured, so the CVM only mounts content
+/// matching the attested root. Empty for a normal deploy.
 pub fn build_app_compose(
     name: &str,
     docker_compose_yaml: &str,
     kms_enabled: bool,
-    verity_volumes: &[(String, String, String)],
+    verity_volumes: &[dstack_types::VerityVolume],
 ) -> String {
     let mut manifest = json!({
         "manifest_version": 2,
@@ -41,10 +40,7 @@ pub fn build_app_compose(
         "secure_time": false,
     });
     if !verity_volumes.is_empty() {
-        manifest["verity_volumes"] = json!(verity_volumes
-            .iter()
-            .map(|(source, root, target)| json!({ "source": source, "verity_root": root, "target": target }))
-            .collect::<Vec<_>>());
+        manifest["verity_volumes"] = json!(verity_volumes);
     }
     // pretty-print via Value's Display (`{:#}`) — infallible, and byte-identical
     // to serde_json::to_string_pretty (avoids an expect on an unfailable Result).

dstack/crates/dstack-cli-core/src/compose.rs

@@ -19,7 +19,9 @@ anyhow.workspace = true
 clap.workspace = true
 dstack-cli-core.workspace = true
 dstack-volume.workspace = true
+dstack-types.workspace = true
 fs-err.workspace = true
+hex.workspace = true
 serde_json.workspace = true
 tokio = { workspace = true, features = ["macros", "rt-multi-thread"] }
 tracing.workspace = true

dstack/crates/dstack-cli/Cargo.toml

@@ -11,7 +11,7 @@
 //! `logs`, a global `-j/--json`).
 
 use anyhow::{bail, Context, Result};
-use clap::{Parser, Subcommand};
+use clap::{Parser, Subcommand, ValueEnum};
 use dstack_cli_core::layout::InstallLayout;
 use dstack_cli_core::vmm::{Vmm, DEFAULT_HOST};
 use dstack_cli_core::{compose, ports, rpc};
@@ -77,7 +77,7 @@ enum Command {
         ports: Vec<String>,
         /// attach a verity volume, as printed by `dstack verity`: the file name
         /// in the vmm's volumes_dir, its verity_root, and the target
-        /// (`docker` or a mount path). Repeatable.
+        /// (an absolute mount path). Repeatable.
         #[arg(long = "volume", value_name = "NAME:VERITY_ROOT:TARGET")]
         volumes: Vec<String>,
         /// deploy in non-KMS mode (ephemeral keys; no KMS required).
@@ -124,11 +124,29 @@ enum Command {
         #[arg(long, short = 'o', default_value = "verity.img")]
         output: String,
         /// squashfs compression: `none` (the default), `zstd`, or `gzip`.
-        #[arg(long, default_value = "none")]
-        compress: String,
+        #[arg(long, value_enum, default_value_t, conflicts_with = "fs_image")]
+        compress: CompressionArg,
     },
 }
 
+#[derive(Clone, Copy, Default, ValueEnum)]
+enum CompressionArg {
+    #[default]
+    None,
+    Zstd,
+    Gzip,
+}
+
+impl From<CompressionArg> for dstack_volume::Compression {
+    fn from(value: CompressionArg) -> Self {
+        match value {
+            CompressionArg::None => Self::None,
+            CompressionArg::Zstd => Self::Zstd,
+            CompressionArg::Gzip => Self::Gzip,
+        }
+    }
+}
+
 #[tokio::main]
 async fn main() -> Result<()> {
     // progress (e.g. `verity` pulling layers) goes to stderr so it never mixes
@@ -223,37 +241,22 @@ async fn main() -> Result<()> {
             fs_image,
             output,
             compress,
-        } => {
-            cmd_verity(
-                dir.as_deref(),
-                fs_image.as_deref(),
-                &output,
-                &compress,
-                json,
-            )
-            .await
-        }
+        } => cmd_verity(dir.as_deref(), fs_image.as_deref(), &output, compress, json).await,
     }
 }
 
 async fn cmd_verity(
     dir: Option<&str>,
     fs_image: Option<&str>,
     output: &str,
-    compress: &str,
+    compress: CompressionArg,
     json: bool,
 ) -> Result<()> {
-    let compress = match compress {
-        "none" => dstack_volume::Compression::None,
-        "zstd" => dstack_volume::Compression::Zstd,
-        "gzip" => dstack_volume::Compression::Gzip,
-        other => bail!("unknown --compress '{other}' (use none|zstd|gzip)"),
-    };
     let result = dstack_volume::verity(dstack_volume::VerityOptions {
         dir: dir.map(std::path::PathBuf::from),
         fs_image: fs_image.map(std::path::PathBuf::from),
         output: output.into(),
-        compress,
+        compress: compress.into(),
     })
     .await?;
 
@@ -290,13 +293,6 @@ async fn cmd_verity(
     Ok(())
 }
 
-/// A parsed `--volume` spec. All fields become part of measured app-compose.
-struct VolumeSpec {
-    source: String,
-    verity_root: String,
-    target: String,
-}
-
 /// Parse a `--volume` spec `NAME:VERITY_ROOT:TARGET`.
 ///
 /// `NAME` is the volume file in the vmm's volumes_dir. `VERITY_ROOT` and `TARGET`
@@ -305,7 +301,7 @@ struct VolumeSpec {
 /// spec to paste.
 ///
 /// `TARGET` is an absolute read-only mount path in the guest.
-fn parse_volume(spec: &str) -> Result<VolumeSpec> {
+fn parse_volume(spec: &str) -> Result<dstack_types::VerityVolume> {
     let mut parts = spec.splitn(3, ':');
     let name = parts.next().unwrap_or_default();
     let (root, target) = match (parts.next(), parts.next()) {
@@ -326,10 +322,12 @@ fn parse_volume(spec: &str) -> Result<VolumeSpec> {
     if !target.starts_with('/') {
         bail!("target '{target}' must be an absolute path");
     }
-    Ok(VolumeSpec {
+    let mut verity_root = [0; 32];
+    hex::decode_to_slice(root, &mut verity_root).context("decoding verity_root")?;
+    Ok(dstack_types::VerityVolume {
         source: name.to_string(),
-        verity_root: root.to_string(),
-        target: target.to_string(),
+        verity_root,
+        target: target.into(),
     })
 }
 
@@ -425,11 +423,7 @@ async fn cmd_deploy(
 
     // each --volume declares a measured verity_volumes entry, so the built
     // app-compose (and thus app_id) binds the attested roots.
-    let verity_volumes: Vec<(String, String, String)> = parsed_volumes
-        .iter()
-        .map(|v| (v.source.clone(), v.verity_root.clone(), v.target.clone()))
-        .collect();
-    let app_compose = compose::build_app_compose(name, &yaml, !no_kms, &verity_volumes);
+    let app_compose = compose::build_app_compose(name, &yaml, !no_kms, &parsed_volumes);
 
     let mut cfg = rpc::VmConfiguration {
         name: name.to_string(),
@@ -669,8 +663,8 @@ mod tests {
         let root = "a".repeat(64);
         let data = parse_volume(&format!("weights.img:{root}:/models/llama")).unwrap();
         assert_eq!(data.source, "weights.img");
-        assert_eq!(data.verity_root, root);
-        assert_eq!(data.target, "/models/llama");
+        assert_eq!(data.verity_root, [0xaa; 32]);
+        assert_eq!(data.target, std::path::Path::new("/models/llama"));
 
         assert!(parse_volume("weights.img").is_err()); // missing verity_root:target
         assert!(parse_volume(&format!("weights.img:{root}")).is_err()); // missing target
@@ -734,5 +728,23 @@ mod tests {
             "rootfs.ext4"
         ])
         .is_err());
+        assert!(Cli::try_parse_from([
+            "dstack",
+            "verity",
+            "--fs-image",
+            "rootfs.ext4",
+            "--compress",
+            "zstd"
+        ])
+        .is_err());
+        assert!(Cli::try_parse_from([
+            "dstack",
+            "verity",
+            "--dir",
+            "data",
+            "--compress",
+            "invalid"
+        ])
+        .is_err());
     }
 }

dstack/crates/dstack-cli/src/main.rs

@@ -9,7 +9,6 @@
 //! table. Everything read from a disk is untrusted: kind handlers use the
 //! measured app compose as their source of policy and cryptographic identity.
 
-use std::collections::HashSet;
 use std::io::Read;
 use std::os::unix::ffi::OsStrExt;
 use std::os::unix::fs::MetadataExt;
@@ -96,9 +95,8 @@ fn mount_all(compose_path: PathBuf) -> Result<()> {
         requested = compose.verity_volumes.len(),
         "discovered dstack volumes"
     );
-    let mut used = HashSet::new();
     for (index, requested) in compose.verity_volumes.iter().enumerate() {
-        activate_requested(index, requested, &volumes, &mut used).with_context(|| {
+        activate_requested(index, requested, &volumes).with_context(|| {
             format!(
                 "failed to activate required volume {index} at {}",
                 requested.target.display()
@@ -276,24 +274,23 @@ fn activate_requested(
     index: usize,
     requested: &RequestedVolume,
     volumes: &[VerityVolume],
-    used: &mut HashSet<usize>,
 ) -> Result<()> {
-    let (candidate_index, candidate) = volumes
+    let candidate = volumes
         .iter()
-        .enumerate()
-        .find(|(candidate_index, volume)| {
-            !used.contains(candidate_index) && volume.root_hash == requested.verity_root
-        })
+        .find(|volume| volume.root_hash == requested.verity_root)
         .context("no attached volume advertises the measured root")?;
 
     let mapper_name = format!("dstack-verity{index}");
     let mapped = PathBuf::from(format!("/dev/mapper/{mapper_name}"));
     let expected_root = hex::encode(requested.verity_root);
     if mapped.exists() {
         if mapping_root(&mapper_name)?.eq_ignore_ascii_case(&expected_root) {
-            verify_first_block(&mapped)?;
-            mount_volume(requested, &mapped)?;
-            used.insert(candidate_index);
+            if let Err(err) =
+                verify_first_block(&mapped).and_then(|_| mount_volume(requested, &mapped))
+            {
+                let _ = run_cmd!(veritysetup close $mapper_name);
+                return Err(err);
+            }
             info!(mapper = mapper_name, "reused active verity mapping");
             return Ok(());
         }
@@ -310,7 +307,6 @@ fn activate_requested(
         let _ = run_cmd!(veritysetup close $mapper_name);
         return Err(err);
     }
-    used.insert(candidate_index);
     Ok(())
 }
 

dstack/crates/dstack-volume/src/bin/dstack-volume.rs

@@ -142,7 +142,8 @@ fn seal_data_image(
     let uuid = uuid_from_data(data_path, data_size)?;
     let out = run_fun!(
         veritysetup format --salt $salt_hex --uuid $uuid
-            --data-block-size 4096 --hash-block-size 4096 $data_path $hash_path
+            --hash sha256 --format 1 --data-block-size 4096 --hash-block-size 4096
+            $data_path $hash_path
     )
     .context("running veritysetup format")?;
     let verity_root =

dstack/crates/dstack-volume/src/volume.rs

@@ -15,7 +15,7 @@ DSTACK_ROOTFS_SRC ?= "${DSTACK_MONOREPO_ROOT}/os/common/rootfs"
 S = "${UNPACKDIR}/repo/dstack"
 DSTACK_ROOTFS_FILES = "${UNPACKDIR}/repo/os/common/rootfs"
 
-RDEPENDS:${PN} += "bash"
+RDEPENDS:${PN} += "bash cryptsetup util-linux-blkid util-linux-mount"
 
 DEPENDS += "rsync-native tpm2-tss"
 DEPENDS += "cmake-native"