petri: hyperv watchdogs and screenshots (#1826)

Take screenshots and collect inspect info using a watchdog when using
the Hyper-V backend as well as OpenVMM

Author: tjones60

petri/src/openhcl_diag.rs

@@ -92,12 +92,16 @@ impl OpenHclDiagHandler {
         client.crash(pid).await
     }
 
-    pub async fn test_inspect(&self) -> anyhow::Result<()> {
+    pub async fn inspect(
+        &self,
+        path: impl Into<String>,
+        depth: Option<usize>,
+        timeout: Option<std::time::Duration>,
+    ) -> anyhow::Result<inspect::Node> {
         self.diag_client()
             .await?
-            .inspect("", None, None)
+            .inspect(path, depth, timeout)
             .await
-            .map(|_| ())
     }
 
     pub async fn kmsg(&self) -> anyhow::Result<KmsgStream> {

petri/src/vm/hyperv/hyperv.psm1

@@ -401,3 +401,35 @@ function Restart-OpenHCL
 
     $result | Trace-CimMethodExecution -CimInstance $guestManagementService -MethodName "ReloadManagementVtl" -TimeoutSeconds $TimeoutHintSeconds
 }
+
+function Get-VmScreenshot
+{
+    [CmdletBinding()]
+    Param(
+        [Parameter(Position = 0, Mandatory = $true, ValueFromPipeline = $true)]
+        [System.Object]
+        $Vm,
+
+        [Parameter(Mandatory = $true)]
+        [string] $Path
+    )
+
+    $vmms = Get-Vmms
+    $vmcs = Get-MsvmComputerSystem $Vm
+
+    # Get the resolution of the screen at the moment
+    $videoHead = $vmcs | Get-CimAssociatedInstance -ResultClassName "Msvm_VideoHead"
+    $x = $videoHead.CurrentHorizontalResolution
+    $y = $videoHead.CurrentVerticalResolution
+
+    # Get screenshot
+    $image = $vmms | Invoke-CimMethod -MethodName "GetVirtualSystemThumbnailImage" -Arguments @{
+        TargetSystem = $vmcs
+        WidthPixels = $x
+        HeightPixels = $y
+    }
+
+    [IO.File]::WriteAllBytes($Path, $image.ImageData)
+
+    return "$x,$y"
+}

petri/src/vm/hyperv/mod.rs

@@ -10,15 +10,18 @@ use vmsocket::VmSocket;
 use super::ProcessorTopology;
 use crate::Firmware;
 use crate::IsolationType;
+use crate::NoPetriVmInspector;
 use crate::OpenHclConfig;
 use crate::OpenHclServicingFlags;
 use crate::PetriVmConfig;
+use crate::PetriVmFramebufferAccess;
 use crate::PetriVmResources;
 use crate::PetriVmRuntime;
 use crate::PetriVmmBackend;
 use crate::SecureBootTemplate;
 use crate::ShutdownKind;
 use crate::UefiConfig;
+use crate::VmScreenshotMeta;
 use crate::disk_image::AgentImage;
 use crate::hyperv::powershell::HyperVSecureBootTemplate;
 use crate::openhcl_diag::OpenHclDiagHandler;
@@ -42,6 +45,8 @@ use pipette_client::PipetteClient;
 use std::fs;
 use std::io::Write;
 use std::path::Path;
+use std::sync::Arc;
+use std::sync::Weak;
 use std::time::Duration;
 use vm::HyperVVM;
 use vmm_core_defs::HaltReason;
@@ -51,10 +56,10 @@ pub struct HyperVPetriBackend {}
 
 /// Resources needed at runtime for a Hyper-V Petri VM
 pub struct HyperVPetriRuntime {
-    vm: HyperVVM,
+    vm: Arc<HyperVVM>,
     log_tasks: Vec<Task<anyhow::Result<()>>>,
     temp_dir: tempfile::TempDir,
-    openhcl_diag_handler: Option<OpenHclDiagHandler>,
+    is_openhcl: bool,
     driver: DefaultDriver,
 }
 
@@ -299,7 +304,7 @@ impl PetriVmmBackend for HyperVPetriBackend {
             }
         }
 
-        let openhcl_diag_handler = if let Some((
+        if let Some((
             src_igvm_file,
             OpenHclConfig {
                 vtl2_nvme_boot: _, // TODO, see #1649.
@@ -365,13 +370,7 @@ impl PetriVmmBackend for HyperVPetriBackend {
                     }
                 }
             }));
-
-            Some(OpenHclDiagHandler::new(
-                diag_client::DiagClient::from_hyperv_id(driver.clone(), *vm.vmid()),
-            ))
-        } else {
-            None
-        };
+        }
 
         let serial_pipe_path = vm.set_vm_com_port(1)?;
         let serial_log_file = log_source.log_file("guest")?;
