petri/vm/openvmm/

runtime.rs

// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.

//! Methods to interact with a running [`PetriVmOpenVmm`].

use super::PetriVmResourcesOpenVmm;
use crate::OpenHclServicingFlags;
use crate::PetriHaltReason;
use crate::PetriHaltReasonDetail;
use crate::PetriVmFramebufferAccess;
use crate::PetriVmInspector;
use crate::PetriVmRuntime;
use crate::ShutdownKind;
use crate::VmScreenshotMeta;
use crate::openhcl_diag::OpenHclDiagHandler;
use crate::worker::Worker;
use anyhow::Context;
use async_trait::async_trait;
use framebuffer::View;
use futures::FutureExt;
use futures_concurrency::future::Race;
use get_resources::ged::FirmwareEvent;
use get_resources::ged::IpmiSelEvent;
use hyperv_ic_resources::shutdown::ShutdownRpc;
use mesh::CancelContext;
use mesh::Receiver;
use mesh::RecvError;
use mesh::rpc::RpcError;
use mesh::rpc::RpcSend;
use mesh_process::Mesh;
use openvmm_defs::rpc::PulseSaveRestoreError;
use pal_async::socket::PolledSocket;
use petri_artifacts_core::ResolvedArtifact;
#[cfg(target_os = "linux")]
use pipette_client::PIPETTE_PORT;
use pipette_client::PipetteClient;
use std::future::Future;
use std::path::Path;
use std::sync::Arc;
use std::time::Duration;
use vmm_core_defs::HaltReason;
#[cfg(target_os = "linux")]
use vmsocket::VmAddress;
#[cfg(target_os = "linux")]
use vmsocket::VmSocket;
use vtl2_settings_proto::Vtl2Settings;

/// A running VM that tests can interact with.
// DEVNOTE: Really the PetriVmInner is the actual VM and channels that we interact
// with. This struct exists as a wrapper to provide error handling, such as not
// hanging indefinitely when waiting on certain channels if the VM crashes.
pub struct PetriVmOpenVmm {
    inner: PetriVmInner,
    halt: PetriVmHaltReceiver,
}

#[async_trait]
impl PetriVmRuntime for PetriVmOpenVmm {
    type VmInspector = OpenVmmInspector;
    type VmFramebufferAccess = OpenVmmFramebufferAccess;

    async fn teardown(self) -> anyhow::Result<()> {
        tracing::info!("waiting for worker");
        let worker = Arc::into_inner(self.inner.worker)
            .context("all references to the OpenVMM worker have not been closed")?;
        worker.shutdown().await?;

        tracing::info!("Worker quit, waiting for mesh");
        self.inner.mesh.shutdown().await;

        tracing::info!("Mesh shutdown, waiting for logging tasks");
        for t in self.inner.resources.log_stream_tasks {
            t.await?;
        }

        Ok(())
    }

    async fn wait_for_halt(&mut self, allow_reset: bool) -> anyhow::Result<PetriHaltReasonDetail> {
        let halt_reason = if let Some(already) = self.halt.already_received.take() {
            already.map_err(anyhow::Error::from)
        } else {
            self.halt
                .halt_notif
                .recv()
                .await
                .context("Failed to get halt reason")
        }?;

        tracing::info!(?halt_reason, "Got halt reason");

        let reason = match halt_reason {
            HaltReason::PowerOff => PetriHaltReason::PowerOff,
            HaltReason::Reset => PetriHaltReason::Reset,
            HaltReason::Hibernate => PetriHaltReason::Hibernate,
            HaltReason::TripleFault { .. } => PetriHaltReason::TripleFault,
            _ => PetriHaltReason::Other,
        };

        if allow_reset && reason == PetriHaltReason::Reset {
            self.reset().await?
        }

        Ok(PetriHaltReasonDetail {
            reason,
            detail: format!("{halt_reason:?}"),
        })
    }

