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

Add --uki flag to enable UEFI disk boot measurement path:
  OVMF → systemd-boot → UKI → vmlinuz (3 EFI app measurements)

When --uki is set, RTMR calculations use the UEFI disk boot event model
instead of the -kernel direct boot model:

  RTMR[0]: Same TD HOB/CFV/EFI vars/ACPI, but BootOrder has 2 entries
           and Boot0001 (disk device) is added. Boot variable digests
           use SHA384(variable_data_bytes).

  RTMR[1]: 8 events — Calling EFI App, separator, GPT event (non-empty
           partition entries only), systemd-boot/UKI/vmlinuz Authenticode
           hashes, ExitBootServices.

  RTMR[2]: 2 events — cmdline as UTF-16LE (Linux EFI stub behavior),
           initrd data hash. No "initrd=initrd" append (unlike -kernel).

Verified against TDX hardware: MRTD, RTMR[1], RTMR[2] all match.
RTMR[0] matches except ACPI tables (dstack-acpi-tables QEMU version
difference — needs matching QEMU_ACPI_COMPAT_VER).

Usage:
  dstack-mr measure metadata.json -c 2 -m 4G \
    --uki dstack-k8s.efi \
    --bootloader BOOTX64.EFI \
    --disk dstack-k8s.raw

Author: kvinwang

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/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,
+        })
+    }
 }

dstack-mr/src/tdvf.rs

@@ -304,6 +304,106 @@ impl<'a> Tdvf<'a> {
         ))
     }
 
