feat(dstack-mr): add UEFI disk boot (UKI) measurement support

dstack-mr/cli/src/main.rs

@@ -81,6 +81,20 @@ struct MachineConfig {
     /// Output JSON
     #[arg(long)]
     json: bool,
+
+    // --- UEFI disk boot (UKI mode) ---
+    /// Path to UKI EFI binary. When set, uses UEFI disk boot measurement path
+    /// instead of -kernel mode. Boot chain: OVMF → systemd-boot → UKI → vmlinuz.
+    #[arg(long)]
+    uki: Option<PathBuf>,
+
+    /// Path to systemd-boot EFI binary (BOOTX64.EFI). Required with --uki.
+    #[arg(long)]
+    bootloader: Option<PathBuf>,
+
+    /// Path to raw disk image (for GPT event measurement). Required with --uki.
+    #[arg(long)]
+    disk: Option<PathBuf>,
 }
 
 fn main() -> Result<()> {
@@ -97,7 +111,18 @@ fn main() -> Result<()> {
             let firmware_path = parent_dir.join(&image_info.bios).display().to_string();
             let kernel_path = parent_dir.join(&image_info.kernel).display().to_string();
             let initrd_path = parent_dir.join(&image_info.initrd).display().to_string();
-            let cmdline = image_info.cmdline + " initrd=initrd";
+            // In -kernel mode, QEMU appends " initrd=initrd" to cmdline.
+            // In UKI mode, cmdline is embedded in the UKI as-is (no append).
+            let cmdline = if config.uki.is_some() {
+                image_info.cmdline
+            } else {
+                image_info.cmdline + " initrd=initrd"
+            };
+
+            // Resolve UKI-related paths
+            let uki_path = config.uki.as_ref().map(|p| p.display().to_string());
+            let bootloader_path = config.bootloader.as_ref().map(|p| p.display().to_string());
+            let disk_path = config.disk.as_ref().map(|p| p.display().to_string());
 
             let machine = Machine::builder()
                 .cpu_count(config.cpu)
@@ -116,6 +141,9 @@ fn main() -> Result<()> {
                 .hotplug_off(config.hotplug_off)
                 .root_verity(config.root_verity)
                 .maybe_qemu_version(config.qemu_version.clone())
+                .maybe_uki(uki_path.as_deref())
+                .maybe_bootloader(bootloader_path.as_deref())
+                .maybe_disk(disk_path.as_deref())
                 .build();
 
             let measurements = machine

dstack-mr/src/acpi.rs

@@ -23,6 +23,101 @@ pub struct Tables {
 
 impl Machine<'_> {
     fn create_tables(&self) -> Result<Vec<u8>> {
+        if self.is_uki_mode() {
+            return self.create_tables_uefi_disk();
+        }
+        self.create_tables_direct_kernel()
+    }
+
+    /// Generate ACPI tables for UEFI disk boot (UKI mode).
+    /// Device config: single disk (hd0) + net. No vsock, no 9p, no -kernel.
+    fn create_tables_uefi_disk(&self) -> Result<Vec<u8>> {
+        if self.cpu_count == 0 {
+            bail!("cpuCount must be greater than 0");
+        }
+        let mem_size_mb = self.memory_size / (1024 * 1024);
+        let dummy_disk = "/bin/sh";
+
+        let mut cmd = std::process::Command::new("dstack-acpi-tables");
+        cmd.args([
+            "-cpu",
+            "qemu64",
+            "-smp",
+            &self.cpu_count.to_string(),
+            "-m",
+            &format!("{mem_size_mb}M"),
+            "-nographic",
+            "-nodefaults",
+            "-serial",
+            "stdio",
+            "-bios",
+            self.firmware,
+            // No -kernel/-initrd: UEFI disk boot loads from ESP
+            "-drive",
+            &format!("file={dummy_disk},if=none,id=hd0,format=raw,readonly=on"),
+            "-device",
+            "virtio-blk-pci,drive=hd0",
+            // Second disk (data disk) — separate from OS disk for RTMR stability
+            "-drive",
+            &format!("file={dummy_disk},if=none,id=hd1,format=raw,readonly=on"),
+            "-device",
+            "virtio-blk-pci,drive=hd1",
+            "-netdev",
+            "user,id=net0",
+            "-device",
+            "virtio-net-pci,netdev=net0",
+            "-object",
+            "tdx-guest,id=tdx",
+            // No vsock, no 9p — UKI mode uses simple device config
+        ]);
+
+        let mut machine =
+            "q35,kernel-irqchip=split,confidential-guest-support=tdx,hpet=off".to_string();
+        if self.smm {
+            machine.push_str(",smm=on");
+        } else {
+            machine.push_str(",smm=off");
+        }
+        let vopt = self
+            .versioned_options()
+            .context("Failed to get versioned options")?;
+        if vopt.pic {
+            machine.push_str(",pic=on");
+        } else {
+            machine.push_str(",pic=off");
+        }
+        cmd.args(["-machine", &machine]);
+
+        if self.hotplug_off {
+            cmd.args([
+                "-global",
+                "ICH9-LPC.acpi-pci-hotplug-with-bridge-support=off",
+            ]);
+        }
+        if let Some(pci_hole64_size) = self.pci_hole64_size {
+            cmd.args([
+                "-global",
+                &format!("q35-pcihost.pci-hole64-size=0x{:x}", pci_hole64_size),
+            ]);
+        }
+
+        let ver = vopt.version;
+        let output = cmd
+            .env(
+                "QEMU_ACPI_COMPAT_VER",
+                format!("{}.{}.{}", ver.0, ver.1, ver.2),
+            )
+            .output()
+            .context("failed to execute dstack-acpi-tables")?;
+
+        if !output.status.success() {
+            let stderr = String::from_utf8_lossy(&output.stderr);
+            bail!("dstack-acpi-tables failed: {stderr}");
+        }
+        Ok(output.stdout)
+    }
+
+    fn create_tables_direct_kernel(&self) -> Result<Vec<u8>> {
         if self.cpu_count == 0 {
             bail!("cpuCount must be greater than 0");
         }

dstack-mr/src/kernel.rs

@@ -7,8 +7,9 @@ use anyhow::{bail, Context, Result};
 use object::pe;
 use sha2::{Digest, Sha384};
 
-/// Calculates the Authenticode hash of a PE/COFF file
-fn authenticode_sha384_hash(data: &[u8]) -> Result<Vec<u8>> {
+/// Calculates the Authenticode hash of a PE/COFF file.
+/// Used by both direct kernel boot (patched kernel) and UEFI disk boot (EFI apps).
+pub(crate) fn authenticode_sha384_hash(data: &[u8]) -> Result<Vec<u8>> {
     let lfanew_offset = 0x3c;
     let lfanew: u32 = read_le(data, lfanew_offset, "DOS header")?;
 

dstack-mr/src/lib.rs

@@ -17,6 +17,7 @@ mod kernel;
 mod machine;
 mod num;
 mod tdvf;
+pub(crate) mod uefi_boot;
 mod util;
 
 /// Contains all the measurement values for TDX.

dstack-mr/src/machine.rs

@@ -32,6 +32,18 @@ pub struct Machine<'a> {
     pub root_verity: bool,
     #[builder(default)]
     pub host_share_mode: String,
+
+    // --- UEFI disk boot (UKI) mode ---
+    // When `uki` is set, use UEFI disk boot measurement path instead of -kernel mode.
+    // In this mode: OVMF → systemd-boot → UKI → vmlinuz (3 EFI app measurements).
+    // The `kernel` field points to vmlinuz, `initrd` to the dracut initrd,
+    // and `kernel_cmdline` to the embedded cmdline — all extracted from the UKI.
+    /// Path to UKI EFI binary. When set, enables UEFI disk boot measurement mode.
+    pub uki: Option<&'a str>,
+    /// Path to systemd-boot EFI binary (BOOTX64.EFI). Required when `uki` is set.
+    pub bootloader: Option<&'a str>,
+    /// Path to raw disk image (for GPT event). Required when `uki` is set.
+    pub disk: Option<&'a str>,
 }
 
 fn parse_version_tuple(v: &str) -> Result<(u32, u32, u32)> {
@@ -93,6 +105,11 @@ impl Machine<'_> {
         self.measure_with_logs().map(|details| details.measurements)
     }
 
+    /// Returns true if this machine is configured for UEFI disk boot (UKI mode).
+    pub fn is_uki_mode(&self) -> bool {
+        self.uki.is_some()
+    }
+
     pub fn measure_with_logs(&self) -> Result<TdxMeasurementDetails> {
         debug!("measuring machine: {self:#?}");
         let fw_data = fs::read(self.firmware)?;
@@ -102,14 +119,29 @@ impl Machine<'_> {
 
         let mrtd = tdvf.mrtd(self).context("Failed to compute MR TD")?;
 
+        if self.is_uki_mode() {
+            self.measure_uefi_disk_boot(&tdvf, &kernel_data, &initrd_data, mrtd)
+        } else {
+            self.measure_direct_kernel_boot(&tdvf, &kernel_data, &initrd_data, mrtd)
+        }
+    }
+
+    /// Direct kernel boot (-kernel mode): original dstack measurement path.
+    fn measure_direct_kernel_boot(
+        &self,
+        tdvf: &Tdvf,
+        kernel_data: &[u8],
+        initrd_data: &[u8],
+        mrtd: Vec<u8>,
+    ) -> Result<TdxMeasurementDetails> {
         let (rtmr0_log, acpi_tables) = tdvf
             .rtmr0_log(self)
             .context("Failed to compute RTMR0 log")?;
         debug_print_log("RTMR0", &rtmr0_log);
         let rtmr0 = measure_log(&rtmr0_log);
 
         let rtmr1_log = kernel::rtmr1_log(
-            &kernel_data,
+            kernel_data,
             initrd_data.len() as u32,
             self.memory_size,
             0x28000,
@@ -119,7 +151,7 @@ impl Machine<'_> {
 
         let rtmr2_log = vec![
             kernel::measure_cmdline(self.kernel_cmdline),
-            measure_sha384(&initrd_data),
+            measure_sha384(initrd_data),
         ];
         debug_print_log("RTMR2", &rtmr2_log);
         let rtmr2 = measure_log(&rtmr2_log);
@@ -135,4 +167,57 @@ impl Machine<'_> {
             acpi_tables,
         })
     }
+
+    /// UEFI disk boot (UKI mode): OVMF → systemd-boot → UKI → vmlinuz.
+    /// Verified against TDX hardware CCEL event log.
+    fn measure_uefi_disk_boot(
+        &self,
+        tdvf: &Tdvf,
+        kernel_data: &[u8],
+        initrd_data: &[u8],
+        mrtd: Vec<u8>,
+    ) -> Result<TdxMeasurementDetails> {
+        let uki_path = self
+            .uki
+            .context("UKI path required for UEFI disk boot mode")?;
+        let bootloader_path = self
+            .bootloader
+            .context("bootloader path required for UEFI disk boot mode")?;
+
+        let uki_data = fs::read(uki_path)?;
+        let bootloader_data = fs::read(bootloader_path)?;
+
+        // RTMR[0]: same structure as direct kernel boot but with UEFI disk boot
+        // differences in boot variables (BootOrder has 2 entries, Boot0001 added).
+        let (rtmr0_log, acpi_tables) = tdvf
+            .rtmr0_log_uefi_disk(self, &bootloader_data)
+            .context("Failed to compute RTMR0 log for UEFI disk boot")?;
+        debug_print_log("RTMR0", &rtmr0_log);
+        let rtmr0 = measure_log(&rtmr0_log);
+
+        // RTMR[1]: EFI application measurements (verified against hardware).
+        // Events: Calling EFI App, separator, GPT, systemd-boot, UKI, vmlinuz,
+        //         Exit Boot Services Invocation, Exit Boot Services Returned.
+        let rtmr1_log =
+            crate::uefi_boot::rtmr1_log(&bootloader_data, &uki_data, kernel_data, self.disk)?;
+        debug_print_log("RTMR1", &rtmr1_log);
+        let rtmr1 = measure_log(&rtmr1_log);
+
+        // RTMR[2]: cmdline (UTF-16LE) + initrd data hash.
+        // Linux EFI stub converts cmdline to UTF-16LE before measuring.
+        let rtmr2_log = crate::uefi_boot::rtmr2_log(self.kernel_cmdline, initrd_data);
+        debug_print_log("RTMR2", &rtmr2_log);
+        let rtmr2 = measure_log(&rtmr2_log);
+
+        Ok(TdxMeasurementDetails {
+            measurements: TdxMeasurements {
+                mrtd,
+                rtmr0,
+                rtmr1,
+                rtmr2,
+            },
+            rtmr_logs: [rtmr0_log, rtmr1_log, rtmr2_log],
+            acpi_tables,
+        })
+    }
 }