Merge pull request #678 from Dstack-TEE/fix/dstack-mr-ovmf-202505-events

dstack-mr: support edk2-stable202505 OVMF event layout

Author: kvinwang

dstack-mr/Cargo.toml

@@ -25,9 +25,9 @@ fs-err.workspace = true
 bon.workspace = true
 log.workspace = true
 scale.workspace = true
+dstack-types.workspace = true
 
 [dev-dependencies]
-dstack-types.workspace = true
 reqwest = { workspace = true, features = ["blocking"] }
 flate2.workspace = true
 tar.workspace = true

dstack-mr/cli/src/main.rs

@@ -4,7 +4,7 @@
 
 use anyhow::{Context, Result};
 use clap::{Parser, Subcommand};
-use dstack_mr::Machine;
+use dstack_mr::{Machine, OvmfVariant, ovmf_variant_for_version};
 use dstack_types::ImageInfo;
 use fs_err as fs;
 use size_parser::parse_memory_size;
@@ -78,6 +78,12 @@ struct MachineConfig {
     #[arg(long)]
     qemu_version: Option<String>,
 
+    /// dstack OS version (MAJOR.MINOR.PATCH), used to pick the OVMF measurement layout.
+    /// 0.5.10 <= ver < 0.6.0 and ver >= 0.6.1 use the edk2-stable202505 layout; everything
+    /// else uses the legacy layout. If omitted, falls back to `image_info.version`.
+    #[arg(long)]
+    dstack_os_version: Option<String>,
+
     /// Output JSON
     #[arg(long)]
     json: bool,
@@ -99,6 +105,21 @@ fn main() -> Result<()> {
             let initrd_path = parent_dir.join(&image_info.initrd).display().to_string();
             let cmdline = image_info.cmdline + " initrd=initrd";
 
+            // CLI flag wins, then the explicit `ovmf_variant` in metadata.json,
+            // and finally the OS version field. Older metadata.json files may
+            // carry neither, in which case fall back to the default.
+            let ovmf_variant = if let Some(v) = config.dstack_os_version.as_deref() {
+                ovmf_variant_for_version(v)
+                    .with_context(|| format!("invalid dstack OS version: {v}"))?
+            } else if let Some(variant) = image_info.ovmf_variant {
+                variant
+            } else if !image_info.version.is_empty() {
+                ovmf_variant_for_version(&image_info.version)
+                    .with_context(|| format!("invalid dstack OS version: {}", image_info.version))?
+            } else {
+                OvmfVariant::default()
+            };
+
             let machine = Machine::builder()
                 .cpu_count(config.cpu)
                 .memory_size(config.memory)
@@ -116,6 +137,7 @@ fn main() -> Result<()> {
                 .hotplug_off(config.hotplug_off)
                 .root_verity(config.root_verity)
                 .maybe_qemu_version(config.qemu_version.clone())
+                .ovmf_variant(ovmf_variant)
                 .build();
 
             let measurements = machine

dstack-mr/src/lib.rs

@@ -2,9 +2,11 @@
 //
 // SPDX-License-Identifier: Apache-2.0
 
+use anyhow::{bail, Context, Result};
 use serde::{Deserialize, Serialize};
 use serde_human_bytes as hex_bytes;
 
+pub use dstack_types::OvmfVariant;
 pub use machine::{Machine, TdxMeasurementDetails};
 
 use util::{measure_log, measure_sha384, utf16_encode};
@@ -17,8 +19,188 @@ mod kernel;
 mod machine;
 mod num;
 mod tdvf;
+mod uefi_var;
 mod util;
 
+/// Pick the OVMF variant for a given dstack OS version string ("MAJOR.MINOR.PATCH").
+///
+/// Treats `0.5.10 <= v < 0.6.0` and `v >= 0.6.1` as `Stable202505`, everything else as
+/// `Pre202505`. Used as a fallback when `VmConfig::ovmf_variant` is absent.
+pub fn ovmf_variant_for_version(version: &str) -> Result<OvmfVariant> {
+    let parts: Vec<u32> = version
+        .split('.')
+        .map(|p| {
+            p.parse::<u32>()
+                .with_context(|| format!("invalid version component: {p}"))
+        })
+        .collect::<Result<Vec<_>, _>>()?;
+    if parts.len() != 3 {
+        bail!("expected MAJOR.MINOR.PATCH, got {version}");
+    }
+    let v = (parts[0], parts[1], parts[2]);
+    let stable = ((0, 5, 10)..(0, 6, 0)).contains(&v) || v >= (0, 6, 1);
+    Ok(if stable {
+        OvmfVariant::Stable202505
+    } else {
+        OvmfVariant::Pre202505
+    })
+}
+
+/// Extract the `MAJOR.MINOR.PATCH` version suffix from a dstack image name.
+///
+/// Recognises any `<prefix>-MAJOR.MINOR.PATCH[.SUFFIX]` shape, e.g.
+/// `dstack-0.5.10`, `dstack-dev-0.5.10`, `dstack-nvidia-0.5.10`,
+/// `dstack-nvidia-dev-0.5.10`, `dstack-0.5.10.rc1`, `dstack-dev-0.6.1.dev`.
+///
+/// The optional `.SUFFIX` is permitted to be non-numeric (pre-release tag,
+/// build label, etc.) and is dropped from the returned slice — only the
+/// numeric `X.Y.Z` is needed to pick the OVMF variant.
+///
+/// Returns `None` when the segment after the last `-` is not at least a valid
+/// `X.Y.Z` triple of non-empty numeric components.
+pub fn extract_version_from_image_name(image: &str) -> Option<&str> {
+    let tail = image.rsplit('-').next()?;
+    let parts: Vec<&str> = tail.split('.').collect();
+    if !(3..=4).contains(&parts.len()) {
+        return None;
+    }
+    let core_numeric = parts[..3]
+        .iter()
+        .all(|p| !p.is_empty() && p.parse::<u32>().is_ok());
+    let suffix_ok = parts.len() == 3 || !parts[3].is_empty();
+    if !(core_numeric && suffix_ok) {
+        return None;
+    }
+    // Slice off the optional `.SUFFIX` so callers get just `X.Y.Z`.
+    let core_len = parts[0].len() + 1 + parts[1].len() + 1 + parts[2].len();
+    Some(&tail[..core_len])
+}
+
+/// Pick the OVMF variant from an image name like `dstack-0.5.10`.
+///
+/// Falls back to `OvmfVariant::default()` (= `Pre202505`) when the image name is
+/// missing or doesn't carry a parseable version suffix. Use this only as a
+/// fallback for images that pre-date `VmConfig::ovmf_variant`.
+pub fn ovmf_variant_for_image(image: Option<&str>) -> OvmfVariant {
+    image
+        .and_then(extract_version_from_image_name)
+        .and_then(|v| ovmf_variant_for_version(v).ok())
+        .unwrap_or_default()
+}
+
+#[cfg(test)]
+mod ovmf_variant_tests {
+    use super::*;
+
+    #[test]
+    fn pre_202505_for_old_versions() {
+        for v in ["0.4.99", "0.5.7", "0.5.8", "0.5.9", "0.6.0"] {
+            assert_eq!(
+                ovmf_variant_for_version(v).unwrap(),
+                OvmfVariant::Pre202505,
+                "{v}"
+            );
+        }
+    }
+
+    #[test]
+    fn stable_202505_for_new_versions() {
+        for v in ["0.5.10", "0.5.99", "0.6.1", "0.6.2", "0.7.0", "1.0.0"] {
+            assert_eq!(
+                ovmf_variant_for_version(v).unwrap(),
+                OvmfVariant::Stable202505,
+                "{v}"
+            );
+        }
+    }
+
+    #[test]
+    fn rejects_malformed_version() {
+        assert!(ovmf_variant_for_version("0.5").is_err());
+        assert!(ovmf_variant_for_version("0.5.10-dev").is_err());
+        assert!(ovmf_variant_for_version("v0.5.10").is_err());
+    }
+
+    #[test]
+    fn parses_version_from_image_name() {
+        // Three-segment plain versions across the known prefix shapes.
+        assert_eq!(
+            extract_version_from_image_name("dstack-0.5.10"),
+            Some("0.5.10")
+        );
+        assert_eq!(
+            extract_version_from_image_name("dstack-dev-0.5.10"),
+            Some("0.5.10")
+        );
+        assert_eq!(
+            extract_version_from_image_name("dstack-nvidia-0.5.10"),
+            Some("0.5.10")
+        );
+        assert_eq!(
+            extract_version_from_image_name("dstack-nvidia-dev-0.6.1"),
+            Some("0.6.1")
+        );
+        // Optional .SUFFIX is allowed and dropped; suffix may be non-numeric.
+        assert_eq!(
+            extract_version_from_image_name("dstack-0.5.10.rc1"),
+            Some("0.5.10")
+        );
+        assert_eq!(
+            extract_version_from_image_name("dstack-dev-0.6.1.dev"),
+            Some("0.6.1")
+        );
+        assert_eq!(
+            extract_version_from_image_name("dstack-nvidia-dev-0.5.10.1"),
+            Some("0.5.10")
+        );
+        // Rejections.
+        assert_eq!(extract_version_from_image_name("dstack"), None);
+        assert_eq!(extract_version_from_image_name("dstack-rc1"), None);
+        assert_eq!(extract_version_from_image_name("dstack-0.5"), None);
+        assert_eq!(extract_version_from_image_name("dstack-0.5.10."), None);
+        assert_eq!(extract_version_from_image_name("dstack-0.5.10.1.2"), None);
+        assert_eq!(extract_version_from_image_name("dstack-0..10"), None);
+        assert_eq!(extract_version_from_image_name("dstack-a.b.c"), None);
+    }
+
+    #[test]
+    fn ovmf_variant_for_image_handles_missing_and_unknown() {
+        assert_eq!(ovmf_variant_for_image(None), OvmfVariant::Pre202505);
+        assert_eq!(
+            ovmf_variant_for_image(Some("dstack")),
+            OvmfVariant::Pre202505
+        );
+        assert_eq!(
+            ovmf_variant_for_image(Some("dstack-0.5.9")),
+            OvmfVariant::Pre202505
+        );
+        assert_eq!(
+            ovmf_variant_for_image(Some("dstack-0.5.10")),
+            OvmfVariant::Stable202505
+        );
+        assert_eq!(
+            ovmf_variant_for_image(Some("dstack-nvidia-dev-0.6.1")),
+            OvmfVariant::Stable202505
+        );
+    }
+
+    #[test]
+    fn serializes_with_snake_case() {
+        assert_eq!(
+            serde_json::to_string(&OvmfVariant::Pre202505).unwrap(),
+            "\"pre202505\""
+        );
+        assert_eq!(
+            serde_json::to_string(&OvmfVariant::Stable202505).unwrap(),
+            "\"stable202505\""
+        );
+        assert_eq!(
+            serde_json::from_str::<OvmfVariant>("\"stable202505\"").unwrap(),
+            OvmfVariant::Stable202505
+        );
+    }
+}
+
 /// Contains all the measurement values for TDX.
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct TdxMeasurements {

dstack-mr/src/machine.rs

@@ -5,7 +5,7 @@
 use crate::acpi::Tables;
 use crate::tdvf::Tdvf;
 use crate::util::debug_print_log;
-use crate::{kernel, RtmrLogs, TdxMeasurements};
+use crate::{kernel, OvmfVariant, RtmrLogs, TdxMeasurements};
 use crate::{measure_log, measure_sha384};
 use anyhow::{bail, Context, Result};
 use fs_err as fs;
@@ -32,6 +32,10 @@ pub struct Machine<'a> {
     pub root_verity: bool,
     #[builder(default)]
     pub host_share_mode: String,
+    /// Selects which OVMF measurement event layout to expect.
+    /// Defaults to the pre-edk2-stable202505 layout for backwards compatibility.
+    #[builder(default)]
+    pub ovmf_variant: OvmfVariant,
 }
 
 fn parse_version_tuple(v: &str) -> Result<(u32, u32, u32)> {