Skip to main content

petri/vm/
mod.rs

1// Copyright (c) Microsoft Corporation.
2// Licensed under the MIT License.
3
4/// Hyper-V VM management
5#[cfg(windows)]
6pub mod hyperv;
7/// OpenVMM VM management
8pub mod openvmm;
9pub mod vtl2_settings;
10
11use crate::PetriLogSource;
12use crate::PetriTestParams;
13use crate::ShutdownKind;
14use crate::disk_image::AgentImage;
15use crate::disk_image::SECTOR_SIZE;
16use crate::openhcl_diag::OpenHclDiagHandler;
17use crate::test::PetriPostTestHook;
18use crate::vtl2_settings::ControllerType;
19use crate::vtl2_settings::Vtl2LunBuilder;
20use crate::vtl2_settings::Vtl2StorageBackingDeviceBuilder;
21use crate::vtl2_settings::Vtl2StorageControllerBuilder;
22use async_trait::async_trait;
23use get_resources::ged::FirmwareEvent;
24use guid::Guid;
25use mesh::CancelContext;
26use openvmm_defs::config::Vtl2BaseAddressType;
27use pal_async::DefaultDriver;
28use pal_async::task::Spawn;
29use pal_async::task::Task;
30use pal_async::timer::PolledTimer;
31use petri_artifacts_common::tags::GuestQuirks;
32use petri_artifacts_common::tags::GuestQuirksInner;
33use petri_artifacts_common::tags::InitialRebootCondition;
34use petri_artifacts_common::tags::IsOpenhclIgvm;
35use petri_artifacts_common::tags::IsTestVmgs;
36use petri_artifacts_common::tags::MachineArch;
37use petri_artifacts_common::tags::OsFlavor;
38use petri_artifacts_core::ArtifactResolver;
39use petri_artifacts_core::ArtifactSource;
40use petri_artifacts_core::ResolvedArtifact;
41use petri_artifacts_core::ResolvedArtifactSource;
42use petri_artifacts_core::ResolvedOptionalArtifact;
43use pipette_client::PipetteClient;
44use std::collections::BTreeMap;
45use std::collections::HashMap;
46use std::collections::hash_map::DefaultHasher;
47use std::fmt::Debug;
48use std::hash::Hash;
49use std::hash::Hasher;
50use std::path::Path;
51use std::path::PathBuf;
52use std::sync::Arc;
53use std::time::Duration;
54use tempfile::TempPath;
55use vmgs_resources::GuestStateEncryptionPolicy;
56use vtl2_settings_proto::StorageController;
57use vtl2_settings_proto::Vtl2Settings;
58
59/// The set of artifacts and resources needed to instantiate a
60/// [`PetriVmBuilder`].
61pub struct PetriVmArtifacts<T: PetriVmmBackend> {
62    /// Artifacts needed to launch the host VMM used for the test
63    pub backend: T,
64    /// Firmware and/or OS to load into the VM and associated settings
65    pub firmware: Firmware,
66    /// The architecture of the VM
67    pub arch: MachineArch,
68    /// Agent to run in the guest
69    pub agent_image: Option<AgentImage>,
70    /// Agent to run in OpenHCL
71    pub openhcl_agent_image: Option<AgentImage>,
72    /// Raw pipette binary path (for embedding in initrd via CPIO append)
73    pub pipette_binary: Option<ResolvedArtifact>,
74}
75
76impl<T: PetriVmmBackend> PetriVmArtifacts<T> {
77    /// Resolves the artifacts needed to instantiate a [`PetriVmBuilder`].
78    ///
79    /// Returns `None` if the supplied configuration is not supported on this platform.
80    pub fn new(
81        resolver: &ArtifactResolver<'_>,
82        firmware: Firmware,
83        arch: MachineArch,
84        with_vtl0_pipette: bool,
85    ) -> Option<Self> {
86        if !T::check_compat(&firmware, arch) {
87            return None;
88        }
89
90        let pipette_binary = if with_vtl0_pipette {
91            Some(Self::resolve_pipette_binary(
92                resolver,
93                firmware.os_flavor(),
94                arch,
95            ))
96        } else {
97            None
98        };
99
100        Some(Self {
101            backend: T::new(resolver),
102            arch,
103            agent_image: Some(if with_vtl0_pipette {
104                AgentImage::new(firmware.os_flavor()).with_pipette(resolver, arch)
105            } else {
106                AgentImage::new(firmware.os_flavor())
107            }),
108            openhcl_agent_image: if firmware.is_openhcl() {
109                Some(AgentImage::new(OsFlavor::Linux).with_pipette(resolver, arch))
110            } else {
111                None
112            },
113            pipette_binary,
114            firmware,
115        })
116    }
117
118    fn resolve_pipette_binary(
119        resolver: &ArtifactResolver<'_>,
120        os_flavor: OsFlavor,
121        arch: MachineArch,
122    ) -> ResolvedArtifact {
123        use petri_artifacts_common::artifacts as common_artifacts;
124        match (os_flavor, arch) {
125            (OsFlavor::Linux, MachineArch::X86_64) => resolver
126                .require(common_artifacts::PIPETTE_LINUX_X64)
127                .erase(),
128            (OsFlavor::Linux, MachineArch::Aarch64) => resolver
129                .require(common_artifacts::PIPETTE_LINUX_AARCH64)
130                .erase(),
131            (OsFlavor::Windows, MachineArch::X86_64) => resolver
132                .require(common_artifacts::PIPETTE_WINDOWS_X64)
133                .erase(),
134            (OsFlavor::Windows, MachineArch::Aarch64) => resolver
135                .require(common_artifacts::PIPETTE_WINDOWS_AARCH64)
136                .erase(),
137            (OsFlavor::FreeBsd | OsFlavor::Uefi, _) => {
138                panic!("No pipette binary for this OS flavor")
139            }
140        }
141    }
142}
143
144/// Petri VM builder
145pub struct PetriVmBuilder<T: PetriVmmBackend> {
146    /// Artifacts needed to launch the host VMM used for the test
147    backend: T,
148    /// VM configuration
149    config: PetriVmConfig,
150    /// Function to modify the VMM-specific configuration
151    modify_vmm_config: Option<ModifyFn<T::VmmConfig>>,
152    /// VMM-agnostic resources
153    resources: PetriVmResources,
154
155    // VMM-specific quirks for the configured firmware
156    guest_quirks: GuestQuirksInner,
157    vmm_quirks: VmmQuirks,
158
159    // Test-specific boot behavior expectations.
160    // Defaults to expected behavior for firmware configuration.
161    expected_boot_event: Option<FirmwareEvent>,
162    override_expect_reset: bool,
163
164    // Config that is used to modify the `PetriVmConfig` before it is passed
165    // to the VMM backend.
166    /// Agent to run in the guest
167    agent_image: Option<AgentImage>,
168    /// Agent to run in OpenHCL
169    openhcl_agent_image: Option<AgentImage>,
170    /// The boot device type for the VM
171    boot_device_type: BootDeviceType,
172    /// Override for the PCIe root port the boot NVMe controller is placed on
173    /// when [`BootDeviceType::PcieNvme`] is used. Defaults to `s0rc0rp0`.
174    pcie_boot_port: Option<String>,
175
176    // Minimal mode: skip default devices, serial, save/restore.
177    minimal_mode: bool,
178    // Raw pipette binary path (for CPIO embedding in initrd).
179    pipette_binary: Option<ResolvedArtifact>,
180    // Enable serial output even in minimal mode (for diagnostics).
181    enable_serial: bool,
182    // Enable periodic framebuffer screenshots.
183    enable_screenshots: bool,
184    // Pre-built initrd with pipette already injected (skips runtime injection).
185    prebuilt_initrd: Option<PathBuf>,
186    // Use virtio vsock instead of VMBus-based hvsocket for guest communication.
187    use_virtio_vsock: bool,
188    // Use the Linux kernel vhost-vsock backend with this guest CID.
189    #[cfg(target_os = "linux")]
190    vhost_vsock_guest_cid: Option<u32>,
191    // Disable VMBus entirely (no vmbus server, no vmbus storage controllers).
192    no_vmbus: bool,
193    // Disable the hypervisor (HV#1) enlightenments. Implies `no_vmbus`.
194    no_hv: bool,
195}
196
197impl<T: PetriVmmBackend> Debug for PetriVmBuilder<T> {
198    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
199        f.debug_struct("PetriVmBuilder")
200            .field("backend", &self.backend)
201            .field("config", &self.config)
202            .field("modify_vmm_config", &self.modify_vmm_config.is_some())
203            .field("resources", &self.resources)
204            .field("guest_quirks", &self.guest_quirks)
205            .field("vmm_quirks", &self.vmm_quirks)
206            .field("expected_boot_event", &self.expected_boot_event)
207            .field("override_expect_reset", &self.override_expect_reset)
208            .field("agent_image", &self.agent_image)
209            .field("openhcl_agent_image", &self.openhcl_agent_image)
210            .field("boot_device_type", &self.boot_device_type)
211            .field("pcie_boot_port", &self.pcie_boot_port)
212            .field("minimal_mode", &self.minimal_mode)
213            .field("enable_serial", &self.enable_serial)
214            .field("enable_screenshots", &self.enable_screenshots)
215            .field("prebuilt_initrd", &self.prebuilt_initrd)
216            .field("use_virtio_vsock", &self.use_virtio_vsock)
217            .field("no_vmbus", &self.no_vmbus)
218            .field("no_hv", &self.no_hv)
219            .finish()
220    }
221}
222
223/// Petri VM configuration
224#[derive(Debug)]
225pub struct PetriVmConfig {
226    /// The name of the VM
227    pub name: String,
228    /// The architecture of the VM
229    pub arch: MachineArch,
230    /// Log levels for the host VMM process.
231    pub host_log_levels: Option<OpenvmmLogConfig>,
232    /// Firmware and/or OS to load into the VM and associated settings
233    pub firmware: Firmware,
234    /// The amount of memory, in bytes, to assign to the VM
235    pub memory: MemoryConfig,
236    /// The processor topology for the VM
237    pub proc_topology: ProcessorTopology,
238    /// VM guest state
239    pub vmgs: PetriVmgsResource,
240    /// TPM configuration
241    pub tpm: Option<TpmConfig>,
242    /// Storage controllers and associated disks
243    pub vmbus_storage_controllers: HashMap<Guid, VmbusStorageController>,
244    /// PCIe NVMe drives.
245    pub pcie_nvme_drives: Vec<PcieNvmeDrive>,
246    /// PCIe virtio-blk drives.
247    pub pcie_virtio_blk_drives: Vec<PcieVirtioBlkDrive>,
248    /// Physical NVMe devices to attach
249    pub physical_nvme_devices: HashMap<Guid, PhysicalNvmeDevice>,
250}
251
252/// PCIe NVMe drive configuration.
253#[derive(Debug)]
254pub struct PcieNvmeDrive {
255    /// PCIe root port name (e.g. "s0rc0rp0").
256    pub port_name: String,
257    /// NVMe namespace ID.
258    pub nsid: u32,
259    /// The drive to attach.
260    pub drive: Drive,
261}
262
263/// PCIe virtio-blk drive configuration.
264#[derive(Debug)]
265pub struct PcieVirtioBlkDrive {
266    /// PCIe root port name (e.g. "s0rc0rp0").
267    pub port_name: String,
268    /// The drive to attach.
269    pub drive: Drive,
270}
271
272/// Physical NVMe device to assign to a VM.
273/// Only used in closed-source HyperV tests
274#[derive(Debug, Clone)]
275pub struct PhysicalNvmeDevice {
276    /// The VTL to assign the physical NVMe device to.
277    pub target_vtl: Vtl,
278    /// NVMe namespace ID.
279    pub nsid: u32,
280    /// Namespace size in MiB
281    pub namespace_size_mib: u64,
282}
283
284/// Static properties about the VM for convenience during contruction and
285/// runtime of a VMM backend
286pub struct PetriVmProperties {
287    /// Whether this VM uses OpenHCL
288    pub is_openhcl: bool,
289    /// Whether this VM is isolated
290    pub is_isolated: bool,
291    /// Whether this VM uses the PCAT BIOS
292    pub is_pcat: bool,
293    /// Whether this VM boots with linux direct
294    pub is_linux_direct: bool,
295    /// Whether this VM is using pipette in VTL0
296    pub using_vtl0_pipette: bool,
297    /// Whether this VM is using VPCI
298    pub using_vpci: bool,
299    /// The OS flavor of the guest in the VM
300    pub os_flavor: OsFlavor,
301    /// Minimal mode: skip default devices, serial, save/restore
302    pub minimal_mode: bool,
303    /// Pipette embeds in initrd as PID 1 (non-OpenHCL Linux direct boot)
304    pub uses_pipette_as_init: bool,
305    /// Enable serial output even in minimal mode
306    pub enable_serial: bool,
307    /// Pre-built initrd path with pipette already injected
308    pub prebuilt_initrd: Option<PathBuf>,
309    /// Whether the VM has a CIDATA agent disk attached
310    pub has_agent_disk: bool,
311    /// Use virtio vsock instead of VMBus-based hvsocket
312    pub use_virtio_vsock: bool,
313    /// Linux kernel vhost-vsock guest CID, when that backend is enabled.
314    #[cfg(target_os = "linux")]
315    pub vhost_vsock_guest_cid: Option<u32>,
316    /// VMBus is entirely disabled
317    pub no_vmbus: bool,
318    /// The hypervisor (HV#1) enlightenments are entirely disabled
319    pub no_hv: bool,
320}
321
322/// VM configuration that can be changed after the VM is created
323pub struct PetriVmRuntimeConfig {
324    /// VTL2 settings
325    pub vtl2_settings: Option<Vtl2Settings>,
326    /// IDE controllers and associated disks
327    pub ide_controllers: Option<[[Option<Drive>; 2]; 2]>,
328    /// Storage controllers and associated disks
329    pub vmbus_storage_controllers: HashMap<Guid, VmbusStorageController>,
330}
331
332/// Resources used by a Petri VM during contruction and runtime
333#[derive(Debug)]
334pub struct PetriVmResources {
335    driver: DefaultDriver,
336    log_source: PetriLogSource,
337}
338
339/// Trait for VMM-specific contruction and runtime resources
340#[async_trait]
341pub trait PetriVmmBackend: Debug {
342    /// VMM-specific configuration
343    type VmmConfig;
344
345    /// Runtime object
346    type VmRuntime: PetriVmRuntime;
347
348    /// Check whether the combination of firmware and architecture is
349    /// supported on the VMM.
350    fn check_compat(firmware: &Firmware, arch: MachineArch) -> bool;
351
352    /// Select backend specific quirks guest and vmm quirks.
353    fn quirks(firmware: &Firmware) -> (GuestQuirksInner, VmmQuirks);
354
355    /// Get the default servicing flags (based on what this backend supports)
356    fn default_servicing_flags() -> OpenHclServicingFlags;
357
358    /// Create a disk for guest crash dumps, and a post-test hook to open the disk
359    /// to allow for reading the dumps.
360    fn create_guest_dump_disk() -> anyhow::Result<
361        Option<(
362            Arc<TempPath>,
363            Box<dyn FnOnce() -> anyhow::Result<Box<dyn fatfs::ReadWriteSeek>>>,
364        )>,
365    >;
366
367    /// Resolve any artifacts needed to use this backend
368    fn new(resolver: &ArtifactResolver<'_>) -> Self;
369
370    /// Create and start VM from the generic config using the VMM backend
371    async fn run(
372        self,
373        config: PetriVmConfig,
374        modify_vmm_config: Option<ModifyFn<Self::VmmConfig>>,
375        resources: &PetriVmResources,
376        properties: PetriVmProperties,
377    ) -> anyhow::Result<(Self::VmRuntime, PetriVmRuntimeConfig)>;
378}
379
380// IDE is only ever offered to VTL0
381pub(crate) const PETRI_IDE_BOOT_CONTROLLER_NUMBER: u32 = 0;
382pub(crate) const PETRI_IDE_BOOT_LUN: u8 = 0;
383pub(crate) const PETRI_IDE_BOOT_CONTROLLER: Guid =
384    guid::guid!("ca56751f-e643-4bef-bf54-f73678e8b7b5");
385
386// SCSI luns used for both VTL0 and VTL2
387pub(crate) const PETRI_SCSI_BOOT_LUN: u32 = 0;
388pub(crate) const PETRI_SCSI_PIPETTE_LUN: u32 = 1;
389pub(crate) const PETRI_SCSI_CRASH_LUN: u32 = 2;
390/// VTL0 SCSI controller instance guid used by Petri
391pub(crate) const PETRI_SCSI_VTL0_CONTROLLER: Guid =
392    guid::guid!("27b553e8-8b39-411b-a55f-839971a7884f");
393/// VTL2 SCSI controller instance guid used by Petri
394pub(crate) const PETRI_SCSI_VTL2_CONTROLLER: Guid =
395    guid::guid!("766e96f8-2ceb-437e-afe3-a93169e48a7c");
396/// SCSI controller instance guid offered to VTL0 by VTL2
397pub(crate) const PETRI_SCSI_VTL0_VIA_VTL2_CONTROLLER: Guid =
398    guid::guid!("6c474f47-ed39-49e6-bbb9-142177a1da6e");
399
400/// The namespace ID used by Petri for the boot disk
401pub(crate) const PETRI_NVME_BOOT_NSID: u32 = 37;
402/// VTL0 NVMe controller instance guid used by Petri
403pub(crate) const PETRI_NVME_BOOT_VTL0_CONTROLLER: Guid =
404    guid::guid!("e23a04e2-90f5-4852-bc9d-e7ac691b756c");
405/// VTL2 NVMe controller instance guid used by Petri
406pub(crate) const PETRI_NVME_BOOT_VTL2_CONTROLLER: Guid =
407    guid::guid!("92bc8346-718b-449a-8751-edbf3dcd27e4");
408
409/// PCIe root port used by Petri for the agent/cidata disk (no-vmbus mode)
410pub(crate) const PETRI_PCIE_NVME_AGENT_PORT: &str = "s0rc0rp1";
411/// NVMe namespace ID used by Petri for the agent/cidata disk (no-vmbus mode)
412pub(crate) const PETRI_PCIE_NVME_AGENT_NSID: u32 = 1;
413
414/// A constructed Petri VM
415pub struct PetriVm<T: PetriVmmBackend> {
416    resources: PetriVmResources,
417    runtime: T::VmRuntime,
418    watchdog_tasks: Vec<Task<()>>,
419    openhcl_diag_handler: Option<OpenHclDiagHandler>,
420
421    arch: MachineArch,
422    guest_quirks: GuestQuirksInner,
423    vmm_quirks: VmmQuirks,
424    expected_boot_event: Option<FirmwareEvent>,
425
426    config: PetriVmRuntimeConfig,
427}
428
429impl<T: PetriVmmBackend> PetriVmBuilder<T> {
430    /// Create a new VM configuration.
431    pub fn new(
432        params: PetriTestParams<'_>,
433        artifacts: PetriVmArtifacts<T>,
434        driver: &DefaultDriver,
435    ) -> anyhow::Result<Self> {
436        let (guest_quirks, vmm_quirks) = T::quirks(&artifacts.firmware);
437        let expected_boot_event = artifacts.firmware.expected_boot_event();
438        let boot_device_type = match artifacts.firmware {
439            Firmware::LinuxDirect { .. } => BootDeviceType::None,
440            Firmware::OpenhclLinuxDirect { .. } => BootDeviceType::None,
441            Firmware::Pcat { .. } => BootDeviceType::Ide,
442            Firmware::OpenhclPcat { .. } => BootDeviceType::IdeViaScsi,
443            Firmware::Uefi {
444                guest: UefiGuest::None,
445                ..
446            }
447            | Firmware::OpenhclUefi {
448                guest: UefiGuest::None,
449                ..
450            } => BootDeviceType::None,
451            Firmware::Uefi { .. } | Firmware::OpenhclUefi { .. } => BootDeviceType::Scsi,
452        };
453
454        Ok(Self {
455            backend: artifacts.backend,
456            config: PetriVmConfig {
457                name: make_vm_safe_name(params.test_name),
458                arch: artifacts.arch,
459                host_log_levels: None,
460                firmware: artifacts.firmware,
461                memory: Default::default(),
462                proc_topology: Default::default(),
463
464                vmgs: PetriVmgsResource::Ephemeral,
465                tpm: None,
466                vmbus_storage_controllers: HashMap::new(),
467                pcie_nvme_drives: Vec::new(),
468                pcie_virtio_blk_drives: Vec::new(),
469                physical_nvme_devices: HashMap::new(),
470            },
471            modify_vmm_config: None,
472            resources: PetriVmResources {
473                driver: driver.clone(),
474                log_source: params.logger.clone(),
475            },
476
477            guest_quirks,
478            vmm_quirks,
479            expected_boot_event,
480            override_expect_reset: false,
481
482            agent_image: artifacts.agent_image,
483            openhcl_agent_image: artifacts.openhcl_agent_image,
484            boot_device_type,
485            pcie_boot_port: None,
486
487            minimal_mode: false,
488            pipette_binary: artifacts.pipette_binary,
489            enable_serial: true,
490            enable_screenshots: true,
491            prebuilt_initrd: None,
492            use_virtio_vsock: false,
493            #[cfg(target_os = "linux")]
494            vhost_vsock_guest_cid: None,
495            no_vmbus: false,
496            no_hv: false,
497        }
498        .add_petri_scsi_controllers()
499        .add_guest_crash_disk(params.post_test_hooks))
500    }
501
502    /// Create a minimal VM builder with only the bare minimum device set.
503    ///
504    /// Unlike [`new()`](Self::new), this constructor:
505    /// - Does not add default VMBus devices (shutdown IC, KVP, etc.)
506    /// - Does not add serial ports
507    /// - Does not add SCSI controllers or crash dump disks
508    /// - Does not verify save/restore on boot
509    ///
510    /// Use builder methods to opt in to specific devices. Intended for
511    /// performance tests where minimal overhead is critical.
512    pub fn minimal(
513        params: PetriTestParams<'_>,
514        artifacts: PetriVmArtifacts<T>,
515        driver: &DefaultDriver,
516    ) -> anyhow::Result<Self> {
517        let (guest_quirks, vmm_quirks) = T::quirks(&artifacts.firmware);
518        let expected_boot_event = artifacts.firmware.expected_boot_event();
519        let boot_device_type = match artifacts.firmware {
520            Firmware::LinuxDirect { .. } => BootDeviceType::None,
521            Firmware::OpenhclLinuxDirect { .. } => BootDeviceType::None,
522            Firmware::Pcat { .. } => BootDeviceType::Ide,
523            Firmware::OpenhclPcat { .. } => BootDeviceType::IdeViaScsi,
524            Firmware::Uefi {
525                guest: UefiGuest::None,
526                ..
527            }
528            | Firmware::OpenhclUefi {
529                guest: UefiGuest::None,
530                ..
531            } => BootDeviceType::None,
532            Firmware::Uefi { .. } | Firmware::OpenhclUefi { .. } => BootDeviceType::Scsi,
533        };
534
535        Ok(Self {
536            backend: artifacts.backend,
537            config: PetriVmConfig {
538                name: make_vm_safe_name(params.test_name),
539                arch: artifacts.arch,
540                host_log_levels: None,
541                firmware: artifacts.firmware,
542                memory: Default::default(),
543                proc_topology: Default::default(),
544
545                vmgs: PetriVmgsResource::Ephemeral,
546                tpm: None,
547                vmbus_storage_controllers: HashMap::new(),
548                pcie_nvme_drives: Vec::new(),
549                pcie_virtio_blk_drives: Vec::new(),
550                physical_nvme_devices: HashMap::new(),
551            },
552            modify_vmm_config: None,
553            resources: PetriVmResources {
554                driver: driver.clone(),
555                log_source: params.logger.clone(),
556            },
557
558            guest_quirks,
559            vmm_quirks,
560            expected_boot_event,
561            override_expect_reset: false,
562
563            agent_image: artifacts.agent_image,
564            openhcl_agent_image: artifacts.openhcl_agent_image,
565            boot_device_type,
566            pcie_boot_port: None,
567
568            minimal_mode: true,
569            pipette_binary: artifacts.pipette_binary,
570            enable_serial: false,
571            enable_screenshots: true,
572            prebuilt_initrd: None,
573            use_virtio_vsock: false,
574            #[cfg(target_os = "linux")]
575            vhost_vsock_guest_cid: None,
576            no_vmbus: false,
577            no_hv: false,
578        })
579    }
580
581    /// Whether this builder is in minimal mode.
582    pub fn is_minimal(&self) -> bool {
583        self.minimal_mode
584    }
585
586    /// Supply a pre-built initrd with pipette already injected.
587    ///
588    /// When set, the builder skips the runtime gzip decompress/inject/
589    /// recompress cycle, using this initrd directly. Use
590    /// [`prepare_initrd`](Self::prepare_initrd) to build the initrd
591    /// ahead of time.
592    pub fn with_prebuilt_initrd(mut self, path: PathBuf) -> Self {
593        self.prebuilt_initrd = Some(path);
594        self
595    }
596
597    /// Pre-build the modified initrd with pipette injected.
598    ///
599    /// Reads the original initrd from the firmware artifacts, injects
600    /// the pipette binary via CPIO, and writes the result to a temp file.
601    /// Returns the path to the temp file. The caller must keep the
602    /// `TempPath` alive until after the VM boots.
603    ///
604    /// Call this once before timing, then pass the path to
605    /// [`with_prebuilt_initrd`](Self::with_prebuilt_initrd) for each
606    /// iteration.
607    pub fn prepare_initrd(&self) -> anyhow::Result<TempPath> {
608        use anyhow::Context;
609        use std::io::Write;
610
611        let initrd_path = self
612            .config
613            .firmware
614            .linux_direct_initrd()
615            .context("prepare_initrd requires Linux direct boot with initrd")?;
616        let pipette_path = self
617            .pipette_binary
618            .as_ref()
619            .context("prepare_initrd requires a pipette binary")?;
620
621        let initrd_gz = std::fs::read(initrd_path)
622            .with_context(|| format!("failed to read initrd at {}", initrd_path.display()))?;
623        let pipette_data = std::fs::read(pipette_path.get()).with_context(|| {
624            format!(
625                "failed to read pipette binary at {}",
626                pipette_path.get().display()
627            )
628        })?;
629
630        let merged_gz =
631            initrd_cpio::inject_into_initrd(&initrd_gz, "pipette", &pipette_data, 0o100755)
632                .context("failed to inject pipette into initrd")?;
633
634        let mut tmp = tempfile::NamedTempFile::new()
635            .context("failed to create temp file for pre-built initrd")?;
636        tmp.write_all(&merged_gz)
637            .context("failed to write pre-built initrd")?;
638
639        Ok(tmp.into_temp_path())
640    }
641
642    /// Enable serial port output even in minimal mode.
643    ///
644    /// Useful for diagnostics — the serial device overhead is negligible;
645    /// the cost comes from kernel console output, which is controlled via
646    /// the kernel cmdline (`quiet loglevel=0`).
647    ///
648    /// Note: this currently only affects LinuxDirect boot (kernel cmdline
649    /// and emulated serial backends). UEFI paths are unaffected.
650    pub fn with_serial_output(mut self) -> Self {
651        self.enable_serial = true;
652        self
653    }
654
655    /// Disable serial port output.
656    ///
657    /// Suppresses serial device creation, eliminating the `[uefi]` / `[openhcl]`
658    /// log lines. Useful for performance tests where serial noise is unwanted.
659    pub fn without_serial_output(mut self) -> Self {
660        self.enable_serial = false;
661        self
662    }
663
664    /// Disable periodic framebuffer screenshots.
665    ///
666    /// Suppresses the watchdog task that takes screenshots every 2 seconds,
667    /// eliminating the "No change in framebuffer" debug log lines.
668    pub fn without_screenshots(mut self) -> Self {
669        self.enable_screenshots = false;
670        self
671    }
672
673    /// Use virtio vsock instead of VMBus-based hvsocket for guest communication.
674    /// The virtio-vsock device will use PCIe, so a PCIe root topology must be
675    /// configured.
676    ///
677    /// When enabled, a virtio-vsock device is added to the VM. This device uses
678    /// the same Unix socket relay path that hvsocket would otherwise use, so
679    /// pipette will connect using this.
680    ///
681    /// For Linux direct boot, this also adjusts the kernel command line to
682    /// blacklist hv_sock instead of virtio_vsock.
683    pub fn with_virtio_vsock(mut self) -> Self {
684        self.use_virtio_vsock = true;
685        #[cfg(target_os = "linux")]
686        {
687            self.vhost_vsock_guest_cid = None;
688        }
689        self
690    }
691
692    /// Use the Linux kernel vhost-vsock backend for guest communication.
693    ///
694    /// The host connects directly to the guest's `AF_VSOCK` listener at
695    /// `guest_cid`. The OpenVMM backend automatically uses shared guest memory,
696    /// which is required by kernel vhost.
697    #[cfg(target_os = "linux")]
698    pub fn with_vhost_vsock(mut self, guest_cid: u32) -> Self {
699        assert!(
700            (3..u32::MAX).contains(&guest_cid),
701            "vhost-vsock guest CID must be between 3 and {}",
702            u32::MAX - 1
703        );
704        self.use_virtio_vsock = true;
705        self.vhost_vsock_guest_cid = Some(guest_cid);
706        self
707    }
708
709    /// Disable VMBus entirely.
710    ///
711    /// This removes all VMBus storage controllers. For Linux guests,
712    /// virtio-vsock is used for pipette communication. For Windows guests,
713    /// the caller must also configure TCP pipette transport via
714    /// `modify_backend(|b| b.with_tcp_pipette_nic(port, mac_address))`. The
715    /// guest must boot from a non-VMBus device (e.g. PCIe NVMe).
716    pub fn with_no_vmbus(mut self) -> Self {
717        self.no_vmbus = true;
718        if self.config.firmware.os_flavor() != OsFlavor::Windows {
719            self.use_virtio_vsock = true;
720        }
721        self.config.vmbus_storage_controllers.clear();
722        self
723    }
724
725    /// Disable the hypervisor (HV#1) enlightenments.
726    ///
727    /// This also disables VMBus, since VMBus depends on the hypervisor. On
728    /// aarch64 UEFI this causes the loader to pass the generic SEC platform
729    /// type to the firmware. This mode is not supported on x86_64 UEFI.
730    pub fn with_no_hv(mut self) -> Self {
731        self.no_hv = true;
732        self.with_no_vmbus()
733    }
734
735    fn add_petri_scsi_controllers(self) -> Self {
736        let builder = self.add_vmbus_storage_controller(
737            &PETRI_SCSI_VTL0_CONTROLLER,
738            Vtl::Vtl0,
739            VmbusStorageType::Scsi,
740        );
741
742        if builder.is_openhcl() {
743            builder.add_vmbus_storage_controller(
744                &PETRI_SCSI_VTL2_CONTROLLER,
745                Vtl::Vtl2,
746                VmbusStorageType::Scsi,
747            )
748        } else {
749            builder
750        }
751    }
752
753    fn add_guest_crash_disk(self, post_test_hooks: &mut Vec<PetriPostTestHook>) -> Self {
754        let logger = self.resources.log_source.clone();
755        let (disk, disk_hook) = matches!(
756            self.config.firmware.os_flavor(),
757            OsFlavor::Windows | OsFlavor::Linux
758        )
759        .then(|| T::create_guest_dump_disk().expect("failed to create guest dump disk"))
760        .flatten()
761        .unzip();
762
763        if let Some(disk_hook) = disk_hook {
764            post_test_hooks.push(PetriPostTestHook::new(
765                "extract guest crash dumps".into(),
766                move |test_passed| {
767                    if test_passed {
768                        return Ok(());
769                    }
770                    let mut disk = disk_hook()?;
771                    let gpt = gptman::GPT::read_from(&mut disk, SECTOR_SIZE)?;
772                    let partition = fscommon::StreamSlice::new(
773                        &mut disk,
774                        gpt[1].starting_lba * SECTOR_SIZE,
775                        gpt[1].ending_lba * SECTOR_SIZE,
776                    )?;
777                    let fs = fatfs::FileSystem::new(partition, fatfs::FsOptions::new())?;
778                    for entry in fs.root_dir().iter() {
779                        let Ok(entry) = entry else {
780                            tracing::warn!(?entry, "failed to read entry in guest crash dump disk");
781                            continue;
782                        };
783                        if !entry.is_file() {
784                            tracing::warn!(
785                                ?entry,
786                                "skipping non-file entry in guest crash dump disk"
787                            );
788                            continue;
789                        }
790                        logger.write_attachment(&entry.file_name(), entry.to_file())?;
791                    }
792                    Ok(())
793                },
794            ));
795        }
796
797        if let Some(disk) = disk {
798            self.add_vmbus_drive(
799                Drive::new(Some(Disk::Temporary(disk)), false),
800                &PETRI_SCSI_VTL0_CONTROLLER,
801                Some(PETRI_SCSI_CRASH_LUN),
802            )
803        } else {
804            self
805        }
806    }
807
808    fn add_agent_disks(self) -> Self {
809        self.add_agent_disk_inner(Vtl::Vtl0)
810            .add_agent_disk_inner(Vtl::Vtl2)
811    }
812
813    fn add_agent_disk_inner(mut self, target_vtl: Vtl) -> Self {
814        let (agent_image, controller_id) = match target_vtl {
815            Vtl::Vtl0 => (self.agent_image.as_ref(), PETRI_SCSI_VTL0_CONTROLLER),
816            Vtl::Vtl1 => panic!("no VTL1 agent disk"),
817            Vtl::Vtl2 => (
818                self.openhcl_agent_image.as_ref(),
819                PETRI_SCSI_VTL2_CONTROLLER,
820            ),
821        };
822
823        // When using pipette-as-init, the VTL0 agent disk is only needed
824        // if it carries extra files (pipette itself is in the initrd).
825        if target_vtl == Vtl::Vtl0
826            && self.uses_pipette_as_init()
827            && !agent_image.is_some_and(|i| i.has_extras())
828        {
829            return self;
830        }
831
832        let Some(agent_disk) = agent_image.and_then(|i| {
833            i.build(crate::disk_image::ImageType::Vhd)
834                .expect("failed to build agent image")
835        }) else {
836            return self;
837        };
838
839        // When VMBus is disabled, route the agent disk through PCIe NVMe
840        // instead of VMBus SCSI.
841        if self.no_vmbus {
842            self.config.pcie_nvme_drives.push(PcieNvmeDrive {
843                port_name: PETRI_PCIE_NVME_AGENT_PORT.into(),
844                nsid: PETRI_PCIE_NVME_AGENT_NSID,
845                drive: Drive::new(
846                    Some(Disk::Temporary(Arc::new(agent_disk.into_temp_path()))),
847                    false,
848                ),
849            });
850            return self;
851        }
852
853        // Ensure the storage controller exists (minimal mode doesn't
854        // add controllers upfront).
855        if !self
856            .config
857            .vmbus_storage_controllers
858            .contains_key(&controller_id)
859        {
860            self = self.add_vmbus_storage_controller(
861                &controller_id,
862                target_vtl,
863                VmbusStorageType::Scsi,
864            );
865        }
866
867        self.add_vmbus_drive(
868            Drive::new(
869                Some(Disk::Temporary(Arc::new(agent_disk.into_temp_path()))),
870                false,
871            ),
872            &controller_id,
873            Some(PETRI_SCSI_PIPETTE_LUN),
874        )
875    }
876
877    fn add_boot_disk(mut self) -> Self {
878        if self.boot_device_type.requires_vtl2() && !self.is_openhcl() {
879            panic!("boot device type {:?} requires vtl2", self.boot_device_type);
880        }
881
882        if self.no_vmbus && self.boot_device_type.requires_vmbus() {
883            panic!(
884                "boot device type {:?} requires vmbus, but vmbus is disabled; \
885                 use with_boot_device_type(BootDeviceType::PcieNvme) or similar",
886                self.boot_device_type
887            );
888        }
889
890        if self.boot_device_type.requires_vpci_boot() {
891            self.config
892                .firmware
893                .uefi_config_mut()
894                .expect("vpci boot requires uefi")
895                .enable_vpci_boot = true;
896        }
897
898        if let Some(boot_drive) = self.config.firmware.boot_drive() {
899            match self.boot_device_type {
900                BootDeviceType::None => unreachable!(),
901                BootDeviceType::Ide => self.add_ide_drive(
902                    boot_drive,
903                    PETRI_IDE_BOOT_CONTROLLER_NUMBER,
904                    PETRI_IDE_BOOT_LUN,
905                ),
906                BootDeviceType::IdeViaScsi => self
907                    .add_vmbus_drive(
908                        boot_drive,
909                        &PETRI_SCSI_VTL2_CONTROLLER,
910                        Some(PETRI_SCSI_BOOT_LUN),
911                    )
912                    .add_vtl2_storage_controller(
913                        Vtl2StorageControllerBuilder::new(ControllerType::Ide)
914                            .with_instance_id(PETRI_IDE_BOOT_CONTROLLER)
915                            .add_lun(
916                                Vtl2LunBuilder::disk()
917                                    .with_channel(PETRI_IDE_BOOT_CONTROLLER_NUMBER)
918                                    .with_location(PETRI_IDE_BOOT_LUN as u32)
919                                    .with_physical_device(Vtl2StorageBackingDeviceBuilder::new(
920                                        ControllerType::Scsi,
921                                        PETRI_SCSI_VTL2_CONTROLLER,
922                                        PETRI_SCSI_BOOT_LUN,
923                                    )),
924                            )
925                            .build(),
926                    ),
927                BootDeviceType::IdeViaNvme => todo!(),
928                BootDeviceType::Scsi => self.add_vmbus_drive(
929                    boot_drive,
930                    &PETRI_SCSI_VTL0_CONTROLLER,
931                    Some(PETRI_SCSI_BOOT_LUN),
932                ),
933                BootDeviceType::ScsiViaScsi => self
934                    .add_vmbus_drive(
935                        boot_drive,
936                        &PETRI_SCSI_VTL2_CONTROLLER,
937                        Some(PETRI_SCSI_BOOT_LUN),
938                    )
939                    .add_vtl2_storage_controller(
940                        Vtl2StorageControllerBuilder::new(ControllerType::Scsi)
941                            .with_instance_id(PETRI_SCSI_VTL0_VIA_VTL2_CONTROLLER)
942                            .add_lun(
943                                Vtl2LunBuilder::disk()
944                                    .with_location(PETRI_SCSI_BOOT_LUN)
945                                    .with_physical_device(Vtl2StorageBackingDeviceBuilder::new(
946                                        ControllerType::Scsi,
947                                        PETRI_SCSI_VTL2_CONTROLLER,
948                                        PETRI_SCSI_BOOT_LUN,
949                                    )),
950                            )
951                            .build(),
952                    ),
953                BootDeviceType::ScsiViaNvme => self
954                    .add_vmbus_storage_controller(
955                        &PETRI_NVME_BOOT_VTL2_CONTROLLER,
956                        Vtl::Vtl2,
957                        VmbusStorageType::Nvme,
958                    )
959                    .add_vmbus_drive(
960                        boot_drive,
961                        &PETRI_NVME_BOOT_VTL2_CONTROLLER,
962                        Some(PETRI_NVME_BOOT_NSID),
963                    )
964                    .add_vtl2_storage_controller(
965                        Vtl2StorageControllerBuilder::new(ControllerType::Scsi)
966                            .with_instance_id(PETRI_SCSI_VTL0_VIA_VTL2_CONTROLLER)
967                            .add_lun(
968                                Vtl2LunBuilder::disk()
969                                    .with_location(PETRI_SCSI_BOOT_LUN)
970                                    .with_physical_device(Vtl2StorageBackingDeviceBuilder::new(
971                                        ControllerType::Nvme,
972                                        PETRI_NVME_BOOT_VTL2_CONTROLLER,
973                                        PETRI_NVME_BOOT_NSID,
974                                    )),
975                            )
976                            .build(),
977                    ),
978                BootDeviceType::Nvme => self
979                    .add_vmbus_storage_controller(
980                        &PETRI_NVME_BOOT_VTL0_CONTROLLER,
981                        Vtl::Vtl0,
982                        VmbusStorageType::Nvme,
983                    )
984                    .add_vmbus_drive(
985                        boot_drive,
986                        &PETRI_NVME_BOOT_VTL0_CONTROLLER,
987                        Some(PETRI_NVME_BOOT_NSID),
988                    ),
989                BootDeviceType::NvmeViaScsi => todo!(),
990                BootDeviceType::NvmeViaNvme => todo!(),
991                BootDeviceType::PcieNvme => {
992                    let port_name = self
993                        .pcie_boot_port
994                        .clone()
995                        .unwrap_or_else(|| "s0rc0rp0".into());
996                    self.config.pcie_nvme_drives.push(PcieNvmeDrive {
997                        port_name,
998                        nsid: 1,
999                        drive: boot_drive,
1000                    });
1001                    self
1002                }
1003                BootDeviceType::PcieVirtioBlk => {
1004                    self.config.pcie_virtio_blk_drives.push(PcieVirtioBlkDrive {
1005                        port_name: "s0rc0rp0".into(),
1006                        drive: boot_drive,
1007                    });
1008                    self
1009                }
1010            }
1011        } else {
1012            self
1013        }
1014    }
1015
1016    /// Whether the VTL0 agent disk will actually be added.
1017    ///
1018    /// False when using pipette-as-init with no extra files (pipette is
1019    /// in the initrd, so the CIDATA disk isn't needed).
1020    fn has_agent_disk(&self) -> bool {
1021        if self.uses_pipette_as_init() {
1022            self.agent_image.as_ref().is_some_and(|i| i.has_extras())
1023        } else {
1024            self.agent_image.is_some()
1025        }
1026    }
1027
1028    /// Get properties about the vm for convenience
1029    pub fn properties(&self) -> PetriVmProperties {
1030        PetriVmProperties {
1031            is_openhcl: self.config.firmware.is_openhcl(),
1032            is_isolated: self.config.firmware.isolation().is_some(),
1033            is_pcat: self.config.firmware.is_pcat(),
1034            is_linux_direct: self.config.firmware.is_linux_direct(),
1035            using_vtl0_pipette: self.using_vtl0_pipette(),
1036            using_vpci: self.boot_device_type.requires_vpci_boot(),
1037            os_flavor: self.config.firmware.os_flavor(),
1038            minimal_mode: self.minimal_mode,
1039            uses_pipette_as_init: self.uses_pipette_as_init(),
1040            enable_serial: self.enable_serial,
1041            prebuilt_initrd: self.prebuilt_initrd.clone(),
1042            has_agent_disk: self.has_agent_disk(),
1043            use_virtio_vsock: self.use_virtio_vsock,
1044            #[cfg(target_os = "linux")]
1045            vhost_vsock_guest_cid: self.vhost_vsock_guest_cid,
1046            no_vmbus: self.no_vmbus,
1047            no_hv: self.no_hv,
1048        }
1049    }
1050
1051    /// Whether pipette will run as PID 1 init in the initrd.
1052    ///
1053    /// True for non-OpenHCL Linux direct boot when a pipette binary is
1054    /// available. Pipette is injected into the initrd via CPIO and set
1055    /// as `rdinit=/pipette`.
1056    fn uses_pipette_as_init(&self) -> bool {
1057        self.config.firmware.is_linux_direct()
1058            && !self.config.firmware.is_openhcl()
1059            && self.pipette_binary.is_some()
1060    }
1061
1062    /// Whether this VM is using pipette in VTL0
1063    pub fn using_vtl0_pipette(&self) -> bool {
1064        self.uses_pipette_as_init()
1065            || self
1066                .agent_image
1067                .as_ref()
1068                .is_some_and(|x| x.contains_pipette())
1069    }
1070
1071    /// Build and run the VM, then wait for the VM to emit the expected boot
1072    /// event (if configured). Does not configure and start pipette. Should
1073    /// only be used for testing platforms that pipette does not support.
1074    pub async fn run_without_agent(self) -> anyhow::Result<PetriVm<T>> {
1075        self.run_core().await
1076    }
1077
1078    /// Build and run the VM, then wait for the VM to emit the expected boot
1079    /// event (if configured). Launches pipette and returns a client to it.
1080    pub async fn run(self) -> anyhow::Result<(PetriVm<T>, PipetteClient)> {
1081        assert!(self.using_vtl0_pipette());
1082
1083        let mut vm = self.run_core().await?;
1084        let client = vm.wait_for_agent().await?;
1085        Ok((vm, client))
1086    }
1087
1088    async fn run_core(mut self) -> anyhow::Result<PetriVm<T>> {
1089        // Add the boot disk now to allow the test to modify the boot type
1090        // Add the agent disks now to allow the test to add custom files
1091        self = self.add_boot_disk().add_agent_disks();
1092
1093        // Auto-prepare the initrd with pipette injected if needed.
1094        // This centralizes the injection logic so backends only ever
1095        // receive a prebuilt_initrd path.
1096        let _prepared_initrd_guard =
1097            if self.uses_pipette_as_init() && self.prebuilt_initrd.is_none() {
1098                let tmp = self.prepare_initrd()?;
1099                self.prebuilt_initrd = Some(tmp.to_path_buf());
1100                Some(tmp)
1101            } else {
1102                None
1103            };
1104
1105        tracing::debug!(builder = ?self);
1106
1107        let arch = self.config.arch;
1108        let expect_reset = self.expect_reset();
1109        let properties = self.properties();
1110
1111        let (mut runtime, config) = self
1112            .backend
1113            .run(
1114                self.config,
1115                self.modify_vmm_config,
1116                &self.resources,
1117                properties,
1118            )
1119            .await?;
1120        let openhcl_diag_handler = runtime.openhcl_diag();
1121        let watchdog_tasks =
1122            Self::start_watchdog_tasks(&self.resources, &mut runtime, self.enable_screenshots)?;
1123
1124        let mut vm = PetriVm {
1125            resources: self.resources,
1126            runtime,
1127            watchdog_tasks,
1128            openhcl_diag_handler,
1129
1130            arch,
1131            guest_quirks: self.guest_quirks,
1132            vmm_quirks: self.vmm_quirks,
1133            expected_boot_event: self.expected_boot_event,
1134
1135            config,
1136        };
1137
1138        if expect_reset {
1139            vm.wait_for_reset_core().await?;
1140        }
1141
1142        vm.wait_for_expected_boot_event().await?;
1143
1144        Ok(vm)
1145    }
1146
1147    fn expect_reset(&self) -> bool {
1148        self.override_expect_reset
1149            || matches!(
1150                (
1151                    self.guest_quirks.initial_reboot,
1152                    self.expected_boot_event,
1153                    &self.config.firmware,
1154                    &self.config.tpm,
1155                ),
1156                (
1157                    Some(InitialRebootCondition::Always),
1158                    Some(FirmwareEvent::BootSuccess | FirmwareEvent::BootAttempt),
1159                    _,
1160                    _,
1161                ) | (
1162                    Some(InitialRebootCondition::WithTpm),
1163                    Some(FirmwareEvent::BootSuccess | FirmwareEvent::BootAttempt),
1164                    _,
1165                    Some(_),
1166                )
1167            )
1168    }
1169
1170    fn start_watchdog_tasks(
1171        resources: &PetriVmResources,
1172        runtime: &mut T::VmRuntime,
1173        enable_screenshots: bool,
1174    ) -> anyhow::Result<Vec<Task<()>>> {
1175        let mut tasks = Vec::new();
1176
1177        {
1178            const TIMEOUT_DURATION_MINUTES: u64 = 10;
1179            const TIMER_DURATION: Duration = Duration::from_secs(TIMEOUT_DURATION_MINUTES * 60);
1180            let log_source = resources.log_source.clone();
1181            let inspect_task =
1182                |name,
1183                 driver: &DefaultDriver,
1184                 inspect: std::pin::Pin<Box<dyn Future<Output = _> + Send>>| {
1185                    driver.spawn(format!("petri-watchdog-inspect-{name}"), async move {
1186                        if CancelContext::new()
1187                            .with_timeout(Duration::from_secs(10))
1188                            .until_cancelled(save_inspect(name, inspect, &log_source))
1189                            .await
1190                            .is_err()
1191                        {
1192                            tracing::warn!(name, "Failed to collect inspect data within timeout");
1193                        }
1194                    })
1195                };
1196
1197            let driver = resources.driver.clone();
1198            let vmm_inspector = runtime.inspector();
1199            let openhcl_diag_handler = runtime.openhcl_diag();
1200            tasks.push(resources.driver.spawn("timer-watchdog", async move {
1201                PolledTimer::new(&driver).sleep(TIMER_DURATION).await;
1202                tracing::warn!("Test timeout reached after {TIMEOUT_DURATION_MINUTES} minutes, collecting diagnostics.");
1203                let mut timeout_tasks = Vec::new();
1204                if let Some(inspector) = vmm_inspector {
1205                    timeout_tasks.push(inspect_task.clone()("vmm", &driver, Box::pin(async move { inspector.inspect("").await })) );
1206                }
1207                if let Some(openhcl_diag_handler) = openhcl_diag_handler {
1208                    timeout_tasks.push(inspect_task("openhcl", &driver, Box::pin(async move { openhcl_diag_handler.inspect("", None, None).await })));
1209                }
1210                futures::future::join_all(timeout_tasks).await;
1211                tracing::error!("Test time out diagnostics collection complete, aborting.");
1212                panic!("Test timed out");
1213            }));
1214        }
1215
1216        if enable_screenshots {
1217            if let Some(mut framebuffer_access) = runtime.take_framebuffer_access() {
1218                let mut timer = PolledTimer::new(&resources.driver);
1219                let log_source = resources.log_source.clone();
1220
1221                tasks.push(
1222                    resources
1223                        .driver
1224                        .spawn("petri-watchdog-screenshot", async move {
1225                            let mut image = Vec::new();
1226                            let mut last_image = Vec::new();
1227                            loop {
1228                                timer.sleep(Duration::from_secs(2)).await;
1229                                tracing::trace!("Taking screenshot.");
1230
1231                                let VmScreenshotMeta {
1232                                    color,
1233                                    width,
1234                                    height,
1235                                } = match framebuffer_access.screenshot(&mut image).await {
1236                                    Ok(Some(meta)) => meta,
1237                                    Ok(None) => {
1238                                        tracing::debug!("VM off, skipping screenshot.");
1239                                        continue;
1240                                    }
1241                                    Err(e) => {
1242                                        tracing::error!(?e, "Failed to take screenshot");
1243                                        continue;
1244                                    }
1245                                };
1246
1247                                if image == last_image {
1248                                    tracing::debug!(
1249                                        "No change in framebuffer, skipping screenshot."
1250                                    );
1251                                    continue;
1252                                }
1253
1254                                let r = log_source.create_attachment("screenshot.png").and_then(
1255                                    |mut f| {
1256                                        image::write_buffer_with_format(
1257                                            &mut f,
1258                                            &image,
1259                                            width.into(),
1260                                            height.into(),
1261                                            color,
1262                                            image::ImageFormat::Png,
1263                                        )
1264                                        .map_err(Into::into)
1265                                    },
1266                                );
1267
1268                                if let Err(e) = r {
1269                                    tracing::error!(?e, "Failed to save screenshot");
1270                                } else {
1271                                    tracing::info!("Screenshot saved.");
1272                                }
1273
1274                                std::mem::swap(&mut image, &mut last_image);
1275                            }
1276                        }),
1277                );
1278            }
1279        }
1280
1281        Ok(tasks)
1282    }
1283
1284    /// Configure the test to expect a boot failure from the VM.
1285    /// Useful for negative tests.
1286    pub fn with_expect_boot_failure(mut self) -> Self {
1287        self.expected_boot_event = Some(FirmwareEvent::BootFailed);
1288        self
1289    }
1290
1291    /// Configure the test to not expect any boot event.
1292    /// Useful for tests that do not boot a VTL0 guest.
1293    pub fn with_expect_no_boot_event(mut self) -> Self {
1294        self.expected_boot_event = None;
1295        self
1296    }
1297
1298    /// Allow the VM to reset once at the beginning of the test. Should only be
1299    /// used if you are using a special VM configuration that causes the guest
1300    /// to reboot when it usually wouldn't.
1301    pub fn with_expect_reset(mut self) -> Self {
1302        self.override_expect_reset = true;
1303        self
1304    }
1305
1306    /// Set the VM to enable secure boot and inject the templates per OS flavor.
1307    pub fn with_secure_boot(mut self) -> Self {
1308        self.config
1309            .firmware
1310            .uefi_config_mut()
1311            .expect("Secure boot is only supported for UEFI firmware.")
1312            .secure_boot_enabled = true;
1313
1314        match self.os_flavor() {
1315            OsFlavor::Windows => self.with_windows_secure_boot_template(),
1316            OsFlavor::Linux => self.with_uefi_ca_secure_boot_template(),
1317            _ => panic!(
1318                "Secure boot unsupported for OS flavor {:?}",
1319                self.os_flavor()
1320            ),
1321        }
1322    }
1323
1324    /// Inject Windows secure boot templates into the VM's UEFI.
1325    pub fn with_windows_secure_boot_template(mut self) -> Self {
1326        self.config
1327            .firmware
1328            .uefi_config_mut()
1329            .expect("Secure boot is only supported for UEFI firmware.")
1330            .secure_boot_template = Some(SecureBootTemplate::MicrosoftWindows);
1331        self
1332    }
1333
1334    /// Inject UEFI CA secure boot templates into the VM's UEFI.
1335    pub fn with_uefi_ca_secure_boot_template(mut self) -> Self {
1336        self.config
1337            .firmware
1338            .uefi_config_mut()
1339            .expect("Secure boot is only supported for UEFI firmware.")
1340            .secure_boot_template = Some(SecureBootTemplate::MicrosoftUefiCertificateAuthority);
1341        self
1342    }
1343
1344    /// Apply a custom UEFI variable delta encoded as JSON.
1345    pub fn with_custom_uefi_json(mut self, json: impl Into<Vec<u8>>) -> Self {
1346        self.config
1347            .firmware
1348            .uefi_config_mut()
1349            .expect("Custom UEFI variables are only supported for UEFI firmware.")
1350            .custom_uefi_json = Some(json.into());
1351        self
1352    }
1353
1354    /// Set the VM to use the specified processor topology.
1355    pub fn with_processor_topology(mut self, topology: ProcessorTopology) -> Self {
1356        self.config.proc_topology = topology;
1357        self
1358    }
1359
1360    /// Set the VM to use the specified memory config.
1361    pub fn with_memory(mut self, memory: MemoryConfig) -> Self {
1362        self.config.memory = memory;
1363        self
1364    }
1365
1366    /// Sets a custom OpenHCL IGVM VTL2 address type. This controls the behavior
1367    /// of where VTL2 is placed in address space, and also the total size of memory
1368    /// allocated for VTL2. VTL2 start will fail if `address_type` is specified
1369    /// and leads to the loader allocating less memory than what is in the IGVM file.
1370    pub fn with_vtl2_base_address_type(mut self, address_type: Vtl2BaseAddressType) -> Self {
1371        self.config
1372            .firmware
1373            .openhcl_config_mut()
1374            .expect("OpenHCL firmware is required to set custom VTL2 address type.")
1375            .vtl2_base_address_type = Some(address_type);
1376        self
1377    }
1378
1379    /// Sets a custom OpenHCL IGVM file to use.
1380    pub fn with_custom_openhcl(mut self, artifact: ResolvedArtifact<impl IsOpenhclIgvm>) -> Self {
1381        match &mut self.config.firmware {
1382            Firmware::OpenhclLinuxDirect { igvm_path, .. }
1383            | Firmware::OpenhclPcat { igvm_path, .. }
1384            | Firmware::OpenhclUefi { igvm_path, .. } => {
1385                *igvm_path = artifact.erase();
1386            }
1387            Firmware::LinuxDirect { .. } | Firmware::Uefi { .. } | Firmware::Pcat { .. } => {
1388                panic!("Custom OpenHCL is only supported for OpenHCL firmware.")
1389            }
1390        }
1391        self
1392    }
1393
1394    /// Append additional command line arguments to pass to the paravisor.
1395    pub fn with_openhcl_command_line(mut self, additional_command_line: &str) -> Self {
1396        append_cmdline(
1397            &mut self
1398                .config
1399                .firmware
1400                .openhcl_config_mut()
1401                .expect("OpenHCL command line is only supported for OpenHCL firmware.")
1402                .custom_command_line,
1403            additional_command_line,
1404        );
1405        self
1406    }
1407
1408    /// Enable confidential filtering, even if the VM is not confidential.
1409    pub fn with_confidential_filtering(self) -> Self {
1410        if !self.config.firmware.is_openhcl() {
1411            panic!("Confidential filtering is only supported for OpenHCL");
1412        }
1413        self.with_openhcl_command_line(&format!(
1414            "{}=1 {}=0",
1415            underhill_confidentiality::OPENHCL_CONFIDENTIAL_ENV_VAR_NAME,
1416            underhill_confidentiality::OPENHCL_CONFIDENTIAL_DEBUG_ENV_VAR_NAME
1417        ))
1418    }
1419
1420    /// Sets the command line parameters passed to OpenHCL related to logging.
1421    pub fn with_openhcl_log_levels(mut self, levels: OpenvmmLogConfig) -> Self {
1422        self.config
1423            .firmware
1424            .openhcl_config_mut()
1425            .expect("OpenHCL firmware is required to set custom OpenHCL log levels.")
1426            .log_levels = levels;
1427        self
1428    }
1429
1430    /// Sets the log levels for the host OpenVMM process.
1431    /// DEVNOTE: In the future, this could be generalized for both HyperV and OpenVMM.
1432    /// For now, this is only implemented for OpenVMM.
1433    pub fn with_host_log_levels(mut self, levels: OpenvmmLogConfig) -> Self {
1434        if let OpenvmmLogConfig::Custom(ref custom_levels) = levels {
1435            for key in custom_levels.keys() {
1436                if !["OPENVMM_LOG", "OPENVMM_SHOW_SPANS"].contains(&key.as_str()) {
1437                    panic!("Unsupported OpenVMM log level key: {}", key);
1438                }
1439            }
1440        }
1441
1442        self.config.host_log_levels = Some(levels.clone());
1443        self
1444    }
1445
1446    /// Adds a file to the VM's pipette agent image.
1447    pub fn with_agent_file(mut self, name: &str, artifact: ResolvedArtifact) -> Self {
1448        self.agent_image
1449            .as_mut()
1450            .expect("no guest pipette")
1451            .add_file(name, artifact);
1452        self
1453    }
1454
1455    /// Adds a file to the paravisor's pipette agent image.
1456    pub fn with_openhcl_agent_file(mut self, name: &str, artifact: ResolvedArtifact) -> Self {
1457        self.openhcl_agent_image
1458            .as_mut()
1459            .expect("no openhcl pipette")
1460            .add_file(name, artifact);
1461        self
1462    }
1463
1464    /// Sets whether UEFI frontpage is enabled.
1465    pub fn with_uefi_frontpage(mut self, enable: bool) -> Self {
1466        self.config
1467            .firmware
1468            .uefi_config_mut()
1469            .expect("UEFI frontpage is only supported for UEFI firmware.")
1470            .disable_frontpage = !enable;
1471        self
1472    }
1473
1474    /// Sets the UEFI diagnostics log level filter.
1475    ///
1476    /// By default only ERROR and WARN level entries are forwarded to the
1477    /// host tracing infrastructure. Use this to also surface INFO (or all)
1478    /// entries when a test needs to observe them.
1479    pub fn with_efi_diagnostics_log_level(mut self, level: EfiDiagnosticsLogLevel) -> Self {
1480        self.config
1481            .firmware
1482            .uefi_config_mut()
1483            .expect("EFI diagnostics log level is only supported for UEFI firmware.")
1484            .efi_diagnostics_log_level = level;
1485        self
1486    }
1487
1488    /// Sets the per-period rate-limit override for UEFI diagnostics emission.
1489    ///
1490    /// - Not called: use the built-in defaults.
1491    /// - `0`: disable rate limiting entirely (emit every entry).
1492    /// - `n > 0`: use `n` as the per-period limit.
1493    pub fn with_efi_diagnostics_rate_limit(mut self, limit: u32) -> Self {
1494        self.config
1495            .firmware
1496            .uefi_config_mut()
1497            .expect("EFI diagnostics rate limit is only supported for UEFI firmware.")
1498            .efi_diagnostics_rate_limit = Some(limit);
1499        self
1500    }
1501
1502    /// Sets whether UEFI should always attempt a default boot.
1503    pub fn with_default_boot_always_attempt(mut self, enable: bool) -> Self {
1504        self.config
1505            .firmware
1506            .uefi_config_mut()
1507            .expect("Default boot always attempt is only supported for UEFI firmware.")
1508            .default_boot_always_attempt = enable;
1509        self
1510    }
1511
1512    /// Force UEFI to bounce-buffer all DMA traffic.
1513    pub fn with_uefi_force_dma_bounce(mut self, enable: bool) -> Self {
1514        self.config
1515            .firmware
1516            .uefi_config_mut()
1517            .expect("force DMA bounce is only supported for UEFI firmware.")
1518            .force_dma_bounce = enable;
1519        self
1520    }
1521
1522    /// Run the VM with Enable VMBus relay enabled
1523    pub fn with_vmbus_redirect(mut self, enable: bool) -> Self {
1524        self.config
1525            .firmware
1526            .openhcl_config_mut()
1527            .expect("VMBus redirection is only supported for OpenHCL firmware.")
1528            .vmbus_redirect = enable;
1529        self
1530    }
1531
1532    /// Specify the guest state lifetime for the VM
1533    pub fn with_guest_state_lifetime(
1534        mut self,
1535        guest_state_lifetime: PetriGuestStateLifetime,
1536    ) -> Self {
1537        let disk = match self.config.vmgs {
1538            PetriVmgsResource::Disk(disk)
1539            | PetriVmgsResource::ReprovisionOnFailure(disk)
1540            | PetriVmgsResource::Reprovision(disk) => disk,
1541            PetriVmgsResource::Ephemeral => PetriVmgsDisk::default(),
1542        };
1543        self.config.vmgs = match guest_state_lifetime {
1544            PetriGuestStateLifetime::Disk => PetriVmgsResource::Disk(disk),
1545            PetriGuestStateLifetime::ReprovisionOnFailure => {
1546                PetriVmgsResource::ReprovisionOnFailure(disk)
1547            }
1548            PetriGuestStateLifetime::Reprovision => PetriVmgsResource::Reprovision(disk),
1549            PetriGuestStateLifetime::Ephemeral => {
1550                if !matches!(disk.disk, Disk::Memory(_)) {
1551                    panic!("attempted to use ephemeral guest state after specifying backing vmgs")
1552                }
1553                PetriVmgsResource::Ephemeral
1554            }
1555        };
1556        self
1557    }
1558
1559    /// Specify the guest state encryption policy for the VM
1560    pub fn with_guest_state_encryption(mut self, policy: GuestStateEncryptionPolicy) -> Self {
1561        match &mut self.config.vmgs {
1562            PetriVmgsResource::Disk(vmgs)
1563            | PetriVmgsResource::ReprovisionOnFailure(vmgs)
1564            | PetriVmgsResource::Reprovision(vmgs) => {
1565                vmgs.encryption_policy = policy;
1566            }
1567            PetriVmgsResource::Ephemeral => {
1568                panic!("attempted to encrypt ephemeral guest state")
1569            }
1570        }
1571        self
1572    }
1573
1574    /// Use the specified backing VMGS file
1575    pub fn with_initial_vmgs(self, disk: ResolvedArtifact<impl IsTestVmgs>) -> Self {
1576        self.with_backing_vmgs(Disk::Differencing(DiskPath::Local(disk.into())))
1577    }
1578
1579    /// Use the specified backing VMGS file
1580    pub fn with_persistent_vmgs(self, disk: impl AsRef<Path>) -> Self {
1581        self.with_backing_vmgs(Disk::Persistent(disk.as_ref().to_path_buf()))
1582    }
1583
1584    fn with_backing_vmgs(mut self, disk: Disk) -> Self {
1585        match &mut self.config.vmgs {
1586            PetriVmgsResource::Disk(vmgs)
1587            | PetriVmgsResource::ReprovisionOnFailure(vmgs)
1588            | PetriVmgsResource::Reprovision(vmgs) => {
1589                if !matches!(vmgs.disk, Disk::Memory(_)) {
1590                    panic!("already specified a backing vmgs file");
1591                }
1592                vmgs.disk = disk;
1593            }
1594            PetriVmgsResource::Ephemeral => {
1595                panic!("attempted to specify a backing vmgs with ephemeral guest state")
1596            }
1597        }
1598        self
1599    }
1600
1601    /// Set the boot device type for the VM.
1602    ///
1603    /// This overrides the default, which is determined by the firmware type.
1604    pub fn with_boot_device_type(mut self, boot: BootDeviceType) -> Self {
1605        self.boot_device_type = boot;
1606        self
1607    }
1608
1609    /// Override the PCIe root port that the boot NVMe controller is placed on
1610    /// when using [`BootDeviceType::PcieNvme`].
1611    ///
1612    /// The named port must exist in the PCIe topology added via the backend's
1613    /// `with_pcie_root_topology`. Defaults to `s0rc0rp0`.
1614    pub fn with_pcie_boot_port(mut self, port_name: &str) -> Self {
1615        self.pcie_boot_port = Some(port_name.to_string());
1616        self
1617    }
1618
1619    /// Enable the TPM for the VM.
1620    pub fn with_tpm(mut self, enable: bool) -> Self {
1621        if enable {
1622            self.config.tpm.get_or_insert_default();
1623        } else {
1624            self.config.tpm = None;
1625        }
1626        self
1627    }
1628
1629    /// Enable or disable the TPM state persistence for the VM.
1630    pub fn with_tpm_state_persistence(mut self, tpm_state_persistence: bool) -> Self {
1631        self.config
1632            .tpm
1633            .as_mut()
1634            .expect("TPM persistence requires a TPM")
1635            .no_persistent_secrets = !tpm_state_persistence;
1636        self
1637    }
1638
1639    /// Set the hardware sealing policy for the VM's TPM.
1640    pub fn with_hardware_sealing_policy(mut self, policy: PetriHardwareSealingPolicy) -> Self {
1641        self.config
1642            .tpm
1643            .as_mut()
1644            .expect("hardware sealing policy requires a TPM")
1645            .hardware_sealing_policy = policy;
1646        self
1647    }
1648
1649    /// Add custom VTL 2 settings.
1650    // TODO: At some point we want to replace uses of this with nicer with_disk,
1651    // with_nic, etc. methods.
1652    pub fn with_custom_vtl2_settings(
1653        mut self,
1654        f: impl FnOnce(&mut Vtl2Settings) + 'static + Send + Sync,
1655    ) -> Self {
1656        f(self
1657            .config
1658            .firmware
1659            .vtl2_settings()
1660            .expect("Custom VTL 2 settings are only supported with OpenHCL"));
1661        self
1662    }
1663
1664    /// Add a storage controller to VTL2
1665    pub fn add_vtl2_storage_controller(self, controller: StorageController) -> Self {
1666        self.with_custom_vtl2_settings(move |v| {
1667            v.dynamic
1668                .as_mut()
1669                .unwrap()
1670                .storage_controllers
1671                .push(controller)
1672        })
1673    }
1674
1675    /// Add an additional SCSI controller to the VM.
1676    pub fn add_vmbus_storage_controller(
1677        mut self,
1678        id: &Guid,
1679        target_vtl: Vtl,
1680        controller_type: VmbusStorageType,
1681    ) -> Self {
1682        if self
1683            .config
1684            .vmbus_storage_controllers
1685            .insert(
1686                *id,
1687                VmbusStorageController::new(target_vtl, controller_type),
1688            )
1689            .is_some()
1690        {
1691            panic!("storage controller {id} already existed");
1692        }
1693        self
1694    }
1695
1696    /// Add a VMBus disk drive to the VM
1697    pub fn add_vmbus_drive(
1698        mut self,
1699        drive: Drive,
1700        controller_id: &Guid,
1701        controller_location: Option<u32>,
1702    ) -> Self {
1703        let controller = self
1704            .config
1705            .vmbus_storage_controllers
1706            .get_mut(controller_id)
1707            .unwrap_or_else(|| panic!("storage controller {controller_id} does not exist"));
1708
1709        _ = controller.set_drive(controller_location, drive, false);
1710
1711        self
1712    }
1713
1714    /// Add a VMBus disk drive to the VM
1715    pub fn add_ide_drive(
1716        mut self,
1717        drive: Drive,
1718        controller_number: u32,
1719        controller_location: u8,
1720    ) -> Self {
1721        self.config
1722            .firmware
1723            .ide_controllers_mut()
1724            .expect("Host IDE requires PCAT with no HCL")[controller_number as usize]
1725            [controller_location as usize] = Some(drive);
1726
1727        self
1728    }
1729
1730    /// Add a physical NVMe device to the VM
1731    pub fn add_physical_nvme_device(mut self, vsid: Guid, device: PhysicalNvmeDevice) -> Self {
1732        if self
1733            .config
1734            .physical_nvme_devices
1735            .insert(vsid, device)
1736            .is_some()
1737        {
1738            panic!("physical NVMe device {vsid} already existed");
1739        }
1740        self
1741    }
1742
1743    /// Get VM's guest OS flavor
1744    pub fn os_flavor(&self) -> OsFlavor {
1745        self.config.firmware.os_flavor()
1746    }
1747
1748    /// Get whether the VM will use OpenHCL
1749    pub fn is_openhcl(&self) -> bool {
1750        self.config.firmware.is_openhcl()
1751    }
1752
1753    /// Get the isolation type of the VM
1754    pub fn isolation(&self) -> Option<IsolationType> {
1755        self.config.firmware.isolation()
1756    }
1757
1758    /// Get the machine architecture
1759    pub fn arch(&self) -> MachineArch {
1760        self.config.arch
1761    }
1762
1763    /// Get the log source for creating additional log files.
1764    pub fn log_source(&self) -> &PetriLogSource {
1765        &self.resources.log_source
1766    }
1767
1768    /// Get the default OpenHCL servicing flags for this config
1769    pub fn default_servicing_flags(&self) -> OpenHclServicingFlags {
1770        T::default_servicing_flags()
1771    }
1772
1773    /// Get the backend-specific config builder
1774    pub fn modify_backend(
1775        mut self,
1776        f: impl FnOnce(T::VmmConfig) -> T::VmmConfig + 'static + Send,
1777    ) -> Self {
1778        if self.modify_vmm_config.is_some() {
1779            panic!("only one modify_backend allowed");
1780        }
1781        self.modify_vmm_config = Some(ModifyFn(Box::new(f)));
1782        self
1783    }
1784}
1785
1786impl<T: PetriVmmBackend> PetriVm<T> {
1787    /// Immediately tear down the VM.
1788    pub async fn teardown(self) -> anyhow::Result<()> {
1789        tracing::info!("Tearing down VM...");
1790        self.runtime.teardown().await
1791    }
1792
1793    /// Wait for the VM to halt, returning the reason for the halt.
1794    pub async fn wait_for_halt(&mut self) -> anyhow::Result<PetriHaltReasonDetail> {
1795        tracing::info!("Waiting for VM to halt...");
1796        let halt_reason = self.runtime.wait_for_halt(false).await?;
1797        tracing::info!("VM halted: {halt_reason:?}. Cancelling watchdogs...");
1798        futures::future::join_all(self.watchdog_tasks.drain(..).map(|t| t.cancel())).await;
1799        Ok(halt_reason)
1800    }
1801
1802    /// Wait for the VM to cleanly shutdown.
1803    pub async fn wait_for_clean_shutdown(&mut self) -> anyhow::Result<()> {
1804        let halt_reason = self.wait_for_halt().await?;
1805        if halt_reason.reason != PetriHaltReason::PowerOff {
1806            anyhow::bail!("Expected PowerOff, got {halt_reason:?}");
1807        }
1808        tracing::info!("VM was cleanly powered off and torn down.");
1809        Ok(())
1810    }
1811
1812    /// Wait for the VM to halt, returning the reason for the halt,
1813    /// and tear down the VM.
1814    pub async fn wait_for_teardown(mut self) -> anyhow::Result<PetriHaltReasonDetail> {
1815        let halt_reason = self.wait_for_halt().await?;
1816        self.teardown().await?;
1817        Ok(halt_reason)
1818    }
1819
1820    /// Wait for the VM to cleanly shutdown and tear down the VM.
1821    pub async fn wait_for_clean_teardown(mut self) -> anyhow::Result<()> {
1822        self.wait_for_clean_shutdown().await?;
1823        self.teardown().await
1824    }
1825
1826    /// Wait for the VM to reset. Does not wait for pipette.
1827    pub async fn wait_for_reset_no_agent(&mut self) -> anyhow::Result<()> {
1828        self.wait_for_reset_core().await?;
1829        self.wait_for_expected_boot_event().await?;
1830        Ok(())
1831    }
1832
1833    /// Wait for the VM to reset and pipette to connect.
1834    pub async fn wait_for_reset(&mut self) -> anyhow::Result<PipetteClient> {
1835        self.wait_for_reset_no_agent().await?;
1836        self.wait_for_agent().await
1837    }
1838
1839    async fn wait_for_reset_core(&mut self) -> anyhow::Result<()> {
1840        tracing::info!("Waiting for VM to reset...");
1841        let halt_reason = self.runtime.wait_for_halt(true).await?;
1842        if halt_reason.reason != PetriHaltReason::Reset {
1843            anyhow::bail!("Expected reset, got {halt_reason:?}");
1844        }
1845        tracing::info!("VM reset.");
1846        Ok(())
1847    }
1848
1849    /// Invoke Inspect on the running OpenHCL instance.
1850    ///
1851    /// IMPORTANT: As mentioned in the Guide, inspect output is *not* guaranteed
1852    /// to be stable. Use this to test that components in OpenHCL are working as
1853    /// you would expect. But, if you are adding a test simply to verify that
1854    /// the inspect output as some other tool depends on it, then that is
1855    /// incorrect.
1856    ///
1857    /// - `timeout` is enforced on the client side
1858    /// - `path` and `depth` are passed to the [`inspect::Inspect`] machinery.
1859    pub async fn inspect_openhcl(
1860        &self,
1861        path: impl Into<String>,
1862        depth: Option<usize>,
1863        timeout: Option<Duration>,
1864    ) -> anyhow::Result<inspect::Node> {
1865        self.openhcl_diag()?
1866            .inspect(path.into().as_str(), depth, timeout)
1867            .await
1868    }
1869
1870    /// Invoke Update (Inspect protocol) on the running OpenHCL instance.
1871    ///
1872    /// IMPORTANT: As mentioned in the Guide, inspect output is *not* guaranteed
1873    /// to be stable. Use this to test that components in OpenHCL are working as
1874    /// you would expect. But, if you are adding a test simply to verify that
1875    /// the inspect output as some other tool depends on it, then that is
1876    /// incorrect.
1877    ///
1878    /// - `path` and `value` are passed to the [`inspect::Inspect`] machinery.
1879    pub async fn inspect_update_openhcl(
1880        &self,
1881        path: impl Into<String>,
1882        value: impl Into<String>,
1883    ) -> anyhow::Result<inspect::Value> {
1884        self.openhcl_diag()?
1885            .inspect_update(path.into(), value.into())
1886            .await
1887    }
1888
1889    /// Test that we are able to inspect OpenHCL.
1890    pub async fn test_inspect_openhcl(&mut self) -> anyhow::Result<()> {
1891        self.inspect_openhcl("", None, None).await.map(|_| ())
1892    }
1893
1894    /// Invoke Inspect on the running VMM process itself (e.g. OpenVMM),
1895    /// returning the inspect tree rooted at `path` (pass `""` for the whole
1896    /// tree).
1897    ///
1898    /// Only backends that expose an inspect interface (currently OpenVMM)
1899    /// support this; other backends return an error.
1900    ///
1901    /// IMPORTANT: As mentioned in the Guide, inspect output is *not* guaranteed
1902    /// to be stable. Use this to verify that components are working as you
1903    /// expect, not to assert on output that some other tool depends on.
1904    pub async fn inspect_vmm(&self, path: &str) -> anyhow::Result<inspect::Node> {
1905        use anyhow::Context;
1906
1907        let inspector = self
1908            .runtime
1909            .inspector()
1910            .context("this VMM backend does not support inspect")?;
1911        inspector.inspect(path).await
1912    }
1913
1914    /// Wait for VTL 2 to report that it is ready to respond to commands.
1915    /// Will fail if the VM is not running OpenHCL.
1916    ///
1917    /// This should only be necessary if you're doing something manual. All
1918    /// Petri-provided methods will wait for VTL 2 to be ready automatically.
1919    pub async fn wait_for_vtl2_ready(&mut self) -> anyhow::Result<()> {
1920        self.openhcl_diag()?.wait_for_vtl2().await
1921    }
1922
1923    /// Get the kmsg stream from OpenHCL.
1924    pub async fn kmsg(&self) -> anyhow::Result<diag_client::kmsg_stream::KmsgStream> {
1925        self.openhcl_diag()?.kmsg().await
1926    }
1927
1928    /// Gets a live core dump of the OpenHCL process specified by 'name' and
1929    /// writes it to 'path'
1930    pub async fn openhcl_core_dump(&self, name: &str, path: &Path) -> anyhow::Result<()> {
1931        self.openhcl_diag()?.core_dump(name, path).await
1932    }
1933
1934    /// Crashes the specified openhcl process
1935    pub async fn openhcl_crash(&self, name: &str) -> anyhow::Result<()> {
1936        self.openhcl_diag()?.crash(name).await
1937    }
1938
1939    /// Wait for a connection from a pipette agent running in the guest.
1940    /// Useful if you've rebooted the vm or are otherwise expecting a fresh connection.
1941    async fn wait_for_agent(&mut self) -> anyhow::Result<PipetteClient> {
1942        // As a workaround for #2470 (where the guest crashes when the pipette
1943        // connection timeout expires due to a vmbus bug), wait for the shutdown
1944        // IC to come online first so that we probably won't time out when
1945        // connecting to the agent.
1946        // TODO: remove this once the bug is fixed, since it shouldn't be
1947        // necessary and a guest could in theory support pipette and not the IC
1948        //
1949        // This is a no-op when the shutdown IC is not configured (e.g.,
1950        // no VMBus or minimal mode).
1951        self.runtime.wait_for_enlightened_shutdown_ready().await?;
1952        self.runtime.wait_for_agent(false).await
1953    }
1954
1955    /// Wait for a connection from a pipette agent running in VTL 2.
1956    /// Useful if you've reset VTL 2 or are otherwise expecting a fresh connection.
1957    /// Will fail if the VM is not running OpenHCL.
1958    pub async fn wait_for_vtl2_agent(&mut self) -> anyhow::Result<PipetteClient> {
1959        // VTL 2's pipette doesn't auto launch, only launch it on demand
1960        self.launch_vtl2_pipette().await?;
1961        self.runtime.wait_for_agent(true).await
1962    }
1963
1964    /// Waits for an event emitted by the firmware about its boot status, and
1965    /// verifies that it is the expected success value.
1966    ///
1967    /// * Linux Direct guests do not emit a boot event, so this method immediately returns Ok.
1968    /// * PCAT guests may not emit an event depending on the PCAT version, this
1969    ///   method is best effort for them.
1970    async fn wait_for_expected_boot_event(&mut self) -> anyhow::Result<()> {
1971        if let Some(expected_event) = self.expected_boot_event {
1972            let event = self.wait_for_boot_event().await?;
1973
1974            anyhow::ensure!(
1975                event == expected_event,
1976                "Did not receive expected boot event"
1977            );
1978        } else {
1979            tracing::warn!("Boot event not emitted for configured firmware or manually ignored.");
1980        }
1981
1982        Ok(())
1983    }
1984
1985    /// Waits for an event emitted by the firmware about its boot status, and
1986    /// returns that status.
1987    async fn wait_for_boot_event(&mut self) -> anyhow::Result<FirmwareEvent> {
1988        tracing::info!("Waiting for boot event...");
1989        let boot_event = loop {
1990            if let Some(event) = self
1991                .runtime
1992                .wait_for_boot_event(self.vmm_quirks.flaky_boot)
1993                .await?
1994            {
1995                break event;
1996            }
1997
1998            tracing::error!("Did not get boot event in required time, resetting...");
1999            if let Some(inspector) = self.runtime.inspector() {
2000                save_inspect(
2001                    "vmm",
2002                    Box::pin(async move { inspector.inspect("").await }),
2003                    &self.resources.log_source,
2004                )
2005                .await;
2006            }
2007
2008            self.runtime.reset().await?;
2009        };
2010        tracing::info!("Got boot event: {boot_event:?}");
2011        Ok(boot_event)
2012    }
2013
2014    /// Wait for the Hyper-V shutdown IC to be ready and use it to instruct
2015    /// the guest to shutdown.
2016    pub async fn send_enlightened_shutdown(&mut self, kind: ShutdownKind) -> anyhow::Result<()> {
2017        tracing::info!("Waiting for enlightened shutdown to be ready");
2018        self.runtime.wait_for_enlightened_shutdown_ready().await?;
2019
2020        // all guests used in testing have been observed to intermittently
2021        // drop shutdown requests if they are sent too soon after the shutdown
2022        // ic comes online. give them a little extra time.
2023        // TODO: use a different method of determining whether the VM has booted
2024        // or debug and fix the shutdown IC.
2025        let mut wait_time = Duration::from_secs(10);
2026
2027        // some guests need even more time
2028        if let Some(duration) = self.guest_quirks.hyperv_shutdown_ic_sleep {
2029            wait_time += duration;
2030        }
2031
2032        tracing::info!(
2033            "Shutdown IC reported ready, waiting for an extra {}s",
2034            wait_time.as_secs()
2035        );
2036        PolledTimer::new(&self.resources.driver)
2037            .sleep(wait_time)
2038            .await;
2039
2040        tracing::info!("Sending enlightened shutdown command");
2041        self.runtime.send_enlightened_shutdown(kind).await
2042    }
2043
2044    /// Instruct the OpenHCL to restart the VTL2 paravisor. Will fail if the VM
2045    /// is not running OpenHCL. Will also fail if the VM is not running.
2046    pub async fn restart_openhcl(
2047        &mut self,
2048        new_openhcl: ResolvedArtifact<impl IsOpenhclIgvm>,
2049        flags: OpenHclServicingFlags,
2050    ) -> anyhow::Result<()> {
2051        self.runtime
2052            .restart_openhcl(&new_openhcl.erase(), flags)
2053            .await
2054    }
2055
2056    /// Update the command line parameter of the running VM that will apply on next boot.
2057    /// Will fail if the VM is not using IGVM load mode.
2058    pub async fn update_command_line(&mut self, command_line: &str) -> anyhow::Result<()> {
2059        self.runtime.update_command_line(command_line).await
2060    }
2061
2062    /// Hot-add a PCIe device to a named port at runtime.
2063    pub async fn add_pcie_device(
2064        &mut self,
2065        port_name: String,
2066        resource: vm_resource::Resource<vm_resource::kind::PciDeviceHandleKind>,
2067    ) -> anyhow::Result<()> {
2068        self.runtime.add_pcie_device(port_name, resource).await
2069    }
2070
2071    /// Hot-remove a PCIe device from a named port at runtime.
2072    pub async fn remove_pcie_device(&mut self, port_name: String) -> anyhow::Result<()> {
2073        self.runtime.remove_pcie_device(port_name).await
2074    }
2075
2076    /// Instruct the OpenHCL to save the state of the VTL2 paravisor. Will fail if the VM
2077    /// is not running OpenHCL. Will also fail if the VM is not running or if this is called twice in succession
2078    pub async fn save_openhcl(
2079        &mut self,
2080        new_openhcl: ResolvedArtifact<impl IsOpenhclIgvm>,
2081        flags: OpenHclServicingFlags,
2082    ) -> anyhow::Result<()> {
2083        self.runtime.save_openhcl(&new_openhcl.erase(), flags).await
2084    }
2085
2086    /// Instruct the OpenHCL to restore the state of the VTL2 paravisor. Will fail if the VM
2087    /// is not running OpenHCL. Will also fail if the VM is running or if this is called without prior save
2088    pub async fn restore_openhcl(&mut self) -> anyhow::Result<()> {
2089        self.runtime.restore_openhcl().await
2090    }
2091
2092    /// Get VM's guest OS flavor
2093    pub fn arch(&self) -> MachineArch {
2094        self.arch
2095    }
2096
2097    /// Get the inner runtime backend to make backend-specific calls
2098    pub fn backend(&mut self) -> &mut T::VmRuntime {
2099        &mut self.runtime
2100    }
2101
2102    async fn launch_vtl2_pipette(&self) -> anyhow::Result<()> {
2103        tracing::debug!("Launching VTL 2 pipette...");
2104
2105        // Start pipette through DiagClient
2106        let res = self
2107            .openhcl_diag()?
2108            .run_vtl2_command("sh", &["-c", "mkdir /cidata && mount LABEL=cidata /cidata"])
2109            .await?;
2110
2111        if !res.exit_status.success() {
2112            anyhow::bail!("Failed to mount VTL 2 pipette drive: {:?}", res);
2113        }
2114
2115        let res = self
2116            .openhcl_diag()?
2117            .run_detached_vtl2_command("sh", &["-c", "/cidata/pipette 2>&1 | logger &"])
2118            .await?;
2119
2120        if !res.success() {
2121            anyhow::bail!("Failed to spawn VTL 2 pipette: {:?}", res);
2122        }
2123
2124        Ok(())
2125    }
2126
2127    fn openhcl_diag(&self) -> anyhow::Result<&OpenHclDiagHandler> {
2128        if let Some(ohd) = self.openhcl_diag_handler.as_ref() {
2129            Ok(ohd)
2130        } else {
2131            anyhow::bail!("VM is not configured with OpenHCL")
2132        }
2133    }
2134
2135    /// Get the path to the VM's guest state file
2136    pub async fn get_guest_state_file(&self) -> anyhow::Result<Option<PathBuf>> {
2137        self.runtime.get_guest_state_file().await
2138    }
2139
2140    /// Modify OpenHCL VTL2 settings.
2141    pub async fn modify_vtl2_settings(
2142        &mut self,
2143        f: impl FnOnce(&mut Vtl2Settings),
2144    ) -> anyhow::Result<()> {
2145        if self.openhcl_diag_handler.is_none() {
2146            panic!("Custom VTL 2 settings are only supported with OpenHCL");
2147        }
2148        f(self
2149            .config
2150            .vtl2_settings
2151            .get_or_insert_with(default_vtl2_settings));
2152        self.runtime
2153            .set_vtl2_settings(self.config.vtl2_settings.as_ref().unwrap())
2154            .await
2155    }
2156
2157    /// Get the list of storage controllers added to this VM
2158    pub fn get_vmbus_storage_controllers(&self) -> &HashMap<Guid, VmbusStorageController> {
2159        &self.config.vmbus_storage_controllers
2160    }
2161
2162    /// Add or modify a VMBus disk drive
2163    pub async fn set_vmbus_drive(
2164        &mut self,
2165        drive: Drive,
2166        controller_id: &Guid,
2167        controller_location: Option<u32>,
2168    ) -> anyhow::Result<()> {
2169        let controller = self
2170            .config
2171            .vmbus_storage_controllers
2172            .get_mut(controller_id)
2173            .unwrap_or_else(|| panic!("storage controller {controller_id} does not exist"));
2174
2175        let controller_location = controller.set_drive(controller_location, drive, true);
2176        let disk = controller.drives.get(&controller_location).unwrap();
2177
2178        self.runtime
2179            .set_vmbus_drive(disk, controller_id, controller_location)
2180            .await?;
2181
2182        Ok(())
2183    }
2184}
2185
2186/// A running VM that tests can interact with.
2187#[async_trait]
2188pub trait PetriVmRuntime: Send + Sync + 'static {
2189    /// Interface for inspecting the VM
2190    type VmInspector: PetriVmInspector;
2191    /// Interface for accessing the framebuffer
2192    type VmFramebufferAccess: PetriVmFramebufferAccess;
2193
2194    /// Cleanly tear down the VM immediately.
2195    async fn teardown(self) -> anyhow::Result<()>;
2196    /// Wait for the VM to halt, returning the reason for the halt. The VM
2197    /// should automatically restart the VM on reset if `allow_reset` is true.
2198    async fn wait_for_halt(&mut self, allow_reset: bool) -> anyhow::Result<PetriHaltReasonDetail>;
2199    /// Wait for a connection from a pipette agent
2200    async fn wait_for_agent(&mut self, set_high_vtl: bool) -> anyhow::Result<PipetteClient>;
2201    /// Get an OpenHCL diagnostics handler for the VM
2202    fn openhcl_diag(&self) -> Option<OpenHclDiagHandler>;
2203    /// Waits for an event emitted by the firmware about its boot status, and
2204    /// returns that status. Returns `None` if `timeout` elapsed first.
2205    async fn wait_for_boot_event(
2206        &mut self,
2207        timeout: Option<Duration>,
2208    ) -> anyhow::Result<Option<FirmwareEvent>>;
2209    /// Waits for the Hyper-V shutdown IC to be ready
2210    // TODO: return a receiver that will be closed when it is no longer ready.
2211    async fn wait_for_enlightened_shutdown_ready(&mut self) -> anyhow::Result<()>;
2212    /// Instruct the guest to shutdown via the Hyper-V shutdown IC.
2213    async fn send_enlightened_shutdown(&mut self, kind: ShutdownKind) -> anyhow::Result<()>;
2214    /// Instruct the OpenHCL to restart the VTL2 paravisor. Will fail if the VM
2215    /// is not running OpenHCL. Will also fail if the VM is not running.
2216    async fn restart_openhcl(
2217        &mut self,
2218        new_openhcl: &ResolvedArtifact,
2219        flags: OpenHclServicingFlags,
2220    ) -> anyhow::Result<()>;
2221    /// Instruct the OpenHCL to save the state of the VTL2 paravisor. Will fail if the VM
2222    /// is not running OpenHCL. Will also fail if the VM is not running or if this is called twice in succession
2223    /// without a call to `restore_openhcl`.
2224    async fn save_openhcl(
2225        &mut self,
2226        new_openhcl: &ResolvedArtifact,
2227        flags: OpenHclServicingFlags,
2228    ) -> anyhow::Result<()>;
2229    /// Instruct the OpenHCL to restore the state of the VTL2 paravisor. Will fail if the VM
2230    /// is not running OpenHCL. Will also fail if the VM is running or if this is called without prior save.
2231    async fn restore_openhcl(&mut self) -> anyhow::Result<()>;
2232    /// Update the command line parameter of the running VM that will apply on next boot.
2233    /// Will fail if the VM is not using IGVM load mode.
2234    async fn update_command_line(&mut self, command_line: &str) -> anyhow::Result<()>;
2235    /// If the backend supports it, get an inspect interface
2236    fn inspector(&self) -> Option<Self::VmInspector> {
2237        None
2238    }
2239    /// If the backend supports it, take the screenshot interface
2240    /// (subsequent calls may return None).
2241    fn take_framebuffer_access(&mut self) -> Option<Self::VmFramebufferAccess> {
2242        None
2243    }
2244    /// Issue a hard reset to the VM
2245    async fn reset(&mut self) -> anyhow::Result<()>;
2246    /// Get the path to the VM's guest state file
2247    async fn get_guest_state_file(&self) -> anyhow::Result<Option<PathBuf>> {
2248        Ok(None)
2249    }
2250    /// Set the OpenHCL VTL2 settings
2251    async fn set_vtl2_settings(&mut self, settings: &Vtl2Settings) -> anyhow::Result<()>;
2252    /// Add or modify a VMBus disk drive
2253    async fn set_vmbus_drive(
2254        &mut self,
2255        disk: &Drive,
2256        controller_id: &Guid,
2257        controller_location: u32,
2258    ) -> anyhow::Result<()>;
2259    /// Hot-add a PCIe device to a named port at runtime.
2260    async fn add_pcie_device(
2261        &mut self,
2262        port_name: String,
2263        resource: vm_resource::Resource<vm_resource::kind::PciDeviceHandleKind>,
2264    ) -> anyhow::Result<()> {
2265        let _ = (port_name, resource);
2266        anyhow::bail!("PCIe hotplug not supported by this backend")
2267    }
2268    /// Hot-remove a PCIe device from a named port at runtime.
2269    async fn remove_pcie_device(&mut self, port_name: String) -> anyhow::Result<()> {
2270        let _ = port_name;
2271        anyhow::bail!("PCIe hotplug not supported by this backend")
2272    }
2273}
2274
2275/// Interface for getting information about the state of the VM
2276#[async_trait]
2277pub trait PetriVmInspector: Send + Sync + 'static {
2278    /// Get information about the state of the VM at the given inspect `path`.
2279    /// Pass `""` to inspect the entire tree.
2280    async fn inspect(&self, path: &str) -> anyhow::Result<inspect::Node>;
2281}
2282
2283/// Use this for the associated type if not supported
2284pub struct NoPetriVmInspector;
2285#[async_trait]
2286impl PetriVmInspector for NoPetriVmInspector {
2287    async fn inspect(&self, _path: &str) -> anyhow::Result<inspect::Node> {
2288        unreachable!()
2289    }
2290}
2291
2292/// Raw VM screenshot
2293pub struct VmScreenshotMeta {
2294    /// color encoding used by the image
2295    pub color: image::ExtendedColorType,
2296    /// x dimension
2297    pub width: u16,
2298    /// y dimension
2299    pub height: u16,
2300}
2301
2302/// Interface for getting screenshots of the VM
2303#[async_trait]
2304pub trait PetriVmFramebufferAccess: Send + 'static {
2305    /// Populates the provided buffer with a screenshot of the VM,
2306    /// returning the dimensions and color type.
2307    async fn screenshot(&mut self, image: &mut Vec<u8>)
2308    -> anyhow::Result<Option<VmScreenshotMeta>>;
2309}
2310
2311/// Common processor topology information for the VM.
2312#[derive(Debug)]
2313pub struct ProcessorTopology {
2314    /// The number of virtual processors.
2315    pub vp_count: u32,
2316    /// Whether SMT (hyperthreading) is enabled.
2317    pub enable_smt: Option<bool>,
2318    /// The number of virtual processors per socket.
2319    pub vps_per_socket: Option<u32>,
2320    /// The APIC configuration (x86-64 only).
2321    pub apic_mode: Option<ApicMode>,
2322}
2323
2324impl Default for ProcessorTopology {
2325    fn default() -> Self {
2326        Self {
2327            vp_count: 2,
2328            enable_smt: None,
2329            vps_per_socket: None,
2330            apic_mode: None,
2331        }
2332    }
2333}
2334
2335impl ProcessorTopology {
2336    /// A large number of VPs
2337    pub fn heavy() -> Self {
2338        Self {
2339            vp_count: 16,
2340            vps_per_socket: Some(8),
2341            ..Default::default()
2342        }
2343    }
2344
2345    /// A very large number of VPs
2346    pub fn very_heavy() -> Self {
2347        Self {
2348            vp_count: 32,
2349            vps_per_socket: Some(16),
2350            ..Default::default()
2351        }
2352    }
2353}
2354
2355/// The APIC mode for the VM.
2356#[derive(Debug, Clone, Copy)]
2357pub enum ApicMode {
2358    /// xAPIC mode only.
2359    Xapic,
2360    /// x2APIC mode supported but not enabled at boot.
2361    X2apicSupported,
2362    /// x2APIC mode enabled at boot.
2363    X2apicEnabled,
2364}
2365
2366/// Common memory configuration information for the VM.
2367#[derive(Debug)]
2368pub struct MemoryConfig {
2369    /// Specifies the amount of memory, in bytes, to assign to the
2370    /// virtual machine.
2371    pub startup_bytes: u64,
2372    /// Specifies the minimum and maximum amount of dynamic memory, in bytes.
2373    ///
2374    /// Dynamic memory will be disabled if this is `None`.
2375    pub dynamic_memory_range: Option<(u64, u64)>,
2376    /// Per-NUMA-node memory sizes. When set, RAM is distributed across
2377    /// vNUMA nodes instead of assigning all RAM to node 0.
2378    pub numa_mem_sizes: Option<Vec<u64>>,
2379    /// Whether to back guest RAM with private anonymous memory rather than a
2380    /// shared (file/memfd-backed) memory section.
2381    ///
2382    /// - `None` (the default) uses private memory whenever the configuration
2383    ///   allows it, falling back to shared memory otherwise. Private anonymous
2384    ///   memory is cheaper to set up and eligible for Transparent Huge Pages,
2385    ///   so it is preferred for performance; this lets each VM get the best
2386    ///   backing for its firmware without the test having to know the details.
2387    /// - `Some(true)` explicitly requires private memory. This is an error if
2388    ///   the configuration is incompatible with private memory (OpenHCL, which
2389    ///   shares VTL0 RAM with VTL2 via a remote mapper, and PCAT/Gen1, which
2390    ///   relies on x86 legacy support, both require shared memory), rather than
2391    ///   silently downgrading to shared.
2392    /// - `Some(false)` explicitly requires shared memory, for tests that need
2393    ///   the guest RAM backing to be shareable with another process, such as
2394    ///   vhost-user backends.
2395    ///
2396    /// Only applies to the OpenVMM backend; ignored by Hyper-V.
2397    pub private_memory: Option<bool>,
2398    /// Mark guest RAM as eligible for Transparent Huge Pages (THP),
2399    /// improving performance for large allocations.
2400    ///
2401    /// Defaults to `true`. Applies to private anonymous guest RAM and to
2402    /// shared memfd-backed RAM, on Linux (via `madvise`) and on Windows (via
2403    /// soft large pages). It has no effect on explicit hugetlb/large-page
2404    /// backings (see
2405    /// [`with_hugepages`](crate::openvmm::PetriVmConfigOpenVmm::with_hugepages)),
2406    /// which are already huge.
2407    ///
2408    /// Only applies to the OpenVMM backend; ignored by Hyper-V.
2409    pub transparent_hugepages: bool,
2410}
2411
2412impl Default for MemoryConfig {
2413    fn default() -> Self {
2414        Self {
2415            startup_bytes: 4 * 1024 * 1024 * 1024, // 4 GiB
2416            dynamic_memory_range: None,
2417            numa_mem_sizes: None,
2418            private_memory: None,
2419            transparent_hugepages: true,
2420        }
2421    }
2422}
2423
2424/// UEFI firmware configuration
2425#[derive(Debug)]
2426pub struct UefiConfig {
2427    /// Enable secure boot
2428    pub secure_boot_enabled: bool,
2429    /// Secure boot template
2430    pub secure_boot_template: Option<SecureBootTemplate>,
2431    /// Custom UEFI variable delta JSON
2432    pub custom_uefi_json: Option<Vec<u8>>,
2433    /// Disable the UEFI frontpage which will cause the VM to shutdown instead when unable to boot.
2434    pub disable_frontpage: bool,
2435    /// Always attempt a default boot
2436    pub default_boot_always_attempt: bool,
2437    /// Enable vPCI boot (for NVMe)
2438    pub enable_vpci_boot: bool,
2439    /// Force UEFI to bounce-buffer all DMA traffic
2440    pub force_dma_bounce: bool,
2441    /// EFI diagnostics log level filter
2442    pub efi_diagnostics_log_level: EfiDiagnosticsLogLevel,
2443    /// Per-period rate-limit override for EFI diagnostics emission.
2444    /// See [`PetriVmBuilder::with_efi_diagnostics_rate_limit()`] for more information.
2445    pub efi_diagnostics_rate_limit: Option<u32>,
2446}
2447
2448impl Default for UefiConfig {
2449    fn default() -> Self {
2450        Self {
2451            secure_boot_enabled: false,
2452            secure_boot_template: None,
2453            custom_uefi_json: None,
2454            disable_frontpage: true,
2455            default_boot_always_attempt: false,
2456            enable_vpci_boot: false,
2457            force_dma_bounce: false,
2458            efi_diagnostics_log_level: EfiDiagnosticsLogLevel::Default,
2459            efi_diagnostics_rate_limit: None,
2460        }
2461    }
2462}
2463
2464/// EFI diagnostics log level filter.
2465///
2466/// Controls which UEFI diagnostics log entries are forwarded to the host
2467/// tracing infrastructure (and thus visible via kmsg / test output).
2468#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
2469pub enum EfiDiagnosticsLogLevel {
2470    /// Default log level (ERROR and WARN only).
2471    #[default]
2472    Default,
2473    /// Include INFO logs (ERROR, WARN, and INFO).
2474    Info,
2475    /// All log levels.
2476    Full,
2477}
2478
2479/// Control the logging configuration of OpenVMM/OpenHCL.
2480#[derive(Debug, Clone)]
2481pub enum OpenvmmLogConfig {
2482    /// Use the default log levels used by petri tests. This will forward
2483    /// `OPENVMM_LOG` and `OPENVMM_SHOW_SPANS` from the environment if they are
2484    /// set, otherwise it will use `debug` and `true` respectively
2485    TestDefault,
2486    /// Use the built-in default log levels of OpenHCL/OpenVMM (e.g. don't pass
2487    /// OPENVMM_LOG or OPENVMM_SHOW_SPANS)
2488    BuiltInDefault,
2489    /// Use the provided custom log levels, specified as key/value pairs. At this time,
2490    /// simply uses the already-defined environment variables (e.g.
2491    /// `OPENVMM_LOG=info,disk_nvme=debug OPENVMM_SHOW_SPANS=true`)
2492    ///
2493    /// See the Guide and source code for configuring these logs.
2494    /// - For the host VMM: see `enable_tracing` in `tracing_init.rs` for details on
2495    ///   the accepted keys and values.
2496    /// - For OpenHCL, see `init_tracing_backend` in `openhcl/src/logging/mod.rs` for details on
2497    ///   the accepted keys and values.
2498    Custom(BTreeMap<String, String>),
2499}
2500
2501/// OpenHCL configuration
2502#[derive(Debug)]
2503pub struct OpenHclConfig {
2504    /// Whether to enable VMBus redirection
2505    pub vmbus_redirect: bool,
2506    /// Test-specified command-line parameters to append to the petri generated
2507    /// command line and pass to OpenHCL. VM backends should use
2508    /// [`OpenHclConfig::command_line()`] rather than reading this directly.
2509    pub custom_command_line: Option<String>,
2510    /// Command line parameters that control OpenHCL logging behavior. Separate
2511    /// from `command_line` so that petri can decide to use default log
2512    /// levels.
2513    pub log_levels: OpenvmmLogConfig,
2514    /// How to place VTL2 in address space. If `None`, the backend VMM
2515    /// will decide on default behavior.
2516    pub vtl2_base_address_type: Option<Vtl2BaseAddressType>,
2517    /// VTL2 settings
2518    pub vtl2_settings: Option<Vtl2Settings>,
2519}
2520
2521impl OpenHclConfig {
2522    /// Returns the command line to pass to OpenHCL based on these parameters. Aggregates
2523    /// the command line and log levels.
2524    pub fn command_line(&self) -> String {
2525        let mut cmdline = self.custom_command_line.clone();
2526
2527        // Enable MANA keep-alive by default for all tests
2528        append_cmdline(&mut cmdline, "OPENHCL_MANA_KEEP_ALIVE=host,privatepool");
2529
2530        match &self.log_levels {
2531            OpenvmmLogConfig::TestDefault => {
2532                let default_log_levels = {
2533                    // Forward OPENVMM_LOG and OPENVMM_SHOW_SPANS to OpenHCL if they're set.
2534                    let openhcl_tracing = if let Ok(x) =
2535                        std::env::var("OPENVMM_LOG").or_else(|_| std::env::var("HVLITE_LOG"))
2536                    {
2537                        format!("OPENVMM_LOG={x}")
2538                    } else {
2539                        "OPENVMM_LOG=debug".to_owned()
2540                    };
2541                    let openhcl_show_spans = if let Ok(x) = std::env::var("OPENVMM_SHOW_SPANS") {
2542                        format!("OPENVMM_SHOW_SPANS={x}")
2543                    } else {
2544                        "OPENVMM_SHOW_SPANS=true".to_owned()
2545                    };
2546                    format!("{openhcl_tracing} {openhcl_show_spans}")
2547                };
2548                append_cmdline(&mut cmdline, &default_log_levels);
2549            }
2550            OpenvmmLogConfig::BuiltInDefault => {
2551                // do nothing, use whatever the built-in default is
2552            }
2553            OpenvmmLogConfig::Custom(levels) => {
2554                levels.iter().for_each(|(key, value)| {
2555                    append_cmdline(&mut cmdline, format!("{key}={value}"));
2556                });
2557            }
2558        }
2559
2560        cmdline.unwrap_or_default()
2561    }
2562}
2563
2564impl Default for OpenHclConfig {
2565    fn default() -> Self {
2566        Self {
2567            vmbus_redirect: false,
2568            custom_command_line: None,
2569            log_levels: OpenvmmLogConfig::TestDefault,
2570            vtl2_base_address_type: None,
2571            vtl2_settings: None,
2572        }
2573    }
2574}
2575
2576/// TPM configuration
2577#[derive(Debug)]
2578pub struct TpmConfig {
2579    /// Use ephemeral TPM state (do not persist to VMGS)
2580    pub no_persistent_secrets: bool,
2581    /// Hardware sealing policy for sealed secrets
2582    pub hardware_sealing_policy: PetriHardwareSealingPolicy,
2583}
2584
2585impl Default for TpmConfig {
2586    fn default() -> Self {
2587        Self {
2588            no_persistent_secrets: true,
2589            hardware_sealing_policy: PetriHardwareSealingPolicy::Default,
2590        }
2591    }
2592}
2593
2594/// Hardware sealing policy used by the test infrastructure.
2595///
2596/// Maps to Hyper-V `Set-GuestStateEncryptionPolicy` values and
2597/// underhill's `HardwareSealingPolicy`.
2598#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
2599pub enum PetriHardwareSealingPolicy {
2600    /// No explicit policy — the backend picks its default.
2601    #[default]
2602    Default,
2603    /// Derive the hardware sealing key from measurement hash.
2604    HashPolicy,
2605    /// Derive the hardware sealing key from signer information.
2606    SignerPolicy,
2607}
2608
2609/// Firmware to load into the test VM.
2610// TODO: remove the guests from the firmware enum so that we don't pass them
2611// to the VMM backend after we have already used them generically.
2612#[derive(Debug)]
2613pub enum Firmware {
2614    /// Boot Linux directly, without any firmware.
2615    LinuxDirect {
2616        /// The kernel to boot.
2617        kernel: ResolvedArtifact,
2618        /// The initrd to use.
2619        initrd: ResolvedArtifact,
2620    },
2621    /// Boot Linux directly, without any firmware, with OpenHCL in VTL2.
2622    OpenhclLinuxDirect {
2623        /// The path to the IGVM file to use.
2624        igvm_path: ResolvedArtifact,
2625        /// OpenHCL configuration
2626        openhcl_config: OpenHclConfig,
2627    },
2628    /// Boot a PCAT-based VM.
2629    Pcat {
2630        /// The guest OS the VM will boot into.
2631        guest: PcatGuest,
2632        /// The firmware to use.
2633        bios_firmware: ResolvedOptionalArtifact,
2634        /// The SVGA firmware to use.
2635        svga_firmware: ResolvedOptionalArtifact,
2636        /// IDE controllers and associated disks
2637        ide_controllers: [[Option<Drive>; 2]; 2],
2638    },
2639    /// Boot a PCAT-based VM with OpenHCL in VTL2.
2640    OpenhclPcat {
2641        /// The guest OS the VM will boot into.
2642        guest: PcatGuest,
2643        /// The path to the IGVM file to use.
2644        igvm_path: ResolvedArtifact,
2645        /// The firmware to use.
2646        bios_firmware: ResolvedOptionalArtifact,
2647        /// The SVGA firmware to use.
2648        svga_firmware: ResolvedOptionalArtifact,
2649        /// OpenHCL configuration
2650        openhcl_config: OpenHclConfig,
2651    },
2652    /// Boot a UEFI-based VM.
2653    Uefi {
2654        /// The guest OS the VM will boot into.
2655        guest: UefiGuest,
2656        /// The firmware to use.
2657        uefi_firmware: ResolvedArtifact,
2658        /// UEFI configuration
2659        uefi_config: UefiConfig,
2660    },
2661    /// Boot a UEFI-based VM with OpenHCL in VTL2.
2662    OpenhclUefi {
2663        /// The guest OS the VM will boot into.
2664        guest: UefiGuest,
2665        /// The isolation type of the VM.
2666        isolation: Option<IsolationType>,
2667        /// The path to the IGVM file to use.
2668        igvm_path: ResolvedArtifact,
2669        /// UEFI configuration
2670        uefi_config: UefiConfig,
2671        /// OpenHCL configuration
2672        openhcl_config: OpenHclConfig,
2673    },
2674}
2675
2676/// The boot device type.
2677#[derive(Debug, Copy, Clone, PartialEq, Eq)]
2678pub enum BootDeviceType {
2679    /// Don't initialize a boot device.
2680    None,
2681    /// Boot from IDE.
2682    Ide,
2683    /// Boot from IDE via SCSI to VTL2.
2684    IdeViaScsi,
2685    /// Boot from IDE via NVME to VTL2.
2686    IdeViaNvme,
2687    /// Boot from SCSI.
2688    Scsi,
2689    /// Boot from SCSI via SCSI to VTL2.
2690    ScsiViaScsi,
2691    /// Boot from SCSI via NVME to VTL2.
2692    ScsiViaNvme,
2693    /// Boot from NVMe.
2694    Nvme,
2695    /// Boot from NVMe via SCSI to VTL2.
2696    NvmeViaScsi,
2697    /// Boot from NVMe via NVMe to VTL2.
2698    NvmeViaNvme,
2699    /// Boot from NVMe attached to a PCIe root port.
2700    PcieNvme,
2701    /// Boot from virtio-blk attached to a PCIe root port.
2702    PcieVirtioBlk,
2703}
2704
2705impl BootDeviceType {
2706    fn requires_vtl2(&self) -> bool {
2707        match self {
2708            BootDeviceType::None
2709            | BootDeviceType::Ide
2710            | BootDeviceType::Scsi
2711            | BootDeviceType::Nvme
2712            | BootDeviceType::PcieNvme
2713            | BootDeviceType::PcieVirtioBlk => false,
2714            BootDeviceType::IdeViaScsi
2715            | BootDeviceType::IdeViaNvme
2716            | BootDeviceType::ScsiViaScsi
2717            | BootDeviceType::ScsiViaNvme
2718            | BootDeviceType::NvmeViaScsi
2719            | BootDeviceType::NvmeViaNvme => true,
2720        }
2721    }
2722
2723    fn requires_vpci_boot(&self) -> bool {
2724        matches!(
2725            self,
2726            BootDeviceType::Nvme | BootDeviceType::NvmeViaScsi | BootDeviceType::NvmeViaNvme
2727        )
2728    }
2729
2730    fn requires_vmbus(&self) -> bool {
2731        match self {
2732            BootDeviceType::None
2733            | BootDeviceType::Ide
2734            | BootDeviceType::PcieNvme
2735            | BootDeviceType::PcieVirtioBlk => false,
2736            BootDeviceType::IdeViaScsi
2737            | BootDeviceType::IdeViaNvme
2738            | BootDeviceType::Scsi
2739            | BootDeviceType::ScsiViaScsi
2740            | BootDeviceType::ScsiViaNvme
2741            | BootDeviceType::Nvme
2742            | BootDeviceType::NvmeViaScsi
2743            | BootDeviceType::NvmeViaNvme => true,
2744        }
2745    }
2746}
2747
2748impl Firmware {
2749    /// Constructs a standard [`Firmware::LinuxDirect`] configuration.
2750    pub fn linux_direct(resolver: &ArtifactResolver<'_>, arch: MachineArch) -> Self {
2751        use petri_artifacts_vmm_test::artifacts::loadable::*;
2752        match arch {
2753            MachineArch::X86_64 => Firmware::LinuxDirect {
2754                kernel: resolver.require(LINUX_DIRECT_TEST_KERNEL_X64).erase(),
2755                initrd: resolver.require(LINUX_DIRECT_TEST_INITRD_X64).erase(),
2756            },
2757            MachineArch::Aarch64 => Firmware::LinuxDirect {
2758                kernel: resolver.require(LINUX_DIRECT_TEST_KERNEL_AARCH64).erase(),
2759                initrd: resolver.require(LINUX_DIRECT_TEST_INITRD_AARCH64).erase(),
2760            },
2761        }
2762    }
2763
2764    /// Constructs a [`Firmware::LinuxDirect`] configuration that uses a
2765    /// compressed bzImage kernel instead of an uncompressed ELF.
2766    ///
2767    /// This is x86_64-only, as bzImage is an x86-specific format.
2768    pub fn linux_direct_bzimage(resolver: &ArtifactResolver<'_>) -> Self {
2769        use petri_artifacts_vmm_test::artifacts::loadable::*;
2770        Firmware::LinuxDirect {
2771            kernel: resolver.require(LINUX_DIRECT_TEST_BZIMAGE_X64).erase(),
2772            initrd: resolver.require(LINUX_DIRECT_TEST_INITRD_X64).erase(),
2773        }
2774    }
2775
2776    /// Constructs a standard [`Firmware::OpenhclLinuxDirect`] configuration.
2777    pub fn openhcl_linux_direct(resolver: &ArtifactResolver<'_>, arch: MachineArch) -> Self {
2778        use petri_artifacts_vmm_test::artifacts::openhcl_igvm::*;
2779        match arch {
2780            MachineArch::X86_64 => Firmware::OpenhclLinuxDirect {
2781                igvm_path: resolver.require(LATEST_LINUX_DIRECT_TEST_X64).erase(),
2782                openhcl_config: Default::default(),
2783            },
2784            MachineArch::Aarch64 => todo!("Linux direct not yet supported on aarch64"),
2785        }
2786    }
2787
2788    /// Constructs a standard [`Firmware::Pcat`] configuration.
2789    pub fn pcat(resolver: &ArtifactResolver<'_>, guest: PcatGuest) -> Self {
2790        use petri_artifacts_vmm_test::artifacts::loadable::*;
2791        Firmware::Pcat {
2792            guest,
2793            bios_firmware: resolver.try_require(PCAT_FIRMWARE_X64).erase(),
2794            svga_firmware: resolver.try_require(SVGA_FIRMWARE_X64).erase(),
2795            ide_controllers: [[None, None], [None, None]],
2796        }
2797    }
2798
2799    /// Constructs a standard [`Firmware::OpenhclPcat`] configuration.
2800    pub fn openhcl_pcat(resolver: &ArtifactResolver<'_>, guest: PcatGuest) -> Self {
2801        use petri_artifacts_vmm_test::artifacts::loadable::*;
2802        use petri_artifacts_vmm_test::artifacts::openhcl_igvm::*;
2803        Firmware::OpenhclPcat {
2804            guest,
2805            igvm_path: resolver.require(LATEST_STANDARD_X64).erase(),
2806            bios_firmware: resolver.try_require(PCAT_FIRMWARE_X64).erase(),
2807            svga_firmware: resolver.try_require(SVGA_FIRMWARE_X64).erase(),
2808            openhcl_config: OpenHclConfig {
2809                // VMBUS redirect is necessary for IDE to be provided by VTL2
2810                vmbus_redirect: true,
2811                ..Default::default()
2812            },
2813        }
2814    }
2815
2816    /// Constructs a standard [`Firmware::Uefi`] configuration.
2817    pub fn uefi(resolver: &ArtifactResolver<'_>, arch: MachineArch, guest: UefiGuest) -> Self {
2818        use petri_artifacts_vmm_test::artifacts::loadable::*;
2819        let uefi_firmware = match arch {
2820            MachineArch::X86_64 => resolver.require(UEFI_FIRMWARE_X64).erase(),
2821            MachineArch::Aarch64 => resolver.require(UEFI_FIRMWARE_AARCH64).erase(),
2822        };
2823        Firmware::Uefi {
2824            guest,
2825            uefi_firmware,
2826            uefi_config: Default::default(),
2827        }
2828    }
2829
2830    /// Constructs a standard [`Firmware::OpenhclUefi`] configuration.
2831    pub fn openhcl_uefi(
2832        resolver: &ArtifactResolver<'_>,
2833        arch: MachineArch,
2834        guest: UefiGuest,
2835        isolation: Option<IsolationType>,
2836    ) -> Self {
2837        use petri_artifacts_vmm_test::artifacts::openhcl_igvm::*;
2838        let igvm_path = match arch {
2839            MachineArch::X86_64 if isolation.is_some() => resolver.require(LATEST_CVM_X64).erase(),
2840            MachineArch::X86_64 => resolver.require(LATEST_STANDARD_X64).erase(),
2841            MachineArch::Aarch64 => resolver.require(LATEST_STANDARD_AARCH64).erase(),
2842        };
2843        Firmware::OpenhclUefi {
2844            guest,
2845            isolation,
2846            igvm_path,
2847            uefi_config: Default::default(),
2848            openhcl_config: Default::default(),
2849        }
2850    }
2851
2852    fn is_openhcl(&self) -> bool {
2853        match self {
2854            Firmware::OpenhclLinuxDirect { .. }
2855            | Firmware::OpenhclUefi { .. }
2856            | Firmware::OpenhclPcat { .. } => true,
2857            Firmware::LinuxDirect { .. } | Firmware::Pcat { .. } | Firmware::Uefi { .. } => false,
2858        }
2859    }
2860
2861    fn isolation(&self) -> Option<IsolationType> {
2862        match self {
2863            Firmware::OpenhclUefi { isolation, .. } => *isolation,
2864            Firmware::LinuxDirect { .. }
2865            | Firmware::Pcat { .. }
2866            | Firmware::Uefi { .. }
2867            | Firmware::OpenhclLinuxDirect { .. }
2868            | Firmware::OpenhclPcat { .. } => None,
2869        }
2870    }
2871
2872    fn is_linux_direct(&self) -> bool {
2873        match self {
2874            Firmware::LinuxDirect { .. } | Firmware::OpenhclLinuxDirect { .. } => true,
2875            Firmware::Pcat { .. }
2876            | Firmware::Uefi { .. }
2877            | Firmware::OpenhclUefi { .. }
2878            | Firmware::OpenhclPcat { .. } => false,
2879        }
2880    }
2881
2882    /// Get the initrd path for Linux direct boot firmware.
2883    pub fn linux_direct_initrd(&self) -> Option<&Path> {
2884        match self {
2885            Firmware::LinuxDirect { initrd, .. } => Some(initrd.get()),
2886            _ => None,
2887        }
2888    }
2889
2890    fn is_pcat(&self) -> bool {
2891        match self {
2892            Firmware::Pcat { .. } | Firmware::OpenhclPcat { .. } => true,
2893            Firmware::Uefi { .. }
2894            | Firmware::OpenhclUefi { .. }
2895            | Firmware::LinuxDirect { .. }
2896            | Firmware::OpenhclLinuxDirect { .. } => false,
2897        }
2898    }
2899
2900    fn os_flavor(&self) -> OsFlavor {
2901        match self {
2902            Firmware::LinuxDirect { .. } | Firmware::OpenhclLinuxDirect { .. } => OsFlavor::Linux,
2903            Firmware::Uefi {
2904                guest: UefiGuest::GuestTestUefi { .. } | UefiGuest::None,
2905                ..
2906            }
2907            | Firmware::OpenhclUefi {
2908                guest: UefiGuest::GuestTestUefi { .. } | UefiGuest::None,
2909                ..
2910            } => OsFlavor::Uefi,
2911            Firmware::Pcat {
2912                guest: PcatGuest::Vhd(cfg),
2913                ..
2914            }
2915            | Firmware::OpenhclPcat {
2916                guest: PcatGuest::Vhd(cfg),
2917                ..
2918            }
2919            | Firmware::Uefi {
2920                guest: UefiGuest::Vhd(cfg),
2921                ..
2922            }
2923            | Firmware::OpenhclUefi {
2924                guest: UefiGuest::Vhd(cfg),
2925                ..
2926            } => cfg.os_flavor,
2927            Firmware::Pcat {
2928                guest: PcatGuest::Iso(cfg),
2929                ..
2930            }
2931            | Firmware::OpenhclPcat {
2932                guest: PcatGuest::Iso(cfg),
2933                ..
2934            } => cfg.os_flavor,
2935        }
2936    }
2937
2938    fn quirks(&self) -> GuestQuirks {
2939        match self {
2940            Firmware::Pcat {
2941                guest: PcatGuest::Vhd(cfg),
2942                ..
2943            }
2944            | Firmware::Uefi {
2945                guest: UefiGuest::Vhd(cfg),
2946                ..
2947            }
2948            | Firmware::OpenhclUefi {
2949                guest: UefiGuest::Vhd(cfg),
2950                ..
2951            } => cfg.quirks.clone(),
2952            Firmware::Pcat {
2953                guest: PcatGuest::Iso(cfg),
2954                ..
2955            } => cfg.quirks.clone(),
2956            _ => Default::default(),
2957        }
2958    }
2959
2960    fn expected_boot_event(&self) -> Option<FirmwareEvent> {
2961        match self {
2962            Firmware::LinuxDirect { .. }
2963            | Firmware::OpenhclLinuxDirect { .. }
2964            | Firmware::Uefi {
2965                guest: UefiGuest::GuestTestUefi(_),
2966                ..
2967            }
2968            | Firmware::OpenhclUefi {
2969                guest: UefiGuest::GuestTestUefi(_),
2970                ..
2971            } => None,
2972            Firmware::Pcat { .. } | Firmware::OpenhclPcat { .. } => {
2973                // TODO: Handle older PCAT versions that don't fire the event
2974                Some(FirmwareEvent::BootAttempt)
2975            }
2976            Firmware::Uefi {
2977                guest: UefiGuest::None,
2978                ..
2979            }
2980            | Firmware::OpenhclUefi {
2981                guest: UefiGuest::None,
2982                ..
2983            } => Some(FirmwareEvent::NoBootDevice),
2984            Firmware::Uefi { .. } | Firmware::OpenhclUefi { .. } => {
2985                Some(FirmwareEvent::BootSuccess)
2986            }
2987        }
2988    }
2989
2990    fn openhcl_config(&self) -> Option<&OpenHclConfig> {
2991        match self {
2992            Firmware::OpenhclLinuxDirect { openhcl_config, .. }
2993            | Firmware::OpenhclUefi { openhcl_config, .. }
2994            | Firmware::OpenhclPcat { openhcl_config, .. } => Some(openhcl_config),
2995            Firmware::LinuxDirect { .. } | Firmware::Pcat { .. } | Firmware::Uefi { .. } => None,
2996        }
2997    }
2998
2999    fn openhcl_config_mut(&mut self) -> Option<&mut OpenHclConfig> {
3000        match self {
3001            Firmware::OpenhclLinuxDirect { openhcl_config, .. }
3002            | Firmware::OpenhclUefi { openhcl_config, .. }
3003            | Firmware::OpenhclPcat { openhcl_config, .. } => Some(openhcl_config),
3004            Firmware::LinuxDirect { .. } | Firmware::Pcat { .. } | Firmware::Uefi { .. } => None,
3005        }
3006    }
3007
3008    #[cfg_attr(not(windows), expect(dead_code))]
3009    fn openhcl_firmware(&self) -> Option<&Path> {
3010        match self {
3011            Firmware::OpenhclLinuxDirect { igvm_path, .. }
3012            | Firmware::OpenhclUefi { igvm_path, .. }
3013            | Firmware::OpenhclPcat { igvm_path, .. } => Some(igvm_path.get()),
3014            Firmware::LinuxDirect { .. } | Firmware::Pcat { .. } | Firmware::Uefi { .. } => None,
3015        }
3016    }
3017
3018    fn into_runtime_config(
3019        self,
3020        vmbus_storage_controllers: HashMap<Guid, VmbusStorageController>,
3021    ) -> PetriVmRuntimeConfig {
3022        match self {
3023            Firmware::OpenhclLinuxDirect { openhcl_config, .. }
3024            | Firmware::OpenhclUefi { openhcl_config, .. }
3025            | Firmware::OpenhclPcat { openhcl_config, .. } => PetriVmRuntimeConfig {
3026                vtl2_settings: Some(
3027                    openhcl_config
3028                        .vtl2_settings
3029                        .unwrap_or_else(default_vtl2_settings),
3030                ),
3031                ide_controllers: None,
3032                vmbus_storage_controllers,
3033            },
3034            Firmware::Pcat {
3035                ide_controllers, ..
3036            } => PetriVmRuntimeConfig {
3037                vtl2_settings: None,
3038                ide_controllers: Some(ide_controllers),
3039                vmbus_storage_controllers,
3040            },
3041            Firmware::LinuxDirect { .. } | Firmware::Uefi { .. } => PetriVmRuntimeConfig {
3042                vtl2_settings: None,
3043                ide_controllers: None,
3044                vmbus_storage_controllers,
3045            },
3046        }
3047    }
3048
3049    fn uefi_config(&self) -> Option<&UefiConfig> {
3050        match self {
3051            Firmware::Uefi { uefi_config, .. } | Firmware::OpenhclUefi { uefi_config, .. } => {
3052                Some(uefi_config)
3053            }
3054            Firmware::LinuxDirect { .. }
3055            | Firmware::OpenhclLinuxDirect { .. }
3056            | Firmware::Pcat { .. }
3057            | Firmware::OpenhclPcat { .. } => None,
3058        }
3059    }
3060
3061    fn uefi_config_mut(&mut self) -> Option<&mut UefiConfig> {
3062        match self {
3063            Firmware::Uefi { uefi_config, .. } | Firmware::OpenhclUefi { uefi_config, .. } => {
3064                Some(uefi_config)
3065            }
3066            Firmware::LinuxDirect { .. }
3067            | Firmware::OpenhclLinuxDirect { .. }
3068            | Firmware::Pcat { .. }
3069            | Firmware::OpenhclPcat { .. } => None,
3070        }
3071    }
3072
3073    fn boot_drive(&self) -> Option<Drive> {
3074        match self {
3075            Firmware::LinuxDirect { .. } | Firmware::OpenhclLinuxDirect { .. } => None,
3076            Firmware::Pcat { guest, .. } | Firmware::OpenhclPcat { guest, .. } => {
3077                Some((guest.disk_path(), guest.is_dvd()))
3078            }
3079            Firmware::Uefi { guest, .. } | Firmware::OpenhclUefi { guest, .. } => {
3080                guest.disk_path().map(|dp| (dp, false))
3081            }
3082        }
3083        .map(|(disk_path, is_dvd)| Drive::new(Some(Disk::Differencing(disk_path)), is_dvd))
3084    }
3085
3086    fn vtl2_settings(&mut self) -> Option<&mut Vtl2Settings> {
3087        self.openhcl_config_mut()
3088            .map(|c| c.vtl2_settings.get_or_insert_with(default_vtl2_settings))
3089    }
3090
3091    fn ide_controllers(&self) -> Option<&[[Option<Drive>; 2]; 2]> {
3092        match self {
3093            Firmware::Pcat {
3094                ide_controllers, ..
3095            } => Some(ide_controllers),
3096            _ => None,
3097        }
3098    }
3099
3100    fn ide_controllers_mut(&mut self) -> Option<&mut [[Option<Drive>; 2]; 2]> {
3101        match self {
3102            Firmware::Pcat {
3103                ide_controllers, ..
3104            } => Some(ide_controllers),
3105            _ => None,
3106        }
3107    }
3108}
3109
3110/// The guest the VM will boot into. A boot drive with the chosen setup
3111/// will be automatically configured.
3112#[derive(Debug)]
3113pub enum PcatGuest {
3114    /// Mount a VHD as the boot drive.
3115    Vhd(BootImageConfig<boot_image_type::Vhd>),
3116    /// Mount an ISO as the CD/DVD drive.
3117    Iso(BootImageConfig<boot_image_type::Iso>),
3118}
3119
3120impl PcatGuest {
3121    fn disk_path(&self) -> DiskPath {
3122        match self {
3123            PcatGuest::Vhd(disk) => disk.disk_path(),
3124            PcatGuest::Iso(disk) => disk.disk_path(),
3125        }
3126    }
3127
3128    fn is_dvd(&self) -> bool {
3129        matches!(self, Self::Iso(_))
3130    }
3131}
3132
3133/// The guest the VM will boot into. A boot drive with the chosen setup
3134/// will be automatically configured.
3135#[derive(Debug)]
3136pub enum UefiGuest {
3137    /// Mount a VHD as the boot drive.
3138    Vhd(BootImageConfig<boot_image_type::Vhd>),
3139    /// The UEFI test image produced by our guest-test infrastructure.
3140    GuestTestUefi(ResolvedArtifact),
3141    /// No guest, just the firmware.
3142    None,
3143}
3144
3145impl UefiGuest {
3146    /// Construct a standard [`UefiGuest::GuestTestUefi`] configuration.
3147    pub fn guest_test_uefi(resolver: &ArtifactResolver<'_>, arch: MachineArch) -> Self {
3148        use petri_artifacts_vmm_test::artifacts::test_vhd::*;
3149        let artifact = match arch {
3150            MachineArch::X86_64 => resolver.require(GUEST_TEST_UEFI_X64).erase(),
3151            MachineArch::Aarch64 => resolver.require(GUEST_TEST_UEFI_AARCH64).erase(),
3152        };
3153        UefiGuest::GuestTestUefi(artifact)
3154    }
3155
3156    fn disk_path(&self) -> Option<DiskPath> {
3157        match self {
3158            UefiGuest::Vhd(vhd) => Some(vhd.disk_path()),
3159            UefiGuest::GuestTestUefi(p) => Some(DiskPath::Local(p.get().to_path_buf())),
3160            UefiGuest::None => None,
3161        }
3162    }
3163}
3164
3165/// Type-tags for [`BootImageConfig`](super::BootImageConfig)
3166pub mod boot_image_type {
3167    mod private {
3168        pub trait Sealed {}
3169        impl Sealed for super::Vhd {}
3170        impl Sealed for super::Iso {}
3171    }
3172
3173    /// Private trait use to seal the set of artifact types BootImageType
3174    /// supports.
3175    pub trait BootImageType: private::Sealed {}
3176
3177    /// BootImageConfig for a VHD file
3178    #[derive(Debug)]
3179    pub enum Vhd {}
3180
3181    /// BootImageConfig for an ISO file
3182    #[derive(Debug)]
3183    pub enum Iso {}
3184
3185    impl BootImageType for Vhd {}
3186    impl BootImageType for Iso {}
3187}
3188
3189/// Configuration information for the boot drive of the VM.
3190#[derive(Debug)]
3191pub struct BootImageConfig<T: boot_image_type::BootImageType> {
3192    /// Artifact source corresponding to the boot media (local or remote).
3193    artifact: ResolvedArtifactSource,
3194    /// The OS flavor.
3195    os_flavor: OsFlavor,
3196    /// Any quirks needed to boot the guest.
3197    ///
3198    /// Most guests should not need any quirks, and can use `Default`.
3199    quirks: GuestQuirks,
3200    /// Marker denoting what type of media `artifact` corresponds to
3201    _type: core::marker::PhantomData<T>,
3202}
3203
3204impl<T: boot_image_type::BootImageType> BootImageConfig<T> {
3205    /// Get a [`DiskPath`] from the artifact source.
3206    fn disk_path(&self) -> DiskPath {
3207        match self.artifact.get() {
3208            ArtifactSource::Local(p) => DiskPath::Local(p.clone()),
3209            ArtifactSource::Remote { url } => DiskPath::Remote { url: url.clone() },
3210        }
3211    }
3212}
3213
3214impl BootImageConfig<boot_image_type::Vhd> {
3215    /// Create a new BootImageConfig from a VHD artifact source
3216    pub fn from_vhd<A>(artifact: ResolvedArtifactSource<A>) -> Self
3217    where
3218        A: petri_artifacts_common::tags::IsTestVhd,
3219    {
3220        BootImageConfig {
3221            artifact: artifact.erase(),
3222            os_flavor: A::OS_FLAVOR,
3223            quirks: A::quirks(),
3224            _type: std::marker::PhantomData,
3225        }
3226    }
3227}
3228
3229impl BootImageConfig<boot_image_type::Iso> {
3230    /// Create a new BootImageConfig from an ISO artifact source
3231    pub fn from_iso<A>(artifact: ResolvedArtifactSource<A>) -> Self
3232    where
3233        A: petri_artifacts_common::tags::IsTestIso,
3234    {
3235        BootImageConfig {
3236            artifact: artifact.erase(),
3237            os_flavor: A::OS_FLAVOR,
3238            quirks: A::quirks(),
3239            _type: std::marker::PhantomData,
3240        }
3241    }
3242}
3243
3244/// Isolation type
3245#[derive(Debug, Clone, Copy)]
3246pub enum IsolationType {
3247    /// VBS
3248    Vbs,
3249    /// SNP
3250    Snp,
3251    /// TDX
3252    Tdx,
3253}
3254
3255/// Flags controlling servicing behavior.
3256#[derive(Debug, Clone, Copy)]
3257pub struct OpenHclServicingFlags {
3258    /// Preserve DMA memory for NVMe devices if supported.
3259    /// Defaults to `true`.
3260    pub enable_nvme_keepalive: bool,
3261    /// Preserve DMA memory for MANA devices if supported.
3262    pub enable_mana_keepalive: bool,
3263    /// Skip any logic that the vmm may have to ignore servicing updates if the supplied igvm file version is not different than the one currently running.
3264    pub override_version_checks: bool,
3265    /// Hint to the OpenHCL runtime how much time to wait when stopping / saving the OpenHCL.
3266    pub stop_timeout_hint_secs: Option<u16>,
3267}
3268
3269/// Where a disk image is located.
3270#[derive(Debug, Clone)]
3271pub enum DiskPath {
3272    /// A local file path.
3273    Local(PathBuf),
3274    /// A remote URL (fetched on demand via HTTP Range requests).
3275    Remote {
3276        /// The URL where the disk can be fetched.
3277        url: String,
3278    },
3279}
3280
3281impl From<PathBuf> for DiskPath {
3282    fn from(path: PathBuf) -> Self {
3283        DiskPath::Local(path)
3284    }
3285}
3286
3287/// Petri disk
3288#[derive(Debug, Clone)]
3289pub enum Disk {
3290    /// Memory backed with specified size
3291    Memory(u64),
3292    /// Memory differencing disk backed by a VHD (local or remote)
3293    Differencing(DiskPath),
3294    /// Persistent VHD
3295    Persistent(PathBuf),
3296    /// Disk backed by a temporary VHD
3297    Temporary(Arc<TempPath>),
3298}
3299
3300/// Petri VMGS disk
3301#[derive(Debug, Clone)]
3302pub struct PetriVmgsDisk {
3303    /// Backing disk
3304    pub disk: Disk,
3305    /// Guest state encryption policy
3306    pub encryption_policy: GuestStateEncryptionPolicy,
3307}
3308
3309impl Default for PetriVmgsDisk {
3310    fn default() -> Self {
3311        PetriVmgsDisk {
3312            disk: Disk::Memory(vmgs_format::VMGS_DEFAULT_CAPACITY),
3313            // TODO: make this strict once we can set it in OpenHCL on Hyper-V
3314            encryption_policy: GuestStateEncryptionPolicy::None(false),
3315        }
3316    }
3317}
3318
3319/// Petri VM guest state resource
3320#[derive(Debug, Clone)]
3321pub enum PetriVmgsResource {
3322    /// Use disk to store guest state
3323    Disk(PetriVmgsDisk),
3324    /// Use disk to store guest state, reformatting if corrupted.
3325    ReprovisionOnFailure(PetriVmgsDisk),
3326    /// Format and use disk to store guest state
3327    Reprovision(PetriVmgsDisk),
3328    /// Store guest state in memory
3329    Ephemeral,
3330}
3331
3332impl PetriVmgsResource {
3333    /// get the inner vmgs disk if one exists
3334    pub fn vmgs(&self) -> Option<&PetriVmgsDisk> {
3335        match self {
3336            PetriVmgsResource::Disk(vmgs)
3337            | PetriVmgsResource::ReprovisionOnFailure(vmgs)
3338            | PetriVmgsResource::Reprovision(vmgs) => Some(vmgs),
3339            PetriVmgsResource::Ephemeral => None,
3340        }
3341    }
3342
3343    /// get the inner disk if one exists
3344    pub fn disk(&self) -> Option<&Disk> {
3345        self.vmgs().map(|vmgs| &vmgs.disk)
3346    }
3347
3348    /// get the encryption policy of the vmgs
3349    pub fn encryption_policy(&self) -> Option<GuestStateEncryptionPolicy> {
3350        self.vmgs().map(|vmgs| vmgs.encryption_policy)
3351    }
3352}
3353
3354/// Petri VM guest state lifetime
3355#[derive(Debug, Clone, Copy)]
3356pub enum PetriGuestStateLifetime {
3357    /// Use a differencing disk backed by a blank, tempory VMGS file
3358    /// or other artifact if one is provided
3359    Disk,
3360    /// Same as default, except reformat the backing disk if corrupted
3361    ReprovisionOnFailure,
3362    /// Same as default, except reformat the backing disk
3363    Reprovision,
3364    /// Store guest state in memory (no backing disk)
3365    Ephemeral,
3366}
3367
3368/// UEFI secure boot template
3369#[derive(Debug, Clone, Copy)]
3370pub enum SecureBootTemplate {
3371    /// The Microsoft Windows template.
3372    MicrosoftWindows,
3373    /// The Microsoft UEFI certificate authority template.
3374    MicrosoftUefiCertificateAuthority,
3375}
3376
3377/// Quirks to workaround certain bugs that only manifest when using a
3378/// particular VMM, and do not depend on which guest is running.
3379#[derive(Default, Debug, Clone)]
3380pub struct VmmQuirks {
3381    /// Automatically reset the VM if we did not recieve a boot event in the
3382    /// specified amount of time.
3383    pub flaky_boot: Option<Duration>,
3384}
3385
3386/// Creates a VM-safe name that respects platform limitations.
3387///
3388/// Hyper-V limits VM names to 100 characters. For names that exceed this limit,
3389/// this function truncates to 96 characters and appends a 4-character hash
3390/// to ensure uniqueness while staying within the limit.
3391fn make_vm_safe_name(name: &str) -> String {
3392    const MAX_VM_NAME_LENGTH: usize = 100;
3393    const HASH_LENGTH: usize = 4;
3394    const MAX_PREFIX_LENGTH: usize = MAX_VM_NAME_LENGTH - HASH_LENGTH;
3395
3396    if name.len() <= MAX_VM_NAME_LENGTH {
3397        name.to_owned()
3398    } else {
3399        // Create a hash of the full name for uniqueness
3400        let mut hasher = DefaultHasher::new();
3401        name.hash(&mut hasher);
3402        let hash = hasher.finish();
3403
3404        // Format hash as a 4-character hex string
3405        let hash_suffix = format!("{:04x}", hash & 0xFFFF);
3406
3407        // Truncate the name and append the hash
3408        let truncated = &name[..MAX_PREFIX_LENGTH];
3409        tracing::debug!(
3410            "VM name too long ({}), truncating '{}' to '{}{}'",
3411            name.len(),
3412            name,
3413            truncated,
3414            hash_suffix
3415        );
3416
3417        format!("{}{}", truncated, hash_suffix)
3418    }
3419}
3420
3421/// The reason that the VM halted
3422#[derive(Debug, Clone, Copy, Eq, PartialEq)]
3423pub enum PetriHaltReason {
3424    /// The vm powered off
3425    PowerOff,
3426    /// The vm reset
3427    Reset,
3428    /// The vm hibernated
3429    Hibernate,
3430    /// The vm triple faulted
3431    TripleFault,
3432    /// The vm halted for some other reason
3433    Other,
3434}
3435
3436impl PetriHaltReason {
3437    /// Construct a halt reason with detailed debug info
3438    pub fn with_detail(self, detail: String) -> PetriHaltReasonDetail {
3439        PetriHaltReasonDetail {
3440            reason: self,
3441            detail,
3442        }
3443    }
3444}
3445
3446/// The reason that the VM halted, with optional addition debug details
3447#[derive(Debug, Clone)]
3448pub struct PetriHaltReasonDetail {
3449    /// The reason for the halt
3450    pub reason: PetriHaltReason,
3451    /// More details about the halt
3452    pub detail: String,
3453}
3454
3455fn append_cmdline(cmd: &mut Option<String>, add_cmd: impl AsRef<str>) {
3456    if let Some(cmd) = cmd.as_mut() {
3457        cmd.push(' ');
3458        cmd.push_str(add_cmd.as_ref());
3459    } else {
3460        *cmd = Some(add_cmd.as_ref().to_string());
3461    }
3462}
3463
3464async fn save_inspect(
3465    name: &str,
3466    inspect: std::pin::Pin<Box<dyn Future<Output = anyhow::Result<inspect::Node>> + Send>>,
3467    log_source: &PetriLogSource,
3468) {
3469    tracing::info!("Collecting {name} inspect details.");
3470    let node = match inspect.await {
3471        Ok(n) => n,
3472        Err(e) => {
3473            tracing::error!(?e, "Failed to get {name}");
3474            return;
3475        }
3476    };
3477    if let Err(e) = log_source.write_attachment(
3478        &format!("timeout_inspect_{name}.log"),
3479        format!("{node:#}").as_bytes(),
3480    ) {
3481        tracing::error!(?e, "Failed to save {name} inspect log");
3482        return;
3483    }
3484    tracing::info!("{name} inspect task finished.");
3485}
3486
3487/// Wrapper for modification functions with stubbed out debug impl
3488pub struct ModifyFn<T>(pub Box<dyn FnOnce(T) -> T + Send>);
3489
3490impl<T> Debug for ModifyFn<T> {
3491    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
3492        write!(f, "_")
3493    }
3494}
3495
3496/// Default VTL 2 settings used by petri
3497fn default_vtl2_settings() -> Vtl2Settings {
3498    Vtl2Settings {
3499        version: vtl2_settings_proto::vtl2_settings_base::Version::V1.into(),
3500        fixed: None,
3501        dynamic: Some(Default::default()),
3502        namespace_settings: Default::default(),
3503    }
3504}
3505
3506/// Virtual trust level
3507#[derive(Debug, Copy, Clone, PartialEq, Eq)]
3508pub enum Vtl {
3509    /// VTL 0
3510    Vtl0 = 0,
3511    /// VTL 1
3512    Vtl1 = 1,
3513    /// VTL 2
3514    Vtl2 = 2,
3515}
3516
3517/// The VMBus storage device type.
3518#[derive(Debug, Copy, Clone, PartialEq, Eq)]
3519pub enum VmbusStorageType {
3520    /// SCSI
3521    Scsi,
3522    /// NVMe
3523    Nvme,
3524    /// Virtio block device
3525    VirtioBlk,
3526}
3527
3528/// VM disk drive
3529#[derive(Debug, Clone)]
3530pub struct Drive {
3531    /// Backing disk
3532    pub disk: Option<Disk>,
3533    /// Whether this is a DVD
3534    pub is_dvd: bool,
3535}
3536
3537impl Drive {
3538    /// Create a new disk
3539    pub fn new(disk: Option<Disk>, is_dvd: bool) -> Self {
3540        Self { disk, is_dvd }
3541    }
3542}
3543
3544/// VMBus storage controller
3545#[derive(Debug, Clone)]
3546pub struct VmbusStorageController {
3547    /// The VTL to assign the storage controller to
3548    pub target_vtl: Vtl,
3549    /// The storage device type
3550    pub controller_type: VmbusStorageType,
3551    /// Drives (with any inserted disks) attached to this storage controller
3552    pub drives: HashMap<u32, Drive>,
3553}
3554
3555impl VmbusStorageController {
3556    /// Create a new storage controller
3557    pub fn new(target_vtl: Vtl, controller_type: VmbusStorageType) -> Self {
3558        Self {
3559            target_vtl,
3560            controller_type,
3561            drives: HashMap::new(),
3562        }
3563    }
3564
3565    /// Add a disk to the storage controller
3566    pub fn set_drive(
3567        &mut self,
3568        lun: Option<u32>,
3569        drive: Drive,
3570        allow_modify_existing: bool,
3571    ) -> u32 {
3572        let lun = lun.unwrap_or_else(|| {
3573            // find the first available lun
3574            let mut lun = None;
3575            for x in 0..u8::MAX as u32 {
3576                if !self.drives.contains_key(&x) {
3577                    lun = Some(x);
3578                    break;
3579                }
3580            }
3581            lun.expect("all locations on this controller are in use")
3582        });
3583
3584        if self.drives.insert(lun, drive).is_some() && !allow_modify_existing {
3585            panic!("a disk with lun {lun} already existed on this controller");
3586        }
3587
3588        lun
3589    }
3590}
3591
3592/// Returns the cache directory for lazy-fetched disk artifacts.
3593pub(crate) fn petri_disk_cache_dir() -> String {
3594    if let Ok(dir) = std::env::var("PETRI_CACHE_DIR") {
3595        return dir;
3596    }
3597
3598    #[cfg(target_os = "macos")]
3599    {
3600        if let Ok(home) = std::env::var("HOME") {
3601            return format!("{home}/Library/Caches/petri");
3602        }
3603    }
3604
3605    #[cfg(windows)]
3606    {
3607        if let Ok(local) = std::env::var("LOCALAPPDATA") {
3608            return format!("{local}\\petri\\cache");
3609        }
3610    }
3611
3612    // Linux / fallback: XDG
3613    if let Ok(xdg) = std::env::var("XDG_CACHE_HOME") {
3614        return format!("{xdg}/petri");
3615    }
3616    if let Ok(home) = std::env::var("HOME") {
3617        return format!("{home}/.cache/petri");
3618    }
3619
3620    ".cache/petri".to_string()
3621}
3622
3623#[cfg(test)]
3624mod tests {
3625    use super::make_vm_safe_name;
3626    use crate::Drive;
3627    use crate::VmbusStorageController;
3628    use crate::VmbusStorageType;
3629    use crate::Vtl;
3630
3631    #[test]
3632    fn test_short_names_unchanged() {
3633        let short_name = "short_test_name";
3634        assert_eq!(make_vm_safe_name(short_name), short_name);
3635    }
3636
3637    #[test]
3638    fn test_exactly_100_chars_unchanged() {
3639        let name_100 = "a".repeat(100);
3640        assert_eq!(make_vm_safe_name(&name_100), name_100);
3641    }
3642
3643    #[test]
3644    fn test_long_name_truncated() {
3645        let long_name = "multiarch::openhcl_servicing::hyperv_openhcl_uefi_aarch64_ubuntu_2404_server_aarch64_openhcl_servicing";
3646        let result = make_vm_safe_name(long_name);
3647
3648        // Should be exactly 100 characters
3649        assert_eq!(result.len(), 100);
3650
3651        // Should start with the truncated prefix
3652        assert!(result.starts_with("multiarch::openhcl_servicing::hyperv_openhcl_uefi_aarch64_ubuntu_2404_server_aarch64_ope"));
3653
3654        // Should end with a 4-character hash
3655        let suffix = &result[96..];
3656        assert_eq!(suffix.len(), 4);
3657        // Should be valid hex
3658        assert!(u16::from_str_radix(suffix, 16).is_ok());
3659    }
3660
3661    #[test]
3662    fn test_deterministic_results() {
3663        let long_name = "very_long_test_name_that_exceeds_the_100_character_limit_and_should_be_truncated_consistently_every_time";
3664        let result1 = make_vm_safe_name(long_name);
3665        let result2 = make_vm_safe_name(long_name);
3666
3667        assert_eq!(result1, result2);
3668        assert_eq!(result1.len(), 100);
3669    }
3670
3671    #[test]
3672    fn test_different_names_different_hashes() {
3673        let name1 = "very_long_test_name_that_definitely_exceeds_the_100_character_limit_and_should_be_truncated_by_the_function_version_1";
3674        let name2 = "very_long_test_name_that_definitely_exceeds_the_100_character_limit_and_should_be_truncated_by_the_function_version_2";
3675
3676        let result1 = make_vm_safe_name(name1);
3677        let result2 = make_vm_safe_name(name2);
3678
3679        // Both should be 100 chars
3680        assert_eq!(result1.len(), 100);
3681        assert_eq!(result2.len(), 100);
3682
3683        // Should have different suffixes since the full names are different
3684        assert_ne!(result1, result2);
3685        assert_ne!(&result1[96..], &result2[96..]);
3686    }
3687
3688    #[test]
3689    fn test_vmbus_storage_controller() {
3690        let mut controller = VmbusStorageController::new(Vtl::Vtl0, VmbusStorageType::Scsi);
3691        assert_eq!(
3692            controller.set_drive(Some(1), Drive::new(None, false), false),
3693            1
3694        );
3695        assert!(controller.drives.contains_key(&1));
3696        assert_eq!(
3697            controller.set_drive(None, Drive::new(None, false), false),
3698            0
3699        );
3700        assert!(controller.drives.contains_key(&0));
3701        assert_eq!(
3702            controller.set_drive(None, Drive::new(None, false), false),
3703            2
3704        );
3705        assert!(controller.drives.contains_key(&2));
3706        assert_eq!(
3707            controller.set_drive(Some(0), Drive::new(None, false), true),
3708            0
3709        );
3710        assert!(controller.drives.contains_key(&0));
3711    }
3712}