// Copyright © 2020, Oracle and/or its affiliates.
//
// Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
//
// Portions Copyright 2017 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE-BSD-3-Clause file.
//
// Copyright © 2019 Intel Corporation
//
// SPDX-License-Identifier: Apache-2.0 AND BSD-3-Clause
//

use std::collections::{BTreeMap, BTreeSet, HashMap};
#[cfg(feature = "fw_cfg")]
use std::ffi;
use std::fs::{File, OpenOptions};
use std::io::{self, Seek, SeekFrom, Write};
use std::num::Wrapping;
use std::ops::Deref;
use std::os::unix::net::UnixStream;
use std::sync::{Arc, Mutex};
#[cfg(not(target_arch = "riscv64"))]
use std::time::Instant;
use std::{any, cmp, result, str, thread};

use anyhow::{Context, anyhow};
#[cfg(any(target_arch = "aarch64", target_arch = "riscv64"))]
use arch::PciSpaceInfo;
#[cfg(target_arch = "x86_64")]
use arch::layout::{KVM_IDENTITY_MAP_START, KVM_TSS_START};
#[cfg(not(target_arch = "x86_64"))]
use arch::uefi;
use arch::{EntryPoint, NumaNode, NumaNodes, get_host_cpu_phys_bits, layout};
use devices::AcpiNotificationFlags;
#[cfg(target_arch = "aarch64")]
use devices::interrupt_controller;
#[cfg(feature = "fw_cfg")]
use devices::legacy::fw_cfg;
#[cfg(feature = "fw_cfg")]
use devices::legacy::fw_cfg::FwCfgItem;
use event_monitor::event;
#[cfg(all(target_arch = "aarch64", feature = "guest_debug"))]
use gdbstub_arch::aarch64::reg::AArch64CoreRegs as CoreRegs;
#[cfg(all(target_arch = "x86_64", feature = "guest_debug"))]
use gdbstub_arch::x86::reg::X86_64CoreRegs as CoreRegs;
#[cfg(target_arch = "aarch64")]
use hypervisor::arch::aarch64::regs::AARCH64_PMU_IRQ;
#[cfg(all(feature = "kvm", target_arch = "x86_64"))]
use hypervisor::arch::x86;
#[cfg(all(feature = "kvm", feature = "sev_snp"))]
use hypervisor::kvm::{
    BOOTLOADER_SIZE, BOOTLOADER_START, KVM_VMSA_PAGE_ADDRESS, KVM_VMSA_PAGE_SIZE,
};
use hypervisor::{HypervisorVmConfig, HypervisorVmError, VmOps};
#[cfg(feature = "igvm")]
use igvm::IgvmFile;
#[cfg(feature = "sev_snp")]
use igvm_defs::SnpPolicy;
use libc::{SIGWINCH, termios};
#[cfg(feature = "tdx")]
use linux_loader::bootparam;
use linux_loader::cmdline::Cmdline;
#[cfg(all(target_arch = "x86_64", feature = "guest_debug"))]
use linux_loader::elf;
use linux_loader::loader::KernelLoader;
#[cfg(target_arch = "x86_64")]
use linux_loader::loader::bzimage::BzImage;
#[cfg(target_arch = "x86_64")]
use linux_loader::loader::elf::PvhBootCapability::PvhEntryPresent;
#[cfg(any(target_arch = "aarch64", target_arch = "riscv64"))]
use linux_loader::loader::pe::Error::InvalidImageMagicNumber;
use linux_loader::{cmdline, loader};
use log::{error, info, warn};
use seccompiler::SeccompAction;
use serde::{Deserialize, Serialize};
use thiserror::Error;
use tracer::trace_scoped;
use vm_device::Bus;
#[cfg(feature = "tdx")]
use vm_memory::GuestMemory;
#[cfg(feature = "tdx")]
use vm_memory::{Address, ByteValued, GuestMemoryRegion, ReadVolatile};
use vm_memory::{Bytes, GuestAddress, GuestAddressSpace, GuestMemoryAtomic};
use vm_migration::protocol::{MemoryRangeTable, Request, Response};
use vm_migration::{
    Migratable, MigratableError, Pausable, Snapshot, Snapshottable, Transportable, snapshot_from_id,
};
use vmm_sys_util::eventfd::EventFd;
use vmm_sys_util::sock_ctrl_msg::ScmSocket;

use crate::config::{MemoryRestoreMode, ValidationError, add_to_config};
use crate::console_devices::{ConsoleDeviceError, ConsoleInfo};
#[cfg(all(target_arch = "x86_64", feature = "guest_debug"))]
use crate::coredump::{
    CpuElf64Writable, DumpState, Elf64Writable, GuestDebuggable, GuestDebuggableError, NoteDescType,
};
use crate::device_manager::{DeviceManager, DeviceManagerError};
use crate::device_tree::DeviceTree;
#[cfg(feature = "guest_debug")]
use crate::gdb::{Debuggable, DebuggableError, GdbRequestPayload, GdbResponsePayload};
#[cfg(feature = "igvm")]
use crate::igvm::{igvm_loader, parse_igvm};
use crate::landlock::LandlockError;
#[cfg(feature = "tdx")]
use crate::memory_manager;
use crate::memory_manager::{
    Error as MemoryManagerError, MemoryManager, MemoryManagerSnapshotData,
};
#[cfg(all(target_arch = "x86_64", feature = "guest_debug"))]
use crate::migration::url_to_file;
use crate::migration::{SNAPSHOT_CONFIG_FILE, SNAPSHOT_STATE_FILE, get_vm_snapshot, url_to_path};
#[cfg(all(feature = "kvm", feature = "sev_snp", feature = "fw_cfg"))]
use crate::sev::MeasuredBootInfo;
#[cfg(feature = "fw_cfg")]
use crate::vm_config::FwCfgConfig;
use crate::vm_config::{
    DeviceConfig, DiskConfig, FsConfig, GenericVhostUserConfig, HotplugMethod, NetConfig,
    NumaConfig, PayloadConfig, PmemConfig, UserDeviceConfig, VdpaConfig, VmConfig, VsockConfig,
};
use crate::{
    CPU_MANAGER_SNAPSHOT_ID, DEVICE_MANAGER_SNAPSHOT_ID, GuestMemoryMmap,
    MEMORY_MANAGER_SNAPSHOT_ID, PciDeviceInfo, acpi, cpu,
};