@@ -390,22 +389,25 @@ impl PetriVmmBackend for HyperVPetriBackend {
         vm.start()?;
 
         Ok(HyperVPetriRuntime {
-            vm,
+            vm: Arc::new(vm),
             log_tasks,
             temp_dir,
-            openhcl_diag_handler,
+            is_openhcl: openhcl_config.is_some(),
             driver: driver.clone(),
         })
     }
 }
 
 #[async_trait]
 impl PetriVmRuntime for HyperVPetriRuntime {
-    async fn teardown(self) -> anyhow::Result<()> {
-        for t in self.log_tasks {
-            _ = t.cancel();
-        }
-        self.vm.remove()
+    type VmInspector = NoPetriVmInspector;
+    type VmFramebufferAccess = HyperVFramebufferAccess;
+
+    async fn teardown(mut self) -> anyhow::Result<()> {
+        futures::future::join_all(self.log_tasks.into_iter().map(|t| t.cancel())).await;
+        Arc::into_inner(self.vm)
+            .context("all references to the Hyper-V VM object have not been closed")?
+            .remove()
     }
 
     async fn wait_for_halt(&mut self) -> anyhow::Result<HaltReason> {
@@ -450,8 +452,13 @@ impl PetriVmRuntime for HyperVPetriRuntime {
             .context("failed to connect to pipette")
     }
 
-    fn openhcl_diag(&self) -> Option<&OpenHclDiagHandler> {
-        self.openhcl_diag_handler.as_ref()
+    fn openhcl_diag(&self) -> Option<OpenHclDiagHandler> {
+        self.is_openhcl.then(|| {
+            OpenHclDiagHandler::new(diag_client::DiagClient::from_hyperv_id(
+                self.driver.clone(),
+                *self.vm.vmid(),
+            ))
+        })
     }
 
     async fn wait_for_successful_boot_event(&mut self) -> anyhow::Result<()> {
@@ -483,6 +490,26 @@ impl PetriVmRuntime for HyperVPetriRuntime {
         // TODO: Updating the file causes failure ... self.vm.set_openhcl_firmware(new_openhcl.get(), false)?;
         self.vm.restart_openhcl(flags).await
     }
+
+    fn take_framebuffer_access(&mut self) -> Option<HyperVFramebufferAccess> {
+        Some(HyperVFramebufferAccess {
+            vm: Arc::downgrade(&self.vm),
+        })
+    }
+}
+
+/// Interface to the Hyper-V framebuffer
+pub struct HyperVFramebufferAccess {
+    vm: Weak<HyperVVM>,
+}
+
+impl PetriVmFramebufferAccess for HyperVFramebufferAccess {
+    fn screenshot(&mut self, image: &mut Vec<u8>) -> anyhow::Result<VmScreenshotMeta> {
+        self.vm
+            .upgrade()
+            .context("VM no longer exists")?
+            .screenshot(image)
+    }
 }
 
 fn acl_read_for_vm(path: &Path, id: Option<guid::Guid>) -> anyhow::Result<()> {

petri/src/vm/hyperv/powershell.rs

@@ -538,7 +538,7 @@ pub fn run_set_vm_com_port(vmid: &Guid, port: u8, path: &Path) -> anyhow::Result
 }
 
 /// Run Set-VMBusRelay commandlet
-pub fn set_vmbus_redirect(vmid: &Guid, ps_mod: &Path, enable: bool) -> anyhow::Result<()> {
+pub fn run_set_vmbus_redirect(vmid: &Guid, ps_mod: &Path, enable: bool) -> anyhow::Result<()> {
     run_cmd(
         PowerShellBuilder::new()
             .cmdlet("Import-Module")
@@ -831,3 +831,30 @@ pub fn run_remove_vm_scsi_controller(vmid: &Guid, controller_number: u32) -> any
     .map(|_| ())
     .context("remove_vm_scsi_controller")
 }
+
+/// Run Get-VmScreenshot commandlet
+pub fn run_get_vm_screenshot(
+    vmid: &Guid,
+    ps_mod: &Path,
+    path: &Path,
+) -> anyhow::Result<(u16, u16)> {
+    let output = run_cmd(
+        PowerShellBuilder::new()
+            .cmdlet("Import-Module")
+            .positional(ps_mod)
+            .next()
+            .cmdlet("Get-VM")
+            .arg("Id", vmid)
+            .pipeline()
+            .cmdlet("Get-VmScreenshot")
+            .arg("Path", path)
+            .finish()
+            .build(),
+    )
+    .context("get_vm_screenshot")?;
+    let (x, y) = output.split_once(',').context("invalid dimensions")?;
+    Ok((
+        x.parse().context("invalid x dimension")?,
+        y.parse().context("invalid y dimension")?,
+    ))
+}

petri/src/vm/hyperv/vm.rs

@@ -8,6 +8,7 @@ use super::hvc::VmState;
 use super::powershell;
 use crate::OpenHclServicingFlags;
 use crate::PetriLogFile;
+use crate::VmScreenshotMeta;
 use anyhow::Context;
 use get_resources::ged::FirmwareEvent;
 use guid::Guid;
@@ -30,7 +31,7 @@ pub struct HyperVVM {
     name: String,
     vmid: Guid,
     destroyed: bool,
-    _temp_dir: TempDir,
+    temp_dir: TempDir,
     ps_mod: PathBuf,
     create_time: Timestamp,
     log_file: PetriLogFile,
@@ -98,7 +99,7 @@ impl HyperVVM {
             name,
             vmid,
             destroyed: false,
-            _temp_dir: temp_dir,
+            temp_dir,
             ps_mod,
             create_time,
             log_file,
@@ -167,7 +168,7 @@ impl HyperVVM {
 
     /// Waits for an event emitted by the firmware about its boot status, and
     /// verifies that it is the expected success value.
-    pub async fn wait_for_successful_boot_event(&mut self) -> anyhow::Result<()> {
+    pub async fn wait_for_successful_boot_event(&self) -> anyhow::Result<()> {
         if let Some(expected_boot_event) = self.expected_boot_event {
             self.wait_for(Self::boot_event, Some(expected_boot_event), 240.seconds())
                 .await
@@ -181,7 +182,7 @@ impl HyperVVM {
 
     /// Waits for an event emitted by the firmware about its boot status, and
     /// returns that status.
-    pub async fn wait_for_boot_event(&mut self) -> anyhow::Result<FirmwareEvent> {
+    pub async fn wait_for_boot_event(&self) -> anyhow::Result<FirmwareEvent> {
         self.wait_for_some(Self::boot_event, 240.seconds()).await
     }
 
@@ -436,13 +437,52 @@ impl HyperVVM {
 
     /// Enable VMBusRelay
     pub fn set_vmbus_redirect(&self, enable: bool) -> anyhow::Result<()> {
-        powershell::set_vmbus_redirect(&self.vmid, &self.ps_mod, enable)
+        powershell::run_set_vmbus_redirect(&self.vmid, &self.ps_mod, enable)
     }
 
     /// Perform an OpenHCL servicing operation.
     pub async fn restart_openhcl(&self, flags: OpenHclServicingFlags) -> anyhow::Result<()> {
         powershell::run_restart_openhcl(&self.vmid, &self.ps_mod, flags)
     }
+
+    /// Take a screenshot of the VM
+    pub fn screenshot(&self, image: &mut Vec<u8>) -> anyhow::Result<VmScreenshotMeta> {
+        const IN_BYTES_PER_PIXEL: usize = 2;
+        const OUT_BYTES_PER_PIXEL: usize = 3;
+        let temp_bin_path = self.temp_dir.path().join("screenshot.bin");
+        let (width, height) =
+            powershell::run_get_vm_screenshot(&self.vmid, &self.ps_mod, &temp_bin_path)?;
+        let (widthsize, heightsize) = (width as usize, height as usize);
+        let in_len = widthsize * heightsize * IN_BYTES_PER_PIXEL;
+        let out_len = widthsize * heightsize * OUT_BYTES_PER_PIXEL;
+        let mut image_rgb565 = fs_err::read(temp_bin_path)?;
+        image_rgb565.truncate(in_len);
+        if image_rgb565.len() != in_len {
+            anyhow::bail!("did not get enough bytes for screenshot");
+        }
+
+        image.resize(out_len, 0);
+        for (out_pixel, in_pixel) in image
+            .chunks_exact_mut(OUT_BYTES_PER_PIXEL)
+            .zip(image_rgb565.chunks_exact(IN_BYTES_PER_PIXEL))
+        {
+            // convert from rgb565 ( gggbbbbb rrrrrggg )
+            // to rgb888 ( rrrrrrrr gggggggg bbbbbbbb )
+
+            // red
+            out_pixel[0] = in_pixel[1] & 0b11111000;
+            // green
+            out_pixel[1] = ((in_pixel[1] & 0b00000111) << 5) + ((in_pixel[0] & 0b11100000) >> 3);
+            // blue
+            out_pixel[2] = in_pixel[0] << 3;
+        }
+
+        Ok(VmScreenshotMeta {
+            color: image::ExtendedColorType::Rgb8,
+            width,
+            height,
+        })
+    }
 }
 
 impl Drop for HyperVVM {