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 async_trait::async_trait;
18use get_resources::ged::FirmwareEvent;
19use mesh::CancelContext;
20use openvmm_defs::config::Vtl2BaseAddressType;
21use pal_async::DefaultDriver;
22use pal_async::task::Spawn;
23use pal_async::task::Task;
24use pal_async::timer::PolledTimer;
25use petri_artifacts_common::tags::GuestQuirks;
26use petri_artifacts_common::tags::GuestQuirksInner;
27use petri_artifacts_common::tags::InitialRebootCondition;
28use petri_artifacts_common::tags::IsOpenhclIgvm;
29use petri_artifacts_common::tags::IsTestVmgs;
30use petri_artifacts_common::tags::MachineArch;
31use petri_artifacts_common::tags::OsFlavor;
32use petri_artifacts_core::ArtifactResolver;
33use petri_artifacts_core::ResolvedArtifact;
34use petri_artifacts_core::ResolvedOptionalArtifact;
35use pipette_client::PipetteClient;
36use std::collections::BTreeMap;
37use std::collections::hash_map::DefaultHasher;
38use std::hash::Hash;
39use std::hash::Hasher;
40use std::path::Path;
41use std::path::PathBuf;
42use std::sync::Arc;
43use std::time::Duration;
44use tempfile::TempPath;
45use vmgs_resources::GuestStateEncryptionPolicy;
46use vtl2_settings_proto::Vtl2Settings;
47
48/// The set of artifacts and resources needed to instantiate a
49/// [`PetriVmBuilder`].
50pub struct PetriVmArtifacts<T: PetriVmmBackend> {
51    /// Artifacts needed to launch the host VMM used for the test
52    pub backend: T,
53    /// Firmware and/or OS to load into the VM and associated settings
54    pub firmware: Firmware,
55    /// The architecture of the VM
56    pub arch: MachineArch,
57    /// Agent to run in the guest
58    pub agent_image: Option<AgentImage>,
59    /// Agent to run in OpenHCL
60    pub openhcl_agent_image: Option<AgentImage>,
61}
62
63impl<T: PetriVmmBackend> PetriVmArtifacts<T> {
64    /// Resolves the artifacts needed to instantiate a [`PetriVmBuilder`].
65    ///
66    /// Returns `None` if the supplied configuration is not supported on this platform.
67    pub fn new(
68        resolver: &ArtifactResolver<'_>,
69        firmware: Firmware,
70        arch: MachineArch,
71        with_vtl0_pipette: bool,
72    ) -> Option<Self> {
73        if !T::check_compat(&firmware, arch) {
74            return None;
75        }
76
77        Some(Self {
78            backend: T::new(resolver),
79            arch,
80            agent_image: Some(if with_vtl0_pipette {
81                AgentImage::new(firmware.os_flavor()).with_pipette(resolver, arch)
82            } else {
83                AgentImage::new(firmware.os_flavor())
84            }),
85            openhcl_agent_image: if firmware.is_openhcl() {
86                Some(AgentImage::new(OsFlavor::Linux).with_pipette(resolver, arch))
87            } else {
88                None
89            },
90            firmware,
91        })
92    }
93}
94
95/// Petri VM builder
96pub struct PetriVmBuilder<T: PetriVmmBackend> {
97    /// Artifacts needed to launch the host VMM used for the test
98    backend: T,
99    /// VM configuration
100    config: PetriVmConfig,
101    /// Function to modify the VMM-specific configuration
102    modify_vmm_config: Option<Box<dyn FnOnce(T::VmmConfig) -> T::VmmConfig + Send>>,
103    /// VMM-agnostic resources
104    resources: PetriVmResources,
105
106    // VMM-specific quirks for the configured firmware
107    guest_quirks: GuestQuirksInner,
108    vmm_quirks: VmmQuirks,
109
110    // Test-specific boot behavior expectations.
111    // Defaults to expected behavior for firmware configuration.
112    expected_boot_event: Option<FirmwareEvent>,
113    override_expect_reset: bool,
114}
115
116/// Petri VM configuration
117pub struct PetriVmConfig {
118    /// The name of the VM
119    pub name: String,
120    /// The architecture of the VM
121    pub arch: MachineArch,
122    /// Log levels for the host VMM process.
123    pub host_log_levels: Option<OpenvmmLogConfig>,
124    /// Firmware and/or OS to load into the VM and associated settings
125    pub firmware: Firmware,
126    /// The amount of memory, in bytes, to assign to the VM
127    pub memory: MemoryConfig,
128    /// The processor tology for the VM
129    pub proc_topology: ProcessorTopology,
130    /// Agent to run in the guest
131    pub agent_image: Option<AgentImage>,
132    /// Agent to run in OpenHCL
133    pub openhcl_agent_image: Option<AgentImage>,
134    /// Disk to use for guest crash dumps
135    pub guest_crash_disk: Option<Arc<TempPath>>,
136    /// VM guest state
137    pub vmgs: PetriVmgsResource,
138    /// The boot device type for the VM
139    pub boot_device_type: BootDeviceType,
140    /// TPM configuration
141    pub tpm: Option<TpmConfig>,
142}
143
144/// VM configuration that can be changed after the VM is created
145pub struct PetriVmRuntimeConfig {
146    /// VTL2 settings
147    pub vtl2_settings: Option<Vtl2Settings>,
148}
149
150/// Resources used by a Petri VM during contruction and runtime
151pub struct PetriVmResources {
152    driver: DefaultDriver,
153    log_source: PetriLogSource,
154}
155
156/// Trait for VMM-specific contruction and runtime resources
157#[async_trait]
158pub trait PetriVmmBackend {
159    /// VMM-specific configuration
160    type VmmConfig;
161
162    /// Runtime object
163    type VmRuntime: PetriVmRuntime;
164
165    /// Check whether the combination of firmware and architecture is
166    /// supported on the VMM.
167    fn check_compat(firmware: &Firmware, arch: MachineArch) -> bool;
168
169    /// Select backend specific quirks guest and vmm quirks.
170    fn quirks(firmware: &Firmware) -> (GuestQuirksInner, VmmQuirks);
171
172    /// Get the default servicing flags (based on what this backend supports)
173    fn default_servicing_flags() -> OpenHclServicingFlags;
174
175    /// Create a disk for guest crash dumps, and a post-test hook to open the disk
176    /// to allow for reading the dumps.
177    fn create_guest_dump_disk() -> anyhow::Result<
178        Option<(
179            Arc<TempPath>,
180            Box<dyn FnOnce() -> anyhow::Result<Box<dyn fatfs::ReadWriteSeek>>>,
181        )>,
182    >;
183
184    /// Resolve any artifacts needed to use this backend
185    fn new(resolver: &ArtifactResolver<'_>) -> Self;
186
187    /// Create and start VM from the generic config using the VMM backend
188    async fn run(
189        self,
190        config: PetriVmConfig,
191        modify_vmm_config: Option<impl FnOnce(Self::VmmConfig) -> Self::VmmConfig + Send>,
192        resources: &PetriVmResources,
193    ) -> anyhow::Result<(Self::VmRuntime, PetriVmRuntimeConfig)>;
194}
195
196pub(crate) const PETRI_VTL0_SCSI_BOOT_LUN: u8 = 0;
197pub(crate) const PETRI_VTL0_SCSI_PIPETTE_LUN: u8 = 1;
198pub(crate) const PETRI_VTL0_SCSI_CRASH_LUN: u8 = 2;
199
200/// A constructed Petri VM
201pub struct PetriVm<T: PetriVmmBackend> {
202    resources: PetriVmResources,
203    runtime: T::VmRuntime,
204    watchdog_tasks: Vec<Task<()>>,
205    openhcl_diag_handler: Option<OpenHclDiagHandler>,
206
207    arch: MachineArch,
208    guest_quirks: GuestQuirksInner,
209    vmm_quirks: VmmQuirks,
210    expected_boot_event: Option<FirmwareEvent>,
211
212    config: PetriVmRuntimeConfig,
213}
214
215impl<T: PetriVmmBackend> PetriVmBuilder<T> {
216    /// Create a new VM configuration.
217    pub fn new(
218        params: PetriTestParams<'_>,
219        artifacts: PetriVmArtifacts<T>,
220        driver: &DefaultDriver,
221    ) -> anyhow::Result<Self> {
222        let (guest_quirks, vmm_quirks) = T::quirks(&artifacts.firmware);
223        let expected_boot_event = artifacts.firmware.expected_boot_event();
224        let boot_device_type = match artifacts.firmware {
225            Firmware::LinuxDirect { .. } => BootDeviceType::None,
226            Firmware::OpenhclLinuxDirect { .. } => BootDeviceType::None,
227            Firmware::Pcat { .. } | Firmware::OpenhclPcat { .. } => BootDeviceType::Ide,
228            Firmware::Uefi {
229                guest: UefiGuest::None,
230                ..
231            }
232            | Firmware::OpenhclUefi {
233                guest: UefiGuest::None,
234                ..
235            } => BootDeviceType::None,
236            Firmware::Uefi { .. } | Firmware::OpenhclUefi { .. } => BootDeviceType::Scsi,
237        };
238
239        let guest_crash_disk = if matches!(
240            artifacts.firmware.os_flavor(),
241            OsFlavor::Windows | OsFlavor::Linux
242        ) {
243            let (guest_crash_disk, guest_dump_disk_hook) = T::create_guest_dump_disk()?.unzip();
244            if let Some(guest_dump_disk_hook) = guest_dump_disk_hook {
245                let logger = params.logger.clone();
246                params
247                    .post_test_hooks
248                    .push(crate::test::PetriPostTestHook::new(
249                        "extract guest crash dumps".into(),
250                        move |test_passed| {
251                            if test_passed {
252                                return Ok(());
253                            }
254                            let mut disk = guest_dump_disk_hook()?;
255                            let gpt = gptman::GPT::read_from(&mut disk, SECTOR_SIZE)?;
256                            let partition = fscommon::StreamSlice::new(
257                                &mut disk,
258                                gpt[1].starting_lba * SECTOR_SIZE,
259                                gpt[1].ending_lba * SECTOR_SIZE,
260                            )?;
261                            let fs = fatfs::FileSystem::new(partition, fatfs::FsOptions::new())?;
262                            for entry in fs.root_dir().iter() {
263                                let Ok(entry) = entry else {
264                                    tracing::warn!(
265                                        ?entry,
266                                        "failed to read entry in guest crash dump disk"
267                                    );
268                                    continue;
269                                };
270                                if !entry.is_file() {
271                                    tracing::warn!(
272                                        ?entry,
273                                        "skipping non-file entry in guest crash dump disk"
274                                    );
275                                    continue;
276                                }
277                                logger.write_attachment(&entry.file_name(), entry.to_file())?;
278                            }
279                            Ok(())
280                        },
281                    ));
282            }
283            guest_crash_disk
284        } else {
285            None
286        };
287
288        Ok(Self {
289            backend: artifacts.backend,
290            config: PetriVmConfig {
291                name: make_vm_safe_name(params.test_name),
292                arch: artifacts.arch,
293                host_log_levels: None,
294                firmware: artifacts.firmware,
295                boot_device_type,
296                memory: Default::default(),
297                proc_topology: Default::default(),
298                agent_image: artifacts.agent_image,
299                openhcl_agent_image: artifacts.openhcl_agent_image,
300                vmgs: PetriVmgsResource::Ephemeral,
301                tpm: None,
302                guest_crash_disk,
303            },
304            modify_vmm_config: None,
305            resources: PetriVmResources {
306                driver: driver.clone(),
307                log_source: params.logger.clone(),
308            },
309
310            guest_quirks,
311            vmm_quirks,
312            expected_boot_event,
313            override_expect_reset: false,
314        })
315    }
316}
317
318impl<T: PetriVmmBackend> PetriVmBuilder<T> {
319    /// Build and run the VM, then wait for the VM to emit the expected boot
320    /// event (if configured). Does not configure and start pipette. Should
321    /// only be used for testing platforms that pipette does not support.
322    pub async fn run_without_agent(self) -> anyhow::Result<PetriVm<T>> {
323        self.run_core().await
324    }
325
326    /// Build and run the VM, then wait for the VM to emit the expected boot
327    /// event (if configured). Launches pipette and returns a client to it.
328    pub async fn run(self) -> anyhow::Result<(PetriVm<T>, PipetteClient)> {
329        assert!(self.config.agent_image.is_some());
330        assert!(self.config.agent_image.as_ref().unwrap().contains_pipette());
331
332        let mut vm = self.run_core().await?;
333        let client = vm.wait_for_agent().await?;
334        Ok((vm, client))
335    }
336
337    async fn run_core(self) -> anyhow::Result<PetriVm<T>> {
338        let arch = self.config.arch;
339        let expect_reset = self.expect_reset();
340
341        let (mut runtime, config) = self
342            .backend
343            .run(self.config, self.modify_vmm_config, &self.resources)
344            .await?;
345        let openhcl_diag_handler = runtime.openhcl_diag();
346        let watchdog_tasks = Self::start_watchdog_tasks(&self.resources, &mut runtime)?;
347
348        let mut vm = PetriVm {
349            resources: self.resources,
350            runtime,
351            watchdog_tasks,
352            openhcl_diag_handler,
353
354            arch,
355            guest_quirks: self.guest_quirks,
356            vmm_quirks: self.vmm_quirks,
357            expected_boot_event: self.expected_boot_event,
358
359            config,
360        };
361
362        if expect_reset {
363            vm.wait_for_reset_core().await?;
364        }
365
366        vm.wait_for_expected_boot_event().await?;
367
368        Ok(vm)
369    }
370
371    fn expect_reset(&self) -> bool {
372        self.override_expect_reset
373            || matches!(
374                (
375                    self.guest_quirks.initial_reboot,
376                    self.expected_boot_event,
377                    &self.config.firmware,
378                    &self.config.tpm,
379                ),
380                (
381                    Some(InitialRebootCondition::Always),
382                    Some(FirmwareEvent::BootSuccess | FirmwareEvent::BootAttempt),
383                    _,
384                    _,
385                ) | (
386                    Some(InitialRebootCondition::WithTpm),
387                    Some(FirmwareEvent::BootSuccess | FirmwareEvent::BootAttempt),
388                    _,
389                    Some(_),
390                )
391            )
392    }
393
394    fn start_watchdog_tasks(
395        resources: &PetriVmResources,
396        runtime: &mut T::VmRuntime,
397    ) -> anyhow::Result<Vec<Task<()>>> {
398        let mut tasks = Vec::new();
399
400        {
401            const TIMEOUT_DURATION_MINUTES: u64 = 10;
402            const TIMER_DURATION: Duration = Duration::from_secs(TIMEOUT_DURATION_MINUTES * 60);
403            let log_source = resources.log_source.clone();
404            let inspect_task =
405                |name,
406                 driver: &DefaultDriver,
407                 inspect: std::pin::Pin<Box<dyn Future<Output = _> + Send>>| {
408                    driver.spawn(format!("petri-watchdog-inspect-{name}"), async move {
409                        if CancelContext::new()
410                            .with_timeout(Duration::from_secs(10))
411                            .until_cancelled(save_inspect(name, inspect, &log_source))
412                            .await
413                            .is_err()
414                        {
415                            tracing::warn!(name, "Failed to collect inspect data within timeout");
416                        }
417                    })
418                };
419
420            let driver = resources.driver.clone();
421            let vmm_inspector = runtime.inspector();
422            let openhcl_diag_handler = runtime.openhcl_diag();
423            tasks.push(resources.driver.spawn("timer-watchdog", async move {
424                PolledTimer::new(&driver).sleep(TIMER_DURATION).await;
425                tracing::warn!("Test timeout reached after {TIMEOUT_DURATION_MINUTES} minutes, collecting diagnostics.");
426                let mut timeout_tasks = Vec::new();
427                if let Some(inspector) = vmm_inspector {
428                    timeout_tasks.push(inspect_task.clone()("vmm", &driver, Box::pin(async move { inspector.inspect_all().await })) );
429                }
430                if let Some(openhcl_diag_handler) = openhcl_diag_handler {
431                    timeout_tasks.push(inspect_task("openhcl", &driver, Box::pin(async move { openhcl_diag_handler.inspect("", None, None).await })));
432                }
433                futures::future::join_all(timeout_tasks).await;
434                tracing::error!("Test time out diagnostics collection complete, aborting.");
435                panic!("Test timed out");
436            }));
437        }
438
439        if let Some(mut framebuffer_access) = runtime.take_framebuffer_access() {
440            let mut timer = PolledTimer::new(&resources.driver);
441            let log_source = resources.log_source.clone();
442
443            tasks.push(
444                resources
445                    .driver
446                    .spawn("petri-watchdog-screenshot", async move {
447                        let mut image = Vec::new();
448                        let mut last_image = Vec::new();
449                        loop {
450                            timer.sleep(Duration::from_secs(2)).await;
451                            tracing::trace!("Taking screenshot.");
452
453                            let VmScreenshotMeta {
454                                color,
455                                width,
456                                height,
457                            } = match framebuffer_access.screenshot(&mut image).await {
458                                Ok(Some(meta)) => meta,
459                                Ok(None) => {
460                                    tracing::debug!("VM off, skipping screenshot.");
461                                    continue;
462                                }
463                                Err(e) => {
464                                    tracing::error!(?e, "Failed to take screenshot");
465                                    continue;
466                                }
467                            };
468
469                            if image == last_image {
470                                tracing::debug!("No change in framebuffer, skipping screenshot.");
471                                continue;
472                            }
473
474                            let r =
475                                log_source
476                                    .create_attachment("screenshot.png")
477                                    .and_then(|mut f| {
478                                        image::write_buffer_with_format(
479                                            &mut f,
480                                            &image,
481                                            width.into(),
482                                            height.into(),
483                                            color,
484                                            image::ImageFormat::Png,
485                                        )
486                                        .map_err(Into::into)
487                                    });
488
489                            if let Err(e) = r {
490                                tracing::error!(?e, "Failed to save screenshot");
491                            } else {
492                                tracing::info!("Screenshot saved.");
493                            }
494
495                            std::mem::swap(&mut image, &mut last_image);
496                        }
497                    }),
498            );
499        }
500
501        Ok(tasks)
502    }
503
504    /// Configure the test to expect a boot failure from the VM.
505    /// Useful for negative tests.
506    pub fn with_expect_boot_failure(mut self) -> Self {
507        self.expected_boot_event = Some(FirmwareEvent::BootFailed);
508        self
509    }
510
511    /// Configure the test to not expect any boot event.
512    /// Useful for tests that do not boot a VTL0 guest.
513    pub fn with_expect_no_boot_event(mut self) -> Self {
514        self.expected_boot_event = None;
515        self
516    }
517
518    /// Allow the VM to reset once at the beginning of the test. Should only be
519    /// used if you are using a special VM configuration that causes the guest
520    /// to reboot when it usually wouldn't.
521    pub fn with_expect_reset(mut self) -> Self {
522        self.override_expect_reset = true;
523        self
524    }
525
526    /// Set the VM to enable secure boot and inject the templates per OS flavor.
527    pub fn with_secure_boot(mut self) -> Self {
528        self.config
529            .firmware
530            .uefi_config_mut()
531            .expect("Secure boot is only supported for UEFI firmware.")
532            .secure_boot_enabled = true;
533
534        match self.os_flavor() {
535            OsFlavor::Windows => self.with_windows_secure_boot_template(),
536            OsFlavor::Linux => self.with_uefi_ca_secure_boot_template(),
537            _ => panic!(
538                "Secure boot unsupported for OS flavor {:?}",
539                self.os_flavor()
540            ),
541        }
542    }
543
544    /// Inject Windows secure boot templates into the VM's UEFI.
545    pub fn with_windows_secure_boot_template(mut self) -> Self {
546        self.config
547            .firmware
548            .uefi_config_mut()
549            .expect("Secure boot is only supported for UEFI firmware.")
550            .secure_boot_template = Some(SecureBootTemplate::MicrosoftWindows);
551        self
552    }
553
554    /// Inject UEFI CA secure boot templates into the VM's UEFI.
555    pub fn with_uefi_ca_secure_boot_template(mut self) -> Self {
556        self.config
557            .firmware
558            .uefi_config_mut()
559            .expect("Secure boot is only supported for UEFI firmware.")
560            .secure_boot_template = Some(SecureBootTemplate::MicrosoftUefiCertificateAuthority);
561        self
562    }
563
564    /// Set the VM to use the specified processor topology.
565    pub fn with_processor_topology(mut self, topology: ProcessorTopology) -> Self {
566        self.config.proc_topology = topology;
567        self
568    }
569
570    /// Set the VM to use the specified processor topology.
571    pub fn with_memory(mut self, memory: MemoryConfig) -> Self {
572        self.config.memory = memory;
573        self
574    }
575
576    /// Sets a custom OpenHCL IGVM VTL2 address type. This controls the behavior
577    /// of where VTL2 is placed in address space, and also the total size of memory
578    /// allocated for VTL2. VTL2 start will fail if `address_type` is specified
579    /// and leads to the loader allocating less memory than what is in the IGVM file.
580    pub fn with_vtl2_base_address_type(mut self, address_type: Vtl2BaseAddressType) -> Self {
581        self.config
582            .firmware
583            .openhcl_config_mut()
584            .expect("OpenHCL firmware is required to set custom VTL2 address type.")
585            .vtl2_base_address_type = Some(address_type);
586        self
587    }
588
589    /// Sets a custom OpenHCL IGVM file to use.
590    pub fn with_custom_openhcl(mut self, artifact: ResolvedArtifact<impl IsOpenhclIgvm>) -> Self {
591        match &mut self.config.firmware {
592            Firmware::OpenhclLinuxDirect { igvm_path, .. }
593            | Firmware::OpenhclPcat { igvm_path, .. }
594            | Firmware::OpenhclUefi { igvm_path, .. } => {
595                *igvm_path = artifact.erase();
596            }
597            Firmware::LinuxDirect { .. } | Firmware::Uefi { .. } | Firmware::Pcat { .. } => {
598                panic!("Custom OpenHCL is only supported for OpenHCL firmware.")
599            }
600        }
601        self
602    }
603
604    /// Append additional command line arguments to pass to the paravisor.
605    pub fn with_openhcl_command_line(mut self, additional_command_line: &str) -> Self {
606        append_cmdline(
607            &mut self
608                .config
609                .firmware
610                .openhcl_config_mut()
611                .expect("OpenHCL command line is only supported for OpenHCL firmware.")
612                .custom_command_line,
613            additional_command_line,
614        );
615        self
616    }
617
618    /// Enable confidential filtering, even if the VM is not confidential.
619    pub fn with_confidential_filtering(self) -> Self {
620        if !self.config.firmware.is_openhcl() {
621            panic!("Confidential filtering is only supported for OpenHCL");
622        }
623        self.with_openhcl_command_line(&format!(
624            "{}=1 {}=0",
625            underhill_confidentiality::OPENHCL_CONFIDENTIAL_ENV_VAR_NAME,
626            underhill_confidentiality::OPENHCL_CONFIDENTIAL_DEBUG_ENV_VAR_NAME
627        ))
628    }
629
630    /// Sets the command line parameters passed to OpenHCL related to logging.
631    pub fn with_openhcl_log_levels(mut self, levels: OpenvmmLogConfig) -> Self {
632        self.config
633            .firmware
634            .openhcl_config_mut()
635            .expect("OpenHCL firmware is required to set custom OpenHCL log levels.")
636            .log_levels = levels;
637        self
638    }
639
640    /// Sets the log levels for the host OpenVMM process.
641    /// DEVNOTE: In the future, this could be generalized for both HyperV and OpenVMM.
642    /// For now, this is only implemented for OpenVMM.
643    pub fn with_host_log_levels(mut self, levels: OpenvmmLogConfig) -> Self {
644        if let OpenvmmLogConfig::Custom(ref custom_levels) = levels {
645            for key in custom_levels.keys() {
646                if !["OPENVMM_LOG", "OPENVMM_SHOW_SPANS"].contains(&key.as_str()) {
647                    panic!("Unsupported OpenVMM log level key: {}", key);
648                }
649            }
650        }
651
652        self.config.host_log_levels = Some(levels.clone());
653        self
654    }
655
656    /// Adds a file to the VM's pipette agent image.
657    pub fn with_agent_file(mut self, name: &str, artifact: ResolvedArtifact) -> Self {
658        self.config
659            .agent_image
660            .as_mut()
661            .expect("no guest pipette")
662            .add_file(name, artifact);
663        self
664    }
665
666    /// Adds a file to the paravisor's pipette agent image.
667    pub fn with_openhcl_agent_file(mut self, name: &str, artifact: ResolvedArtifact) -> Self {
668        self.config
669            .openhcl_agent_image
670            .as_mut()
671            .expect("no openhcl pipette")
672            .add_file(name, artifact);
673        self
674    }
675
676    /// Sets whether UEFI frontpage is enabled.
677    pub fn with_uefi_frontpage(mut self, enable: bool) -> Self {
678        self.config
679            .firmware
680            .uefi_config_mut()
681            .expect("UEFI frontpage is only supported for UEFI firmware.")
682            .disable_frontpage = !enable;
683        self
684    }
685
686    /// Sets whether UEFI should always attempt a default boot.
687    pub fn with_default_boot_always_attempt(mut self, enable: bool) -> Self {
688        self.config
689            .firmware
690            .uefi_config_mut()
691            .expect("Default boot always attempt is only supported for UEFI firmware.")
692            .default_boot_always_attempt = enable;
693        self
694    }
695
696    /// Run the VM with Enable VMBus relay enabled
697    pub fn with_vmbus_redirect(mut self, enable: bool) -> Self {
698        self.config
699            .firmware
700            .openhcl_config_mut()
701            .expect("VMBus redirection is only supported for OpenHCL firmware.")
702            .vmbus_redirect = enable;
703        self
704    }
705
706    /// Specify the guest state lifetime for the VM
707    pub fn with_guest_state_lifetime(
708        mut self,
709        guest_state_lifetime: PetriGuestStateLifetime,
710    ) -> Self {
711        let disk = match self.config.vmgs {
712            PetriVmgsResource::Disk(disk)
713            | PetriVmgsResource::ReprovisionOnFailure(disk)
714            | PetriVmgsResource::Reprovision(disk) => disk,
715            PetriVmgsResource::Ephemeral => PetriVmgsDisk::default(),
716        };
717        self.config.vmgs = match guest_state_lifetime {
718            PetriGuestStateLifetime::Disk => PetriVmgsResource::Disk(disk),
719            PetriGuestStateLifetime::ReprovisionOnFailure => {
720                PetriVmgsResource::ReprovisionOnFailure(disk)
721            }
722            PetriGuestStateLifetime::Reprovision => PetriVmgsResource::Reprovision(disk),
723            PetriGuestStateLifetime::Ephemeral => {
724                if !matches!(disk.disk, PetriDiskType::Memory) {
725                    panic!("attempted to use ephemeral guest state after specifying backing vmgs")
726                }
727                PetriVmgsResource::Ephemeral
728            }
729        };
730        self
731    }
732
733    /// Specify the guest state encryption policy for the VM
734    pub fn with_guest_state_encryption(mut self, policy: GuestStateEncryptionPolicy) -> Self {
735        match &mut self.config.vmgs {
736            PetriVmgsResource::Disk(vmgs)
737            | PetriVmgsResource::ReprovisionOnFailure(vmgs)
738            | PetriVmgsResource::Reprovision(vmgs) => {
739                vmgs.encryption_policy = policy;
740            }
741            PetriVmgsResource::Ephemeral => {
742                panic!("attempted to encrypt ephemeral guest state")
743            }
744        }
745        self
746    }
747
748    /// Use the specified backing VMGS file
749    pub fn with_initial_vmgs(self, disk: ResolvedArtifact<impl IsTestVmgs>) -> Self {
750        self.with_backing_vmgs(PetriDiskType::Differencing(disk.into()))
751    }
752
753    /// Use the specified backing VMGS file
754    pub fn with_persistent_vmgs(self, disk: impl AsRef<Path>) -> Self {
755        self.with_backing_vmgs(PetriDiskType::Persistent(disk.as_ref().to_path_buf()))
756    }
757
758    fn with_backing_vmgs(mut self, disk: PetriDiskType) -> Self {
759        match &mut self.config.vmgs {
760            PetriVmgsResource::Disk(vmgs)
761            | PetriVmgsResource::ReprovisionOnFailure(vmgs)
762            | PetriVmgsResource::Reprovision(vmgs) => {
763                if !matches!(vmgs.disk, PetriDiskType::Memory) {
764                    panic!("already specified a backing vmgs file");
765                }
766                vmgs.disk = disk;
767            }
768            PetriVmgsResource::Ephemeral => {
769                panic!("attempted to specify a backing vmgs with ephemeral guest state")
770            }
771        }
772        self
773    }
774
775    /// Set the boot device type for the VM.
776    ///
777    /// This overrides the default, which is determined by the firmware type.
778    pub fn with_boot_device_type(mut self, boot: BootDeviceType) -> Self {
779        self.config.boot_device_type = boot;
780        self
781    }
782
783    /// Enable the TPM for the VM.
784    pub fn with_tpm(mut self, enable: bool) -> Self {
785        if enable {
786            self.config.tpm.get_or_insert_default();
787        } else {
788            self.config.tpm = None;
789        }
790        self
791    }
792
793    /// Enable or disable the TPM state persistence for the VM.
794    pub fn with_tpm_state_persistence(mut self, tpm_state_persistence: bool) -> Self {
795        self.config
796            .tpm
797            .as_mut()
798            .expect("TPM persistence requires a TPM")
799            .no_persistent_secrets = !tpm_state_persistence;
800        self
801    }
802
803    /// Add custom VTL 2 settings.
804    // TODO: At some point we want to replace uses of this with nicer with_disk,
805    // with_nic, etc. methods.
806    pub fn with_custom_vtl2_settings(
807        mut self,
808        f: impl FnOnce(&mut Vtl2Settings) + 'static + Send + Sync,
809    ) -> Self {
810        let openhcl_config = self
811            .config
812            .firmware
813            .openhcl_config_mut()
814            .expect("Custom VTL 2 settings are only supported with OpenHCL");
815        if openhcl_config.modify_vtl2_settings.is_some() {
816            panic!("only one with_custom_vtl2_settings allowed");
817        }
818        openhcl_config.modify_vtl2_settings = Some(ModifyFn(Box::new(f)));
819        self
820    }
821
822    /// Get VM's guest OS flavor
823    pub fn os_flavor(&self) -> OsFlavor {
824        self.config.firmware.os_flavor()
825    }
826
827    /// Get whether the VM will use OpenHCL
828    pub fn is_openhcl(&self) -> bool {
829        self.config.firmware.is_openhcl()
830    }
831
832    /// Get the isolation type of the VM
833    pub fn isolation(&self) -> Option<IsolationType> {
834        self.config.firmware.isolation()
835    }
836
837    /// Get the machine architecture
838    pub fn arch(&self) -> MachineArch {
839        self.config.arch
840    }
841
842    /// Get the default OpenHCL servicing flags for this config
843    pub fn default_servicing_flags(&self) -> OpenHclServicingFlags {
844        T::default_servicing_flags()
845    }
846
847    /// Get the backend-specific config builder
848    pub fn modify_backend(
849        mut self,
850        f: impl FnOnce(T::VmmConfig) -> T::VmmConfig + 'static + Send,
851    ) -> Self {
852        if self.modify_vmm_config.is_some() {
853            panic!("only one modify_backend allowed");
854        }
855        self.modify_vmm_config = Some(Box::new(f));
856        self
857    }
858}
859
860impl<T: PetriVmmBackend> PetriVm<T> {
861    /// Immediately tear down the VM.
862    pub async fn teardown(self) -> anyhow::Result<()> {
863        tracing::info!("Tearing down VM...");
864        self.runtime.teardown().await
865    }
866
867    /// Wait for the VM to halt, returning the reason for the halt.
868    pub async fn wait_for_halt(&mut self) -> anyhow::Result<PetriHaltReason> {
869        tracing::info!("Waiting for VM to halt...");
870        let halt_reason = self.runtime.wait_for_halt(false).await?;
871        tracing::info!("VM halted: {halt_reason:?}. Cancelling watchdogs...");
872        futures::future::join_all(self.watchdog_tasks.drain(..).map(|t| t.cancel())).await;
873        Ok(halt_reason)
874    }
875
876    /// Wait for the VM to cleanly shutdown.
877    pub async fn wait_for_clean_shutdown(&mut self) -> anyhow::Result<()> {
878        let halt_reason = self.wait_for_halt().await?;
879        if halt_reason != PetriHaltReason::PowerOff {
880            anyhow::bail!("Expected PowerOff, got {halt_reason:?}");
881        }
882        tracing::info!("VM was cleanly powered off and torn down.");
883        Ok(())
884    }
885
886    /// Wait for the VM to halt, returning the reason for the halt,
887    /// and tear down the VM.
888    pub async fn wait_for_teardown(mut self) -> anyhow::Result<PetriHaltReason> {
889        let halt_reason = self.wait_for_halt().await?;
890        self.teardown().await?;
891        Ok(halt_reason)
892    }
893
894    /// Wait for the VM to cleanly shutdown and tear down the VM.
895    pub async fn wait_for_clean_teardown(mut self) -> anyhow::Result<()> {
896        self.wait_for_clean_shutdown().await?;
897        self.teardown().await
898    }
899
900    /// Wait for the VM to reset. Does not wait for pipette.
901    pub async fn wait_for_reset_no_agent(&mut self) -> anyhow::Result<()> {
902        self.wait_for_reset_core().await?;
903        self.wait_for_expected_boot_event().await?;
904        Ok(())
905    }
906
907    /// Wait for the VM to reset and pipette to connect.
908    pub async fn wait_for_reset(&mut self) -> anyhow::Result<PipetteClient> {
909        self.wait_for_reset_no_agent().await?;
910        self.wait_for_agent().await
911    }
912
913    async fn wait_for_reset_core(&mut self) -> anyhow::Result<()> {
914        tracing::info!("Waiting for VM to reset...");
915        let halt_reason = self.runtime.wait_for_halt(true).await?;
916        if halt_reason != PetriHaltReason::Reset {
917            anyhow::bail!("Expected reset, got {halt_reason:?}");
918        }
919        tracing::info!("VM reset.");
920        Ok(())
921    }
922
923    /// Invoke Inspect on the running OpenHCL instance.
924    ///
925    /// IMPORTANT: As mentioned in the Guide, inspect output is *not* guaranteed
926    /// to be stable. Use this to test that components in OpenHCL are working as
927    /// you would expect. But, if you are adding a test simply to verify that
928    /// the inspect output as some other tool depends on it, then that is
929    /// incorrect.
930    ///
931    /// - `timeout` is enforced on the client side
932    /// - `path` and `depth` are passed to the [`inspect::Inspect`] machinery.
933    pub async fn inspect_openhcl(
934        &self,
935        path: impl Into<String>,
936        depth: Option<usize>,
937        timeout: Option<Duration>,
938    ) -> anyhow::Result<inspect::Node> {
939        self.openhcl_diag()?
940            .inspect(path.into().as_str(), depth, timeout)
941            .await
942    }
943
944    /// Test that we are able to inspect OpenHCL.
945    pub async fn test_inspect_openhcl(&mut self) -> anyhow::Result<()> {
946        self.inspect_openhcl("", None, None).await.map(|_| ())
947    }
948
949    /// Wait for VTL 2 to report that it is ready to respond to commands.
950    /// Will fail if the VM is not running OpenHCL.
951    ///
952    /// This should only be necessary if you're doing something manual. All
953    /// Petri-provided methods will wait for VTL 2 to be ready automatically.
954    pub async fn wait_for_vtl2_ready(&mut self) -> anyhow::Result<()> {
955        self.openhcl_diag()?.wait_for_vtl2().await
956    }
957
958    /// Get the kmsg stream from OpenHCL.
959    pub async fn kmsg(&self) -> anyhow::Result<diag_client::kmsg_stream::KmsgStream> {
960        self.openhcl_diag()?.kmsg().await
961    }
962
963    /// Gets a live core dump of the OpenHCL process specified by 'name' and
964    /// writes it to 'path'
965    pub async fn openhcl_core_dump(&self, name: &str, path: &Path) -> anyhow::Result<()> {
966        self.openhcl_diag()?.core_dump(name, path).await
967    }
968
969    /// Crashes the specified openhcl process
970    pub async fn openhcl_crash(&self, name: &str) -> anyhow::Result<()> {
971        self.openhcl_diag()?.crash(name).await
972    }
973
974    /// Wait for a connection from a pipette agent running in the guest.
975    /// Useful if you've rebooted the vm or are otherwise expecting a fresh connection.
976    async fn wait_for_agent(&mut self) -> anyhow::Result<PipetteClient> {
977        // As a workaround for #2470 (where the guest crashes when the pipette
978        // connection timeout expires due to a vmbus bug), wait for the shutdown
979        // IC to come online first so that we probably won't time out when
980        // connecting to the agent.
981        // TODO: remove this once the bug is fixed, since it shouldn't be
982        // necessary and a guest could in theory support pipette and not the IC
983        self.runtime.wait_for_enlightened_shutdown_ready().await?;
984        self.runtime.wait_for_agent(false).await
985    }
986
987    /// Wait for a connection from a pipette agent running in VTL 2.
988    /// Useful if you've reset VTL 2 or are otherwise expecting a fresh connection.
989    /// Will fail if the VM is not running OpenHCL.
990    pub async fn wait_for_vtl2_agent(&mut self) -> anyhow::Result<PipetteClient> {
991        // VTL 2's pipette doesn't auto launch, only launch it on demand
992        self.launch_vtl2_pipette().await?;
993        self.runtime.wait_for_agent(true).await
994    }
995
996    /// Waits for an event emitted by the firmware about its boot status, and
997    /// verifies that it is the expected success value.
998    ///
999    /// * Linux Direct guests do not emit a boot event, so this method immediately returns Ok.
1000    /// * PCAT guests may not emit an event depending on the PCAT version, this
1001    ///   method is best effort for them.
1002    async fn wait_for_expected_boot_event(&mut self) -> anyhow::Result<()> {
1003        if let Some(expected_event) = self.expected_boot_event {
1004            let event = self.wait_for_boot_event().await?;
1005
1006            anyhow::ensure!(
1007                event == expected_event,
1008                "Did not receive expected boot event"
1009            );
1010        } else {
1011            tracing::warn!("Boot event not emitted for configured firmware or manually ignored.");
1012        }
1013
1014        Ok(())
1015    }
1016
1017    /// Waits for an event emitted by the firmware about its boot status, and
1018    /// returns that status.
1019    async fn wait_for_boot_event(&mut self) -> anyhow::Result<FirmwareEvent> {
1020        tracing::info!("Waiting for boot event...");
1021        let boot_event = loop {
1022            match CancelContext::new()
1023                .with_timeout(self.vmm_quirks.flaky_boot.unwrap_or(Duration::MAX))
1024                .until_cancelled(self.runtime.wait_for_boot_event())
1025                .await
1026            {
1027                Ok(res) => break res?,
1028                Err(_) => {
1029                    tracing::error!("Did not get boot event in required time, resetting...");
1030                    if let Some(inspector) = self.runtime.inspector() {
1031                        save_inspect(
1032                            "vmm",
1033                            Box::pin(async move { inspector.inspect_all().await }),
1034                            &self.resources.log_source,
1035                        )
1036                        .await;
1037                    }
1038
1039                    self.runtime.reset().await?;
1040                    continue;
1041                }
1042            }
1043        };
1044        tracing::info!("Got boot event: {boot_event:?}");
1045        Ok(boot_event)
1046    }
1047
1048    /// Wait for the Hyper-V shutdown IC to be ready and use it to instruct
1049    /// the guest to shutdown.
1050    pub async fn send_enlightened_shutdown(&mut self, kind: ShutdownKind) -> anyhow::Result<()> {
1051        tracing::info!("Waiting for enlightened shutdown to be ready");
1052        self.runtime.wait_for_enlightened_shutdown_ready().await?;
1053
1054        // all guests used in testing have been observed to intermittently
1055        // drop shutdown requests if they are sent too soon after the shutdown
1056        // ic comes online. give them a little extra time.
1057        // TODO: use a different method of determining whether the VM has booted
1058        // or debug and fix the shutdown IC.
1059        let mut wait_time = Duration::from_secs(10);
1060
1061        // some guests need even more time
1062        if let Some(duration) = self.guest_quirks.hyperv_shutdown_ic_sleep {
1063            wait_time += duration;
1064        }
1065
1066        tracing::info!(
1067            "Shutdown IC reported ready, waiting for an extra {}s",
1068            wait_time.as_secs()
1069        );
1070        PolledTimer::new(&self.resources.driver)
1071            .sleep(wait_time)
1072            .await;
1073
1074        tracing::info!("Sending enlightened shutdown command");
1075        self.runtime.send_enlightened_shutdown(kind).await
1076    }
1077
1078    /// Instruct the OpenHCL to restart the VTL2 paravisor. Will fail if the VM
1079    /// is not running OpenHCL. Will also fail if the VM is not running.
1080    pub async fn restart_openhcl(
1081        &mut self,
1082        new_openhcl: ResolvedArtifact<impl IsOpenhclIgvm>,
1083        flags: OpenHclServicingFlags,
1084    ) -> anyhow::Result<()> {
1085        self.runtime
1086            .restart_openhcl(&new_openhcl.erase(), flags)
1087            .await
1088    }
1089
1090    /// Update the command line parameter of the running VM that will apply on next boot.
1091    /// Will fail if the VM is not using IGVM load mode.
1092    pub async fn update_command_line(&mut self, command_line: &str) -> anyhow::Result<()> {
1093        self.runtime.update_command_line(command_line).await
1094    }
1095
1096    /// Instruct the OpenHCL to save the state of the VTL2 paravisor. Will fail if the VM
1097    /// is not running OpenHCL. Will also fail if the VM is not running or if this is called twice in succession
1098    pub async fn save_openhcl(
1099        &mut self,
1100        new_openhcl: ResolvedArtifact<impl IsOpenhclIgvm>,
1101        flags: OpenHclServicingFlags,
1102    ) -> anyhow::Result<()> {
1103        self.runtime.save_openhcl(&new_openhcl.erase(), flags).await
1104    }
1105
1106    /// Instruct the OpenHCL to restore the state of the VTL2 paravisor. Will fail if the VM
1107    /// is not running OpenHCL. Will also fail if the VM is running or if this is called without prior save
1108    pub async fn restore_openhcl(&mut self) -> anyhow::Result<()> {
1109        self.runtime.restore_openhcl().await
1110    }
1111
1112    /// Get VM's guest OS flavor
1113    pub fn arch(&self) -> MachineArch {
1114        self.arch
1115    }
1116
1117    /// Get the inner runtime backend to make backend-specific calls
1118    pub fn backend(&mut self) -> &mut T::VmRuntime {
1119        &mut self.runtime
1120    }
1121
1122    async fn launch_vtl2_pipette(&self) -> anyhow::Result<()> {
1123        tracing::debug!("Launching VTL 2 pipette...");
1124
1125        // Start pipette through DiagClient
1126        let res = self
1127            .openhcl_diag()?
1128            .run_vtl2_command("sh", &["-c", "mkdir /cidata && mount LABEL=cidata /cidata"])
1129            .await?;
1130
1131        if !res.exit_status.success() {
1132            anyhow::bail!("Failed to mount VTL 2 pipette drive: {:?}", res);
1133        }
1134
1135        let res = self
1136            .openhcl_diag()?
1137            .run_detached_vtl2_command("sh", &["-c", "/cidata/pipette 2>&1 | logger &"])
1138            .await?;
1139
1140        if !res.success() {
1141            anyhow::bail!("Failed to spawn VTL 2 pipette: {:?}", res);
1142        }
1143
1144        Ok(())
1145    }
1146
1147    fn openhcl_diag(&self) -> anyhow::Result<&OpenHclDiagHandler> {
1148        if let Some(ohd) = self.openhcl_diag_handler.as_ref() {
1149            Ok(ohd)
1150        } else {
1151            anyhow::bail!("VM is not configured with OpenHCL")
1152        }
1153    }
1154
1155    /// Get the path to the VM's guest state file
1156    pub async fn get_guest_state_file(&self) -> anyhow::Result<Option<PathBuf>> {
1157        self.runtime.get_guest_state_file().await
1158    }
1159
1160    /// Modify OpenHCL VTL2 settings.
1161    pub async fn modify_vtl2_settings(
1162        &mut self,
1163        f: impl FnOnce(&mut Vtl2Settings),
1164    ) -> anyhow::Result<()> {
1165        if self.openhcl_diag_handler.is_none() {
1166            panic!("Custom VTL 2 settings are only supported with OpenHCL");
1167        }
1168        f(self
1169            .config
1170            .vtl2_settings
1171            .get_or_insert_with(default_vtl2_settings));
1172        self.runtime
1173            .set_vtl2_settings(self.config.vtl2_settings.as_ref().unwrap())
1174            .await
1175    }
1176}
1177
1178/// A running VM that tests can interact with.
1179#[async_trait]
1180pub trait PetriVmRuntime: Send + Sync + 'static {
1181    /// Interface for inspecting the VM
1182    type VmInspector: PetriVmInspector;
1183    /// Interface for accessing the framebuffer
1184    type VmFramebufferAccess: PetriVmFramebufferAccess;
1185
1186    /// Cleanly tear down the VM immediately.
1187    async fn teardown(self) -> anyhow::Result<()>;
1188    /// Wait for the VM to halt, returning the reason for the halt. The VM
1189    /// should automatically restart the VM on reset if `allow_reset` is true.
1190    async fn wait_for_halt(&mut self, allow_reset: bool) -> anyhow::Result<PetriHaltReason>;
1191    /// Wait for a connection from a pipette agent
1192    async fn wait_for_agent(&mut self, set_high_vtl: bool) -> anyhow::Result<PipetteClient>;
1193    /// Get an OpenHCL diagnostics handler for the VM
1194    fn openhcl_diag(&self) -> Option<OpenHclDiagHandler>;
1195    /// Waits for an event emitted by the firmware about its boot status, and
1196    /// returns that status.
1197    async fn wait_for_boot_event(&mut self) -> anyhow::Result<FirmwareEvent>;
1198    /// Waits for the Hyper-V shutdown IC to be ready
1199    // TODO: return a receiver that will be closed when it is no longer ready.
1200    async fn wait_for_enlightened_shutdown_ready(&mut self) -> anyhow::Result<()>;
1201    /// Instruct the guest to shutdown via the Hyper-V shutdown IC.
1202    async fn send_enlightened_shutdown(&mut self, kind: ShutdownKind) -> anyhow::Result<()>;
1203    /// Instruct the OpenHCL to restart the VTL2 paravisor. Will fail if the VM
1204    /// is not running OpenHCL. Will also fail if the VM is not running.
1205    async fn restart_openhcl(
1206        &mut self,
1207        new_openhcl: &ResolvedArtifact,
1208        flags: OpenHclServicingFlags,
1209    ) -> anyhow::Result<()>;
1210    /// Instruct the OpenHCL to save the state of the VTL2 paravisor. Will fail if the VM
1211    /// is not running OpenHCL. Will also fail if the VM is not running or if this is called twice in succession
1212    /// without a call to `restore_openhcl`.
1213    async fn save_openhcl(
1214        &mut self,
1215        new_openhcl: &ResolvedArtifact,
1216        flags: OpenHclServicingFlags,
1217    ) -> anyhow::Result<()>;
1218    /// Instruct the OpenHCL to restore the state of the VTL2 paravisor. Will fail if the VM
1219    /// is not running OpenHCL. Will also fail if the VM is running or if this is called without prior save.
1220    async fn restore_openhcl(&mut self) -> anyhow::Result<()>;
1221    /// Update the command line parameter of the running VM that will apply on next boot.
1222    /// Will fail if the VM is not using IGVM load mode.
1223    async fn update_command_line(&mut self, command_line: &str) -> anyhow::Result<()>;
1224    /// If the backend supports it, get an inspect interface
1225    fn inspector(&self) -> Option<Self::VmInspector> {
1226        None
1227    }
1228    /// If the backend supports it, take the screenshot interface
1229    /// (subsequent calls may return None).
1230    fn take_framebuffer_access(&mut self) -> Option<Self::VmFramebufferAccess> {
1231        None
1232    }
1233    /// Issue a hard reset to the VM
1234    async fn reset(&mut self) -> anyhow::Result<()>;
1235    /// Get the path to the VM's guest state file
1236    async fn get_guest_state_file(&self) -> anyhow::Result<Option<PathBuf>> {
1237        Ok(None)
1238    }
1239    /// Set the OpenHCL VTL2 settings.
1240    async fn set_vtl2_settings(&mut self, settings: &Vtl2Settings) -> anyhow::Result<()>;
1241}
1242
1243/// Interface for getting information about the state of the VM
1244#[async_trait]
1245pub trait PetriVmInspector: Send + Sync + 'static {
1246    /// Get information about the state of the VM
1247    async fn inspect_all(&self) -> anyhow::Result<inspect::Node>;
1248}
1249
1250/// Use this for the associated type if not supported
1251pub struct NoPetriVmInspector;
1252#[async_trait]
1253impl PetriVmInspector for NoPetriVmInspector {
1254    async fn inspect_all(&self) -> anyhow::Result<inspect::Node> {
1255        unreachable!()
1256    }
1257}
1258
1259/// Raw VM screenshot
1260pub struct VmScreenshotMeta {
1261    /// color encoding used by the image
1262    pub color: image::ExtendedColorType,
1263    /// x dimension
1264    pub width: u16,
1265    /// y dimension
1266    pub height: u16,
1267}
1268
1269/// Interface for getting screenshots of the VM
1270#[async_trait]
1271pub trait PetriVmFramebufferAccess: Send + 'static {
1272    /// Populates the provided buffer with a screenshot of the VM,
1273    /// returning the dimensions and color type.
1274    async fn screenshot(&mut self, image: &mut Vec<u8>)
1275    -> anyhow::Result<Option<VmScreenshotMeta>>;
1276}
1277
1278/// Use this for the associated type if not supported
1279pub struct NoPetriVmFramebufferAccess;
1280#[async_trait]
1281impl PetriVmFramebufferAccess for NoPetriVmFramebufferAccess {
1282    async fn screenshot(
1283        &mut self,
1284        _image: &mut Vec<u8>,
1285    ) -> anyhow::Result<Option<VmScreenshotMeta>> {
1286        unreachable!()
1287    }
1288}
1289
1290/// Common processor topology information for the VM.
1291#[derive(Debug)]
1292pub struct ProcessorTopology {
1293    /// The number of virtual processors.
1294    pub vp_count: u32,
1295    /// Whether SMT (hyperthreading) is enabled.
1296    pub enable_smt: Option<bool>,
1297    /// The number of virtual processors per socket.
1298    pub vps_per_socket: Option<u32>,
1299    /// The APIC configuration (x86-64 only).
1300    pub apic_mode: Option<ApicMode>,
1301}
1302
1303impl Default for ProcessorTopology {
1304    fn default() -> Self {
1305        Self {
1306            vp_count: 2,
1307            enable_smt: None,
1308            vps_per_socket: None,
1309            apic_mode: None,
1310        }
1311    }
1312}
1313
1314/// The APIC mode for the VM.
1315#[derive(Debug, Clone, Copy)]
1316pub enum ApicMode {
1317    /// xAPIC mode only.
1318    Xapic,
1319    /// x2APIC mode supported but not enabled at boot.
1320    X2apicSupported,
1321    /// x2APIC mode enabled at boot.
1322    X2apicEnabled,
1323}
1324
1325/// Common memory configuration information for the VM.
1326pub struct MemoryConfig {
1327    /// Specifies the amount of memory, in bytes, to assign to the
1328    /// virtual machine.
1329    pub startup_bytes: u64,
1330    /// Specifies the minimum and maximum amount of dynamic memory, in bytes.
1331    ///
1332    /// Dynamic memory will be disabled if this is `None`.
1333    pub dynamic_memory_range: Option<(u64, u64)>,
1334}
1335
1336impl Default for MemoryConfig {
1337    fn default() -> Self {
1338        Self {
1339            startup_bytes: 0x1_0000_0000,
1340            dynamic_memory_range: None,
1341        }
1342    }
1343}
1344
1345/// UEFI firmware configuration
1346#[derive(Debug)]
1347pub struct UefiConfig {
1348    /// Enable secure boot
1349    pub secure_boot_enabled: bool,
1350    /// Secure boot template
1351    pub secure_boot_template: Option<SecureBootTemplate>,
1352    /// Disable the UEFI frontpage which will cause the VM to shutdown instead when unable to boot.
1353    pub disable_frontpage: bool,
1354    /// Always attempt a default boot
1355    pub default_boot_always_attempt: bool,
1356}
1357
1358impl Default for UefiConfig {
1359    fn default() -> Self {
1360        Self {
1361            secure_boot_enabled: false,
1362            secure_boot_template: None,
1363            disable_frontpage: true,
1364            default_boot_always_attempt: false,
1365        }
1366    }
1367}
1368
1369/// Control the logging configuration of OpenVMM/OpenHCL.
1370#[derive(Debug, Clone)]
1371pub enum OpenvmmLogConfig {
1372    /// Use the default log levels used by petri tests. This will forward
1373    /// `OPENVMM_LOG` and `OPENVMM_SHOW_SPANS` from the environment if they are
1374    /// set, otherwise it will use `debug` and `true` respectively
1375    TestDefault,
1376    /// Use the built-in default log levels of OpenHCL/OpenVMM (e.g. don't pass
1377    /// OPENVMM_LOG or OPENVMM_SHOW_SPANS)
1378    BuiltInDefault,
1379    /// Use the provided custom log levels, specified as key/value pairs. At this time,
1380    /// simply uses the already-defined environment variables (e.g.
1381    /// `OPENVMM_LOG=info,disk_nvme=debug OPENVMM_SHOW_SPANS=true`)
1382    ///
1383    /// See the Guide and source code for configuring these logs.
1384    /// - For the host VMM: see `enable_tracing` in `tracing_init.rs` for details on
1385    ///   the accepted keys and values.
1386    /// - For OpenHCL, see `init_tracing_backend` in `openhcl/src/logging/mod.rs` for details on
1387    ///   the accepted keys and values.
1388    Custom(BTreeMap<String, String>),
1389}
1390
1391/// OpenHCL configuration
1392#[derive(Debug)]
1393pub struct OpenHclConfig {
1394    /// Emulate SCSI via NVME to VTL2, with the provided namespace ID on
1395    /// the controller with `BOOT_NVME_INSTANCE`.
1396    pub vtl2_nvme_boot: bool,
1397    /// Whether to enable VMBus redirection
1398    pub vmbus_redirect: bool,
1399    /// Test-specified command-line parameters to append to the petri generated
1400    /// command line and pass to OpenHCL. VM backends should use
1401    /// [`OpenHclConfig::command_line()`] rather than reading this directly.
1402    pub custom_command_line: Option<String>,
1403    /// Command line parameters that control OpenHCL logging behavior. Separate
1404    /// from `command_line` so that petri can decide to use default log
1405    /// levels.
1406    pub log_levels: OpenvmmLogConfig,
1407    /// How to place VTL2 in address space. If `None`, the backend VMM
1408    /// will decide on default behavior.
1409    pub vtl2_base_address_type: Option<Vtl2BaseAddressType>,
1410    /// Optional manual modification of the VTL2 settings
1411    pub modify_vtl2_settings: Option<ModifyFn<Vtl2Settings>>,
1412}
1413
1414impl OpenHclConfig {
1415    /// Returns the command line to pass to OpenHCL based on these parameters. Aggregates
1416    /// the command line and log levels.
1417    pub fn command_line(&self) -> String {
1418        let mut cmdline = self.custom_command_line.clone();
1419
1420        // Enable MANA keep-alive by default for all tests
1421        append_cmdline(&mut cmdline, "OPENHCL_MANA_KEEP_ALIVE=host,privatepool");
1422
1423        match &self.log_levels {
1424            OpenvmmLogConfig::TestDefault => {
1425                let default_log_levels = {
1426                    // Forward OPENVMM_LOG and OPENVMM_SHOW_SPANS to OpenHCL if they're set.
1427                    let openhcl_tracing = if let Ok(x) =
1428                        std::env::var("OPENVMM_LOG").or_else(|_| std::env::var("HVLITE_LOG"))
1429                    {
1430                        format!("OPENVMM_LOG={x}")
1431                    } else {
1432                        "OPENVMM_LOG=debug".to_owned()
1433                    };
1434                    let openhcl_show_spans = if let Ok(x) = std::env::var("OPENVMM_SHOW_SPANS") {
1435                        format!("OPENVMM_SHOW_SPANS={x}")
1436                    } else {
1437                        "OPENVMM_SHOW_SPANS=true".to_owned()
1438                    };
1439                    format!("{openhcl_tracing} {openhcl_show_spans}")
1440                };
1441                append_cmdline(&mut cmdline, &default_log_levels);
1442            }
1443            OpenvmmLogConfig::BuiltInDefault => {
1444                // do nothing, use whatever the built-in default is
1445            }
1446            OpenvmmLogConfig::Custom(levels) => {
1447                levels.iter().for_each(|(key, value)| {
1448                    append_cmdline(&mut cmdline, format!("{key}={value}"));
1449                });
1450            }
1451        }
1452
1453        cmdline.unwrap_or_default()
1454    }
1455}
1456
1457impl Default for OpenHclConfig {
1458    fn default() -> Self {
1459        Self {
1460            vtl2_nvme_boot: false,
1461            vmbus_redirect: false,
1462            custom_command_line: None,
1463            log_levels: OpenvmmLogConfig::TestDefault,
1464            vtl2_base_address_type: None,
1465            modify_vtl2_settings: None,
1466        }
1467    }
1468}
1469
1470/// TPM configuration
1471#[derive(Debug)]
1472pub struct TpmConfig {
1473    /// Use ephemeral TPM state (do not persist to VMGS)
1474    pub no_persistent_secrets: bool,
1475}
1476
1477impl Default for TpmConfig {
1478    fn default() -> Self {
1479        Self {
1480            no_persistent_secrets: true,
1481        }
1482    }
1483}
1484
1485/// Firmware to load into the test VM.
1486#[derive(Debug)]
1487pub enum Firmware {
1488    /// Boot Linux directly, without any firmware.
1489    LinuxDirect {
1490        /// The kernel to boot.
1491        kernel: ResolvedArtifact,
1492        /// The initrd to use.
1493        initrd: ResolvedArtifact,
1494    },
1495    /// Boot Linux directly, without any firmware, with OpenHCL in VTL2.
1496    OpenhclLinuxDirect {
1497        /// The path to the IGVM file to use.
1498        igvm_path: ResolvedArtifact,
1499        /// OpenHCL configuration
1500        openhcl_config: OpenHclConfig,
1501    },
1502    /// Boot a PCAT-based VM.
1503    Pcat {
1504        /// The guest OS the VM will boot into.
1505        guest: PcatGuest,
1506        /// The firmware to use.
1507        bios_firmware: ResolvedOptionalArtifact,
1508        /// The SVGA firmware to use.
1509        svga_firmware: ResolvedOptionalArtifact,
1510    },
1511    /// Boot a PCAT-based VM with OpenHCL in VTL2.
1512    OpenhclPcat {
1513        /// The guest OS the VM will boot into.
1514        guest: PcatGuest,
1515        /// The path to the IGVM file to use.
1516        igvm_path: ResolvedArtifact,
1517        /// The firmware to use.
1518        bios_firmware: ResolvedOptionalArtifact,
1519        /// The SVGA firmware to use.
1520        svga_firmware: ResolvedOptionalArtifact,
1521        /// OpenHCL configuration
1522        openhcl_config: OpenHclConfig,
1523    },
1524    /// Boot a UEFI-based VM.
1525    Uefi {
1526        /// The guest OS the VM will boot into.
1527        guest: UefiGuest,
1528        /// The firmware to use.
1529        uefi_firmware: ResolvedArtifact,
1530        /// UEFI configuration
1531        uefi_config: UefiConfig,
1532    },
1533    /// Boot a UEFI-based VM with OpenHCL in VTL2.
1534    OpenhclUefi {
1535        /// The guest OS the VM will boot into.
1536        guest: UefiGuest,
1537        /// The isolation type of the VM.
1538        isolation: Option<IsolationType>,
1539        /// The path to the IGVM file to use.
1540        igvm_path: ResolvedArtifact,
1541        /// UEFI configuration
1542        uefi_config: UefiConfig,
1543        /// OpenHCL configuration
1544        openhcl_config: OpenHclConfig,
1545    },
1546}
1547
1548/// The boot device type.
1549#[derive(Debug, Copy, Clone, PartialEq, Eq)]
1550pub enum BootDeviceType {
1551    /// Don't initialize a boot device.
1552    None,
1553    /// Boot from IDE.
1554    Ide,
1555    /// Boot from SCSI.
1556    Scsi,
1557    /// Boot from an NVMe controller.
1558    Nvme,
1559}
1560
1561impl Firmware {
1562    /// Constructs a standard [`Firmware::LinuxDirect`] configuration.
1563    pub fn linux_direct(resolver: &ArtifactResolver<'_>, arch: MachineArch) -> Self {
1564        use petri_artifacts_vmm_test::artifacts::loadable::*;
1565        match arch {
1566            MachineArch::X86_64 => Firmware::LinuxDirect {
1567                kernel: resolver.require(LINUX_DIRECT_TEST_KERNEL_X64).erase(),
1568                initrd: resolver.require(LINUX_DIRECT_TEST_INITRD_X64).erase(),
1569            },
1570            MachineArch::Aarch64 => Firmware::LinuxDirect {
1571                kernel: resolver.require(LINUX_DIRECT_TEST_KERNEL_AARCH64).erase(),
1572                initrd: resolver.require(LINUX_DIRECT_TEST_INITRD_AARCH64).erase(),
1573            },
1574        }
1575    }
1576
1577    /// Constructs a standard [`Firmware::OpenhclLinuxDirect`] configuration.
1578    pub fn openhcl_linux_direct(resolver: &ArtifactResolver<'_>, arch: MachineArch) -> Self {
1579        use petri_artifacts_vmm_test::artifacts::openhcl_igvm::*;
1580        match arch {
1581            MachineArch::X86_64 => Firmware::OpenhclLinuxDirect {
1582                igvm_path: resolver.require(LATEST_LINUX_DIRECT_TEST_X64).erase(),
1583                openhcl_config: Default::default(),
1584            },
1585            MachineArch::Aarch64 => todo!("Linux direct not yet supported on aarch64"),
1586        }
1587    }
1588
1589    /// Constructs a standard [`Firmware::Pcat`] configuration.
1590    pub fn pcat(resolver: &ArtifactResolver<'_>, guest: PcatGuest) -> Self {
1591        use petri_artifacts_vmm_test::artifacts::loadable::*;
1592        Firmware::Pcat {
1593            guest,
1594            bios_firmware: resolver.try_require(PCAT_FIRMWARE_X64).erase(),
1595            svga_firmware: resolver.try_require(SVGA_FIRMWARE_X64).erase(),
1596        }
1597    }
1598
1599    /// Constructs a standard [`Firmware::Uefi`] configuration.
1600    pub fn uefi(resolver: &ArtifactResolver<'_>, arch: MachineArch, guest: UefiGuest) -> Self {
1601        use petri_artifacts_vmm_test::artifacts::loadable::*;
1602        let uefi_firmware = match arch {
1603            MachineArch::X86_64 => resolver.require(UEFI_FIRMWARE_X64).erase(),
1604            MachineArch::Aarch64 => resolver.require(UEFI_FIRMWARE_AARCH64).erase(),
1605        };
1606        Firmware::Uefi {
1607            guest,
1608            uefi_firmware,
1609            uefi_config: Default::default(),
1610        }
1611    }
1612
1613    /// Constructs a standard [`Firmware::OpenhclUefi`] configuration.
1614    pub fn openhcl_uefi(
1615        resolver: &ArtifactResolver<'_>,
1616        arch: MachineArch,
1617        guest: UefiGuest,
1618        isolation: Option<IsolationType>,
1619        vtl2_nvme_boot: bool,
1620    ) -> Self {
1621        use petri_artifacts_vmm_test::artifacts::openhcl_igvm::*;
1622        let igvm_path = match arch {
1623            MachineArch::X86_64 if isolation.is_some() => resolver.require(LATEST_CVM_X64).erase(),
1624            MachineArch::X86_64 => resolver.require(LATEST_STANDARD_X64).erase(),
1625            MachineArch::Aarch64 => resolver.require(LATEST_STANDARD_AARCH64).erase(),
1626        };
1627        Firmware::OpenhclUefi {
1628            guest,
1629            isolation,
1630            igvm_path,
1631            uefi_config: Default::default(),
1632            openhcl_config: OpenHclConfig {
1633                vtl2_nvme_boot,
1634                ..Default::default()
1635            },
1636        }
1637    }
1638
1639    fn is_openhcl(&self) -> bool {
1640        match self {
1641            Firmware::OpenhclLinuxDirect { .. }
1642            | Firmware::OpenhclUefi { .. }
1643            | Firmware::OpenhclPcat { .. } => true,
1644            Firmware::LinuxDirect { .. } | Firmware::Pcat { .. } | Firmware::Uefi { .. } => false,
1645        }
1646    }
1647
1648    fn isolation(&self) -> Option<IsolationType> {
1649        match self {
1650            Firmware::OpenhclUefi { isolation, .. } => *isolation,
1651            Firmware::LinuxDirect { .. }
1652            | Firmware::Pcat { .. }
1653            | Firmware::Uefi { .. }
1654            | Firmware::OpenhclLinuxDirect { .. }
1655            | Firmware::OpenhclPcat { .. } => None,
1656        }
1657    }
1658
1659    fn is_linux_direct(&self) -> bool {
1660        match self {
1661            Firmware::LinuxDirect { .. } | Firmware::OpenhclLinuxDirect { .. } => true,
1662            Firmware::Pcat { .. }
1663            | Firmware::Uefi { .. }
1664            | Firmware::OpenhclUefi { .. }
1665            | Firmware::OpenhclPcat { .. } => false,
1666        }
1667    }
1668
1669    fn is_pcat(&self) -> bool {
1670        match self {
1671            Firmware::Pcat { .. } | Firmware::OpenhclPcat { .. } => true,
1672            Firmware::Uefi { .. }
1673            | Firmware::OpenhclUefi { .. }
1674            | Firmware::LinuxDirect { .. }
1675            | Firmware::OpenhclLinuxDirect { .. } => false,
1676        }
1677    }
1678
1679    fn os_flavor(&self) -> OsFlavor {
1680        match self {
1681            Firmware::LinuxDirect { .. } | Firmware::OpenhclLinuxDirect { .. } => OsFlavor::Linux,
1682            Firmware::Uefi {
1683                guest: UefiGuest::GuestTestUefi { .. } | UefiGuest::None,
1684                ..
1685            }
1686            | Firmware::OpenhclUefi {
1687                guest: UefiGuest::GuestTestUefi { .. } | UefiGuest::None,
1688                ..
1689            } => OsFlavor::Uefi,
1690            Firmware::Pcat {
1691                guest: PcatGuest::Vhd(cfg),
1692                ..
1693            }
1694            | Firmware::OpenhclPcat {
1695                guest: PcatGuest::Vhd(cfg),
1696                ..
1697            }
1698            | Firmware::Uefi {
1699                guest: UefiGuest::Vhd(cfg),
1700                ..
1701            }
1702            | Firmware::OpenhclUefi {
1703                guest: UefiGuest::Vhd(cfg),
1704                ..
1705            } => cfg.os_flavor,
1706            Firmware::Pcat {
1707                guest: PcatGuest::Iso(cfg),
1708                ..
1709            }
1710            | Firmware::OpenhclPcat {
1711                guest: PcatGuest::Iso(cfg),
1712                ..
1713            } => cfg.os_flavor,
1714        }
1715    }
1716
1717    fn quirks(&self) -> GuestQuirks {
1718        match self {
1719            Firmware::Pcat {
1720                guest: PcatGuest::Vhd(cfg),
1721                ..
1722            }
1723            | Firmware::Uefi {
1724                guest: UefiGuest::Vhd(cfg),
1725                ..
1726            }
1727            | Firmware::OpenhclUefi {
1728                guest: UefiGuest::Vhd(cfg),
1729                ..
1730            } => cfg.quirks.clone(),
1731            Firmware::Pcat {
1732                guest: PcatGuest::Iso(cfg),
1733                ..
1734            } => cfg.quirks.clone(),
1735            _ => Default::default(),
1736        }
1737    }
1738
1739    fn expected_boot_event(&self) -> Option<FirmwareEvent> {
1740        match self {
1741            Firmware::LinuxDirect { .. }
1742            | Firmware::OpenhclLinuxDirect { .. }
1743            | Firmware::Uefi {
1744                guest: UefiGuest::GuestTestUefi(_),
1745                ..
1746            }
1747            | Firmware::OpenhclUefi {
1748                guest: UefiGuest::GuestTestUefi(_),
1749                ..
1750            } => None,
1751            Firmware::Pcat { .. } | Firmware::OpenhclPcat { .. } => {
1752                // TODO: Handle older PCAT versions that don't fire the event
1753                Some(FirmwareEvent::BootAttempt)
1754            }
1755            Firmware::Uefi {
1756                guest: UefiGuest::None,
1757                ..
1758            }
1759            | Firmware::OpenhclUefi {
1760                guest: UefiGuest::None,
1761                ..
1762            } => Some(FirmwareEvent::NoBootDevice),
1763            Firmware::Uefi { .. } | Firmware::OpenhclUefi { .. } => {
1764                Some(FirmwareEvent::BootSuccess)
1765            }
1766        }
1767    }
1768
1769    fn openhcl_config(&self) -> Option<&OpenHclConfig> {
1770        match self {
1771            Firmware::OpenhclLinuxDirect { openhcl_config, .. }
1772            | Firmware::OpenhclUefi { openhcl_config, .. }
1773            | Firmware::OpenhclPcat { openhcl_config, .. } => Some(openhcl_config),
1774            Firmware::LinuxDirect { .. } | Firmware::Pcat { .. } | Firmware::Uefi { .. } => None,
1775        }
1776    }
1777
1778    fn openhcl_config_mut(&mut self) -> Option<&mut OpenHclConfig> {
1779        match self {
1780            Firmware::OpenhclLinuxDirect { openhcl_config, .. }
1781            | Firmware::OpenhclUefi { openhcl_config, .. }
1782            | Firmware::OpenhclPcat { openhcl_config, .. } => Some(openhcl_config),
1783            Firmware::LinuxDirect { .. } | Firmware::Pcat { .. } | Firmware::Uefi { .. } => None,
1784        }
1785    }
1786
1787    fn into_openhcl_config(self) -> Option<OpenHclConfig> {
1788        match self {
1789            Firmware::OpenhclLinuxDirect { openhcl_config, .. }
1790            | Firmware::OpenhclUefi { openhcl_config, .. }
1791            | Firmware::OpenhclPcat { openhcl_config, .. } => Some(openhcl_config),
1792            Firmware::LinuxDirect { .. } | Firmware::Pcat { .. } | Firmware::Uefi { .. } => None,
1793        }
1794    }
1795
1796    fn uefi_config(&self) -> Option<&UefiConfig> {
1797        match self {
1798            Firmware::Uefi { uefi_config, .. } | Firmware::OpenhclUefi { uefi_config, .. } => {
1799                Some(uefi_config)
1800            }
1801            Firmware::LinuxDirect { .. }
1802            | Firmware::OpenhclLinuxDirect { .. }
1803            | Firmware::Pcat { .. }
1804            | Firmware::OpenhclPcat { .. } => None,
1805        }
1806    }
1807
1808    fn uefi_config_mut(&mut self) -> Option<&mut UefiConfig> {
1809        match self {
1810            Firmware::Uefi { uefi_config, .. } | Firmware::OpenhclUefi { uefi_config, .. } => {
1811                Some(uefi_config)
1812            }
1813            Firmware::LinuxDirect { .. }
1814            | Firmware::OpenhclLinuxDirect { .. }
1815            | Firmware::Pcat { .. }
1816            | Firmware::OpenhclPcat { .. } => None,
1817        }
1818    }
1819}
1820
1821/// The guest the VM will boot into. A boot drive with the chosen setup
1822/// will be automatically configured.
1823#[derive(Debug)]
1824pub enum PcatGuest {
1825    /// Mount a VHD as the boot drive.
1826    Vhd(BootImageConfig<boot_image_type::Vhd>),
1827    /// Mount an ISO as the CD/DVD drive.
1828    Iso(BootImageConfig<boot_image_type::Iso>),
1829}
1830
1831impl PcatGuest {
1832    fn artifact(&self) -> &ResolvedArtifact {
1833        match self {
1834            PcatGuest::Vhd(disk) => &disk.artifact,
1835            PcatGuest::Iso(disk) => &disk.artifact,
1836        }
1837    }
1838}
1839
1840/// The guest the VM will boot into. A boot drive with the chosen setup
1841/// will be automatically configured.
1842#[derive(Debug)]
1843pub enum UefiGuest {
1844    /// Mount a VHD as the boot drive.
1845    Vhd(BootImageConfig<boot_image_type::Vhd>),
1846    /// The UEFI test image produced by our guest-test infrastructure.
1847    GuestTestUefi(ResolvedArtifact),
1848    /// No guest, just the firmware.
1849    None,
1850}
1851
1852impl UefiGuest {
1853    /// Construct a standard [`UefiGuest::GuestTestUefi`] configuration.
1854    pub fn guest_test_uefi(resolver: &ArtifactResolver<'_>, arch: MachineArch) -> Self {
1855        use petri_artifacts_vmm_test::artifacts::test_vhd::*;
1856        let artifact = match arch {
1857            MachineArch::X86_64 => resolver.require(GUEST_TEST_UEFI_X64).erase(),
1858            MachineArch::Aarch64 => resolver.require(GUEST_TEST_UEFI_AARCH64).erase(),
1859        };
1860        UefiGuest::GuestTestUefi(artifact)
1861    }
1862
1863    fn artifact(&self) -> Option<&ResolvedArtifact> {
1864        match self {
1865            UefiGuest::Vhd(vhd) => Some(&vhd.artifact),
1866            UefiGuest::GuestTestUefi(p) => Some(p),
1867            UefiGuest::None => None,
1868        }
1869    }
1870}
1871
1872/// Type-tags for [`BootImageConfig`](super::BootImageConfig)
1873pub mod boot_image_type {
1874    mod private {
1875        pub trait Sealed {}
1876        impl Sealed for super::Vhd {}
1877        impl Sealed for super::Iso {}
1878    }
1879
1880    /// Private trait use to seal the set of artifact types BootImageType
1881    /// supports.
1882    pub trait BootImageType: private::Sealed {}
1883
1884    /// BootImageConfig for a VHD file
1885    #[derive(Debug)]
1886    pub enum Vhd {}
1887
1888    /// BootImageConfig for an ISO file
1889    #[derive(Debug)]
1890    pub enum Iso {}
1891
1892    impl BootImageType for Vhd {}
1893    impl BootImageType for Iso {}
1894}
1895
1896/// Configuration information for the boot drive of the VM.
1897#[derive(Debug)]
1898pub struct BootImageConfig<T: boot_image_type::BootImageType> {
1899    /// Artifact handle corresponding to the boot media.
1900    artifact: ResolvedArtifact,
1901    /// The OS flavor.
1902    os_flavor: OsFlavor,
1903    /// Any quirks needed to boot the guest.
1904    ///
1905    /// Most guests should not need any quirks, and can use `Default`.
1906    quirks: GuestQuirks,
1907    /// Marker denoting what type of media `artifact` corresponds to
1908    _type: core::marker::PhantomData<T>,
1909}
1910
1911impl BootImageConfig<boot_image_type::Vhd> {
1912    /// Create a new BootImageConfig from a VHD artifact handle
1913    pub fn from_vhd<A>(artifact: ResolvedArtifact<A>) -> Self
1914    where
1915        A: petri_artifacts_common::tags::IsTestVhd,
1916    {
1917        BootImageConfig {
1918            artifact: artifact.erase(),
1919            os_flavor: A::OS_FLAVOR,
1920            quirks: A::quirks(),
1921            _type: std::marker::PhantomData,
1922        }
1923    }
1924}
1925
1926impl BootImageConfig<boot_image_type::Iso> {
1927    /// Create a new BootImageConfig from an ISO artifact handle
1928    pub fn from_iso<A>(artifact: ResolvedArtifact<A>) -> Self
1929    where
1930        A: petri_artifacts_common::tags::IsTestIso,
1931    {
1932        BootImageConfig {
1933            artifact: artifact.erase(),
1934            os_flavor: A::OS_FLAVOR,
1935            quirks: A::quirks(),
1936            _type: std::marker::PhantomData,
1937        }
1938    }
1939}
1940
1941/// Isolation type
1942#[derive(Debug, Clone, Copy)]
1943pub enum IsolationType {
1944    /// VBS
1945    Vbs,
1946    /// SNP
1947    Snp,
1948    /// TDX
1949    Tdx,
1950}
1951
1952/// Flags controlling servicing behavior.
1953#[derive(Debug, Clone, Copy)]
1954pub struct OpenHclServicingFlags {
1955    /// Preserve DMA memory for NVMe devices if supported.
1956    /// Defaults to `true`.
1957    pub enable_nvme_keepalive: bool,
1958    /// Preserve DMA memory for MANA devices if supported.
1959    pub enable_mana_keepalive: bool,
1960    /// 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.
1961    pub override_version_checks: bool,
1962    /// Hint to the OpenHCL runtime how much time to wait when stopping / saving the OpenHCL.
1963    pub stop_timeout_hint_secs: Option<u16>,
1964}
1965
1966/// Petri disk type
1967#[derive(Debug, Clone)]
1968pub enum PetriDiskType {
1969    /// Memory backed
1970    Memory,
1971    /// Memory differencing disk backed by a file
1972    Differencing(PathBuf),
1973    /// Persistent disk
1974    Persistent(PathBuf),
1975}
1976
1977/// Petri VMGS disk
1978#[derive(Debug, Clone)]
1979pub struct PetriVmgsDisk {
1980    /// Backing disk
1981    pub disk: PetriDiskType,
1982    /// Guest state encryption policy
1983    pub encryption_policy: GuestStateEncryptionPolicy,
1984}
1985
1986impl Default for PetriVmgsDisk {
1987    fn default() -> Self {
1988        PetriVmgsDisk {
1989            disk: PetriDiskType::Memory,
1990            // TODO: make this strict once we can set it in OpenHCL on Hyper-V
1991            encryption_policy: GuestStateEncryptionPolicy::None(false),
1992        }
1993    }
1994}
1995
1996/// Petri VM guest state resource
1997#[derive(Debug, Clone)]
1998pub enum PetriVmgsResource {
1999    /// Use disk to store guest state
2000    Disk(PetriVmgsDisk),
2001    /// Use disk to store guest state, reformatting if corrupted.
2002    ReprovisionOnFailure(PetriVmgsDisk),
2003    /// Format and use disk to store guest state
2004    Reprovision(PetriVmgsDisk),
2005    /// Store guest state in memory
2006    Ephemeral,
2007}
2008
2009impl PetriVmgsResource {
2010    /// get the inner vmgs disk if one exists
2011    pub fn disk(&self) -> Option<&PetriVmgsDisk> {
2012        match self {
2013            PetriVmgsResource::Disk(vmgs)
2014            | PetriVmgsResource::ReprovisionOnFailure(vmgs)
2015            | PetriVmgsResource::Reprovision(vmgs) => Some(vmgs),
2016            PetriVmgsResource::Ephemeral => None,
2017        }
2018    }
2019}
2020
2021/// Petri VM guest state lifetime
2022#[derive(Debug, Clone, Copy)]
2023pub enum PetriGuestStateLifetime {
2024    /// Use a differencing disk backed by a blank, tempory VMGS file
2025    /// or other artifact if one is provided
2026    Disk,
2027    /// Same as default, except reformat the backing disk if corrupted
2028    ReprovisionOnFailure,
2029    /// Same as default, except reformat the backing disk
2030    Reprovision,
2031    /// Store guest state in memory (no backing disk)
2032    Ephemeral,
2033}
2034
2035/// UEFI secure boot template
2036#[derive(Debug, Clone, Copy)]
2037pub enum SecureBootTemplate {
2038    /// The Microsoft Windows template.
2039    MicrosoftWindows,
2040    /// The Microsoft UEFI certificate authority template.
2041    MicrosoftUefiCertificateAuthority,
2042}
2043
2044/// Quirks to workaround certain bugs that only manifest when using a
2045/// particular VMM, and do not depend on which guest is running.
2046#[derive(Default, Debug, Clone)]
2047pub struct VmmQuirks {
2048    /// Automatically reset the VM if we did not recieve a boot event in the
2049    /// specified amount of time.
2050    pub flaky_boot: Option<Duration>,
2051}
2052
2053/// Creates a VM-safe name that respects platform limitations.
2054///
2055/// Hyper-V limits VM names to 100 characters. For names that exceed this limit,
2056/// this function truncates to 96 characters and appends a 4-character hash
2057/// to ensure uniqueness while staying within the limit.
2058fn make_vm_safe_name(name: &str) -> String {
2059    const MAX_VM_NAME_LENGTH: usize = 100;
2060    const HASH_LENGTH: usize = 4;
2061    const MAX_PREFIX_LENGTH: usize = MAX_VM_NAME_LENGTH - HASH_LENGTH;
2062
2063    if name.len() <= MAX_VM_NAME_LENGTH {
2064        name.to_owned()
2065    } else {
2066        // Create a hash of the full name for uniqueness
2067        let mut hasher = DefaultHasher::new();
2068        name.hash(&mut hasher);
2069        let hash = hasher.finish();
2070
2071        // Format hash as a 4-character hex string
2072        let hash_suffix = format!("{:04x}", hash & 0xFFFF);
2073
2074        // Truncate the name and append the hash
2075        let truncated = &name[..MAX_PREFIX_LENGTH];
2076        tracing::debug!(
2077            "VM name too long ({}), truncating '{}' to '{}{}'",
2078            name.len(),
2079            name,
2080            truncated,
2081            hash_suffix
2082        );
2083
2084        format!("{}{}", truncated, hash_suffix)
2085    }
2086}
2087
2088/// The reason that the VM halted
2089#[derive(Debug, Clone, Copy, Eq, PartialEq)]
2090pub enum PetriHaltReason {
2091    /// The vm powered off
2092    PowerOff,
2093    /// The vm reset
2094    Reset,
2095    /// The vm hibernated
2096    Hibernate,
2097    /// The vm triple faulted
2098    TripleFault,
2099    /// The vm halted for some other reason
2100    Other,
2101}
2102
2103fn append_cmdline(cmd: &mut Option<String>, add_cmd: impl AsRef<str>) {
2104    if let Some(cmd) = cmd.as_mut() {
2105        cmd.push(' ');
2106        cmd.push_str(add_cmd.as_ref());
2107    } else {
2108        *cmd = Some(add_cmd.as_ref().to_string());
2109    }
2110}
2111
2112async fn save_inspect(
2113    name: &str,
2114    inspect: std::pin::Pin<Box<dyn Future<Output = anyhow::Result<inspect::Node>> + Send>>,
2115    log_source: &PetriLogSource,
2116) {
2117    tracing::info!("Collecting {name} inspect details.");
2118    let node = match inspect.await {
2119        Ok(n) => n,
2120        Err(e) => {
2121            tracing::error!(?e, "Failed to get {name}");
2122            return;
2123        }
2124    };
2125    if let Err(e) = log_source.write_attachment(
2126        &format!("timeout_inspect_{name}.log"),
2127        format!("{node:#}").as_bytes(),
2128    ) {
2129        tracing::error!(?e, "Failed to save {name} inspect log");
2130        return;
2131    }
2132    tracing::info!("{name} inspect task finished.");
2133}
2134
2135/// Wrapper for modification functions with stubbed out debug impl
2136pub struct ModifyFn<T>(pub Box<dyn FnOnce(&mut T) + Send + Sync>);
2137
2138impl<T> std::fmt::Debug for ModifyFn<T> {
2139    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
2140        write!(f, "_")
2141    }
2142}
2143
2144/// Default VTL 2 settings used by petri
2145fn default_vtl2_settings() -> Vtl2Settings {
2146    Vtl2Settings {
2147        version: vtl2_settings_proto::vtl2_settings_base::Version::V1.into(),
2148        fixed: None,
2149        dynamic: Some(Default::default()),
2150        namespace_settings: Default::default(),
2151    }
2152}
2153
2154#[cfg(test)]
2155mod tests {
2156    use super::make_vm_safe_name;
2157
2158    #[test]
2159    fn test_short_names_unchanged() {
2160        let short_name = "short_test_name";
2161        assert_eq!(make_vm_safe_name(short_name), short_name);
2162    }
2163
2164    #[test]
2165    fn test_exactly_100_chars_unchanged() {
2166        let name_100 = "a".repeat(100);
2167        assert_eq!(make_vm_safe_name(&name_100), name_100);
2168    }
2169
2170    #[test]
2171    fn test_long_name_truncated() {
2172        let long_name = "multiarch::openhcl_servicing::hyperv_openhcl_uefi_aarch64_ubuntu_2404_server_aarch64_openhcl_servicing";
2173        let result = make_vm_safe_name(long_name);
2174
2175        // Should be exactly 100 characters
2176        assert_eq!(result.len(), 100);
2177
2178        // Should start with the truncated prefix
2179        assert!(result.starts_with("multiarch::openhcl_servicing::hyperv_openhcl_uefi_aarch64_ubuntu_2404_server_aarch64_ope"));
2180
2181        // Should end with a 4-character hash
2182        let suffix = &result[96..];
2183        assert_eq!(suffix.len(), 4);
2184        // Should be valid hex
2185        assert!(u16::from_str_radix(suffix, 16).is_ok());
2186    }
2187
2188    #[test]
2189    fn test_deterministic_results() {
2190        let long_name = "very_long_test_name_that_exceeds_the_100_character_limit_and_should_be_truncated_consistently_every_time";
2191        let result1 = make_vm_safe_name(long_name);
2192        let result2 = make_vm_safe_name(long_name);
2193
2194        assert_eq!(result1, result2);
2195        assert_eq!(result1.len(), 100);
2196    }
2197
2198    #[test]
2199    fn test_different_names_different_hashes() {
2200        let name1 = "very_long_test_name_that_definitely_exceeds_the_100_character_limit_and_should_be_truncated_by_the_function_version_1";
2201        let name2 = "very_long_test_name_that_definitely_exceeds_the_100_character_limit_and_should_be_truncated_by_the_function_version_2";
2202
2203        let result1 = make_vm_safe_name(name1);
2204        let result2 = make_vm_safe_name(name2);
2205
2206        // Both should be 100 chars
2207        assert_eq!(result1.len(), 100);
2208        assert_eq!(result2.len(), 100);
2209
2210        // Should have different suffixes since the full names are different
2211        assert_ne!(result1, result2);
2212        assert_ne!(&result1[96..], &result2[96..]);
2213    }
2214}