    async fn wait_for_agent(&mut self, set_high_vtl: bool) -> anyhow::Result<PipetteClient> {
        Self::wait_for_agent(self, set_high_vtl).await
    }

    fn openhcl_diag(&self) -> Option<OpenHclDiagHandler> {
        self.inner.resources.vtl2_vsock_path.as_ref().map(|path| {
            OpenHclDiagHandler::new(diag_client::DiagClient::from_hybrid_vsock(
                self.inner.resources.driver.clone(),
                path,
            ))
        })
    }

    async fn wait_for_boot_event(
        &mut self,
        timeout: Option<Duration>,
    ) -> anyhow::Result<Option<FirmwareEvent>> {
        // The event arrives on a channel, so cancelling the wait can't discard
        // one that has already been delivered.
        CancelContext::new()
            .with_timeout(timeout.unwrap_or(Duration::MAX))
            .until_cancelled(Self::wait_for_boot_event(self))
            .await
            .ok()
            .transpose()
    }

    async fn wait_for_enlightened_shutdown_ready(&mut self) -> anyhow::Result<()> {
        Self::wait_for_enlightened_shutdown_ready(self)
            .await
            .map(|_| ())
    }

    async fn send_enlightened_shutdown(&mut self, kind: ShutdownKind) -> anyhow::Result<()> {
        Self::send_enlightened_shutdown(self, kind).await
    }

    async fn restart_openhcl(
        &mut self,
        new_openhcl: &ResolvedArtifact,
        flags: OpenHclServicingFlags,
    ) -> anyhow::Result<()> {
        Self::save_openhcl(self, new_openhcl, flags).await?;
        Self::restore_openhcl(self).await
    }

    async fn save_openhcl(
        &mut self,
        new_openhcl: &ResolvedArtifact,
        flags: OpenHclServicingFlags,
    ) -> anyhow::Result<()> {
        Self::save_openhcl(self, new_openhcl, flags).await
    }

    async fn restore_openhcl(&mut self) -> anyhow::Result<()> {
        Self::restore_openhcl(self).await
    }

    async fn update_command_line(&mut self, command_line: &str) -> anyhow::Result<()> {
        Self::update_command_line(self, command_line).await
    }

    fn inspector(&self) -> Option<OpenVmmInspector> {
        Some(OpenVmmInspector {
            worker: self.inner.worker.clone(),
        })
    }

    fn take_framebuffer_access(&mut self) -> Option<OpenVmmFramebufferAccess> {
        self.inner
            .framebuffer_view
            .take()
            .map(|view| OpenVmmFramebufferAccess { view })
    }

    async fn reset(&mut self) -> anyhow::Result<()> {
        Self::reset(self).await
    }

    async fn set_vtl2_settings(&mut self, settings: &Vtl2Settings) -> anyhow::Result<()> {
        Self::set_vtl2_settings(self, settings).await
    }

    async fn set_vmbus_drive(
        &mut self,
        _disk: &crate::Drive,
        _controller_id: &guid::Guid,
        _controller_location: u32,
    ) -> anyhow::Result<()> {
        todo!("openvmm set vmbus drive")
    }

    async fn add_pcie_device(
        &mut self,
        port_name: String,
        resource: vm_resource::Resource<vm_resource::kind::PciDeviceHandleKind>,
    ) -> anyhow::Result<()> {
        Self::add_pcie_device(self, port_name, resource).await
    }

    async fn remove_pcie_device(&mut self, port_name: String) -> anyhow::Result<()> {
        Self::remove_pcie_device(self, port_name).await
    }
}

pub(super) struct PetriVmInner {
    pub(super) resources: PetriVmResourcesOpenVmm,
    pub(super) mesh: Mesh,
    pub(super) worker: Arc<Worker>,
    pub(super) framebuffer_view: Option<View>,
    /// Whether CIDATA has already been mounted inside the guest.
    /// Used to skip re-mounting after save/restore (where guest state is
    /// preserved) while still mounting after a full reset/reboot.
    pub(super) cidata_mounted: bool,
    /// Resolved TCP pipette port for no-vmbus Windows guests. Set once
    /// during startup and reused across reconnections (e.g. after reset).
    pub(super) tcp_pipette_port: Option<u16>,
    pub(super) pid: i32,
}

