refactor(verity): simplify volume envelope

Author: kvinwang

dstack/Cargo.lock

@@ -143,7 +143,7 @@ flate2 = "1.1"
 borsh = { version = "1.5.7", default-features = false, features = ["derive"] }
 bon = { version = "3.4.0", default-features = false }
 base64 = "0.22.1"
-hex = { version = "0.4.3", default-features = false, features = ["serde"] }
+hex = { version = "0.4.3", default-features = false }
 hex_fmt = "0.3.0"
 hex-literal = "1.0.0"
 prost = "0.13.5"

dstack/Cargo.toml

@@ -532,7 +532,8 @@ mod tests {
         img.seek(SeekFrom::Start(metadata.first_lba * SECTOR))?;
         img.read_exact(&mut header)?;
         assert_eq!(&header[..16], b"DSTACK_VOLUME\0\0\0");
-        assert_eq!(&header[32..64], &hex::decode(root_hash)?);
+        let decoded = DstackVolumeHeader::decode(&header)?;
+        assert_eq!(decoded.root_hash.as_slice(), hex::decode(root_hash)?);
         let mut buf = vec![0; data.len()];
         img.seek(SeekFrom::Start(data_partition.first_lba * SECTOR))?;
         img.read_exact(&mut buf)?;

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

@@ -133,7 +133,7 @@ pub struct AppCompose {
 pub struct VerityVolume {
     /// dm-verity root hash (hex): the volume's content identity and integrity
     /// check. The guest matches attached devices against it.
-    #[serde(with = "hex::serde")]
+    #[serde(with = "hex_bytes")]
     pub verity_root: [u8; 32],
     /// `"docker"` (seed the docker overlay2 image store), or an absolute path
     /// where the volume's filesystem is mounted (e.g. model weights).

dstack/dstack-types/src/lib.rs

@@ -10,34 +10,21 @@ use binrw::{binrw, BinRead, BinWrite};
 
 pub const DSTACK_VOLUME_MAGIC: &[u8; 16] = b"DSTACK_VOLUME\0\0\0";
 pub const DSTACK_VOLUME_HEADER_SIZE: usize = 4096;
-pub const DSTACK_VOLUME_FORMAT_VERSION: u16 = 1;
 pub const DSTACK_VOLUME_KIND_VERITY: u32 = 1;
 
 #[binrw]
 #[brw(little, magic = b"DSTACK_VOLUME\0\0\0")]
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct DstackVolumeHeader {
-    pub format_version: u16,
-    pub header_size: u16,
     pub kind: u32,
-    pub kind_version: u32,
-    pub flags: u32,
     pub root_hash: [u8; 32],
-    pub data_block_size: u32,
-    pub hash_block_size: u32,
 }
 
 impl DstackVolumeHeader {
     pub fn new_verity(root_hash: [u8; 32]) -> Self {
         Self {
-            format_version: DSTACK_VOLUME_FORMAT_VERSION,
-            header_size: DSTACK_VOLUME_HEADER_SIZE as u16,
             kind: DSTACK_VOLUME_KIND_VERITY,
-            kind_version: 1,
-            flags: 0,
             root_hash,
-            data_block_size: 4096,
-            hash_block_size: 4096,
         }
     }
 
@@ -57,23 +44,7 @@ impl DstackVolumeHeader {
                 ),
             });
         }
-        let header = Self::read(&mut Cursor::new(block))?;
-        if header.format_version != DSTACK_VOLUME_FORMAT_VERSION {
-            return Err(binrw::Error::AssertFail {
-                pos: 16,
-                message: format!(
-                    "unsupported volume format version {}",
-                    header.format_version
-                ),
-            });
-        }
-        if header.header_size as usize != DSTACK_VOLUME_HEADER_SIZE {
-            return Err(binrw::Error::AssertFail {
-                pos: 18,
-                message: format!("invalid volume header size {}", header.header_size),
-            });
-        }
-        Ok(header)
+        Self::read(&mut Cursor::new(block))
     }
 }
 
@@ -91,14 +62,18 @@ mod tests {
     }
 
     #[test]
-    fn verity_volume_validates_root_and_target_during_deserialization() {
+    fn verity_volume_validates_root_and_target_during_deserialization(
+    ) -> Result<(), serde_json::Error> {
         let volume: VerityVolume = serde_json::from_value(serde_json::json!({
             "verity_root": "5a".repeat(32),
             "target": "/run/models"
-        }))
-        .unwrap();
+        }))?;
         assert_eq!(volume.verity_root, [0x5a; 32]);
         assert_eq!(volume.target, VolumeTarget::Mount("/run/models".into()));