/// Errors associated with VM management
#[derive(Debug, Error)]
pub enum Error {
    #[error("Cannot open kernel file")]
    KernelFile(#[source] io::Error),

    #[error("Cannot open initramfs file")]
    InitramfsFile(#[source] io::Error),

    #[error("Cannot load the kernel into memory")]
    KernelLoad(#[source] loader::Error),

    #[cfg(target_arch = "aarch64")]
    #[error("Cannot load the UEFI binary in memory")]
    UefiLoad(#[source] arch::aarch64::uefi::Error),

    #[cfg(target_arch = "riscv64")]
    #[error("Cannot load the UEFI binary in memory")]
    UefiLoad(#[source] arch::riscv64::uefi::Error),

    #[error("Cannot load the initramfs into memory")]
    InitramfsLoad(#[source] io::Error),

    #[error("Cannot determine initramfs load address")]
    InitramfsAddress(#[source] arch::Error),

    #[error("Cannot read initramfs into guest memory")]
    InitramfsRead(#[source] vm_memory::GuestMemoryError),

    #[error("Cannot load the kernel command line in memory")]
    LoadCmdLine(#[source] loader::Error),

    #[error("Failed to apply landlock config during vm_create")]
    ApplyLandlock(#[source] LandlockError),

    #[error("Cannot modify the kernel command line")]
    CmdLineInsertStr(#[source] cmdline::Error),

    #[error("Cannot create the kernel command line")]
    CmdLineCreate(#[source] cmdline::Error),

    #[error("Cannot configure system")]
    ConfigureSystem(#[source] arch::Error),

    #[cfg(target_arch = "aarch64")]
    #[error("Cannot enable interrupt controller")]
    EnableInterruptController(#[source] interrupt_controller::Error),

    #[error("Error from device manager")]
    DeviceManager(#[source] DeviceManagerError),

    #[error("Error initializing VM")]
    InitializeVm(#[source] hypervisor::HypervisorVmError),

    #[error("No device with id {0:?} to remove")]
    NoDeviceToRemove(String),

    #[error("Cannot spawn a signal handler thread")]
    SignalHandlerSpawn(#[source] io::Error),

    #[error("Failed to join on threads: {0:?}")]
    ThreadCleanup(Box<dyn any::Any + Send>),

    #[error("VM config is missing")]
    VmMissingConfig,

    #[error("VM is not created")]
    VmNotCreated,

    #[error("VM is already created")]
    VmAlreadyCreated,

    #[error("VM is not running")]
    VmNotRunning,

    #[error("VM is currently migrating and can't be modified")]
    VmMigrating,

    #[error("VM on-demand memory restore is still in progress")]
    VmRestoring,

    #[error("Cannot clone EventFd")]
    EventFdClone(#[source] io::Error),

    #[error("invalid VM state transition: {0:?} to {1:?}")]
    InvalidStateTransition(VmState, VmState),

    #[error("Error from CPU manager")]
    CpuManager(#[source] cpu::Error),

    #[error("Cannot pause devices")]
    PauseDevices(#[source] MigratableError),

    #[error("Cannot resume devices")]
    ResumeDevices(#[source] MigratableError),

    #[error("Cannot pause CPUs")]
    PauseCpus(#[source] MigratableError),

    #[error("Cannot resume cpus")]
    ResumeCpus(#[source] MigratableError),

    #[error("Cannot pause VM")]
    Pause(#[source] MigratableError),

    #[error("Cannot resume VM")]
    Resume(#[source] MigratableError),

    #[error("Memory manager error")]
    MemoryManager(#[source] MemoryManagerError),

    #[error("Eventfd write error")]
    EventfdError(#[source] io::Error),

    #[error("Cannot snapshot VM")]
    Snapshot(#[source] MigratableError),

    #[error("Cannot restore VM")]
    Restore(#[source] MigratableError),

    #[error("Cannot send VM snapshot")]
    SnapshotSend(#[source] MigratableError),

    #[error("Invalid restore source URL")]
    InvalidRestoreSourceUrl,

    #[error("Failed to validate config")]
    ConfigValidation(#[source] ValidationError),

    #[error("Too many virtio-vsock devices")]
    TooManyVsockDevices,

    #[error("Failed serializing into JSON")]
    SerializeJson(#[source] serde_json::Error),

    #[error("Invalid NUMA configuration")]
    InvalidNumaConfig,

    #[error("Cannot create seccomp filter")]
    CreateSeccompFilter(#[source] seccompiler::Error),

    #[error("Cannot apply seccomp filter")]
    ApplySeccompFilter(#[source] seccompiler::Error),

    #[error("Failed resizing a memory zone")]
    ResizeZone,

    #[error("Failed resizing a disk image")]
    ResizeDisk,

    #[error("Cannot activate virtio devices")]
    ActivateVirtioDevices(#[source] DeviceManagerError),

    #[error("Error triggering power button")]
    PowerButton(#[source] DeviceManagerError),

    #[error("Kernel lacks PVH header")]
    KernelMissingPvhHeader,

    #[error("Failed to allocate firmware RAM")]
    AllocateFirmwareMemory(#[source] MemoryManagerError),

    #[error("Error manipulating firmware file")]
    FirmwareFile(#[source] io::Error),

    #[error("Firmware too big")]
    FirmwareTooLarge,

    #[error("Failed to copy firmware to memory")]
    FirmwareLoad(#[source] vm_memory::GuestMemoryError),

    #[cfg(feature = "sev_snp")]
    #[error("Error enabling SEV-SNP VM")]
    InitializeSevSnpVm(#[source] hypervisor::HypervisorVmError),

    #[cfg(feature = "tdx")]
    #[error("Error performing I/O on TDX firmware file")]
    LoadTdvf(#[source] io::Error),

    #[cfg(feature = "tdx")]
    #[error("Error performing I/O on the TDX payload file")]
    LoadPayload(#[source] io::Error),

    #[cfg(feature = "tdx")]
    #[error("Error parsing TDVF")]
    ParseTdvf(#[source] arch::x86_64::tdx::TdvfError),

    #[cfg(feature = "tdx")]
    #[error("Error populating TDX HOB")]
    PopulateHob(#[source] arch::x86_64::tdx::TdvfError),

    #[cfg(feature = "tdx")]
    #[error("Error allocating TDVF memory")]
    AllocatingTdvfMemory(#[source] memory_manager::Error),

    #[cfg(feature = "tdx")]
    #[error("Error enabling TDX VM")]
    InitializeTdxVm(#[source] hypervisor::HypervisorVmError),

    #[cfg(feature = "tdx")]
    #[error("Error enabling TDX memory region")]
    InitializeTdxMemoryRegion(#[source] hypervisor::HypervisorVmError),

    #[cfg(feature = "tdx")]
    #[error("Error finalizing TDX VM")]
    FinalizeTdx(#[source] hypervisor::HypervisorVmError),

    #[cfg(feature = "tdx")]
    #[error("TDX firmware missing")]
    TdxFirmwareMissing,

    #[cfg(feature = "tdx")]
    #[error("Invalid TDX payload type")]
    InvalidPayloadType,

    #[cfg(feature = "guest_debug")]
    #[error("Error debugging VM")]
    Debug(#[source] DebuggableError),

    #[error("Error spawning kernel loading thread")]
    KernelLoadThreadSpawn(#[source] io::Error),

    #[error("Error joining kernel loading thread")]
    KernelLoadThreadJoin(Box<dyn any::Any + Send>),

    #[cfg(all(target_arch = "x86_64", feature = "guest_debug"))]
    #[error("Error coredumping VM")]
    Coredump(#[source] GuestDebuggableError),

    #[cfg(feature = "igvm")]
    #[error("Cannot load the igvm into memory")]
    IgvmLoad(#[source] igvm_loader::Error),

    #[error("Error injecting NMI")]
    ErrorNmi(#[source] cpu::Error),

    #[error("Error resuming the VM")]
    ResumeVm(#[source] hypervisor::HypervisorVmError),

    #[error("Error creating console devices")]
    CreateConsoleDevices(#[source] ConsoleDeviceError),

    #[error("Error locking disk images: Another instance likely holds a lock")]
    LockingError(#[source] DeviceManagerError),

    #[cfg(feature = "fw_cfg")]
    #[error("Fw Cfg missing kernel")]
    MissingFwCfgKernelFile(#[source] io::Error),

    #[cfg(feature = "fw_cfg")]
    #[error("Fw Cfg missing initramfs")]
    MissingFwCfgInitramfs(#[source] io::Error),

    #[cfg(feature = "fw_cfg")]
    #[error("Fw Cfg missing kernel cmdline")]
    MissingFwCfgCmdline,

    #[cfg(feature = "fw_cfg")]
    #[error("Error creating e820 map")]
    CreatingE820Map(#[source] io::Error),

    #[cfg(feature = "fw_cfg")]
    #[error("Error creating acpi tables")]
    CreatingAcpiTables(#[source] io::Error),

    #[cfg(feature = "fw_cfg")]
    #[error("Error adding fw_cfg item")]
    AddingFwCfgItem(#[source] io::Error),

    #[cfg(feature = "fw_cfg")]
    #[error("Error populating fw_cfg")]
    ErrorPopulatingFwCfg(#[source] io::Error),

    #[cfg(feature = "fw_cfg")]
    #[error("Error using fw_cfg while disabled")]
    FwCfgDisabled,
}
pub type Result<T> = result::Result<T, Error>;

#[derive(Clone, Copy, Debug, Deserialize, Serialize, PartialEq, Eq)]
pub enum VmState {
    Created,
    Running,
    Shutdown,
    Paused,
    BreakPoint,
}

impl VmState {
    fn valid_transition(self, new_state: VmState) -> Result<()> {
        match self {
            VmState::Created => match new_state {
                VmState::Created => Err(Error::InvalidStateTransition(self, new_state)),
                VmState::Running | VmState::Paused | VmState::BreakPoint | VmState::Shutdown => {
                    Ok(())
                }
            },

            VmState::Running => match new_state {
                VmState::Created | VmState::Running => {
                    Err(Error::InvalidStateTransition(self, new_state))
                }
                VmState::Paused | VmState::Shutdown | VmState::BreakPoint => Ok(()),
            },

            VmState::Shutdown => match new_state {
                VmState::Paused | VmState::Created | VmState::Shutdown | VmState::BreakPoint => {
                    Err(Error::InvalidStateTransition(self, new_state))
                }
                VmState::Running => Ok(()),
            },

            VmState::Paused => match new_state {
                VmState::Created | VmState::Paused | VmState::BreakPoint => {
                    Err(Error::InvalidStateTransition(self, new_state))
                }
                VmState::Running | VmState::Shutdown => Ok(()),
            },
            VmState::BreakPoint => match new_state {
                VmState::Created | VmState::Running => Ok(()),
                _ => Err(Error::InvalidStateTransition(self, new_state)),
            },
        }
    }
}

struct VmOpsHandler {
    memory: GuestMemoryAtomic<GuestMemoryMmap>,
    #[cfg(target_arch = "x86_64")]
    io_bus: Arc<Bus>,
    mmio_bus: Arc<Bus>,
}

impl VmOps for VmOpsHandler {
    fn guest_mem_write(&self, gpa: u64, buf: &[u8]) -> result::Result<usize, HypervisorVmError> {
        self.memory
            .memory()
            .write(buf, GuestAddress(gpa))
            .map_err(|e| HypervisorVmError::GuestMemWrite(e.into()))
    }

    fn guest_mem_read(&self, gpa: u64, buf: &mut [u8]) -> result::Result<usize, HypervisorVmError> {
        self.memory
            .memory()
            .read(buf, GuestAddress(gpa))
            .map_err(|e| HypervisorVmError::GuestMemRead(e.into()))
    }

    fn mmio_read(&self, gpa: u64, data: &mut [u8]) -> result::Result<(), HypervisorVmError> {
        if let Err(vm_device::BusError::MissingAddressRange) = self.mmio_bus.read(gpa, data) {
            info!("Guest MMIO read from unregistered address 0x{gpa:x}");
            data.fill(0xff); // 0xff is sentinel value for invalid reads
        }
        Ok(())
    }

    fn mmio_write(&self, gpa: u64, data: &[u8]) -> result::Result<(), HypervisorVmError> {
        match self.mmio_bus.write(gpa, data) {
            Err(vm_device::BusError::MissingAddressRange) => {
                info!("Guest MMIO write to unregistered address 0x{gpa:x}");
            }
            Ok(Some(barrier)) => {
                info!("Waiting for barrier");
                barrier.wait();
                info!("Barrier released");
            }
            _ => {}
        }
        Ok(())
    }

    #[cfg(target_arch = "x86_64")]
    fn pio_read(&self, port: u64, data: &mut [u8]) -> result::Result<(), HypervisorVmError> {
        if let Err(vm_device::BusError::MissingAddressRange) = self.io_bus.read(port, data) {
            info!("Guest PIO read from unregistered address 0x{port:x}");
            data.fill(0xff); // 0xff is sentinel value for invalid reads
        }
        Ok(())
    }

    #[cfg(target_arch = "x86_64")]
    fn pio_write(&self, port: u64, data: &[u8]) -> result::Result<(), HypervisorVmError> {
        match self.io_bus.write(port, data) {
            Err(vm_device::BusError::MissingAddressRange) => {
                info!("Guest PIO write to unregistered address 0x{port:x}");
            }
            Ok(Some(barrier)) => {
                info!("Waiting for barrier");
                barrier.wait();
                info!("Barrier released");
            }
            _ => {}
        }
        Ok(())
    }
}

pub fn physical_bits(hypervisor: &dyn hypervisor::Hypervisor, max_phys_bits: u8) -> u8 {
    let host_phys_bits = get_host_cpu_phys_bits(hypervisor);

    cmp::min(host_phys_bits, max_phys_bits)
}

/// Guest clock baseline captured for snapshot/restore, plus how it must be
/// re-established on the next resume. `mode` is runtime-only (`SnapshotRestore`
/// after restore/migration-receive, `SameHostResume` for a live same-host capture)
/// and is not serialized: it is re-derived from "was this VM built from a snapshot".
#[derive(Clone, Copy)]
struct SavedClock {
    mode: hypervisor::ClockRestoreMode,
    state: hypervisor::ClockState,
}

pub struct Vm {
    #[cfg(feature = "tdx")]
    kernel: Option<File>,
    initramfs: Option<File>,
    threads: Vec<thread::JoinHandle<()>>,
    device_manager: Arc<Mutex<DeviceManager>>,
    config: Arc<Mutex<VmConfig>>,
    state: VmState,
    cpu_manager: Arc<Mutex<cpu::CpuManager>>,
    memory_manager: Arc<Mutex<MemoryManager>>,
    #[cfg_attr(any(not(feature = "kvm"), target_arch = "aarch64"), allow(dead_code))]
    // The hypervisor abstracted virtual machine.
    vm: Arc<dyn hypervisor::Vm>,
    saved_clock: Option<SavedClock>,
    #[cfg(not(target_arch = "riscv64"))]
    numa_nodes: NumaNodes,
    #[cfg_attr(any(not(feature = "kvm"), target_arch = "aarch64"), allow(dead_code))]
    #[cfg(not(target_arch = "riscv64"))]
    hypervisor: Arc<dyn hypervisor::Hypervisor>,
    stop_on_boot: bool,
    load_payload_handle: Option<thread::JoinHandle<Result<EntryPoint>>>,
}

impl Vm {
    pub const HANDLED_SIGNALS: [i32; 1] = [SIGWINCH];

    #[cfg(feature = "sev_snp")]
    pub fn get_default_sev_snp_guest_policy() -> SnpPolicy {
        SnpPolicy::new()
            .with_abi_minor(0)
            .with_abi_major(0)
            // SMT permitted: allows the guest to run on an SMT-enabled host.
            // This is the permissive default; future work can expose this as a
            // configurable platform option.
            .with_smt(1)
            .with_reserved_must_be_one(1)
            .with_migrate_ma(0)
    }

    #[expect(clippy::needless_pass_by_value)]
    #[expect(clippy::too_many_arguments)]
    pub fn new_from_memory_manager(
        config: Arc<Mutex<VmConfig>>,
        memory_manager: Arc<Mutex<MemoryManager>>,
        vm: Arc<dyn hypervisor::Vm>,
        exit_evt: EventFd,
        reset_evt: EventFd,
        guest_exit_evt: EventFd,
        #[cfg(feature = "guest_debug")] vm_debug_evt: EventFd,
        seccomp_action: &SeccompAction,
        hypervisor: Arc<dyn hypervisor::Hypervisor>,
        activate_evt: EventFd,
        #[cfg(not(target_arch = "riscv64"))] timestamp: Instant,
        console_info: Option<ConsoleInfo>,
        console_resize_pipe: Option<Arc<File>>,
        original_termios: Arc<Mutex<Option<termios>>>,
        snapshot: Option<&Snapshot>,
        #[cfg(feature = "igvm")] igvm_file: Option<IgvmFile>,
    ) -> Result<Self> {
        trace_scoped!("Vm::new_from_memory_manager");

        let boot_id_list = config
            .lock()
            .unwrap()
            .validate()
            .map_err(Error::ConfigValidation)?;

        info!("Booting VM from config: {config:?}");

        // Create NUMA nodes based on NumaConfig.
        let numa_nodes =
            Self::create_numa_nodes(config.lock().unwrap().numa.as_deref(), &memory_manager)?;

        // Determine if VIRTIO_F_ACCESS_PLATFORM should be forced (e.g. for TDX/SEV-SNP)
        let force_access_platform = Self::should_force_access_platform(&config);

        let stop_on_boot = Self::should_stop_on_boot(&config);

        let memory = memory_manager.lock().unwrap().guest_memory();
        let io_bus = Arc::new(Bus::new());
        let mmio_bus = Arc::new(Bus::new());

        let vm_ops: Arc<dyn VmOps> = Arc::new(VmOpsHandler {
            memory,
            #[cfg(target_arch = "x86_64")]
            io_bus: io_bus.clone(),
            mmio_bus: mmio_bus.clone(),
        });

        // Create CPU manager
        let cpu_manager = Self::create_cpu_manager(
            &config,
            vm.clone(),
            exit_evt.try_clone().map_err(Error::EventFdClone)?,
            reset_evt.try_clone().map_err(Error::EventFdClone)?,
            #[cfg(feature = "guest_debug")]
            vm_debug_evt,
            &hypervisor,
            seccomp_action.clone(),
            vm_ops,
            &numa_nodes,
        )?;

        // Perform hypervisor-specific TDX initialization if enabled
        #[cfg(feature = "tdx")]
        Self::init_tdx_if_enabled(&config, &vm, &cpu_manager)?;

        // Create device manager
        let device_manager = Self::create_device_manager(
            io_bus,
            mmio_bus,
            vm.clone(),
            config.clone(),
            memory_manager.clone(),
            cpu_manager.clone(),
            exit_evt.try_clone().map_err(Error::EventFdClone)?,
            reset_evt,
            guest_exit_evt,
            seccomp_action.clone(),
            numa_nodes.clone(),
            &activate_evt,
            force_access_platform,
            boot_id_list,
            #[cfg(not(target_arch = "riscv64"))]
            timestamp,
            snapshot,
        )?;

        // Perform hypervisor-specific initialization
        let load_payload_handle = Self::hypervisor_specific_init(
            &vm,
            &memory_manager,
            &cpu_manager,
            &device_manager,
            &config,
            &hypervisor,
            console_info.as_ref(),
            console_resize_pipe.as_ref(),
            &original_termios,
            snapshot,
            #[cfg(feature = "igvm")]
            igvm_file,
        )?;

        // Load kernel and initramfs files
        #[cfg(feature = "tdx")]
        let kernel = config
            .lock()
            .unwrap()
            .payload
            .as_ref()
            .map(|p| p.kernel.as_ref().map(File::open))
            .unwrap_or_default()
            .transpose()
            .map_err(Error::KernelFile)?;

        let initramfs = config
            .lock()
            .unwrap()
            .payload
            .as_ref()
            .map(|p| p.initramfs.as_ref().map(File::open))
            .unwrap_or_default()
            .transpose()
            .map_err(Error::InitramfsFile)?;

        let saved_clock = if let Some(snapshot) = snapshot.as_ref() {
            let vm_snapshot = get_vm_snapshot(snapshot).map_err(Error::Restore)?;
            // Restored or migrated in: the guest clock must catch up to wall time
            // on resume (the counter was reset to the saved value).
            vm_snapshot.clock.map(|state| SavedClock {
                mode: hypervisor::ClockRestoreMode::SnapshotRestore,
                state,
            })
        } else {
            None
        };

        let state = if snapshot.is_some() {
            VmState::Paused
        } else {
            VmState::Created
        };

        Ok(Vm {
            #[cfg(feature = "tdx")]
            kernel,
            initramfs,
            device_manager,
            config,
            threads: Vec::with_capacity(1),
            state,
            cpu_manager,
            memory_manager,
            vm,
            saved_clock,
            #[cfg(not(target_arch = "riscv64"))]
            numa_nodes,
            #[cfg(not(target_arch = "riscv64"))]
            hypervisor,
            stop_on_boot,
            load_payload_handle,
        })
    }

    /// Determine if VIRTIO_F_ACCESS_PLATFORM should be forced based on
    /// confidential computing features.
    fn should_force_access_platform(_config: &Arc<Mutex<VmConfig>>) -> bool {
        #[cfg(feature = "tdx")]
        if _config.lock().unwrap().is_tdx_enabled() {
            return true;
        }
        #[cfg(feature = "sev_snp")]
        if _config.lock().unwrap().is_sev_snp_enabled() {
            return true;
        }
        false
    }

    /// Determine if VM should stop on boot (for debugging).
    fn should_stop_on_boot(config: &Arc<Mutex<VmConfig>>) -> bool {
        #[cfg(feature = "guest_debug")]
        {
            config.lock().unwrap().gdb
        }
        #[cfg(not(feature = "guest_debug"))]
        {
            let _ = config;
            false
        }
    }

    /// Create and configure the CPU manager.
    #[expect(clippy::too_many_arguments)]
    fn create_cpu_manager(
        config: &Arc<Mutex<VmConfig>>,
        vm: Arc<dyn hypervisor::Vm>,
        exit_evt: EventFd,
        reset_evt: EventFd,
        #[cfg(feature = "guest_debug")] vm_debug_evt: EventFd,
        hypervisor: &Arc<dyn hypervisor::Hypervisor>,
        seccomp_action: SeccompAction,
        vm_ops: Arc<dyn VmOps>,
        numa_nodes: &NumaNodes,
    ) -> Result<Arc<Mutex<cpu::CpuManager>>> {
        #[cfg(feature = "tdx")]
        let tdx_enabled = config.lock().unwrap().is_tdx_enabled();
        #[cfg(feature = "sev_snp")]
        let sev_snp_enabled = config.lock().unwrap().is_sev_snp_enabled();
        #[cfg(feature = "igvm")]
        let igvm_enabled = config
            .lock()
            .unwrap()
            .payload
            .as_ref()
            .and_then(|p| p.igvm.as_ref())
            .is_some();

        let cpus_config = config.lock().unwrap().cpus.clone();
        let cpu_manager = cpu::CpuManager::new(
            &cpus_config,
            vm,
            exit_evt,
            reset_evt,
            #[cfg(feature = "guest_debug")]
            vm_debug_evt,
            hypervisor.clone(),
            seccomp_action,
            vm_ops,
            #[cfg(feature = "tdx")]
            tdx_enabled,
            numa_nodes,
            #[cfg(feature = "sev_snp")]
            sev_snp_enabled,
            #[cfg(feature = "igvm")]
            igvm_enabled,
        )
        .map_err(Error::CpuManager)?;

        #[cfg(target_arch = "x86_64")]
        cpu_manager
            .lock()
            .unwrap()
            .populate_cpuid(
                hypervisor.as_ref(),
                #[cfg(feature = "tdx")]
                tdx_enabled,
            )
            .map_err(Error::CpuManager)?;

        Ok(cpu_manager)
    }

    /// Initialize TDX if enabled.
    #[cfg(feature = "tdx")]
    fn init_tdx_if_enabled(
        config: &Arc<Mutex<VmConfig>>,
        vm: &Arc<dyn hypervisor::Vm>,
        cpu_manager: &Arc<Mutex<cpu::CpuManager>>,
    ) -> Result<()> {
        if config.lock().unwrap().is_tdx_enabled() {
            let cpuid = cpu_manager.lock().unwrap().common_cpuid();
            let max_vcpus = cpu_manager.lock().unwrap().max_vcpus();
            vm.tdx_init(&cpuid, max_vcpus)
                .map_err(Error::InitializeTdxVm)?;
        }
        Ok(())
    }

    /// Create and configure the device manager.
    #[expect(clippy::too_many_arguments)]
    fn create_device_manager(
        io_bus: Arc<Bus>,
        mmio_bus: Arc<Bus>,
        vm: Arc<dyn hypervisor::Vm>,
        config: Arc<Mutex<VmConfig>>,
        memory_manager: Arc<Mutex<MemoryManager>>,
        cpu_manager: Arc<Mutex<cpu::CpuManager>>,
        exit_evt: EventFd,
        reset_evt: EventFd,
        guest_exit_evt: EventFd,
        seccomp_action: SeccompAction,
        numa_nodes: NumaNodes,
        activate_evt: &EventFd,
        force_access_platform: bool,
        boot_id_list: BTreeSet<String>,
        #[cfg(not(target_arch = "riscv64"))] timestamp: Instant,
        snapshot: Option<&Snapshot>,
    ) -> Result<Arc<Mutex<DeviceManager>>> {
        #[cfg(feature = "tdx")]
        let dynamic = !config.lock().unwrap().is_tdx_enabled();
        #[cfg(not(feature = "tdx"))]
        let dynamic = true;

        DeviceManager::new(
            io_bus,
            mmio_bus,
            vm,
            config,
            memory_manager,
            cpu_manager,
            exit_evt,
            reset_evt,
            guest_exit_evt,
            seccomp_action,
            numa_nodes,
            activate_evt,
            force_access_platform,
            boot_id_list,
            #[cfg(not(target_arch = "riscv64"))]
            timestamp,
            snapshot_from_id(snapshot, DEVICE_MANAGER_SNAPSHOT_ID),
            dynamic,
        )
        .map_err(Error::DeviceManager)
    }

    /// Perform hypervisor-specific initialization.
    ///
    /// This handles the different initialization sequences required by:
    /// - KVM (x86_64, aarch64, riscv64)
    /// - MSHV (x86_64, aarch64)
    /// - SEV-SNP (MSHV with confidential computing)
    #[expect(clippy::too_many_arguments)]
    fn hypervisor_specific_init(
        vm: &Arc<dyn hypervisor::Vm>,
        memory_manager: &Arc<Mutex<MemoryManager>>,
        cpu_manager: &Arc<Mutex<cpu::CpuManager>>,
        device_manager: &Arc<Mutex<DeviceManager>>,
        config: &Arc<Mutex<VmConfig>>,
        hypervisor: &Arc<dyn hypervisor::Hypervisor>,
        console_info: Option<&ConsoleInfo>,
        console_resize_pipe: Option<&Arc<File>>,
        original_termios: &Arc<Mutex<Option<termios>>>,
        snapshot: Option<&Snapshot>,
        #[cfg(feature = "igvm")] igvm_file: Option<IgvmFile>,
    ) -> Result<Option<thread::JoinHandle<Result<EntryPoint>>>> {
        #[cfg(feature = "mshv")]
        let is_mshv = matches!(
            hypervisor.hypervisor_type(),
            hypervisor::HypervisorType::Mshv
        );
        #[cfg(feature = "kvm")]
        let is_kvm = matches!(
            hypervisor.hypervisor_type(),
            hypervisor::HypervisorType::Kvm
        );

        #[cfg(feature = "sev_snp")]
        let sev_snp_enabled = config.lock().unwrap().is_sev_snp_enabled();

        // MSHV-specific initialization (non-aarch64)
        #[cfg(all(feature = "mshv", not(target_arch = "aarch64")))]
        if is_mshv {
            vm.init().map_err(Error::InitializeVm)?;
        }

        // SEV-SNP specific initialization
        #[cfg(feature = "sev_snp")]
        if sev_snp_enabled {
            return Self::init_sev_snp(
                vm,
                memory_manager,
                cpu_manager,
                device_manager,
                config,
                console_info,
                console_resize_pipe,
                original_termios,
                snapshot,
                igvm_file,
            );
        }

        // MSHV initialization (create interrupt controller and devices)
        #[cfg(feature = "mshv")]
        if is_mshv {
            Self::init_mshv(
                vm,
                device_manager,
                console_info,
                console_resize_pipe,
                original_termios,
                snapshot,
            )?;
        }

        // Allocate address space for non-SEV-SNP guests
        memory_manager
            .lock()
            .unwrap()
            .allocate_address_space()
            .map_err(Error::MemoryManager)?;

        // Load payload asynchronously
        let load_payload_handle = if snapshot.is_none() {
            Self::load_payload_async(
                memory_manager,
                config,
                #[cfg(feature = "igvm")]
                cpu_manager,
                #[cfg(feature = "igvm")]
                igvm_file,
            )?
        } else {
            None
        };

        // Create boot vCPUs
        cpu_manager
            .lock()
            .unwrap()
            .create_boot_vcpus(snapshot_from_id(snapshot, CPU_MANAGER_SNAPSHOT_ID))
            .map_err(Error::CpuManager)?;

        // KVM-specific initialization
        #[cfg(feature = "kvm")]
        if is_kvm {
            Self::init_kvm(
                vm,
                device_manager,
                console_info.cloned(),
                console_resize_pipe.cloned(),
                original_termios.clone(),
                snapshot,
            )?;
        }

        // Create fw_cfg device if configured
        #[cfg(feature = "fw_cfg")]
        Self::create_fw_cfg_if_enabled(config, device_manager)?;

        Ok(load_payload_handle)
    }

    /// Initialize SEV-SNP specific components.
    #[cfg(feature = "sev_snp")]
    #[expect(clippy::too_many_arguments)]
    fn init_sev_snp(
        vm: &Arc<dyn hypervisor::Vm>,
        memory_manager: &Arc<Mutex<MemoryManager>>,
        cpu_manager: &Arc<Mutex<cpu::CpuManager>>,
        device_manager: &Arc<Mutex<DeviceManager>>,
        config: &Arc<Mutex<VmConfig>>,
        console_info: Option<&ConsoleInfo>,
        console_resize_pipe: Option<&Arc<File>>,
        original_termios: &Arc<Mutex<Option<termios>>>,
        snapshot: Option<&Snapshot>,
        igvm_file: Option<IgvmFile>,
    ) -> Result<Option<thread::JoinHandle<Result<EntryPoint>>>> {
        // Create boot vCPUs before SEV-SNP initialization
        cpu_manager
            .lock()
            .unwrap()
            .create_boot_vcpus(snapshot_from_id(snapshot, CPU_MANAGER_SNAPSHOT_ID))
            .map_err(Error::CpuManager)?;

        // Extract guest policy from IGVM if available, otherwise use default.
        let guest_policy = igvm_file
            .as_ref()
            .and_then(igvm_loader::extract_guest_policy)
            .unwrap_or_else(Self::get_default_sev_snp_guest_policy);

        vm.sev_snp_init(guest_policy)
            .map_err(Error::InitializeSevSnpVm)?;

        // Load payload for SEV-SNP (IGVM parser needs cpu_manager for cpuid)
        let load_payload_handle = if snapshot.is_none() {
            Self::load_payload_async(memory_manager, config, cpu_manager, igvm_file)?
        } else {
            None
        };

        // Create interrupt controller and devices for MSHV
        let dm_snapshot = snapshot_from_id(snapshot, DEVICE_MANAGER_SNAPSHOT_ID);
        let ic = device_manager
            .lock()
            .unwrap()
            .create_interrupt_controller(dm_snapshot)
            .map_err(Error::DeviceManager)?;

        #[cfg(target_arch = "aarch64")]
        vm.init().map_err(Error::InitializeVm)?;

        device_manager
            .lock()
            .unwrap()
            .create_devices(
                console_info.cloned(),
                console_resize_pipe.cloned(),
                original_termios.clone(),
                ic,
                dm_snapshot,
            )
            .map_err(Error::DeviceManager)?;

        #[cfg(feature = "fw_cfg")]
        Self::create_fw_cfg_if_enabled(config, device_manager)?;

        Ok(load_payload_handle)
    }

    /// Initialize MSHV-specific components.
    #[cfg(feature = "mshv")]
    fn init_mshv(
        _vm: &Arc<dyn hypervisor::Vm>,
        device_manager: &Arc<Mutex<DeviceManager>>,
        console_info: Option<&ConsoleInfo>,
        console_resize_pipe: Option<&Arc<File>>,
        original_termios: &Arc<Mutex<Option<termios>>>,
        snapshot: Option<&Snapshot>,
    ) -> Result<()> {
        let dm_snapshot = snapshot_from_id(snapshot, DEVICE_MANAGER_SNAPSHOT_ID);
        let ic = device_manager
            .lock()
            .unwrap()
            .create_interrupt_controller(dm_snapshot)
            .map_err(Error::DeviceManager)?;

        #[cfg(target_arch = "aarch64")]
        _vm.init().map_err(Error::InitializeVm)?;

        device_manager
            .lock()
            .unwrap()
            .create_devices(
                console_info.cloned(),
                console_resize_pipe.cloned(),
                original_termios.clone(),
                ic,
                dm_snapshot,
            )
            .map_err(Error::DeviceManager)?;

        Ok(())
    }

    /// Initialize KVM-specific components.
    #[cfg(feature = "kvm")]
    fn init_kvm(
        vm: &Arc<dyn hypervisor::Vm>,
        device_manager: &Arc<Mutex<DeviceManager>>,
        console_info: Option<ConsoleInfo>,
        console_resize_pipe: Option<Arc<File>>,
        original_termios: Arc<Mutex<Option<termios>>>,
        snapshot: Option<&Snapshot>,
    ) -> Result<()> {
        // For KVM, create interrupt controller after boot vcpus
        // because GIC state is restored from snapshot during vcpu creation
        let dm_snapshot = snapshot_from_id(snapshot, DEVICE_MANAGER_SNAPSHOT_ID);
        let ic = device_manager
            .lock()
            .unwrap()
            .create_interrupt_controller(dm_snapshot)
            .map_err(Error::DeviceManager)?;

        vm.init().map_err(Error::InitializeVm)?;

        device_manager
            .lock()
            .unwrap()
            .create_devices(
                console_info,
                console_resize_pipe,
                original_termios,
                ic,
                dm_snapshot,
            )
            .map_err(Error::DeviceManager)?;

        Ok(())
    }

    /// Create fw_cfg device if enabled in configuration.
    #[cfg(feature = "fw_cfg")]
    fn create_fw_cfg_if_enabled(
        config: &Arc<Mutex<VmConfig>>,
        device_manager: &Arc<Mutex<DeviceManager>>,
    ) -> Result<()> {
        let fw_cfg_enabled = config
            .lock()
            .unwrap()
            .payload
            .as_ref()
            .is_some_and(|p| p.fw_cfg_config.is_some());

        if fw_cfg_enabled {
            device_manager
                .lock()
                .unwrap()
                .create_fw_cfg_device()
                .map_err(Error::DeviceManager)?;
        }

        Ok(())
    }

    #[cfg(feature = "fw_cfg")]
    fn populate_fw_cfg(
        fw_cfg_config: &FwCfgConfig,
        device_manager: &Arc<Mutex<DeviceManager>>,
        config: &Arc<Mutex<VmConfig>>,
        #[cfg(target_arch = "x86_64")] kvm_sev_snp_enabled: bool,
    ) -> Result<()> {
        let mut e820_option: Option<usize> = None;
        if fw_cfg_config.e820 {
            e820_option = Some(config.lock().unwrap().memory.size as usize);
        }
        let mut kernel_option: Option<File> = None;
        if fw_cfg_config.kernel {
            let kernel = config
                .lock()
                .unwrap()
                .payload
                .as_ref()
                .map(|p| p.kernel.as_ref().map(File::open))
                .unwrap_or_default()
                .transpose()
                .map_err(Error::MissingFwCfgKernelFile)?;
            kernel_option = kernel;
        }
        let mut cmdline_option: Option<ffi::CString> = None;
        if fw_cfg_config.cmdline {
            let cmdline = Vm::generate_cmdline(
                config.lock().unwrap().payload.as_ref().unwrap(),
                #[cfg(target_arch = "aarch64")]
                device_manager,
            )
            .map_err(|_| Error::MissingFwCfgCmdline)?
            .as_cstring()
            .map_err(|_| Error::MissingFwCfgCmdline)?;
            cmdline_option = Some(cmdline);
        }
        let mut initramfs_option: Option<File> = None;
        if fw_cfg_config.initramfs {
            let initramfs = config
                .lock()
                .unwrap()
                .payload
                .as_ref()
                .map(|p| p.initramfs.as_ref().map(File::open))
                .unwrap_or_default()
                .transpose()
                .map_err(Error::MissingFwCfgInitramfs)?;
            // We measure the initramfs when running Oak Containers in SNP mode (initramfs = Stage1)
            // o/w use Stage0 to launch cloud disk images
            initramfs_option = initramfs;
        }
        let mut fw_cfg_item_list_option: Option<Vec<FwCfgItem>> = None;
        if let Some(fw_cfg_items) = &fw_cfg_config.items {
            let mut fw_cfg_item_list = vec![];
            for fw_cfg_item in fw_cfg_items.item_list.clone() {
                let content = match (fw_cfg_item.string, fw_cfg_item.file) {
                    (Some(string_val), None) => {
                        fw_cfg::FwCfgContent::Bytes(string_val.into_bytes())
                    }
                    (None, Some(file_path)) => fw_cfg::FwCfgContent::File(
                        0,
                        File::open(file_path).map_err(Error::AddingFwCfgItem)?,
                    ),
                    _ => unreachable!(
                        "PayloadConfig::validate() ensures either 'file' or 'string' is present"
                    ),
                };
                fw_cfg_item_list.push(FwCfgItem {
                    name: fw_cfg_item.name,
                    content,
                });
            }
            fw_cfg_item_list_option = Some(fw_cfg_item_list);
        }

        let device_manager_binding = device_manager.lock().unwrap();
        let Some(fw_cfg) = device_manager_binding.fw_cfg() else {
            return Err(Error::FwCfgDisabled);
        };

        fw_cfg
            .lock()
            .unwrap()
            .populate_fw_cfg(
                e820_option,
                kernel_option,
                initramfs_option,
                cmdline_option,
                fw_cfg_item_list_option,
                #[cfg(target_arch = "x86_64")]
                kvm_sev_snp_enabled,
            )
            .map_err(Error::ErrorPopulatingFwCfg)?;
        Ok(())
    }

    fn create_numa_nodes(
        configs: Option<&[NumaConfig]>,
        memory_manager: &Arc<Mutex<MemoryManager>>,
    ) -> Result<NumaNodes> {
        let mm = memory_manager.lock().unwrap();
        let mm_zones = mm.memory_zones();
        let mut numa_nodes = BTreeMap::new();

        if let Some(configs) = &configs {
            for config in configs.iter() {
                if numa_nodes.contains_key(&config.guest_numa_id) {
                    error!("Can't define twice the same NUMA node");
                    return Err(Error::InvalidNumaConfig);
                }

                let mut node = NumaNode::default();

                if let Some(memory_zones) = &config.memory_zones {
                    for memory_zone in memory_zones.iter() {
                        if let Some(mm_zone) = mm_zones.get(memory_zone) {
                            node.memory_regions.extend(mm_zone.regions().clone());
                            if let Some(virtiomem_zone) = mm_zone.virtio_mem_zone() {
                                node.hotplug_regions.push(virtiomem_zone.region().clone());
                            }
                            node.memory_zones.push(memory_zone.clone());
                        } else {
                            error!("Unknown memory zone '{memory_zone}'");
                            return Err(Error::InvalidNumaConfig);
                        }
                    }
                }

                if let Some(cpus) = &config.cpus {
                    node.cpus.extend(cpus);
                }

                if let Some(pci_segments) = &config.pci_segments {
                    node.pci_segments.extend(pci_segments);
                }

                if let Some(distances) = &config.distances {
                    for distance in distances.iter() {
                        let dest = distance.destination;
                        let dist = distance.distance;

                        if dest == config.guest_numa_id && dist != 10 {
                            warn!(
                                "Ignoring self-distance {dest}@{dist} (must be 10 per ACPI spec)"
                            );
                        }

                        if !configs.iter().any(|cfg| cfg.guest_numa_id == dest) {
                            error!("Unknown destination NUMA node {dest}");
                            return Err(Error::InvalidNumaConfig);
                        }

                        if node.distances.contains_key(&dest) {
                            error!("Destination NUMA node {dest} has been already set");
                            return Err(Error::InvalidNumaConfig);
                        }

                        node.distances.insert(dest, dist);
                    }
                }

                node.device_id = config.device_id.clone();

                numa_nodes.insert(config.guest_numa_id, node);
            }
        }

        Ok(numa_nodes)
    }

    #[expect(clippy::too_many_arguments)]
    pub fn new(
        vm_config: Arc<Mutex<VmConfig>>,
        exit_evt: EventFd,
        reset_evt: EventFd,
        guest_exit_evt: EventFd,
        #[cfg(feature = "guest_debug")] vm_debug_evt: EventFd,
        seccomp_action: &SeccompAction,
        hypervisor: Arc<dyn hypervisor::Hypervisor>,
        activate_evt: EventFd,
        console_info: Option<ConsoleInfo>,
        console_resize_pipe: Option<Arc<File>>,
        original_termios: Arc<Mutex<Option<termios>>>,
        snapshot: Option<&Snapshot>,
        source_url: Option<&str>,
        prefault: Option<bool>,
        memory_restore_mode: Option<MemoryRestoreMode>,
    ) -> Result<Self> {
        trace_scoped!("Vm::new");

        #[cfg(not(target_arch = "riscv64"))]
        let timestamp = Instant::now();

        #[cfg(feature = "tdx")]
        let tdx_enabled = if snapshot.is_some() {
            false
        } else {
            vm_config.lock().unwrap().is_tdx_enabled()
        };

        #[cfg(feature = "igvm")]
        let igvm_file = {
            let config = vm_config.lock().unwrap();
            config
                .payload
                .as_ref()
                .and_then(|p| p.igvm.as_ref())
                .map(|igvm_path| parse_igvm(igvm_path))
                .transpose()
                .map_err(Error::IgvmLoad)?
        };

        let vm = {
            #[allow(unused_mut)]
            let mut hv_config: hypervisor::HypervisorVmConfig =
                vm_config.as_ref().lock().unwrap().deref().into();
            #[cfg(feature = "sev_snp")]
            if let Some(ref igvm) = igvm_file {
                hv_config.vmsa_features = igvm_loader::extract_sev_features(igvm);
            }
            Self::create_hypervisor_vm(hypervisor.as_ref(), hv_config)?
        };

        #[cfg(all(feature = "kvm", target_arch = "x86_64"))]
        if vm_config.lock().unwrap().max_apic_id() > arch::x86_64::MAX_SUPPORTED_CPUS_LEGACY {
            vm.enable_x2apic_api().unwrap();
        }

        let phys_bits = physical_bits(
            hypervisor.as_ref(),
            vm_config.lock().unwrap().cpus.max_phys_bits,
        );

        let memory_manager =
            if let Some(snapshot) = snapshot_from_id(snapshot, MEMORY_MANAGER_SNAPSHOT_ID) {
                MemoryManager::new_from_snapshot(
                    snapshot,
                    vm.clone(),
                    &vm_config.lock().unwrap().memory.clone(),
                    source_url,
                    prefault.unwrap_or(false),
                    memory_restore_mode.unwrap_or_default(),
                    phys_bits,
                    &exit_evt,
                )
                .map_err(Error::MemoryManager)?
            } else {
                MemoryManager::new(
                    vm.clone(),
                    &vm_config.lock().unwrap().memory.clone(),
                    None,
                    phys_bits,
                    #[cfg(feature = "tdx")]
                    tdx_enabled,
                    None,
                    Default::default(),
                )
                .map_err(Error::MemoryManager)?
            };

        Vm::new_from_memory_manager(
            vm_config,
            memory_manager,
            vm,
            exit_evt,
            reset_evt,
            guest_exit_evt,
            #[cfg(feature = "guest_debug")]
            vm_debug_evt,
            seccomp_action,
            hypervisor,
            activate_evt,
            #[cfg(not(target_arch = "riscv64"))]
            timestamp,
            console_info,
            console_resize_pipe,
            original_termios,
            snapshot,
            #[cfg(feature = "igvm")]
            igvm_file,
        )
    }

    pub fn create_hypervisor_vm(
        hypervisor: &dyn hypervisor::Hypervisor,
        config: HypervisorVmConfig,
    ) -> Result<Arc<dyn hypervisor::Vm>> {
        hypervisor.check_required_extensions().unwrap();

        let vm = hypervisor.create_vm(config).unwrap();

        #[cfg(target_arch = "x86_64")]
        {
            vm.set_identity_map_address(KVM_IDENTITY_MAP_START.0)
                .unwrap();
            vm.set_tss_address(KVM_TSS_START.0 as usize).unwrap();
            vm.enable_split_irq().unwrap();
        }

        Ok(vm)
    }

    fn load_initramfs(&mut self, guest_mem: &GuestMemoryMmap) -> Result<arch::InitramfsConfig> {
        let initramfs = self.initramfs.as_mut().unwrap();
        let size: usize = initramfs
            .seek(SeekFrom::End(0))
            .map_err(Error::InitramfsLoad)?
            .try_into()
            .unwrap();
        initramfs.rewind().map_err(Error::InitramfsLoad)?;

        let address =
            arch::initramfs_load_addr(guest_mem, size).map_err(Error::InitramfsAddress)?;
        let address = GuestAddress(address);

        guest_mem
            .read_volatile_from(address, initramfs, size)
            .map_err(Error::InitramfsRead)?;

        info!("Initramfs loaded: address = 0x{:x}", address.0);
        Ok(arch::InitramfsConfig { address, size })
    }

    pub fn generate_cmdline(
        payload: &PayloadConfig,
        #[cfg(any(target_arch = "aarch64", target_arch = "riscv64"))] device_manager: &Arc<
            Mutex<DeviceManager>,
        >,
    ) -> Result<Cmdline> {
        let mut cmdline = Cmdline::new(arch::CMDLINE_MAX_SIZE).map_err(Error::CmdLineCreate)?;
        if let Some(s) = payload.cmdline.as_ref() {
            cmdline.insert_str(s).map_err(Error::CmdLineInsertStr)?;
        }

        #[cfg(any(target_arch = "aarch64", target_arch = "riscv64"))]
        for entry in device_manager.lock().unwrap().cmdline_additions() {
            cmdline.insert_str(entry).map_err(Error::CmdLineInsertStr)?;
        }
        Ok(cmdline)
    }

    #[expect(clippy::needless_pass_by_value)]
    #[cfg(any(target_arch = "aarch64", target_arch = "riscv64"))]
    fn load_firmware(
        mut firmware: &File,
        memory_manager: Arc<Mutex<MemoryManager>>,
    ) -> Result<EntryPoint> {
        let mut memory_manager = memory_manager.lock().unwrap();
        memory_manager
            .add_uefi_flash()
            .map_err(Error::MemoryManager)?;
        let uefi_flash = memory_manager.uefi_flash();
        let mem = uefi_flash.memory();
        uefi::load_uefi(mem.deref(), layout::UEFI_START, &mut firmware).map_err(Error::UefiLoad)?;
        Ok(EntryPoint {
            entry_addr: layout::UEFI_START,
        })
    }

    #[allow(clippy::needless_pass_by_value)]
    #[cfg(any(target_arch = "aarch64", target_arch = "riscv64"))]
    fn load_kernel(
        mut kernel: File,
        memory_manager: Arc<Mutex<MemoryManager>>,
    ) -> Result<EntryPoint> {
        let guest_memory = memory_manager.lock().as_ref().unwrap().guest_memory();
        let mem = guest_memory.memory();
        let alignment = 0x20_0000;
        // round up
        let aligned_kernel_addr = layout::KERNEL_START.0.div_ceil(alignment) * alignment;
        let entry_addr = {
            match loader::pe::PE::load(
                mem.deref(),
                Some(GuestAddress(aligned_kernel_addr)),
                &mut kernel,
                None,
            ) {
                Ok(entry_addr) => entry_addr.kernel_load,
                // Try to load the binary as kernel PE file at first.
                // If failed, retry to load it as UEFI binary.
                // As the UEFI binary is formatless, it must be the last option to try.
                Err(loader::Error::Pe(InvalidImageMagicNumber)) => {
                    Self::load_firmware(&kernel, memory_manager)?;
                    layout::UEFI_START
                }
                Err(e) => {
                    return Err(Error::KernelLoad(e));
                }
            }
        };

        Ok(EntryPoint { entry_addr })
    }

    #[cfg(all(feature = "kvm", feature = "sev_snp"))]
    fn reserve_bootloader_regions(memory_manager: &Arc<Mutex<MemoryManager>>) -> Result<()> {
        let mut mm = memory_manager.lock().unwrap();
        mm.add_ram_region(BOOTLOADER_START, BOOTLOADER_SIZE)
            .map_err(Error::MemoryManager)?;
        mm.add_ram_region(KVM_VMSA_PAGE_ADDRESS, KVM_VMSA_PAGE_SIZE)
            .map_err(Error::MemoryManager)?;
        Ok(())
    }

    #[cfg(feature = "igvm")]
    #[expect(clippy::needless_pass_by_value)]
    fn load_igvm(
        igvm_file: IgvmFile,
        memory_manager: Arc<Mutex<MemoryManager>>,
        cpu_manager: Arc<Mutex<cpu::CpuManager>>,
        #[cfg(all(feature = "kvm", feature = "sev_snp", feature = "fw_cfg"))] measured_boot: Option<
            MeasuredBootInfo,
        >,
        #[cfg(feature = "sev_snp")] host_data: &Option<String>,
    ) -> Result<EntryPoint> {
        // Only reserve bootloader/VMSA regions for KVM + SEV-SNP; other hypervisors
        // (e.g. MSHV) handle this through their own import path.
        #[cfg(all(feature = "kvm", feature = "sev_snp"))]
        if cpu_manager.lock().unwrap().sev_snp_enabled()
            && cpu_manager.lock().unwrap().hypervisor_type() == hypervisor::HypervisorType::Kvm
        {
            Self::reserve_bootloader_regions(&memory_manager)?;
        }

        let res = igvm_loader::load_igvm(
            igvm_file,
            memory_manager,
            cpu_manager.clone(),
            "",
            #[cfg(all(feature = "kvm", feature = "sev_snp", feature = "fw_cfg"))]
            measured_boot,
            #[cfg(feature = "sev_snp")]
            host_data,
        )
        .map_err(Error::IgvmLoad)?;

        cfg_if::cfg_if! {
            if #[cfg(feature = "sev_snp")] {
                let entry_point = if cpu_manager.lock().unwrap().sev_snp_enabled() {
                    EntryPoint { entry_addr: vm_memory::GuestAddress(res.vmsa_gpa), setup_header: None }
                } else {
                    EntryPoint {entry_addr: vm_memory::GuestAddress(res.vmsa.rip), setup_header: None }
                };
            } else {
               let entry_point = EntryPoint { entry_addr: vm_memory::GuestAddress(res.vmsa.rip), setup_header: None };
            }
        };
        Ok(entry_point)
    }

    /// Loads the kernel or a firmware file.
    ///
    /// For x86_64, the boot path is the same.
    #[cfg(target_arch = "x86_64")]
    #[expect(clippy::needless_pass_by_value)]
    fn load_kernel(
        mut kernel: File,
        cmdline: Option<Cmdline>,
        memory_manager: Arc<Mutex<MemoryManager>>,
    ) -> Result<EntryPoint> {
        info!("Loading kernel");

        let mem = {
            let guest_memory = memory_manager.lock().as_ref().unwrap().guest_memory();
            guest_memory.memory()
        };

        // Try ELF binary with PVH boot.
        let entry_addr =
            loader::elf::Elf::load(mem.deref(), None, &mut kernel, Some(layout::HIGH_RAM_START))
                // Try loading kernel as bzImage.
                .or_else(|_| {
                    BzImage::load(mem.deref(), None, &mut kernel, Some(layout::HIGH_RAM_START))
                })
                .map_err(Error::KernelLoad)?;

        if let Some(cmdline) = cmdline {
            loader::load_cmdline(mem.deref(), layout::CMDLINE_START, &cmdline)
                .map_err(Error::LoadCmdLine)?;
        }

        if let PvhEntryPresent(entry_addr) = entry_addr.pvh_boot_cap {
            // Use the PVH kernel entry point to boot the guest
            info!("PVH kernel loaded: entry_addr = 0x{:x}", entry_addr.0);
            Ok(EntryPoint {
                entry_addr,
                setup_header: None,
            })
        } else if entry_addr.setup_header.is_some() {
            // Use the bzImage 32bit entry point to boot the guest
            info!(
                "bzImage kernel loaded: entry_addr = 0x{:x}",
                entry_addr.kernel_load.0
            );
            Ok(EntryPoint {
                entry_addr: entry_addr.kernel_load,
                setup_header: entry_addr.setup_header,
            })
        } else {
            Err(Error::KernelMissingPvhHeader)
        }
    }

    #[cfg(target_arch = "x86_64")]
    fn load_payload(
        payload: &PayloadConfig,
        memory_manager: Arc<Mutex<MemoryManager>>,
        #[cfg(feature = "igvm")] cpu_manager: Arc<Mutex<cpu::CpuManager>>,
        #[cfg(feature = "igvm")] igvm_file: Option<IgvmFile>,
    ) -> Result<EntryPoint> {
        trace_scoped!("load_payload");
        #[cfg(feature = "igvm")]
        {
            if payload.igvm.is_some() {
                let igvm_file =
                    igvm_file.ok_or(Error::IgvmLoad(igvm_loader::Error::MissingIgvm))?;
                #[cfg(all(feature = "kvm", feature = "sev_snp", feature = "fw_cfg"))]
                let measured_boot = if let (true, Some(kernel_path), Some(_cmdline)) = (
                    payload
                        .fw_cfg_config
                        .as_ref()
                        .is_some_and(|cfg| cfg.kernel && cfg.cmdline),
                    payload.kernel.as_ref(),
                    payload.cmdline.as_ref(),
                ) {
                    let kernel = File::open(kernel_path).map_err(Error::KernelFile)?;
                    let initramfs = if payload
                        .fw_cfg_config
                        .as_ref()
                        .is_some_and(|cfg| cfg.initramfs)
                    {
                        payload
                            .initramfs
                            .as_ref()
                            .map(File::open)
                            .transpose()
                            .map_err(Error::InitramfsFile)?
                    } else {
                        None
                    };

                    // Must byte-match the cmdline that populate_fw_cfg sends to the guest,
                    // otherwise the launch measurement will diverge.
                    let cmdline = Self::generate_cmdline(payload)?
                        .as_cstring()
                        .map_err(Error::CmdLineCreate)?;

                    Some(MeasuredBootInfo {
                        kernel,
                        initramfs,
                        cmdline,
                    })
                } else {
                    None
                };
                return Self::load_igvm(
                    igvm_file,
                    memory_manager,
                    cpu_manager,
                    #[cfg(all(feature = "kvm", feature = "sev_snp", feature = "fw_cfg"))]
                    measured_boot,
                    #[cfg(feature = "sev_snp")]
                    &payload.host_data,
                );
            }
        }
        match (&payload.firmware, &payload.kernel) {
            (Some(firmware), None) => {
                let firmware = File::open(firmware).map_err(Error::FirmwareFile)?;
                Self::load_kernel(firmware, None, memory_manager)
            }
            (None, Some(kernel)) => {
                let kernel = File::open(kernel).map_err(Error::KernelFile)?;
                let cmdline = Self::generate_cmdline(payload)?;
                Self::load_kernel(kernel, Some(cmdline), memory_manager)
            }
            _ => unreachable!(
                "Unsupported boot configuration: programming error from 'PayloadConfigError::validate()'"
            ),
        }
    }

    #[allow(clippy::needless_pass_by_value)]
    #[cfg(any(target_arch = "aarch64", target_arch = "riscv64"))]
    fn load_payload(
        payload: &PayloadConfig,
        memory_manager: Arc<Mutex<MemoryManager>>,
    ) -> Result<EntryPoint> {
        match (&payload.firmware, &payload.kernel) {
            (Some(firmware), None) => {
                let firmware = File::open(firmware).map_err(Error::FirmwareFile)?;
                Self::load_firmware(&firmware, memory_manager)
            }
            (None, Some(kernel)) => {
                let kernel = File::open(kernel).map_err(Error::KernelFile)?;
                Self::load_kernel(kernel, memory_manager)
            }
            _ => unreachable!(
                "Unsupported boot configuration: programming error from 'PayloadConfigError::validate()'"
            ),
        }
    }

    fn load_payload_async(
        memory_manager: &Arc<Mutex<MemoryManager>>,
        config: &Arc<Mutex<VmConfig>>,
        #[cfg(feature = "igvm")] cpu_manager: &Arc<Mutex<cpu::CpuManager>>,
        #[cfg(feature = "igvm")] igvm_file: Option<IgvmFile>,
    ) -> Result<Option<thread::JoinHandle<Result<EntryPoint>>>> {
        // Kernel with TDX is loaded in a different manner
        #[cfg(feature = "tdx")]
        if config.lock().unwrap().is_tdx_enabled() {
            return Ok(None);
        }

        config
            .lock()
            .unwrap()
            .payload
            .as_ref()
            .map(|payload| {
                let memory_manager = memory_manager.clone();
                let payload = payload.clone();
                #[cfg(feature = "igvm")]
                let cpu_manager = cpu_manager.clone();

                thread::Builder::new()
                    .name("payload_loader".into())
                    .spawn(move || {
                        Self::load_payload(
                            &payload,
                            memory_manager,
                            #[cfg(feature = "igvm")]
                            cpu_manager,
                            #[cfg(feature = "igvm")]
                            igvm_file,
                        )
                    })
                    .map_err(Error::KernelLoadThreadSpawn)
            })
            .transpose()
    }

    #[cfg(target_arch = "x86_64")]
    fn configure_system(
        &mut self,
        rsdp_addr: Option<GuestAddress>,
        entry_addr: EntryPoint,
    ) -> Result<()> {
        trace_scoped!("configure_system");
        info!("Configuring system");
        let mem = self.memory_manager.lock().unwrap().boot_guest_memory();

        let initramfs_config = match self.initramfs {
            Some(_) => Some(self.load_initramfs(&mem)?),
            None => None,
        };

        let boot_vcpus = self.cpu_manager.lock().unwrap().boot_vcpus();

        let smbios = self
            .config
            .lock()
            .unwrap()
            .platform
            .as_ref()
            .and_then(|p| p.smbios_config());

        let topology = self.cpu_manager.lock().unwrap().get_vcpu_topology();

        arch::configure_system(
            &mem,
            layout::CMDLINE_START,
            layout::CMDLINE_MAX_SIZE,
            &initramfs_config,
            boot_vcpus,
            entry_addr.setup_header,
            rsdp_addr,
            smbios.as_ref(),
            topology,
        )
        .map_err(Error::ConfigureSystem)?;
        Ok(())
    }

    #[cfg(target_arch = "aarch64")]
    fn configure_system(
        &mut self,
        _rsdp_addr: Option<GuestAddress>,
        _entry_addr: EntryPoint,
    ) -> Result<()> {
        let cmdline = Self::generate_cmdline(
            self.config.lock().unwrap().payload.as_ref().unwrap(),
            &self.device_manager,
        )?;
        let vcpu_mpidrs = self.cpu_manager.lock().unwrap().get_mpidrs();
        let vcpu_topology = self.cpu_manager.lock().unwrap().get_vcpu_topology();
        let mem = self.memory_manager.lock().unwrap().boot_guest_memory();
        let mut pci_space_info: Vec<PciSpaceInfo> = Vec::new();
        let initramfs_config = match self.initramfs {
            Some(_) => Some(self.load_initramfs(&mem)?),
            None => None,
        };

        let device_info = &self
            .device_manager
            .lock()
            .unwrap()
            .get_device_info()
            .clone();

        for pci_segment in self.device_manager.lock().unwrap().pci_segments().iter() {
            let pci_space = PciSpaceInfo {
                pci_segment_id: pci_segment.id,
                mmio_config_address: pci_segment.mmio_config_address,
                pci_device_space_start: pci_segment.start_of_mem64_area,
                pci_device_space_size: pci_segment.end_of_mem64_area
                    - pci_segment.start_of_mem64_area
                    + 1,
            };
            pci_space_info.push(pci_space);
        }

        let virtio_iommu_bdf = self
            .device_manager
            .lock()
            .unwrap()
            .iommu_attached_devices()
            .as_ref()
            .map(|(v, _)| *v);

        let vgic = self
            .device_manager
            .lock()
            .unwrap()
            .get_interrupt_controller()
            .unwrap()
            .lock()
            .unwrap()
            .get_vgic()
            .map_err(|_| {
                Error::ConfigureSystem(arch::Error::PlatformSpecific(
                    arch::aarch64::Error::SetupGic,
                ))
            })?;

        // PMU interrupt sticks to PPI, so need to be added by 16 to get real irq number.
        let pmu_supported = self
            .cpu_manager
            .lock()
            .unwrap()
            .init_pmu(AARCH64_PMU_IRQ + 16)
            .map_err(|_| {
                Error::ConfigureSystem(arch::Error::PlatformSpecific(
                    arch::aarch64::Error::VcpuInitPmu,
                ))
            })?;

        arch::configure_system(
            &mem,
            cmdline.as_cstring().unwrap().to_str().unwrap(),
            &vcpu_mpidrs,
            vcpu_topology,
            device_info,
            &initramfs_config,
            &pci_space_info,
            virtio_iommu_bdf.map(|bdf| bdf.into()),
            &vgic,
            &self.numa_nodes,
            pmu_supported,
        )
        .map_err(Error::ConfigureSystem)?;

        Ok(())
    }

    #[cfg(target_arch = "riscv64")]
    fn configure_system(&mut self) -> Result<()> {
        let cmdline = Self::generate_cmdline(
            self.config.lock().unwrap().payload.as_ref().unwrap(),
            &self.device_manager,
        )?;
        let num_vcpu = self.cpu_manager.lock().unwrap().vcpus().len();
        let mem = self.memory_manager.lock().unwrap().boot_guest_memory();
        let mut pci_space_info: Vec<PciSpaceInfo> = Vec::new();
        let initramfs_config = match self.initramfs {
            Some(_) => Some(self.load_initramfs(&mem)?),
            None => None,
        };

        let device_info = &self
            .device_manager
            .lock()
            .unwrap()
            .get_device_info()
            .clone();

        for pci_segment in self.device_manager.lock().unwrap().pci_segments().iter() {
            let pci_space = PciSpaceInfo {
                pci_segment_id: pci_segment.id,
                mmio_config_address: pci_segment.mmio_config_address,
                pci_device_space_start: pci_segment.start_of_mem64_area,
                pci_device_space_size: pci_segment.end_of_mem64_area
                    - pci_segment.start_of_mem64_area
                    + 1,
            };
            pci_space_info.push(pci_space);
        }

        // TODO: IOMMU for riscv64 is not yet support in kernel.

        let vaia = self
            .device_manager
            .lock()
            .unwrap()
            .get_interrupt_controller()
            .unwrap()
            .lock()
            .unwrap()
            .get_vaia()
            .map_err(|_| {
                Error::ConfigureSystem(arch::Error::PlatformSpecific(
                    arch::riscv64::Error::SetupAia,
                ))
            })?;

        // TODO: PMU support for riscv64 is scheduled to next stage.

        let timebase_frequency = self
            .cpu_manager
            .lock()
            .unwrap()
            .vcpus()
            .first()
            .and_then(|vcpu| vcpu.lock().unwrap().get_timebase_frequency().ok())
            .map(|f| f as u32)
            .unwrap_or(0x989680);

        arch::configure_system(
            &mem,
            cmdline.as_cstring().unwrap().to_str().unwrap(),
            num_vcpu as u32,
            device_info,
            &initramfs_config,
            &pci_space_info,
            &vaia,
            timebase_frequency,
        )
        .map_err(Error::ConfigureSystem)?;

        Ok(())
    }

    pub fn console_resize_pipe(&self) -> Option<Arc<File>> {
        self.device_manager.lock().unwrap().console_resize_pipe()
    }

    pub fn shutdown(&mut self) -> Result<()> {
        let new_state = VmState::Shutdown;

        self.state.valid_transition(new_state)?;

        // Wake up the DeviceManager threads so they will get terminated cleanly
        self.device_manager
            .lock()
            .unwrap()
            .resume()
            .map_err(Error::Resume)?;

        self.cpu_manager
            .lock()
            .unwrap()
            .shutdown()
            .map_err(Error::CpuManager)?;

        // Wait for all the threads to finish
        for thread in self.threads.drain(..) {
            thread.join().map_err(Error::ThreadCleanup)?;
        }
        self.state = new_state;

        Ok(())
    }

    pub fn resize(
        &mut self,
        desired_vcpus: Option<u32>,
        desired_memory: Option<u64>,
        desired_balloon: Option<u64>,
    ) -> Result<()> {
        event!("vm", "resizing");

        if let Some(desired_vcpus) = desired_vcpus {
            if self
                .cpu_manager
                .lock()
                .unwrap()
                .resize(desired_vcpus)
                .map_err(Error::CpuManager)?
            {
                self.device_manager
                    .lock()
                    .unwrap()
                    .notify_hotplug(AcpiNotificationFlags::CPU_DEVICES_CHANGED)
                    .map_err(Error::DeviceManager)?;
            }
            self.config.lock().unwrap().cpus.boot_vcpus = desired_vcpus;
        }

        if let Some(desired_memory) = desired_memory {
            let new_region = self
                .memory_manager
                .lock()
                .unwrap()
                .resize(desired_memory)
                .map_err(Error::MemoryManager)?;

            let memory_config = &mut self.config.lock().unwrap().memory;

            if let Some(new_region) = &new_region {
                self.device_manager
                    .lock()
                    .unwrap()
                    .update_memory(new_region)
                    .map_err(Error::DeviceManager)?;

                match memory_config.hotplug_method {
                    HotplugMethod::Acpi => {
                        self.device_manager
                            .lock()
                            .unwrap()
                            .notify_hotplug(AcpiNotificationFlags::MEMORY_DEVICES_CHANGED)
                            .map_err(Error::DeviceManager)?;
                    }
                    HotplugMethod::VirtioMem => {}
                }
            }

            // We update the VM config regardless of the actual guest resize
            // operation result (happened or not), so that if the VM reboots
            // it will be running with the last configure memory size.
            match memory_config.hotplug_method {
                HotplugMethod::Acpi => memory_config.size = desired_memory,
                HotplugMethod::VirtioMem => {
                    if desired_memory > memory_config.size {
                        memory_config.hotplugged_size = Some(desired_memory - memory_config.size);
                    } else {
                        memory_config.hotplugged_size = None;
                    }
                }
            }
        }

        if let Some(desired_balloon) = desired_balloon {
            self.device_manager
                .lock()
                .unwrap()
                .resize_balloon(desired_balloon)
                .map_err(Error::DeviceManager)?;

            // Update the configuration value for the balloon size to ensure
            // a reboot would use the right value.
            if let Some(balloon_config) = &mut self.config.lock().unwrap().balloon {
                balloon_config.size = desired_balloon;
            }
        }

        event!("vm", "resized");

        Ok(())
    }

    pub fn resize_disk(&mut self, id: &str, desired_size: u64) -> Result<()> {
        self.device_manager
            .lock()
            .unwrap()
            .resize_disk(id, desired_size)
            .map_err(Error::DeviceManager)?;

        Ok(())
    }

    pub fn resize_zone(&mut self, id: &str, desired_memory: u64) -> Result<()> {
        let memory_config = &mut self.config.lock().unwrap().memory;

        if let Some(zones) = &mut memory_config.zones {
            for zone in zones.iter_mut() {
                if zone.id == id {
                    self.memory_manager
                        .lock()
                        .unwrap()
                        .resize_zone(zone, desired_memory)
                        .map_err(Error::MemoryManager)?;
                    // We update the memory zone config regardless of the
                    // actual 'resize-zone' operation result (happened or
                    // not), so that if the VM reboots it will be running
                    // with the last configured memory zone size.
                    zone.hotplugged_size = Some(desired_memory - zone.size);

                    return Ok(());
                }
            }
        }

        error!("Could not find the memory zone {id} for the resize");
        Err(Error::ResizeZone)
    }

    pub fn add_device(&mut self, mut device_cfg: DeviceConfig) -> Result<PciDeviceInfo> {
        let pci_device_info = self
            .device_manager
            .lock()
            .unwrap()
            .add_device(&mut device_cfg)
            .map_err(Error::DeviceManager)?;

        // Update VmConfig by adding the new device. This is important to
        // ensure the device would be created in case of a reboot.
        {
            let mut config = self.config.lock().unwrap();
            add_to_config(&mut config.devices, device_cfg);
        }

        self.device_manager
            .lock()
            .unwrap()
            .notify_hotplug(AcpiNotificationFlags::PCI_DEVICES_CHANGED)
            .map_err(Error::DeviceManager)?;

        Ok(pci_device_info)
    }

    pub fn add_user_device(&mut self, mut device_cfg: UserDeviceConfig) -> Result<PciDeviceInfo> {
        let pci_device_info = self
            .device_manager
            .lock()
            .unwrap()
            .add_user_device(&mut device_cfg)
            .map_err(Error::DeviceManager)?;

        // Update VmConfig by adding the new device. This is important to
        // ensure the device would be created in case of a reboot.
        {
            let mut config = self.config.lock().unwrap();
            add_to_config(&mut config.user_devices, device_cfg);
        }

        self.device_manager
            .lock()
            .unwrap()
            .notify_hotplug(AcpiNotificationFlags::PCI_DEVICES_CHANGED)
            .map_err(Error::DeviceManager)?;

        Ok(pci_device_info)
    }

    pub fn remove_device(&mut self, id: &str) -> Result<()> {
        self.device_manager
            .lock()
            .unwrap()
            .remove_device(id)
            .map_err(Error::DeviceManager)?;

        self.device_manager
            .lock()
            .unwrap()
            .notify_hotplug(AcpiNotificationFlags::PCI_DEVICES_CHANGED)
            .map_err(Error::DeviceManager)?;
        Ok(())
    }

    pub fn add_disk(&mut self, mut disk_cfg: DiskConfig) -> Result<PciDeviceInfo> {
        let pci_device_info = self
            .device_manager
            .lock()
            .unwrap()
            .add_disk(&mut disk_cfg)
            .map_err(Error::DeviceManager)?;

        // Update VmConfig by adding the new device. This is important to
        // ensure the device would be created in case of a reboot.
        {
            let mut config = self.config.lock().unwrap();
            add_to_config(&mut config.disks, disk_cfg);
        }

        self.device_manager
            .lock()
            .unwrap()
            .notify_hotplug(AcpiNotificationFlags::PCI_DEVICES_CHANGED)
            .map_err(Error::DeviceManager)?;

        Ok(pci_device_info)
    }

    pub fn add_fs(&mut self, mut fs_cfg: FsConfig) -> Result<PciDeviceInfo> {
        let pci_device_info = self
            .device_manager
            .lock()
            .unwrap()
            .add_fs(&mut fs_cfg)
            .map_err(Error::DeviceManager)?;

        // Update VmConfig by adding the new device. This is important to
        // ensure the device would be created in case of a reboot.
        {
            let mut config = self.config.lock().unwrap();
            add_to_config(&mut config.fs, fs_cfg);
        }

        self.device_manager
            .lock()
            .unwrap()
            .notify_hotplug(AcpiNotificationFlags::PCI_DEVICES_CHANGED)
            .map_err(Error::DeviceManager)?;

        Ok(pci_device_info)
    }

    pub fn add_generic_vhost_user(
        &mut self,
        mut generic_vhost_user_cfg: GenericVhostUserConfig,
    ) -> Result<PciDeviceInfo> {
        let pci_device_info = self
            .device_manager
            .lock()
            .unwrap()
            .add_generic_vhost_user(&mut generic_vhost_user_cfg)
            .map_err(Error::DeviceManager)?;

        // Update VmConfig by adding the new device. This is important to
        // ensure the device would be created in case of a reboot.
        {
            let mut config = self.config.lock().unwrap();
            add_to_config(&mut config.generic_vhost_user, generic_vhost_user_cfg);
        }

        self.device_manager
            .lock()
            .unwrap()
            .notify_hotplug(AcpiNotificationFlags::PCI_DEVICES_CHANGED)
            .map_err(Error::DeviceManager)?;

        Ok(pci_device_info)
    }

    pub fn add_pmem(&mut self, mut pmem_cfg: PmemConfig) -> Result<PciDeviceInfo> {
        let pci_device_info = self
            .device_manager
            .lock()
            .unwrap()
            .add_pmem(&mut pmem_cfg)
            .map_err(Error::DeviceManager)?;

        // Update VmConfig by adding the new device. This is important to
        // ensure the device would be created in case of a reboot.
        {
            let mut config = self.config.lock().unwrap();
            add_to_config(&mut config.pmem, pmem_cfg);
        }

        self.device_manager
            .lock()
            .unwrap()
            .notify_hotplug(AcpiNotificationFlags::PCI_DEVICES_CHANGED)
            .map_err(Error::DeviceManager)?;

        Ok(pci_device_info)
    }

    pub fn add_net(&mut self, mut net_cfg: NetConfig) -> Result<PciDeviceInfo> {
        let pci_device_info = self
            .device_manager
            .lock()
            .unwrap()
            .add_net(&mut net_cfg)
            .map_err(Error::DeviceManager)?;

        // Update VmConfig by adding the new device. This is important to
        // ensure the device would be created in case of a reboot.
        {
            let mut config = self.config.lock().unwrap();
            add_to_config(&mut config.net, net_cfg);
        }

        self.device_manager
            .lock()
            .unwrap()
            .notify_hotplug(AcpiNotificationFlags::PCI_DEVICES_CHANGED)
            .map_err(Error::DeviceManager)?;

        Ok(pci_device_info)
    }

    pub fn add_vdpa(&mut self, mut vdpa_cfg: VdpaConfig) -> Result<PciDeviceInfo> {
        let pci_device_info = self
            .device_manager
            .lock()
            .unwrap()
            .add_vdpa(&mut vdpa_cfg)
            .map_err(Error::DeviceManager)?;

        // Update VmConfig by adding the new device. This is important to
        // ensure the device would be created in case of a reboot.
        {
            let mut config = self.config.lock().unwrap();
            add_to_config(&mut config.vdpa, vdpa_cfg);
        }

        self.device_manager
            .lock()
            .unwrap()
            .notify_hotplug(AcpiNotificationFlags::PCI_DEVICES_CHANGED)
            .map_err(Error::DeviceManager)?;

        Ok(pci_device_info)
    }

    pub fn add_vsock(&mut self, mut vsock_cfg: VsockConfig) -> Result<PciDeviceInfo> {
        let pci_device_info = self
            .device_manager
            .lock()
            .unwrap()
            .add_vsock(&mut vsock_cfg)
            .map_err(Error::DeviceManager)?;

        // Update VmConfig by adding the new device. This is important to
        // ensure the device would be created in case of a reboot.
        {
            let mut config = self.config.lock().unwrap();
            config.vsock = Some(vsock_cfg);
        }

        self.device_manager
            .lock()
            .unwrap()
            .notify_hotplug(AcpiNotificationFlags::PCI_DEVICES_CHANGED)
            .map_err(Error::DeviceManager)?;

        Ok(pci_device_info)
    }

    pub fn counters(&self) -> Result<HashMap<String, HashMap<&'static str, Wrapping<u64>>>> {
        Ok(self.device_manager.lock().unwrap().counters())
    }

    #[cfg(feature = "tdx")]
    fn extract_tdvf_sections(&mut self) -> Result<(Vec<arch::x86_64::tdx::TdvfSection>, bool)> {
        use arch::x86_64::tdx::*;

        let firmware_path = self
            .config
            .lock()
            .unwrap()
            .payload
            .as_ref()
            .unwrap()
            .firmware
            .clone()
            .ok_or(Error::TdxFirmwareMissing)?;
        // The TDVF file contains a table of section as well as code
        let mut firmware_file = File::open(firmware_path).map_err(Error::LoadTdvf)?;

        // For all the sections allocate some RAM backing them
        parse_tdvf_sections(&mut firmware_file).map_err(Error::ParseTdvf)
    }

    #[cfg(feature = "tdx")]
    fn hob_memory_resources(
        mut sorted_sections: Vec<arch::x86_64::tdx::TdvfSection>,
        guest_memory: &GuestMemoryMmap,
    ) -> Vec<(u64, u64, bool)> {
        let mut list = Vec::new();

        let mut current_section = sorted_sections.pop();

        // RAM regions interleaved with TDVF sections
        let mut next_start_addr = 0;
        for region in guest_memory.iter() {
            let region_start = region.start_addr().0;
            let region_end = region.last_addr().0;
            if region_start > next_start_addr {
                next_start_addr = region_start;
            }

            loop {
                let (start, size, ram) = if let Some(section) = &current_section {
                    if section.address <= next_start_addr {
                        (section.address, section.size, false)
                    } else {
                        let last_addr = cmp::min(section.address - 1, region_end);
                        (next_start_addr, last_addr - next_start_addr + 1, true)
                    }
                } else {
                    (next_start_addr, region_end - next_start_addr + 1, true)
                };

                list.push((start, size, ram));

                if !ram {
                    current_section = sorted_sections.pop();
                }

                next_start_addr = start + size;

                if region_start > next_start_addr {
                    next_start_addr = region_start;
                }

                if next_start_addr > region_end {
                    break;
                }
            }
        }

        // Once all the interleaved sections have been processed, let's simply
        // pull the remaining ones.
        if let Some(section) = current_section {
            list.push((section.address, section.size, false));
        }
        while let Some(section) = sorted_sections.pop() {
            list.push((section.address, section.size, false));
        }

        list
    }

    #[cfg(feature = "tdx")]
    fn populate_tdx_sections(
        &mut self,
        sections: &[arch::x86_64::tdx::TdvfSection],
        guid_found: bool,
    ) -> Result<Option<u64>> {
        use arch::x86_64::tdx::*;
        // Get the memory end *before* we start adding TDVF ram regions
        let boot_guest_memory = self
            .memory_manager
            .lock()
            .as_ref()
            .unwrap()
            .boot_guest_memory();
        for section in sections {
            // No need to allocate if the section falls within guest RAM ranges
            if boot_guest_memory.address_in_range(GuestAddress(section.address)) {
                info!(
                    "Not allocating TDVF Section: {section:x?} since it is already part of guest RAM"
                );
                continue;
            }

            info!("Allocating TDVF Section: {section:x?}");
            self.memory_manager
                .lock()
                .unwrap()
                .add_ram_region(GuestAddress(section.address), section.size as usize)
                .map_err(Error::AllocatingTdvfMemory)?;
        }

        // The TDVF file contains a table of section as well as code
        let firmware_path = self
            .config
            .lock()
            .unwrap()
            .payload
            .as_ref()
            .unwrap()
            .firmware
            .clone()
            .ok_or(Error::TdxFirmwareMissing)?;
        let mut firmware_file = File::open(firmware_path).map_err(Error::LoadTdvf)?;

        // The guest memory at this point now has all the required regions so it
        // is safe to copy from the TDVF file into it.
        let guest_memory = self.memory_manager.lock().as_ref().unwrap().guest_memory();
        let mem = guest_memory.memory();
        let mut payload_info = None;
        let mut hob_offset = None;
        for section in sections {
            info!("Populating TDVF Section: {section:x?}");
            match section.r#type {
                TdvfSectionType::Bfv | TdvfSectionType::Cfv => {
                    info!("Copying section to guest memory");
                    firmware_file
                        .seek(SeekFrom::Start(section.data_offset as u64))
                        .map_err(Error::LoadTdvf)?;
                    mem.read_volatile_from(
                        GuestAddress(section.address),
                        &mut firmware_file,
                        section.data_size as usize,
                    )
                    .unwrap();
                }
                TdvfSectionType::TdHob => {
                    hob_offset = Some(section.address);
                }
                TdvfSectionType::Payload => {
                    info!("Copying payload to guest memory");
                    if let Some(payload_file) = self.kernel.as_mut() {
                        let payload_size = payload_file
                            .seek(SeekFrom::End(0))
                            .map_err(Error::LoadPayload)?;

                        payload_file
                            .seek(SeekFrom::Start(0x1f1))
                            .map_err(Error::LoadPayload)?;

                        let mut payload_header = bootparam::setup_header::default();
                        payload_file
                            .read_volatile(&mut payload_header.as_bytes())
                            .unwrap();

                        if payload_header.header != 0x5372_6448 {
                            return Err(Error::InvalidPayloadType);
                        }

                        if (payload_header.version < 0x0200)
                            || ((payload_header.loadflags & 0x1) == 0x0)
                        {
                            return Err(Error::InvalidPayloadType);
                        }

                        payload_file.rewind().map_err(Error::LoadPayload)?;
                        mem.read_volatile_from(
                            GuestAddress(section.address),
                            payload_file,
                            payload_size as usize,
                        )
                        .unwrap();

                        // Create the payload info that will be inserted into
                        // the HOB.
                        payload_info = Some(PayloadInfo {
                            image_type: PayloadImageType::BzImage,
                            entry_point: section.address,
                        });
                    }
                }
                TdvfSectionType::PayloadParam => {
                    info!("Copying payload parameters to guest memory");
                    let cmdline = Self::generate_cmdline(
                        self.config.lock().unwrap().payload.as_ref().unwrap(),
                    )?;
                    mem.write_slice(
                        cmdline.as_cstring().unwrap().as_bytes_with_nul(),
                        GuestAddress(section.address),
                    )
                    .unwrap();
                }
                _ => {}
            }
        }

        // Generate HOB
        let mut hob = TdHob::start(hob_offset.unwrap());

        let mut sorted_sections = sections.to_vec();
        sorted_sections.retain(|section| matches!(section.r#type, TdvfSectionType::TempMem));

        sorted_sections.sort_by_key(|section| section.address);
        sorted_sections.reverse();

        for (start, size, ram) in Vm::hob_memory_resources(sorted_sections, &boot_guest_memory) {
            hob.add_memory_resource(&mem, start, size, ram, guid_found)
                .map_err(Error::PopulateHob)?;
        }

        // MMIO regions
        hob.add_mmio_resource(
            &mem,
            layout::MEM_32BIT_DEVICES_START.raw_value(),
            layout::APIC_START.raw_value() - layout::MEM_32BIT_DEVICES_START.raw_value(),
        )
        .map_err(Error::PopulateHob)?;
        let start_of_device_area = self
            .memory_manager
            .lock()
            .unwrap()
            .start_of_device_area()
            .raw_value();
        let end_of_device_area = self
            .memory_manager
            .lock()
            .unwrap()
            .end_of_device_area()
            .raw_value();
        hob.add_mmio_resource(
            &mem,
            start_of_device_area,
            end_of_device_area - start_of_device_area,
        )
        .map_err(Error::PopulateHob)?;

        // Loop over the ACPI tables and copy them to the HOB.

        for acpi_table in acpi::create_acpi_tables_tdx(
            &self.device_manager.lock().unwrap(),
            &self.cpu_manager.lock().unwrap(),
            &self.memory_manager.lock().unwrap(),
            &self.numa_nodes,
        ) {
            hob.add_acpi_table(&mem, acpi_table.as_slice())
                .map_err(Error::PopulateHob)?;
        }

        // If a payload info has been created, let's insert it into the HOB.
        if let Some(payload_info) = payload_info {
            hob.add_payload(&mem, payload_info)
                .map_err(Error::PopulateHob)?;
        }

        hob.finish(&mem).map_err(Error::PopulateHob)?;

        Ok(hob_offset)
    }

    #[cfg(feature = "tdx")]
    fn init_tdx_memory(&mut self, sections: &[arch::x86_64::tdx::TdvfSection]) -> Result<()> {
        let guest_memory = self.memory_manager.lock().as_ref().unwrap().guest_memory();
        let mem = guest_memory.memory();

        for section in sections {
            let size = section.size.try_into().unwrap();
            // SAFETY: get_host_address_range does proper bounds checking
            unsafe {
                self.vm.tdx_init_memory_region(
                    virtio_devices::get_host_address_range(
                        &*mem,
                        GuestAddress(section.address),
                        size,
                    )
                    .unwrap(),
                    section.address,
                    size,
                    /* TDVF_SECTION_ATTRIBUTES_EXTENDMR */
                    section.attributes == 1,
                )
            }
            .map_err(Error::InitializeTdxMemoryRegion)?;
        }

        Ok(())
    }

    // Creates ACPI tables
    // In case of TDX being used, this is a no-op since the tables will be
    // created and passed when populating the HOB.

    #[cfg(not(target_arch = "riscv64"))]
    fn create_acpi_tables(&self) -> Option<GuestAddress> {
        #[cfg(feature = "tdx")]
        if self.config.lock().unwrap().is_tdx_enabled() {
            return None;
        }
        let mem = self.memory_manager.lock().unwrap().guest_memory().memory();
        let tpm_enabled = self.config.lock().unwrap().tpm.is_some();
        let rsdp_addr = acpi::create_acpi_tables(
            &mem,
            &self.device_manager.lock().unwrap(),
            &self.cpu_manager.lock().unwrap(),
            &self.memory_manager.lock().unwrap(),
            &self.numa_nodes,
            tpm_enabled,
        );
        info!("Created ACPI tables: rsdp_addr = 0x{:x}", rsdp_addr.0);

        Some(rsdp_addr)
    }

    fn entry_point(&mut self) -> Result<Option<EntryPoint>> {
        trace_scoped!("entry_point");

        self.load_payload_handle
            .take()
            .map(|handle| handle.join().map_err(Error::KernelLoadThreadJoin)?)
            .transpose()
    }

    pub fn boot(&mut self) -> Result<()> {
        trace_scoped!("Vm::boot");
        let current_state = self.state;
        if current_state == VmState::Paused {
            return self.resume().map_err(Error::Resume);
        }

        // We acquire all advisory disk image locks here and not on device creation
        // to enable live-migration without locking issues.
        self.device_manager
            .lock()
            .unwrap()
            .try_lock_disks()
            .map_err(Error::LockingError)?;

        let new_state = if self.stop_on_boot {
            VmState::BreakPoint
        } else {
            VmState::Running
        };

        current_state.valid_transition(new_state)?;

        #[cfg(feature = "fw_cfg")]
        {
            let fw_cfg_enabled = self
                .config
                .lock()
                .unwrap()
                .payload
                .as_ref()
                .is_some_and(|p| p.fw_cfg_config.is_some());
            if fw_cfg_enabled {
                let fw_cfg_config = self
                    .config
                    .lock()
                    .unwrap()
                    .payload
                    .as_ref()
                    .map(|p| p.fw_cfg_config.clone())
                    .unwrap_or_default()
                    .ok_or(Error::VmMissingConfig)?;
                #[cfg(target_arch = "x86_64")]
                let kvm_sev_snp_enabled = {
                    #[cfg(feature = "sev_snp")]
                    {
                        self.config.lock().unwrap().is_sev_snp_enabled()
                            && self.hypervisor.hypervisor_type() == hypervisor::HypervisorType::Kvm
                    }
                    #[cfg(not(feature = "sev_snp"))]
                    {
                        false
                    }
                };

                Self::populate_fw_cfg(
                    &fw_cfg_config,
                    &self.device_manager,
                    &self.config,
                    #[cfg(target_arch = "x86_64")]
                    kvm_sev_snp_enabled,
                )?;

                if fw_cfg_config.acpi_tables {
                    let tpm_enabled = self.config.lock().unwrap().tpm.is_some();
                    acpi::create_acpi_tables_for_fw_cfg(
                        &self.device_manager.lock().unwrap(),
                        &self.cpu_manager.lock().unwrap(),
                        &self.memory_manager.lock().unwrap(),
                        &self.numa_nodes,
                        tpm_enabled,
                    )?;
                }
            }
        }

        // Do earlier to parallelise with loading kernel
        #[cfg(target_arch = "x86_64")]
        cfg_if::cfg_if! {
            if #[cfg(feature = "sev_snp")] {
                let sev_snp_enabled = self.config.lock().unwrap().is_sev_snp_enabled();
                let rsdp_addr = if sev_snp_enabled {
                    // In case of SEV-SNP guest ACPI tables are provided via
                    // IGVM. So skip the creation of ACPI tables and set the
                    // rsdp addr to None.
                    None
                } else {
                    self.create_acpi_tables()
                };
            } else {
                let rsdp_addr = self.create_acpi_tables();
            }
        }

        // Load kernel synchronously or if asynchronous then wait for load to
        // finish.
        let entry_point = self.entry_point()?;

        #[cfg(feature = "tdx")]
        let tdx_enabled = self.config.lock().unwrap().is_tdx_enabled();

        #[cfg(target_arch = "aarch64")]
        let vgic = self
            .device_manager
            .lock()
            .unwrap()
            .get_interrupt_controller()
            .unwrap()
            .lock()
            .unwrap()
            .get_vgic()
            .unwrap();

        #[cfg(target_arch = "aarch64")]
        let redist_addr = vgic.lock().unwrap().device_properties();

        // Configure the vcpus that have been created
        let vcpus = self.cpu_manager.lock().unwrap().vcpus();
        for vcpu in vcpus {
            let guest_memory = &self.memory_manager.lock().as_ref().unwrap().guest_memory();
            let boot_setup = entry_point.map(|e| (e, guest_memory));
            let mut vcpu = vcpu.lock().unwrap();
            self.cpu_manager
                .lock()
                .unwrap()
                .configure_vcpu(&mut vcpu, boot_setup)
                .map_err(Error::CpuManager)?;

            #[cfg(target_arch = "aarch64")]
            vcpu.set_gic_redistributor_addr(redist_addr[2], redist_addr[3])
                .map_err(Error::CpuManager)?;
        }

        #[cfg(feature = "mshv")]
        {
            self.cpu_manager
                .lock()
                .unwrap()
                .set_processors_per_socket_property()
                .ok();
        }

        #[cfg(feature = "tdx")]
        let (sections, guid_found) = if tdx_enabled {
            self.extract_tdvf_sections()?
        } else {
            (Vec::new(), false)
        };

        // Configuring the TDX regions requires that the vCPUs are created.
        #[cfg(feature = "tdx")]
        let hob_address = if tdx_enabled {
            // TDX sections are written to memory.
            self.populate_tdx_sections(&sections, guid_found)?
        } else {
            None
        };

        // On aarch64 the ACPI tables depend on the vCPU mpidr which is only
        // available after they are configured
        #[cfg(target_arch = "aarch64")]
        let rsdp_addr = self.create_acpi_tables();

        #[cfg(not(target_arch = "riscv64"))]
        // Configure shared state based on loaded kernel.
        // Skip for SEV-SNP guests where system configuration is provided via IGVM.
        if rsdp_addr.is_some() {
            entry_point
                .map(|entry_point| self.configure_system(rsdp_addr, entry_point))
                .transpose()?;
        }

        #[cfg(target_arch = "riscv64")]
        self.configure_system().unwrap();

        #[cfg(feature = "tdx")]
        if let Some(hob_address) = hob_address {
            // With the HOB address extracted the vCPUs can have
            // their TDX state configured.
            self.cpu_manager
                .lock()
                .unwrap()
                .initialize_tdx(hob_address)
                .map_err(Error::CpuManager)?;
            // Let the hypervisor know which memory ranges are shared with the
            // guest. This prevents the guest from ignoring/discarding memory
            // regions provided by the host.
            self.init_tdx_memory(&sections)?;
            // With TDX memory and CPU state configured TDX setup is complete
            self.vm.tdx_finalize().map_err(Error::FinalizeTdx)?;
        }

        // Resume the vm for MSHV
        if current_state == VmState::Created {
            self.vm.resume().map_err(Error::ResumeVm)?;
        }

        self.cpu_manager
            .lock()
            .unwrap()
            .start_boot_vcpus(new_state == VmState::BreakPoint)
            .map_err(Error::CpuManager)?;

        self.state = new_state;
        Ok(())
    }

    pub fn restore(&mut self) -> Result<()> {
        event!("vm", "restoring");

        // We acquire all advisory disk image locks again.
        self.device_manager
            .lock()
            .unwrap()
            .try_lock_disks()
            .map_err(Error::LockingError)?;

        // Now we can start all vCPUs from here.
        self.cpu_manager
            .lock()
            .unwrap()
            .start_restored_vcpus()
            .map_err(Error::CpuManager)?;

        event!("vm", "restored");
        Ok(())
    }

    /// Gets a thread-safe reference counted pointer to the VM configuration.
    pub fn get_config(&self) -> Arc<Mutex<VmConfig>> {
        Arc::clone(&self.config)
    }

    /// Get the VM state.
    pub fn get_state(&self) -> VmState {
        self.state
    }

    /// Gets the actual size of the balloon.
    pub fn balloon_size(&self) -> u64 {
        self.device_manager.lock().unwrap().balloon_size()
    }

    /// Get the actual size of the virtio_mem regions
    pub fn virtio_mem_plugged_size(&self) -> u64 {
        self.memory_manager
            .lock()
            .unwrap()
            .virtio_mem_plugged_size()
    }

    pub fn send_memory_fds(
        &mut self,
        socket: &mut UnixStream,
    ) -> result::Result<(), MigratableError> {
        for (slot, fd) in self
            .memory_manager
            .lock()
            .unwrap()
            .memory_slot_fds()
            .drain()
        {
            Request::memory_fd(size_of_val(&slot) as u64)
                .write_to(socket)
                .context("Error sending memory fd request")
                .map_err(MigratableError::MigrateSend)?;
            socket
                .send_with_fd(&slot.to_le_bytes()[..], fd)
                .context("Error sending memory fd")
                .map_err(MigratableError::MigrateSend)?;

            Response::read_from(socket)?.ok_or_fatal_error(MigratableError::MigrateSend(
                anyhow!("Error during memory fd migration"),
            ))?;
        }

        Ok(())
    }

    pub fn memory_range_table(&self) -> result::Result<MemoryRangeTable, MigratableError> {
        self.memory_manager
            .lock()
            .unwrap()
            .memory_range_table(false)
    }

    pub fn guest_memory(&self) -> GuestMemoryAtomic<GuestMemoryMmap> {
        self.memory_manager.lock().unwrap().guest_memory()
    }

    pub fn restoring(&self) -> bool {
        self.memory_manager.lock().unwrap().restoring()
    }

    pub fn device_tree(&self) -> Arc<Mutex<DeviceTree>> {
        self.device_manager.lock().unwrap().device_tree()
    }

    /// Release all advisory locks held for the disk images.
    ///
    /// This should only be called when the VM is stopped and the VMM supposed
    /// to shut down. A new VMM, either after a live migration or a
    /// state save/resume cycle, should then acquire all locks before the VM
    /// starts to run.
    pub fn release_disk_locks(&self) -> Result<()> {
        self.device_manager
            .lock()
            .unwrap()
            .release_disk_locks()
            .map_err(Error::LockingError)?;
        Ok(())
    }

    pub fn activate_virtio_devices(&self) -> Result<()> {
        self.device_manager
            .lock()
            .unwrap()
            .activate_virtio_devices()
            .map_err(Error::ActivateVirtioDevices)
    }

    #[cfg(target_arch = "x86_64")]
    pub fn power_button(&self) -> Result<()> {
        return self
            .device_manager
            .lock()
            .unwrap()
            .notify_power_button()
            .map_err(Error::PowerButton);
    }

    #[cfg(target_arch = "aarch64")]
    pub fn power_button(&self) -> Result<()> {
        self.device_manager
            .lock()
            .unwrap()
            .notify_power_button()
            .map_err(Error::PowerButton)
    }

    #[cfg(target_arch = "riscv64")]
    pub fn power_button(&self) -> Result<()> {
        unimplemented!()
    }

    pub fn memory_manager_data(&self) -> MemoryManagerSnapshotData {
        self.memory_manager.lock().unwrap().snapshot_data()
    }

    #[cfg(feature = "guest_debug")]
    pub fn debug_request(
        &mut self,
        gdb_request: &GdbRequestPayload,
        cpu_id: usize,
    ) -> Result<GdbResponsePayload> {
        use GdbRequestPayload::*;
        match gdb_request {
            SetSingleStep(single_step) => {
                self.set_guest_debug(cpu_id, &[], *single_step)
                    .map_err(Error::Debug)?;
            }
            SetHwBreakPoint(addrs) => {
                self.set_guest_debug(cpu_id, addrs, false)
                    .map_err(Error::Debug)?;
            }
            Pause => {
                self.debug_pause().map_err(Error::Debug)?;
            }
            Resume => {
                self.debug_resume().map_err(Error::Debug)?;
            }
            ReadRegs => {
                let regs = self.read_regs(cpu_id).map_err(Error::Debug)?;
                return Ok(GdbResponsePayload::RegValues(Box::new(regs)));
            }
            WriteRegs(regs) => {
                self.write_regs(cpu_id, regs).map_err(Error::Debug)?;
            }
            ReadMem(vaddr, len) => {
                let guest_memory = self.memory_manager.lock().as_ref().unwrap().guest_memory();
                let mem = self
                    .read_mem(&guest_memory, cpu_id, *vaddr, *len)
                    .map_err(Error::Debug)?;
                return Ok(GdbResponsePayload::MemoryRegion(mem));
            }
            WriteMem(vaddr, data) => {
                let guest_memory = self.memory_manager.lock().as_ref().unwrap().guest_memory();
                self.write_mem(&guest_memory, cpu_id, vaddr, data)
                    .map_err(Error::Debug)?;
            }
            ActiveVcpus => {
                let active_vcpus = self.active_vcpus();
                return Ok(GdbResponsePayload::ActiveVcpus(active_vcpus));
            }
        }
        Ok(GdbResponsePayload::CommandComplete)
    }

    #[cfg(all(target_arch = "x86_64", feature = "guest_debug"))]
    fn get_dump_state(
        &mut self,
        destination_url: &str,
    ) -> result::Result<DumpState, GuestDebuggableError> {
        let nr_cpus = self.config.lock().unwrap().cpus.boot_vcpus;
        let elf_note_size = self.get_note_size(NoteDescType::ElfAndVmm, nr_cpus) as isize;
        let mut elf_phdr_num = 1;
        let elf_sh_info = 0;
        let coredump_file_path = url_to_file(destination_url)?;
        let mapping_num = self.memory_manager.lock().unwrap().num_guest_ram_mappings();

        if mapping_num < UINT16_MAX - 2 {
            elf_phdr_num += mapping_num as u16;
        } else {
            panic!("mapping num beyond 65535 not supported");
        }
        let coredump_file = OpenOptions::new()
            .read(true)
            .write(true)
            .create_new(true)
            .open(coredump_file_path)
            .map_err(|e| GuestDebuggableError::Coredump(e.into()))?;

        let mem_offset = self.coredump_get_mem_offset(elf_phdr_num, elf_note_size);
        let mem_data = self
            .memory_manager
            .lock()
            .unwrap()
            .coredump_memory_regions(mem_offset);

        Ok(DumpState {
            elf_note_size,
            elf_phdr_num,
            elf_sh_info,
            mem_offset,
            mem_info: Some(mem_data),
            file: Some(coredump_file),
        })
    }

    #[cfg(all(target_arch = "x86_64", feature = "guest_debug"))]
    fn coredump_get_mem_offset(&self, phdr_num: u16, note_size: isize) -> u64 {
        size_of::<elf::Elf64_Ehdr>() as u64
            + note_size as u64
            + size_of::<elf::Elf64_Phdr>() as u64 * phdr_num as u64
    }

    pub fn nmi(&self) -> Result<()> {
        return self
            .cpu_manager
            .lock()
            .unwrap()
            .nmi()
            .map_err(Error::ErrorNmi);
    }

    /// Capture the guest clock for a same-host resume (mode `SameHostResume`).
    /// `None` if the VM is not booted or the backend has no guest clock.
    fn capture_guest_clock(&self) -> result::Result<Option<SavedClock>, MigratableError> {
        let Some(boot_vcpu) = self.cpu_manager.lock().unwrap().boot_vcpu() else {
            return Ok(None);
        };
        let boot_vcpu = boot_vcpu.lock().unwrap();
        Ok(self
            .vm
            .snapshot_clock(boot_vcpu.hypervisor_vcpu())
            .map_err(|e| MigratableError::Pause(anyhow!("Could not capture guest clock: {e}")))?
            .map(|state| SavedClock {
                mode: hypervisor::ClockRestoreMode::SameHostResume,
                state,
            }))
    }

    /// Re-establish the guest clock before the vCPUs resume. No-op if none captured.
    fn restore_guest_clock(&self) -> result::Result<(), MigratableError> {
        let Some(saved) = &self.saved_clock else {
            return Ok(());
        };
        let vcpus = self.cpu_manager.lock().unwrap().vcpus();
        let guards: Vec<_> = vcpus.iter().map(|v| v.lock().unwrap()).collect();
        let hv_vcpus: Vec<&dyn hypervisor::Vcpu> =
            guards.iter().map(|g| g.hypervisor_vcpu()).collect();
        self.vm
            .restore_clock(&hv_vcpus, &saved.state, saved.mode)
            .map_err(|e| MigratableError::Resume(anyhow!("Could not restore guest clock: {e}")))
    }

    pub fn device_manager(&self) -> &Arc<Mutex<DeviceManager>> {
        &self.device_manager
    }
}

impl Pausable for Vm {
    fn pause(&mut self) -> result::Result<(), MigratableError> {
        event!("vm", "pausing");
        let new_state = VmState::Paused;
        self.state
            .valid_transition(new_state)
            .map_err(|e| MigratableError::Pause(anyhow!("Invalid transition: {e:?}")))?;

        // Before pausing the vCPUs activate any pending virtio devices that might
        // need activation between starting the pause (or e.g. a migration it's part of)
        self.activate_virtio_devices().map_err(|e| {
            MigratableError::Pause(anyhow!("Error activating pending virtio devices: {e:?}"))
        })?;

        self.cpu_manager.lock().unwrap().pause()?;

        // Capture the guest clock now that the vCPUs are quiesced.
        self.saved_clock = self.capture_guest_clock()?;

        self.device_manager.lock().unwrap().pause()?;

        self.vm
            .pause()
            .map_err(|e| MigratableError::Pause(anyhow!("Could not pause the VM: {e}")))?;

        self.state = new_state;

        event!("vm", "paused");
        Ok(())
    }

    fn resume(&mut self) -> result::Result<(), MigratableError> {
        event!("vm", "resuming");
        let current_state = self.get_state();
        let new_state = VmState::Running;

        self.state
            .valid_transition(new_state)
            .map_err(|e| MigratableError::Resume(anyhow!("Invalid transition: {e:?}")))?;

        // Restore the guest clock before the vCPUs start running.
        self.restore_guest_clock()?;

        if current_state == VmState::Paused {
            self.vm
                .resume()
                .map_err(|e| MigratableError::Resume(anyhow!("Could not resume the VM: {e}")))?;
        }

        self.device_manager.lock().unwrap().resume()?;
        self.cpu_manager.lock().unwrap().resume()?;

        // And we're back to the Running state.
        self.state = new_state;
        event!("vm", "resumed");
        Ok(())
    }
}

#[derive(Serialize, Deserialize)]
pub struct VmSnapshot {
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub clock: Option<hypervisor::ClockState>,
    #[cfg(all(feature = "kvm", target_arch = "x86_64"))]
    pub common_cpuid: Vec<x86::CpuIdEntry>,
}

pub const VM_SNAPSHOT_ID: &str = "vm";
impl Snapshottable for Vm {
    fn id(&self) -> String {
        VM_SNAPSHOT_ID.to_string()
    }

    fn snapshot(&mut self) -> result::Result<Snapshot, MigratableError> {
        event!("vm", "snapshotting");

        #[cfg(feature = "tdx")]
        {
            if self.config.lock().unwrap().is_tdx_enabled() {
                return Err(MigratableError::Snapshot(anyhow!(
                    "Snapshot not possible with TDX VM"
                )));
            }
        }

        if self.get_state() != VmState::Paused {
            return Err(MigratableError::Snapshot(anyhow!(
                "Trying to snapshot while VM is running"
            )));
        }

        #[cfg(all(feature = "kvm", target_arch = "x86_64"))]
        let common_cpuid = {
            let (amx, max_phys_bits, kvm_hyperv, profile) = {
                let guard = self.config.lock().unwrap();
                let VmConfig { cpus, .. } = &*guard;
                (
                    cpus.features.amx,
                    cpus.max_phys_bits,
                    cpus.kvm_hyperv,
                    cpus.profile,
                )
            };

            let phys_bits = physical_bits(self.hypervisor.as_ref(), max_phys_bits);

            arch::generate_common_cpuid(
                self.hypervisor.as_ref(),
                &arch::CpuidConfig {
                    phys_bits,
                    kvm_hyperv,
                    #[cfg(feature = "tdx")]
                    tdx: false,
                    amx,
                    profile,
                },
            )
            .context("Error generating common cpuid")
            .map_err(MigratableError::MigrateReceive)?
        };

        let vm_snapshot_state = VmSnapshot {
            clock: self.saved_clock.map(|saved| saved.state),
            #[cfg(all(feature = "kvm", target_arch = "x86_64"))]
            common_cpuid,
        };

        let mut vm_snapshot = Snapshot::new_from_state(&vm_snapshot_state)?;

        let (id, snapshot) = {
            let mut cpu_manager = self.cpu_manager.lock().unwrap();
            (cpu_manager.id(), cpu_manager.snapshot()?)
        };
        vm_snapshot.add_snapshot(id, snapshot);
        let (id, snapshot) = {
            let mut memory_manager = self.memory_manager.lock().unwrap();
            (memory_manager.id(), memory_manager.snapshot()?)
        };
        vm_snapshot.add_snapshot(id, snapshot);
        let (id, snapshot) = {
            let mut device_manager = self.device_manager.lock().unwrap();
            (device_manager.id(), device_manager.snapshot()?)
        };
        vm_snapshot.add_snapshot(id, snapshot);

        event!("vm", "snapshotted");
        Ok(vm_snapshot)
    }
}

impl Transportable for Vm {
    fn send(
        &self,
        snapshot: &Snapshot,
        destination_url: &str,
    ) -> result::Result<(), MigratableError> {
        let mut snapshot_config_path = url_to_path(destination_url)?;
        snapshot_config_path.push(SNAPSHOT_CONFIG_FILE);

        // Create the snapshot config file
        let mut snapshot_config_file = OpenOptions::new()
            .read(true)
            .write(true)
            .create_new(true)
            .open(snapshot_config_path)
            .context("Error creating VM config snapshot file")
            .map_err(MigratableError::MigrateSend)?;

        // Serialize and write the snapshot config
        let vm_config = serde_json::to_string(self.config.lock().unwrap().deref())
            .context("Error serializing VM config snapshot")
            .map_err(MigratableError::MigrateSend)?;

        snapshot_config_file
            .write(vm_config.as_bytes())
            .context("Error writing VM config snapshot")
            .map_err(MigratableError::MigrateSend)?;

        let mut snapshot_state_path = url_to_path(destination_url)?;
        snapshot_state_path.push(SNAPSHOT_STATE_FILE);

        // Create the snapshot state file
        let mut snapshot_state_file = OpenOptions::new()
            .read(true)
            .write(true)
            .create_new(true)
            .open(snapshot_state_path)
            .context("Error creating VM state snapshot file")
            .map_err(MigratableError::MigrateSend)?;

        // Serialize and write the snapshot state
        let vm_state = serde_json::to_vec(snapshot)
            .context("Error serializing VM state snapshot")
            .map_err(MigratableError::MigrateSend)?;

        snapshot_state_file
            .write(&vm_state)
            .context("Error writing VM state snapshot")
            .map_err(MigratableError::MigrateSend)?;

        // Tell the memory manager to also send/write its own snapshot.
        if let Some(memory_manager_snapshot) = snapshot.snapshots.get(MEMORY_MANAGER_SNAPSHOT_ID) {
            self.memory_manager
                .lock()
                .unwrap()
                .send(&memory_manager_snapshot.clone(), destination_url)?;
        } else {
            return Err(MigratableError::Restore(anyhow!(
                "Missing memory manager snapshot"
            )));
        }

        Ok(())
    }
}

impl Migratable for Vm {
    fn start_dirty_log(&mut self) -> result::Result<(), MigratableError> {
        self.memory_manager.lock().unwrap().start_dirty_log()?;
        self.device_manager.lock().unwrap().start_dirty_log()
    }

    fn stop_dirty_log(&mut self) -> result::Result<(), MigratableError> {
        self.memory_manager.lock().unwrap().stop_dirty_log()?;
        self.device_manager.lock().unwrap().stop_dirty_log()
    }

    fn dirty_log(&mut self) -> result::Result<MemoryRangeTable, MigratableError> {
        Ok(MemoryRangeTable::new_from_tables(vec![
            self.memory_manager.lock().unwrap().dirty_log()?,
            self.device_manager.lock().unwrap().dirty_log()?,
        ]))
    }

    fn start_migration(&mut self) -> result::Result<(), MigratableError> {
        self.memory_manager.lock().unwrap().start_migration()?;
        self.device_manager.lock().unwrap().start_migration()
    }

    fn complete_migration(&mut self) -> result::Result<(), MigratableError> {
        self.memory_manager.lock().unwrap().complete_migration()?;
        self.device_manager.lock().unwrap().complete_migration()
    }
}

#[cfg(feature = "guest_debug")]
impl Debuggable for Vm {
    fn set_guest_debug(
        &self,
        cpu_id: usize,
        addrs: &[GuestAddress],
        singlestep: bool,
    ) -> result::Result<(), DebuggableError> {
        self.cpu_manager
            .lock()
            .unwrap()
            .set_guest_debug(cpu_id, addrs, singlestep)
    }

    fn debug_pause(&mut self) -> result::Result<(), DebuggableError> {
        if self.state == VmState::Running {
            self.pause().map_err(DebuggableError::Pause)?;
        }

        self.state = VmState::BreakPoint;
        Ok(())
    }

    fn debug_resume(&mut self) -> result::Result<(), DebuggableError> {
        if self.state == VmState::BreakPoint {
            self.resume().map_err(DebuggableError::Pause)?;
        }

        Ok(())
    }

    fn read_regs(&self, cpu_id: usize) -> result::Result<CoreRegs, DebuggableError> {
        self.cpu_manager.lock().unwrap().read_regs(cpu_id)
    }

    fn write_regs(&self, cpu_id: usize, regs: &CoreRegs) -> result::Result<(), DebuggableError> {
        self.cpu_manager.lock().unwrap().write_regs(cpu_id, regs)
    }

    fn read_mem(
        &self,
        guest_memory: &GuestMemoryAtomic<GuestMemoryMmap>,
        cpu_id: usize,
        vaddr: GuestAddress,
        len: usize,
    ) -> result::Result<Vec<u8>, DebuggableError> {
        self.cpu_manager
            .lock()
            .unwrap()
            .read_mem(guest_memory, cpu_id, vaddr, len)
    }

    fn write_mem(
        &self,
        guest_memory: &GuestMemoryAtomic<GuestMemoryMmap>,
        cpu_id: usize,
        vaddr: &GuestAddress,
        data: &[u8],
    ) -> result::Result<(), DebuggableError> {
        self.cpu_manager
            .lock()
            .unwrap()
            .write_mem(guest_memory, cpu_id, vaddr, data)
    }

    fn active_vcpus(&self) -> usize {
        let active_vcpus = self.cpu_manager.lock().unwrap().active_vcpus();
        if active_vcpus > 0 {
            active_vcpus
        } else {
            // The VM is not booted yet. Report boot_vcpus() instead.
            self.cpu_manager.lock().unwrap().boot_vcpus() as usize
        }
    }
}

#[cfg(feature = "guest_debug")]
pub const UINT16_MAX: u32 = 65535;

#[cfg(all(target_arch = "x86_64", feature = "guest_debug"))]
impl Elf64Writable for Vm {}

#[cfg(all(target_arch = "x86_64", feature = "guest_debug"))]
impl GuestDebuggable for Vm {
    fn coredump(&mut self, destination_url: &str) -> result::Result<(), GuestDebuggableError> {
        event!("vm", "coredumping");

        let mut resume = false;

        #[cfg(feature = "tdx")]
        {
            if let Some(ref platform) = self.config.lock().unwrap().platform
                && platform.tdx
            {
                return Err(GuestDebuggableError::Coredump(anyhow!(
                    "Coredump not possible with TDX VM"
                )));
            }
        }

        match self.get_state() {
            VmState::Running => {
                self.pause().map_err(GuestDebuggableError::Pause)?;
                resume = true;
            }
            VmState::Paused => {}
            _ => {
                return Err(GuestDebuggableError::Coredump(anyhow!(
                    "Trying to coredump while VM is not running or paused"
                )));
            }
        }

        let coredump_state = self.get_dump_state(destination_url)?;

        self.write_header(&coredump_state)?;
        self.write_note(&coredump_state)?;
        self.write_loads(&coredump_state)?;

        self.cpu_manager
            .lock()
            .unwrap()
            .cpu_write_elf64_note(&coredump_state)?;
        self.cpu_manager
            .lock()
            .unwrap()
            .cpu_write_vmm_note(&coredump_state)?;

        self.memory_manager
            .lock()
            .unwrap()
            .coredump_iterate_save_mem(&coredump_state)?;

        if resume {
            self.resume().map_err(GuestDebuggableError::Resume)?;
        }

        Ok(())
    }
}

#[cfg(all(feature = "kvm", target_arch = "x86_64"))]
#[cfg(test)]
mod unit_tests {
    use super::*;

    fn test_vm_state_transitions(state: VmState) {
        match state {
            VmState::Created => {
                // Check the transitions from Created
                state.valid_transition(VmState::Created).unwrap_err();
                state.valid_transition(VmState::Running).unwrap();
                state.valid_transition(VmState::Shutdown).unwrap();
                state.valid_transition(VmState::Paused).unwrap();
                state.valid_transition(VmState::BreakPoint).unwrap();
            }
            VmState::Running => {
                // Check the transitions from Running
                state.valid_transition(VmState::Created).unwrap_err();
                state.valid_transition(VmState::Running).unwrap_err();
                state.valid_transition(VmState::Shutdown).unwrap();
                state.valid_transition(VmState::Paused).unwrap();
                state.valid_transition(VmState::BreakPoint).unwrap();
            }
            VmState::Shutdown => {
                // Check the transitions from Shutdown
                state.valid_transition(VmState::Created).unwrap_err();
                state.valid_transition(VmState::Running).unwrap();
                state.valid_transition(VmState::Shutdown).unwrap_err();
                state.valid_transition(VmState::Paused).unwrap_err();
                state.valid_transition(VmState::BreakPoint).unwrap_err();
            }
            VmState::Paused => {
                // Check the transitions from Paused
                state.valid_transition(VmState::Created).unwrap_err();
                state.valid_transition(VmState::Running).unwrap();
                state.valid_transition(VmState::Shutdown).unwrap();
                state.valid_transition(VmState::Paused).unwrap_err();
                state.valid_transition(VmState::BreakPoint).unwrap_err();
            }
            VmState::BreakPoint => {
                // Check the transitions from Breakpoint
                state.valid_transition(VmState::Created).unwrap();
                state.valid_transition(VmState::Running).unwrap();
                state.valid_transition(VmState::Shutdown).unwrap_err();
                state.valid_transition(VmState::Paused).unwrap_err();
                state.valid_transition(VmState::BreakPoint).unwrap_err();
            }
        }
    }

    #[test]
    fn test_vm_created_transitions() {
        test_vm_state_transitions(VmState::Created);
    }

    #[test]
    fn test_vm_running_transitions() {
        test_vm_state_transitions(VmState::Running);
    }

    #[test]
    fn test_vm_shutdown_transitions() {
        test_vm_state_transitions(VmState::Shutdown);
    }

    #[test]
    fn test_vm_paused_transitions() {
        test_vm_state_transitions(VmState::Paused);
    }

    #[cfg(feature = "tdx")]
    #[test]
    fn test_hob_memory_resources() {
        // Case 1: Two TDVF sections in the middle of the RAM
        let sections = vec![
            arch::x86_64::tdx::TdvfSection {
                address: 0xc000,
                size: 0x1000,
                ..Default::default()
            },
            arch::x86_64::tdx::TdvfSection {
                address: 0x1000,
                size: 0x4000,
                ..Default::default()
            },
        ];
        let guest_ranges: Vec<(GuestAddress, usize)> = vec![(GuestAddress(0), 0x1000_0000)];
        let expected = vec![
            (0, 0x1000, true),
            (0x1000, 0x4000, false),
            (0x5000, 0x7000, true),
            (0xc000, 0x1000, false),
            (0xd000, 0x0fff_3000, true),
        ];
        assert_eq!(
            expected,
            Vm::hob_memory_resources(
                sections,
                &GuestMemoryMmap::from_ranges(&guest_ranges).unwrap()
            )
        );

        // Case 2: Two TDVF sections with no conflict with the RAM
        let sections = vec![
            arch::x86_64::tdx::TdvfSection {
                address: 0x1000_1000,
                size: 0x1000,
                ..Default::default()
            },
            arch::x86_64::tdx::TdvfSection {
                address: 0,
                size: 0x1000,
                ..Default::default()
            },
        ];
        let guest_ranges: Vec<(GuestAddress, usize)> = vec![(GuestAddress(0x1000), 0x1000_0000)];
        let expected = vec![
            (0, 0x1000, false),
            (0x1000, 0x1000_0000, true),
            (0x1000_1000, 0x1000, false),
        ];
        assert_eq!(
            expected,
            Vm::hob_memory_resources(
                sections,
                &GuestMemoryMmap::from_ranges(&guest_ranges).unwrap()
            )
        );

        // Case 3: Two TDVF sections with partial conflicts with the RAM
        let sections = vec![
            arch::x86_64::tdx::TdvfSection {
                address: 0x1000_0000,
                size: 0x2000,
                ..Default::default()
            },
            arch::x86_64::tdx::TdvfSection {
                address: 0,
                size: 0x2000,
                ..Default::default()
            },
        ];
        let guest_ranges: Vec<(GuestAddress, usize)> = vec![(GuestAddress(0x1000), 0x1000_0000)];
        let expected = vec![
            (0, 0x2000, false),
            (0x2000, 0x0fff_e000, true),
            (0x1000_0000, 0x2000, false),
        ];
        assert_eq!(
            expected,
            Vm::hob_memory_resources(
                sections,
                &GuestMemoryMmap::from_ranges(&guest_ranges).unwrap()
            )
        );

        // Case 4: Two TDVF sections with no conflict before the RAM and two
        // more additional sections with no conflict after the RAM.
        let sections = vec![
            arch::x86_64::tdx::TdvfSection {
                address: 0x2000_1000,
                size: 0x1000,
                ..Default::default()
            },
            arch::x86_64::tdx::TdvfSection {
                address: 0x2000_0000,
                size: 0x1000,
                ..Default::default()
            },
            arch::x86_64::tdx::TdvfSection {
                address: 0x1000,
                size: 0x1000,
                ..Default::default()
            },
            arch::x86_64::tdx::TdvfSection {
                address: 0,
                size: 0x1000,
                ..Default::default()
            },
        ];
        let guest_ranges: Vec<(GuestAddress, usize)> = vec![(GuestAddress(0x4000), 0x1000_0000)];
        let expected = vec![
            (0, 0x1000, false),
            (0x1000, 0x1000, false),
            (0x4000, 0x1000_0000, true),
            (0x2000_0000, 0x1000, false),
            (0x2000_1000, 0x1000, false),
        ];
        assert_eq!(
            expected,
            Vm::hob_memory_resources(
                sections,
                &GuestMemoryMmap::from_ranges(&guest_ranges).unwrap()
            )
        );

        // Case 5: One TDVF section overriding the entire RAM
        let sections = vec![arch::x86_64::tdx::TdvfSection {
            address: 0,
            size: 0x2000_0000,
            ..Default::default()
        }];
        let guest_ranges: Vec<(GuestAddress, usize)> = vec![(GuestAddress(0x1000), 0x1000_0000)];
        let expected = vec![(0, 0x2000_0000, false)];
        assert_eq!(
            expected,
            Vm::hob_memory_resources(
                sections,
                &GuestMemoryMmap::from_ranges(&guest_ranges).unwrap()
            )
        );

        // Case 6: Two TDVF sections with no conflict with 2 RAM regions
        let sections = vec![
            arch::x86_64::tdx::TdvfSection {
                address: 0x1000_2000,
                size: 0x2000,
                ..Default::default()
            },
            arch::x86_64::tdx::TdvfSection {
                address: 0,
                size: 0x2000,
                ..Default::default()
            },
        ];
        let guest_ranges: Vec<(GuestAddress, usize)> = vec![
            (GuestAddress(0x2000), 0x1000_0000),
            (GuestAddress(0x1000_4000), 0x1000_0000),
        ];
        let expected = vec![
            (0, 0x2000, false),
            (0x2000, 0x1000_0000, true),
            (0x1000_2000, 0x2000, false),
            (0x1000_4000, 0x1000_0000, true),
        ];
        assert_eq!(
            expected,
            Vm::hob_memory_resources(
                sections,
                &GuestMemoryMmap::from_ranges(&guest_ranges).unwrap()
            )
        );

        // Case 7: Two TDVF sections with partial conflicts with 2 RAM regions
        let sections = vec![
            arch::x86_64::tdx::TdvfSection {
                address: 0x1000_0000,
                size: 0x4000,
                ..Default::default()
            },
            arch::x86_64::tdx::TdvfSection {
                address: 0,
                size: 0x4000,
                ..Default::default()
            },
        ];
        let guest_ranges: Vec<(GuestAddress, usize)> = vec![
            (GuestAddress(0x1000), 0x1000_0000),
            (GuestAddress(0x1000_3000), 0x1000_0000),
        ];
        let expected = vec![
            (0, 0x4000, false),
            (0x4000, 0x0fff_c000, true),
            (0x1000_0000, 0x4000, false),
            (0x1000_4000, 0x0fff_f000, true),
        ];
        assert_eq!(
            expected,
            Vm::hob_memory_resources(
                sections,
                &GuestMemoryMmap::from_ranges(&guest_ranges).unwrap()
            )
        );
    }

    #[test]
    pub fn test_vm() {
        use hypervisor::VmExit;
        use vm_memory::{Address, GuestMemory, GuestMemoryRegion};
        // This example based on https://lwn.net/Articles/658511/
        let code = [
            0xba, 0xf8, 0x03, /* mov $0x3f8, %dx */
            0x00, 0xd8, /* add %bl, %al */
            0x04, b'0', /* add $'0', %al */
            0xee, /* out %al, (%dx) */
            0xb0, b'\n', /* mov $'\n', %al */
            0xee,  /* out %al, (%dx) */
            0xf4,  /* hlt */
        ];

        let mem_size = 0x1000;
        let load_addr = GuestAddress(0x1000);
        let mem = GuestMemoryMmap::from_ranges(&[(load_addr, mem_size)]).unwrap();

        let hv = hypervisor::new().unwrap();
        let vm = hv
            .create_vm(HypervisorVmConfig::default())
            .expect("new VM creation failed");

        for (index, region) in mem.iter().enumerate() {
            // SAFETY: inputs are valid
            unsafe {
                vm.create_user_memory_region(
                    index as u32,
                    region.start_addr().raw_value(),
                    region.len().try_into().unwrap(),
                    region.as_ptr(),
                    false,
                    false,
                )
                .expect("Cannot configure guest memory");
            }
        }
        mem.write_slice(&code, load_addr)
            .expect("Writing code to memory failed");

        let mut vcpu = vm.create_vcpu(0, None).expect("new Vcpu failed");

        let mut vcpu_sregs = vcpu.get_sregs().expect("get sregs failed");
        vcpu_sregs.cs.base = 0;
        vcpu_sregs.cs.selector = 0;
        vcpu.set_sregs(&vcpu_sregs).expect("set sregs failed");

        let mut vcpu_regs = vcpu.get_regs().expect("get regs failed");
        vcpu_regs.set_rip(0x1000);
        vcpu_regs.set_rax(2);
        vcpu_regs.set_rbx(3);
        vcpu_regs.set_rflags(2);
        vcpu.set_regs(&vcpu_regs).expect("set regs failed");

        loop {
            match vcpu.run().expect("run failed") {
                VmExit::Reset => {
                    println!("HLT");
                    break;
                }
                VmExit::Ignore => {}
                r => panic!("unexpected exit reason: {r:?}"),
            }
        }
    }
}

#[cfg(target_arch = "aarch64")]
#[cfg(test)]
mod unit_tests {
    use arch::{DeviceType, MmioDeviceInfo};
    use devices::gic::Gic;

    use super::*;

    const LEN: u64 = 4096;

    #[test]
    fn test_create_fdt_with_devices() {
        let regions = vec![(
            arch::aarch64::layout::RAM_START,
            (arch::aarch64::layout::FDT_MAX_SIZE + 0x1000) as usize,
        )];
        let mem = GuestMemoryMmap::from_ranges(&regions).expect("Cannot initialize memory");

        let dev_info: HashMap<(DeviceType, String), MmioDeviceInfo> = [
            (
                (DeviceType::Serial, DeviceType::Serial.to_string()),
                MmioDeviceInfo {
                    addr: 0x00,
                    len: LEN,
                    irq: 33,
                },
            ),
            (
                (DeviceType::Virtio(1), "virtio".to_string()),
                MmioDeviceInfo {
                    addr: LEN,
                    len: LEN,
                    irq: 34,
                },
            ),
            (
                (DeviceType::Rtc, "rtc".to_string()),
                MmioDeviceInfo {
                    addr: 2 * LEN,
                    len: LEN,
                    irq: 35,
                },
            ),
        ]
        .iter()
        .cloned()
        .collect();

        let hv = hypervisor::new().unwrap();
        let vm = hv.create_vm(HypervisorVmConfig::default()).unwrap();
        let vgic_config = Gic::create_default_config(1);
        let gic = vm.create_vgic(&vgic_config).expect("Cannot create gic");
        arch::aarch64::fdt::create_fdt(
            &mem,
            "console=tty0",
            &[0],
            Some((0, 0, 0, 0)),
            &dev_info,
            &gic,
            &None,
            &Vec::new(),
            &BTreeMap::new(),
            None,
            true,
        )
        .unwrap();
    }
}

#[cfg(all(feature = "kvm", target_arch = "x86_64"))]
#[test]
pub fn test_vm() {
    use hypervisor::VmExit;
    use vm_memory::{Address, GuestMemory, GuestMemoryRegion};
    // This example based on https://lwn.net/Articles/658511/
    let code = [
        0xba, 0xf8, 0x03, /* mov $0x3f8, %dx */
        0x00, 0xd8, /* add %bl, %al */
        0x04, b'0', /* add $'0', %al */
        0xee, /* out %al, (%dx) */
        0xb0, b'\n', /* mov $'\n', %al */
        0xee,  /* out %al, (%dx) */
        0xf4,  /* hlt */
    ];

    let mem_size = 0x1000;
    let load_addr = GuestAddress(0x1000);
    let mem = GuestMemoryMmap::from_ranges(&[(load_addr, mem_size)]).unwrap();

    let hv = hypervisor::new().unwrap();
    let vm = hv
        .create_vm(HypervisorVmConfig::default())
        .expect("new VM creation failed");

    for (index, region) in mem.iter().enumerate() {
        // SAFETY: parameters are correct
        unsafe {
            vm.create_user_memory_region(
                index as u32,
                region.start_addr().raw_value(),
                region.len().try_into().unwrap(),
                region.as_ptr().cast(),
                false,
                false,
            )
            .expect("Cannot configure guest memory");
        }
    }
    mem.write_slice(&code, load_addr)
        .expect("Writing code to memory failed");

    let mut vcpu = vm.create_vcpu(0, None).expect("new Vcpu failed");

    let mut vcpu_sregs = vcpu.get_sregs().expect("get sregs failed");
    vcpu_sregs.cs.base = 0;
    vcpu_sregs.cs.selector = 0;
    vcpu.set_sregs(&vcpu_sregs).expect("set sregs failed");

    let mut vcpu_regs = vcpu.get_regs().expect("get regs failed");
    vcpu_regs.set_rip(0x1000);
    vcpu_regs.set_rax(2);
    vcpu_regs.set_rbx(3);
    vcpu_regs.set_rflags(2);
    vcpu.set_regs(&vcpu_regs).expect("set regs failed");

    loop {
        match vcpu.run().expect("run failed") {
            VmExit::Reset => {
                println!("HLT");
                break;
            }
            VmExit::Ignore => {}
            r => panic!("unexpected exit reason: {r:?}"),
        }
    }
}