+    /// RTMR[0] event log for UEFI disk boot (UKI mode).
+    ///
+    /// Differences from direct kernel boot (-kernel mode):
+    /// - EFI variables have different values (SecureBoot not enforced, empty PK/KEK/db/dbx)
+    ///   but the digest calculation is the same (measure_tdx_efi_variable with data_len=0)
+    /// - BootOrder contains 2 entries [0x0000, 0x0001] instead of 1
+    ///   Digest = SHA384(variable_data_bytes), not SHA384(UEFI_VARIABLE_DATA struct)
+    /// - Boot0001 added: UEFI Misc Device (virtio-blk disk)
+    ///   Digest = SHA384(EFI_LOAD_OPTION bytes)
+    /// - Boot0000 is the same (UiApp from OVMF)
+    /// - Second separator at the end
+    ///
+    /// Verified against TDX hardware CCEL event log.
+    pub fn rtmr0_log_uefi_disk(
+        &self,
+        machine: &Machine,
+        _bootloader_data: &[u8],
+    ) -> Result<(RtmrLog, Tables)> {
+        let td_hob_hash = self.measure_td_hob(machine.memory_size)?;
+        let cfv_image_hash = hex!(
+            "344BC51C980BA621AAA00DA3ED7436F7D6E549197DFE699515DFA2C6583D95E6412AF21C097D473155875FFD561D6790"
+        );
+        // Boot0000: UiApp from OVMF (fixed, same in both boot modes).
+        // Digest = SHA384(EFI_LOAD_OPTION variable data).
+        let boot000_hash = hex!(
+            "23ADA07F5261F12F34A0BD8E46760962D6B4D576A416F1FEA1C64BC656B1D28EACF7047AE6E967C58FD2A98BFA74C298"
+        );
+
+        let tables = machine.build_tables()?;
+        let acpi_tables_hash = measure_sha384(&tables.tables);
+        let acpi_rsdp_hash = measure_sha384(&tables.rsdp);
+        let acpi_loader_hash = measure_sha384(&tables.loader);
+
+        // BootOrder for UEFI disk boot: [0x0000, 0x0001] (2 boot entries).
+        // Digest = SHA384(variable_data_bytes), NOT SHA384(UEFI_VARIABLE_DATA struct).
+        let boot_order_data: [u8; 4] = [0x00, 0x00, 0x01, 0x00]; // LE u16: 0, 1
+        let boot_order_hash = measure_sha384(&boot_order_data);
+
+        // Boot0001: "UEFI Misc Device" at PciRoot(0x0)/Pci(0x1,0x0) — first virtio-blk.
+        // This is the EFI_LOAD_OPTION structure for the disk boot entry.
+        // The device path encodes PciRoot(0x0)/Pci(0x1,0x0)/End.
+        // Digest = SHA384(EFI_LOAD_OPTION variable data).
+        //
+        // EFI_LOAD_OPTION layout:
+        //   Attributes(4): 0x00000001 (LOAD_OPTION_ACTIVE)
+        //   FilePathListLength(2): 0x0016 (22 bytes)
+        //   Description(UTF-16LE+null): "UEFI Misc Device\0"
+        //   FilePathList: PciRoot(0x0)/Pci(0x1,0x0)/End
+        //   OptionalData: VenMedia GUID (virtio device signature)
+        let boot0001_data: Vec<u8> = {
+            let mut d = Vec::new();
+            // Attributes: LOAD_OPTION_ACTIVE
+            d.extend_from_slice(&0x00000001u32.to_le_bytes());
+            // FilePathListLength
+            d.extend_from_slice(&0x0016u16.to_le_bytes());
+            // Description: "UEFI Misc Device" in UTF-16LE + null terminator
+            for c in "UEFI Misc Device".encode_utf16() {
+                d.extend_from_slice(&c.to_le_bytes());
+            }
+            d.extend_from_slice(&[0x00, 0x00]); // null terminator
+                                                // FilePathList: PciRoot(0x0)/Pci(0x1,0x0)/End
+                                                // ACPI device path: type=0x02, subtype=0x01, length=0x0c, HID=0x0a0341d0(PNP0A03), UID=0
+            d.extend_from_slice(&[0x02, 0x01, 0x0c, 0x00]);
+            d.extend_from_slice(&0x0a0341d0u32.to_le_bytes()); // PNP0A03 (PCI root)
+            d.extend_from_slice(&0x00000000u32.to_le_bytes()); // UID=0
+                                                               // PCI device path: type=0x01, subtype=0x01, length=0x06, Function=0, Device=1
+            d.extend_from_slice(&[0x01, 0x01, 0x06, 0x00, 0x00, 0x01]);
+            // End device path: type=0x7f, subtype=0xff, length=0x04
+            d.extend_from_slice(&[0x7f, 0xff, 0x04, 0x00]);
+            // OptionalData: VenMedia GUID for virtio device signature
+            // 4e ac 08 81 11 9f 59 4d 85 0e e2 1a 52 2c 59 b2
+            d.extend_from_slice(&[
+                0x4e, 0xac, 0x08, 0x81, 0x11, 0x9f, 0x59, 0x4d, 0x85, 0x0e, 0xe2, 0x1a, 0x52, 0x2c,
+                0x59, 0xb2,
+            ]);
+            d
+        };
+        let boot0001_hash = measure_sha384(&boot0001_data);
+
+        Ok((
+            vec![
+                td_hob_hash,
+                cfv_image_hash.to_vec(),
+                measure_tdx_efi_variable("8BE4DF61-93CA-11D2-AA0D-00E098032B8C", "SecureBoot")?,
+                measure_tdx_efi_variable("8BE4DF61-93CA-11D2-AA0D-00E098032B8C", "PK")?,
+                measure_tdx_efi_variable("8BE4DF61-93CA-11D2-AA0D-00E098032B8C", "KEK")?,
+                measure_tdx_efi_variable("D719B2CB-3D3A-4596-A3BC-DAD00E67656F", "db")?,
+                measure_tdx_efi_variable("D719B2CB-3D3A-4596-A3BC-DAD00E67656F", "dbx")?,
+                measure_sha384(&[0x00, 0x00, 0x00, 0x00]), // Separator
+                acpi_loader_hash,
+                acpi_rsdp_hash,
+                acpi_tables_hash,
+                boot_order_hash,
+                boot000_hash.to_vec(),
+                boot0001_hash,
+            ],
+            tables,
+        ))
+    }
+
     fn measure_td_hob(&self, memory_size: u64) -> Result<Vec<u8>> {
         let mut memory_acceptor = MemoryAcceptor::new(0, memory_size);
         let mut td_hob = Vec::new();