+        assert_eq!(
+            serde_json::to_value(&volume)?["verity_root"],
+            "5a".repeat(32)
+        );
 
         assert!(serde_json::from_value::<VerityVolume>(serde_json::json!({
             "verity_root": "abcd",
@@ -110,5 +85,6 @@ mod tests {
             "target": "relative/path"
         }))
         .is_err());
+        Ok(())
     }
 }

dstack/dstack-types/src/volume.rs

@@ -15,9 +15,9 @@ use std::io::Read;
 use std::os::unix::ffi::OsStrExt;
 use std::os::unix::fs::MetadataExt;
 use std::path::{Path, PathBuf};
-use std::process::{Command, Output};
 
 use anyhow::{bail, Context, Result};
+use cmd_lib::{run_cmd, run_fun};
 use dstack_types::volume::{
     DstackVolumeHeader, DSTACK_VOLUME_HEADER_SIZE, DSTACK_VOLUME_KIND_VERITY, DSTACK_VOLUME_MAGIC,
 };
@@ -43,7 +43,7 @@ struct VerityVolume {
 }
 
 fn main() -> Result<()> {
-    tracing_subscriber::fmt().with_target(false).init();
+    tracing_subscriber::fmt().init();
     let compose_path = std::env::args_os()
         .nth(1)
         .map(PathBuf::from)
@@ -56,8 +56,8 @@ fn main() -> Result<()> {
         return Ok(());
     }
 
-    let _ = run(Command::new("modprobe").arg("dm-verity"));
-    let _ = run(Command::new("udevadm").args(["settle", "--timeout=5"]));
+    let _ = run_cmd!(modprobe dm-verity);
+    let _ = run_cmd!(udevadm settle --timeout=5);
 
     let volumes = discover_volumes()?;
     info!(
@@ -194,19 +194,6 @@ fn parse_header(bytes: &[u8]) -> Result<Option<DstackVolumeHeader>> {
 }
 
 fn resolve_verity(disk: &BlockDisk, header: DstackVolumeHeader) -> Result<VerityVolume> {
-    if header.kind_version != 1 {
-        bail!("unsupported verity version {}", header.kind_version);
-    }
-    if header.flags != 0 {
-        bail!("unsupported verity flags {:#x}", header.flags);
-    }
-    if header.data_block_size != 4096 || header.hash_block_size != 4096 {
-        bail!(
-            "unsupported verity block sizes {}/{}",
-            header.data_block_size,
-            header.hash_block_size
-        );
-    }
     if disk.partitions.is_empty() {
         bail!("raw verity layout is not defined by kind version 1");
     }
@@ -249,27 +236,19 @@ fn activate_requested(
             info!(mapper = mapper_name, "reused active verity mapping");
             return Ok(());
         }
-        checked(
-            Command::new("veritysetup").args(["close", &mapper_name]),
-            "closing stale verity mapping",
-        )?;
+        run_cmd!(veritysetup close $mapper_name).context("closing stale verity mapping")?;
     }
 
     // The on-disk root only selected a candidate. Pass the root from the
     // measured compose to veritysetup, which is the actual trust decision.
-    checked(
-        Command::new("veritysetup")
-            .arg("open")
-            .arg(&candidate.data)
-            .arg(&mapper_name)
-            .arg(&candidate.hash)
-            .arg(&expected_root),
-        "opening dm-verity volume",
-    )?;
+    let data = &candidate.data;
+    let hash = &candidate.hash;
+    run_cmd!(veritysetup open $data $mapper_name $hash $expected_root)
+        .context("opening dm-verity volume")?;
     if let Err(err) =
         verify_first_block(&mapped).and_then(|_| mount_volume(index, requested, &mapped))
     {
-        let _ = run(Command::new("veritysetup").args(["close", &mapper_name]));
+        let _ = run_cmd!(veritysetup close $mapper_name);
         return Err(err);
     }
     used.insert(candidate_index);
@@ -284,12 +263,7 @@ fn verify_first_block(path: &Path) -> Result<()> {
 }
 
 fn mount_volume(index: usize, requested: &RequestedVolume, mapped: &Path) -> Result<()> {
-    let fs_type = command_stdout(
-        Command::new("blkid")
-            .args(["-o", "value", "-s", "TYPE"])
-            .arg(mapped),
-    )
-    .unwrap_or_default();
+    let fs_type = run_fun!(blkid -o value -s TYPE $mapped).unwrap_or_default();
     match &requested.target {
         VolumeTarget::DockerSeed => {
             let mountpoint = PathBuf::from(format!("/run/dstack-verity/{index}"));
@@ -300,7 +274,7 @@ fn mount_volume(index: usize, requested: &RequestedVolume, mapped: &Path) -> Res
                 mount_read_only(mapped, &mountpoint, fs_type.trim())?;
             }
             if let Err(err) = seed_docker(&mountpoint) {
-                let _ = run(Command::new("umount").arg(&mountpoint));
+                let _ = run_cmd!(umount $mountpoint);
                 return Err(err);
             }
             info!(root = %hex::encode(requested.verity_root), "seeded docker from verity volume");
@@ -331,14 +305,12 @@ fn mount_read_only(device: &Path, target: &Path, fs_type: &str) -> Result<()> {
     } else {
         "ro"
     };
-    let mut command = Command::new("mount");
-    if !fs_type.is_empty() {
-        command.args(["-t", fs_type]);
-    }
-    checked(
-        command.arg("-o").arg(options).arg(device).arg(target),
-        "mounting verity volume",
-    )?;
+    if fs_type.is_empty() {
+        run_cmd!(mount -o $options $device $target).context("mounting verity volume")?;
+    } else {
+        run_cmd!(mount -t $fs_type -o $options $device $target)
+            .context("mounting verity volume")?;
+    }
     Ok(())
 }
 
@@ -360,23 +332,15 @@ fn seed_docker(volume: &Path) -> Result<()> {
         let target = store.join("overlay2").join(layer_id).join("diff");
         fs::create_dir_all(&target)?;
         if !is_mountpoint(&target)? {
-            if let Err(err) = checked(
-                Command::new("mount")
-                    .args(["--bind"])
-                    .arg(&source)
-                    .arg(&target),
-                "binding docker layer",
-            ) {
+            if let Err(err) = run_cmd!(mount --bind $source $target).context("binding docker layer")
+            {
                 unwind_binds(&bound);
                 return Err(err);
             }
             // A bind inherits neither the intended policy nor all mount flags.
-            if let Err(err) = checked(
-                Command::new("mount")
-                    .args(["-o", "remount,bind,ro"])
-                    .arg(&target),
-                "making docker layer read-only",
-            ) {
+            if let Err(err) =
+                run_cmd!(mount -o remount,bind,ro $target).context("making docker layer read-only")
+            {
                 bound.push(target);
                 unwind_binds(&bound);
                 return Err(err);
@@ -419,13 +383,8 @@ fn child_directories(path: &Path) -> Result<Vec<PathBuf>> {
 
 fn copy_contents(source: &Path, target: &Path) -> Result<()> {
     fs::create_dir_all(target)?;
-    checked(
-        Command::new("cp")
-            .args(["-a"])
-            .arg(source.join("."))
-            .arg(target),
-        "copying docker metadata",
-    )?;
+    let source_contents = source.join(".");
+    run_cmd!(cp -a $source_contents $target).context("copying docker metadata")?;
     Ok(())
 }
 
@@ -474,10 +433,7 @@ fn copy_layer_metadata(volume: &Path, store: &Path) -> Result<()> {
             if source.file_name() == Some(OsStr::new("diff")) {
                 continue;
             }
-            checked(
-                Command::new("cp").args(["-a"]).arg(&source).arg(&target),
-                "copying docker layer metadata",
-            )?;
+            run_cmd!(cp -a $source $target).context("copying docker layer metadata")?;
         }
     }
     Ok(())
@@ -549,20 +505,12 @@ fn unescape_mountinfo(value: &[u8]) -> Vec<u8> {
 
 fn unwind_binds(paths: &[PathBuf]) {
     for path in paths.iter().rev() {
-        let _ = run(Command::new("umount").arg(path));
-    }
-}
-
-fn command_stdout(command: &mut Command) -> Result<String> {
-    let output = command.output()?;
-    if !output.status.success() {
-        bail!("command failed with {}", output.status);
+        let _ = run_cmd!(umount $path);
     }
-    Ok(String::from_utf8(output.stdout)?)
 }
 
 fn mapping_root(mapper_name: &str) -> Result<String> {
-    let status = command_stdout(Command::new("veritysetup").args(["status", mapper_name]))?;
+    let status = run_fun!(veritysetup status $mapper_name)?;
     status
         .lines()
         .find_map(|line| {
@@ -574,22 +522,6 @@ fn mapping_root(mapper_name: &str) -> Result<String> {
         .context("verity mapping status has no root hash")
 }
 
-fn run(command: &mut Command) -> Result<Output> {
-    command.output().context("running command")
-}
-
-fn checked(command: &mut Command, operation: &str) -> Result<Output> {
-    let output = command.output().with_context(|| operation.to_string())?;
-    if !output.status.success() {
-        bail!(
-            "{operation} failed with {}: {}",
-            output.status,
-            String::from_utf8_lossy(&output.stderr).trim()
-        );
-    }
-    Ok(output)
-}
-
 #[cfg(test)]
 mod tests {
     use super::*;