struct PetriVmHaltReceiver {
    halt_notif: Receiver<HaltReason>,
    already_received: Option<Result<HaltReason, RecvError>>,
}

// Wrap a PetriVmInner function in [`PetriVmOpenVmm::wait_for_halt_or_internal`] to
// provide better error handling.
macro_rules! petri_vm_fn {
    ($(#[$($attrss:tt)*])* $vis:vis async fn $fn_name:ident (&mut self $(,$arg:ident: $ty:ty)*) $(-> $ret:ty)?) => {
        $(#[$($attrss)*])*
        $vis async fn $fn_name(&mut self, $($arg:$ty,)*) $(-> $ret)? {
            Self::wait_for_halt_or_internal(&mut self.halt, self.inner.$fn_name($($arg,)*)).await
        }
    };
}

// TODO: Add all runtime functions that are not backend specific
// to the `PetriVmRuntime` trait
impl PetriVmOpenVmm {
    pub(super) fn new(inner: PetriVmInner, halt_notif: Receiver<HaltReason>) -> Self {
        Self {
            inner,
            halt: PetriVmHaltReceiver {
                halt_notif,
                already_received: None,
            },
        }
    }

    /// Get the path to the VTL 2 vsock socket, if the VM is configured with OpenHCL.
    pub fn vtl2_vsock_path(&self) -> anyhow::Result<&Path> {
        self.inner
            .resources
            .vtl2_vsock_path
            .as_deref()
            .context("VM is not configured with OpenHCL")
    }

    /// Get the PID of the openvmm child process.
    pub fn pid(&self) -> i32 {
        self.inner.pid
    }

    petri_vm_fn!(
        /// 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>
    );
    petri_vm_fn!(
        /// Waits for an IPMI SEL notification received from OpenHCL.
        pub async fn wait_for_ipmi_sel(&mut self) -> anyhow::Result<IpmiSelEvent>
    );
    petri_vm_fn!(
        /// Waits for the Hyper-V shutdown IC to be ready, returning a receiver
        /// that will be closed when it is no longer ready. Returns `None` if
        /// the shutdown IC is not configured.
        pub async fn wait_for_enlightened_shutdown_ready(&mut self) -> anyhow::Result<Option<mesh::OneshotReceiver<()>>>
    );
    petri_vm_fn!(
        /// Instruct the guest to shutdown via the Hyper-V shutdown IC.
        pub async fn send_enlightened_shutdown(&mut self, kind: ShutdownKind) -> anyhow::Result<()>
    );
    petri_vm_fn!(
        /// Waits for the KVP IC to be ready, returning a sender that can be used
        /// to send requests to it.
        pub async fn wait_for_kvp(&mut self) -> anyhow::Result<mesh::Sender<hyperv_ic_resources::kvp::KvpRpc>>
    );
    petri_vm_fn!(
        /// Stages the new OpenHCL file and saves the existing state.
        pub async fn save_openhcl(
            &mut self,
            new_openhcl: &ResolvedArtifact,
            flags: OpenHclServicingFlags
        ) -> anyhow::Result<()>
    );
    petri_vm_fn!(
        /// Restores OpenHCL from a previously saved state.
        pub async fn restore_openhcl(
            &mut self
        ) -> anyhow::Result<()>
    );
    petri_vm_fn!(
        /// Updates the command line parameters of the running VM.
        pub async fn update_command_line(
            &mut self,
            command_line: &str
        ) -> anyhow::Result<()>
    );

    petri_vm_fn!(
        /// Hot-add a PCIe device to a named port at runtime.
        pub async fn add_pcie_device(
            &mut self,
            port_name: String,
            resource: vm_resource::Resource<vm_resource::kind::PciDeviceHandleKind>
        ) -> anyhow::Result<()>
    );
    petri_vm_fn!(
        /// Hot-remove a PCIe device from a named port at runtime.
        pub async fn remove_pcie_device(
            &mut self,
            port_name: String
        ) -> anyhow::Result<()>
    );
    petri_vm_fn!(
        /// Resets the hardware state of the VM, simulating a power cycle.
        pub async fn reset(&mut self) -> anyhow::Result<()>
    );
    petri_vm_fn!(
        /// Dumps the VM's processor and memory state to a `.vmrs` file at `path`.
        pub async fn dump_state(&mut self, path: &Path) -> anyhow::Result<()>
    );
    petri_vm_fn!(
        /// Wait for a connection from a pipette agent
        pub async fn wait_for_agent(&mut self, set_high_vtl: bool) -> anyhow::Result<PipetteClient>
    );
    petri_vm_fn!(
        /// Set the OpenHCL VTL2 settings.
        pub async fn set_vtl2_settings(&mut self, settings: &Vtl2Settings) -> anyhow::Result<()>
    );

    petri_vm_fn!(
        /// Pause the VM. Call [`resume`](Self::resume) to continue execution.
        pub async fn pause(&mut self) -> anyhow::Result<()>
    );
    petri_vm_fn!(
        /// Save the VM's device and processor state, returning the serialized
        /// bytes. The VM should be paused before calling this.
        pub async fn save_state(&mut self) -> anyhow::Result<Vec<u8>>
    );
    petri_vm_fn!(
        /// Resume a paused VM.
        pub async fn resume(&mut self) -> anyhow::Result<()>
    );
    petri_vm_fn!(
        /// Perform a pulse save/restore cycle: pause the VM, save all state,
        /// reset, restore, and resume. Useful for verifying that device state
        /// survives a save/restore round-trip.
        pub async fn verify_save_restore(&mut self) -> anyhow::Result<()>
    );
    petri_vm_fn!(pub(crate) async fn launch_linux_direct_pipette(&mut self) -> anyhow::Result<()>);

    /// Wrap the provided future in a race with the worker process's halt
    /// notification channel. This is useful for preventing a future from
    /// waiting indefinitely if the VM dies for any reason. If the worker
    /// process crashes the halt notification channel will return an error, and
    /// if the VM halts for any other reason the future will complete with that
    /// reason.
    pub async fn wait_for_halt_or<T, F: Future<Output = anyhow::Result<T>>>(
        &mut self,
        future: F,
    ) -> anyhow::Result<T> {
        Self::wait_for_halt_or_internal(&mut self.halt, future).await
    }

    async fn wait_for_halt_or_internal<T, F: Future<Output = anyhow::Result<T>>>(
        halt: &mut PetriVmHaltReceiver,
        future: F,
    ) -> anyhow::Result<T> {
        let future = &mut std::pin::pin!(future);
        enum Either<T> {
            Future(anyhow::Result<T>),
            Halt(Result<HaltReason, RecvError>),
        }
        let res = (
            future.map(Either::Future),
            halt.halt_notif.recv().map(Either::Halt),
        )
            .race()
            .await;

        match res {
            Either::Future(Ok(success)) => Ok(success),
            Either::Future(Err(e)) => {
                tracing::warn!(
                    ?e,
                    "Future returned with an error, sleeping for 5 seconds to let outstanding work finish"
                );
                let mut c = CancelContext::new().with_timeout(Duration::from_secs(5));
                c.cancelled().await;
                Err(e)
            }
            Either::Halt(halt_result) => {
                tracing::warn!(
                    halt_result = format_args!("{:x?}", halt_result),
                    "Halt channel returned while waiting for other future, sleeping for 5 seconds to let outstanding work finish"
                );
                let mut c = CancelContext::new().with_timeout(Duration::from_secs(5));
                let try_again = c.until_cancelled(future).await;

                match try_again {
                    Ok(fut_result) => {
                        halt.already_received = Some(halt_result);
                        if let Err(e) = &fut_result {
                            tracing::warn!(
                                ?e,
                                "Future returned with an error, sleeping for 5 seconds to let outstanding work finish"
                            );
                            let mut c = CancelContext::new().with_timeout(Duration::from_secs(5));
                            c.cancelled().await;
                        }
                        fut_result
                    }
                    Err(_cancel) => match halt_result {
                        Ok(halt_reason) => Err(anyhow::anyhow!("VM halted: {:x?}", halt_reason)),
                        Err(e) => Err(e).context("VM disappeared"),
                    },
                }
            }
        }
    }
}

impl PetriVmInner {
    async fn wait_for_boot_event(&mut self) -> anyhow::Result<FirmwareEvent> {
        self.resources
            .firmware_event_recv
            .recv()
            .await
            .context("Failed to get firmware boot event")
    }

    async fn wait_for_ipmi_sel(&mut self) -> anyhow::Result<IpmiSelEvent> {
        CancelContext::new()
            .with_timeout(Duration::from_secs(30))
            .until_cancelled(self.resources.ipmi_sel_event_recv.recv())
            .await
            .context("timed out waiting for an IPMI SEL host notification")?
            .context("IPMI SEL host notification channel closed")
    }

    async fn wait_for_enlightened_shutdown_ready(
        &mut self,
    ) -> anyhow::Result<Option<mesh::OneshotReceiver<()>>> {
        let Some(send) = self.resources.shutdown_ic_send.as_ref() else {
            return Ok(None);
        };
        let recv = send
            .call(ShutdownRpc::WaitReady, ())
            .await
            .context("waiting for shutdown IC to be ready")?;
        Ok(Some(recv))
    }

    async fn send_enlightened_shutdown(&mut self, kind: ShutdownKind) -> anyhow::Result<()> {
        let send = self
            .resources
            .shutdown_ic_send
            .as_ref()
            .context("shutdown IC not configured")?;
        let shutdown_result = send
            .call(
                ShutdownRpc::Shutdown,
                hyperv_ic_resources::shutdown::ShutdownParams {
                    shutdown_type: match kind {
                        ShutdownKind::Shutdown => {
                            hyperv_ic_resources::shutdown::ShutdownType::PowerOff
                        }
                        ShutdownKind::Reboot => hyperv_ic_resources::shutdown::ShutdownType::Reboot,
                        ShutdownKind::Hibernate => {
                            hyperv_ic_resources::shutdown::ShutdownType::Hibernate
                        }
                    },
                    force: false,
                },
            )
            .await?;

        tracing::info!(?shutdown_result, "Shutdown sent");
        anyhow::ensure!(
            shutdown_result == hyperv_ic_resources::shutdown::ShutdownResult::Ok,
            "Got non-Ok shutdown response"
        );

        Ok(())
    }

    async fn wait_for_kvp(
        &mut self,
    ) -> anyhow::Result<mesh::Sender<hyperv_ic_resources::kvp::KvpRpc>> {
        tracing::info!("Waiting for KVP IC");
        let send = self
            .resources
            .kvp_ic_send
            .as_ref()
            .context("KVP IC not configured")?;
        let (send, _) = send
            .call_failable(hyperv_ic_resources::kvp::KvpConnectRpc::WaitForGuest, ())
            .await
            .context("failed to connect to KVP IC")?;

        Ok(send)
    }

    async fn save_openhcl(
        &self,
        new_openhcl: &ResolvedArtifact,
        flags: OpenHclServicingFlags,
    ) -> anyhow::Result<()> {
        let ged_send = self
            .resources
            .ged_send
            .as_ref()
            .context("openhcl not configured")?;

        let igvm_file = fs_err::File::open(new_openhcl).context("failed to open igvm file")?;
        self.worker
            .save_openhcl(ged_send, flags, igvm_file.into())
            .await
    }

    async fn update_command_line(&mut self, command_line: &str) -> anyhow::Result<()> {
        self.worker.update_command_line(command_line).await
    }

    async fn add_pcie_device(
        &mut self,
        port_name: String,
        resource: vm_resource::Resource<vm_resource::kind::PciDeviceHandleKind>,
    ) -> anyhow::Result<()> {
        self.worker.add_pcie_device(port_name, resource).await
    }

    async fn remove_pcie_device(&mut self, port_name: String) -> anyhow::Result<()> {
        self.worker.remove_pcie_device(port_name).await
    }
    async fn dump_state(&mut self, path: &Path) -> anyhow::Result<()> {
        self.worker.dump_state(path).await
    }
    async fn restore_openhcl(&self) -> anyhow::Result<()> {
        let ged_send = self
            .resources
            .ged_send
            .as_ref()
            .context("openhcl not configured")?;

        self.worker.restore_openhcl(ged_send).await
    }

    async fn set_vtl2_settings(&self, settings: &Vtl2Settings) -> anyhow::Result<()> {
        let ged_send = self
            .resources
            .ged_send
            .as_ref()
            .context("openhcl not configured")?;

        ged_send
            .call_failable(
                get_resources::ged::GuestEmulationRequest::ModifyVtl2Settings,
                prost::Message::encode_to_vec(settings),
            )
            .await?;

        Ok(())
    }

    async fn reset(&mut self) -> anyhow::Result<()> {
        tracing::info!("Resetting VM");
        self.worker.reset().await?;
        // Discard any firmware events from the boot that was abandoned, so
        // that they aren't mistaken for the results of the new boot.
        while self.resources.firmware_event_recv.try_recv().is_ok() {}
        // Guest state is lost on reset, so CIDATA needs to be remounted.
        self.cidata_mounted = false;
        // On linux direct, pipette won't auto-start unless it is the init
        // process. When it isn't, restart it over serial. (When pipette runs
        // as PID 1 via rdinit=/pipette, linux_direct_serial_agent is None, so
        // this block is skipped and pipette restarts automatically on reboot.)
        if let Some(agent) = self.resources.linux_direct_serial_agent.as_mut() {
            agent.reset();

            if self.resources.properties.using_vtl0_pipette {
                self.launch_linux_direct_pipette().await?;
            }
        }
        Ok(())
    }

    async fn wait_for_agent(&mut self, set_high_vtl: bool) -> anyhow::Result<PipetteClient> {
        #[cfg(target_os = "linux")]
        if let Some(guest_cid) = self.resources.properties.vhost_vsock_guest_cid {
            assert!(
                !set_high_vtl,
                "kernel vhost-vsock pipette transport does not support VTL2"
            );
            return self.wait_for_agent_vhost_vsock(guest_cid).await;
        }

        // Use TCP transport if configured (Windows no-vmbus guests).
        if let Some(port) = self.tcp_pipette_port {
            assert!(!set_high_vtl, "TCP pipette transport does not support VTL2");
            return self.wait_for_agent_tcp(port).await;
        }

        let listener = if set_high_vtl {
            self.resources
                .vtl2_pipette_listener
                .as_mut()
                .context("VM is not configured with VTL 2")?
        } else {
            &mut self.resources.pipette_listener
        };

        tracing::info!(set_high_vtl, "listening for pipette connection");
        let client = loop {
            let (conn, _) = listener
                .accept()
                .await
                .context("failed to accept pipette connection")?;
            tracing::info!(set_high_vtl, "handshaking with pipette");
            let socket = PolledSocket::new(&self.resources.driver, conn)?;
            match PipetteClient::new(&self.resources.driver, socket, &self.resources.output_dir)
                .await
            {
                Ok(client) => break client,
                Err(e) => {
                    // During save/restore cycles, stale connections from
                    // previous hvsock relay sessions can accumulate in the
                    // listener backlog. These are already-closed sockets
                    // that fail during the mesh handshake. Drain them and
                    // retry until we get a live connection.
                    tracing::warn!(
                        error = e.as_ref() as &dyn std::error::Error,
                        "pipette connection not live, retrying"
                    );
                }
            }
        };
        tracing::info!(set_high_vtl, "completed pipette handshake");

        // When pipette runs as PID 1 init and a CIDATA agent disk is
        // attached, mount it so test files are available at /cidata.
        // Skip if already mounted (e.g. reconnecting after save/restore
        // where guest state is preserved).
        if !set_high_vtl
            && self.resources.properties.uses_pipette_as_init
            && self.resources.properties.has_agent_disk
            && !self.cidata_mounted
        {
            tracing::info!("mounting CIDATA agent disk via pipette");
            client
                .unix_shell()
                .cmd("mkdir")
                .arg("-p")
                .arg("/cidata")
                .run()
                .await
                .context("failed to create /cidata mount point")?;
            client
                .unix_shell()
                .cmd("mount")
                .arg("LABEL=cidata")
                .arg("/cidata")
                .run()
                .await
                .context("failed to mount CIDATA disk")?;
            self.cidata_mounted = true;
        }

        Ok(client)
    }

    /// Connect to pipette directly through the host's AF_VSOCK namespace.
    #[cfg(target_os = "linux")]
    async fn wait_for_agent_vhost_vsock(
        &mut self,
        guest_cid: u32,
    ) -> anyhow::Result<PipetteClient> {
        tracing::info!(
            guest_cid,
            port = PIPETTE_PORT,
            "connecting to pipette via kernel vhost-vsock"
        );
        let socket = loop {
            let connect = async {
                let socket = VmSocket::new().context("failed to create AF_VSOCK socket")?;
                socket
                    .set_connect_timeout(Duration::from_secs(5))
                    .context("failed to set AF_VSOCK connect timeout")?;
                let mut socket = PolledSocket::new(&self.resources.driver, socket)
                    .context("failed to create polled AF_VSOCK socket")?
                    .convert();
                socket
                    .connect(&VmAddress::vsock(guest_cid, PIPETTE_PORT).into())
                    .await
                    .context("failed to connect to guest AF_VSOCK listener")?;
                Ok::<_, anyhow::Error>(socket)
            };

            match connect.await {
                Ok(socket) => break socket,
                Err(error) => {
                    tracing::trace!(
                        error = error.as_ref() as &dyn std::error::Error,
                        "AF_VSOCK connect failed, guest not ready yet"
                    );
                }
            }

            pal_async::timer::PolledTimer::new(&self.resources.driver)
                .sleep(Duration::from_secs(1))
                .await;
        };
        tracing::info!("AF_VSOCK connected, handshaking with pipette");
        let client = PipetteClient::new(&self.resources.driver, socket, &self.resources.output_dir)
            .await
            .context("pipette AF_VSOCK handshake failed")?;
        tracing::info!("completed pipette AF_VSOCK handshake");
        Ok(client)
    }

    /// Connect to pipette via TCP through consomme port forwarding.
    ///
    /// The guest pipette agent listens on `0.0.0.0:{port}` and consomme
    /// forwards connections from `localhost:{port}` on the host into the
    /// guest. We retry until the guest's network stack and pipette are up.
    async fn wait_for_agent_tcp(&mut self, port: u16) -> anyhow::Result<PipetteClient> {
        tracing::info!(port, "connecting to pipette via TCP");
        let addr = std::net::SocketAddr::from((std::net::Ipv4Addr::LOCALHOST, port));
        let client = loop {
            match PolledSocket::connect_tcp(&self.resources.driver, addr).await {
                Ok(socket) => {
                    socket
                        .get()
                        .set_nodelay(true)
                        .context("failed to set TCP_NODELAY")?;
                    tracing::info!("TCP connected, handshaking with pipette");
                    match PipetteClient::new(
                        &self.resources.driver,
                        socket,
                        &self.resources.output_dir,
                    )
                    .await
                    {
                        Ok(client) => break client,
                        Err(e) => {
                            tracing::warn!(
                                error = e.as_ref() as &dyn std::error::Error,
                                "pipette TCP connection failed, retrying"
                            );
                        }
                    }
                }
                Err(e) => {
                    tracing::trace!(
                        error = &e as &dyn std::error::Error,
                        "TCP connect failed, guest not ready yet"
                    );
                }
            }
            // Wait before retrying — guest network stack may not be up yet.
            pal_async::timer::PolledTimer::new(&self.resources.driver)
                .sleep(Duration::from_secs(1))
                .await;
        };
        tracing::info!("completed pipette TCP handshake");
        Ok(client)
    }

    async fn pause(&self) -> anyhow::Result<()> {
        self.worker.pause().await?;
        Ok(())
    }

    async fn save_state(&self) -> anyhow::Result<Vec<u8>> {
        let state_msg = self.worker.save().await?;
        Ok(mesh::payload::encode(state_msg))
    }

    async fn resume(&self) -> anyhow::Result<()> {
        self.worker.resume().await?;
        Ok(())
    }

    async fn verify_save_restore(&self) -> anyhow::Result<()> {
        for i in 0..2 {
            let result = self.worker.pulse_save_restore().await;
            match result {
                Ok(()) => {}
                Err(RpcError::Channel(err)) => return Err(err.into()),
                Err(RpcError::Call(PulseSaveRestoreError::ResetNotSupported)) => {
                    tracing::warn!("Reset not supported, could not test save + restore.");
                    break;
                }
                Err(RpcError::Call(PulseSaveRestoreError::Other(err))) => {
                    return Err(anyhow::Error::from(err))
                        .context(format!("Save + restore {i} failed."));
                }
            }
        }

        Ok(())
    }

    async fn launch_linux_direct_pipette(&mut self) -> anyhow::Result<()> {
        // Start pipette through serial on linux direct.
        self.resources
            .linux_direct_serial_agent
            .as_mut()
            .unwrap()
            .run_command("mkdir /cidata && mount LABEL=cidata /cidata && sh -c '/cidata/pipette &'")
            .await?;
        Ok(())
    }
}

/// Interface for inspecting OpenVMM
pub struct OpenVmmInspector {
    worker: Arc<Worker>,
}

#[async_trait]
impl PetriVmInspector for OpenVmmInspector {
    async fn inspect(&self, path: &str) -> anyhow::Result<inspect::Node> {
        Ok(self.worker.inspect(path).await)
    }
}

/// Interface to the OpenVMM framebuffer
pub struct OpenVmmFramebufferAccess {
    view: View,
}

#[async_trait]
impl PetriVmFramebufferAccess for OpenVmmFramebufferAccess {
    async fn screenshot(
        &mut self,
        image: &mut Vec<u8>,
    ) -> anyhow::Result<Option<VmScreenshotMeta>> {
        // Our framebuffer uses 4 bytes per pixel, approximating an
        // BGRA image, however it only actually contains BGR data.
        // The fourth byte is effectively noise. We can set the 'alpha'
        // value to 0xFF to make the image opaque.
        const BYTES_PER_PIXEL: usize = 4;
        let (width, height) = self.view.resolution();
        let (widthsize, heightsize) = (width as usize, height as usize);
        let len = widthsize * heightsize * BYTES_PER_PIXEL;

        image.resize(len, 0);
        for (i, line) in (0..height).zip(image.chunks_exact_mut(widthsize * BYTES_PER_PIXEL)) {
            self.view.read_line(i, line);
            for pixel in line.chunks_exact_mut(BYTES_PER_PIXEL) {
                pixel.swap(0, 2);
                pixel[3] = 0xFF;
            }
        }

        Ok(Some(VmScreenshotMeta {
            color: image::ExtendedColorType::Rgba8,
            width,
            height,
        }))
    }
}