axvm 0.5.26

// Copyright 2025 The Axvisor Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Runtime configuration structures for an AxVM instance.

use std::{string::String, sync::Arc, vec::Vec};

use axdevice::{NullSerialBackendFactory, SerialBackendFactory};
use axvm_types::InterruptTriggerMode;
pub use axvm_types::{
    AddressSpacePolicy, GuestPhysAddr, HostAddressAssignment, HostDeviceAssignment,
    HostPortAssignment, ReservedAddressConfig, VMBootProtocol, VmMemConfig, VmMemMappingType,
};
use axvmconfig::VirtualDeviceRequest;

use crate::{arch::*, machine::*};

/// Policy used by AxVM when deriving runtime guest boot image addresses.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum GuestBootPolicy {
    /// Keep the load addresses exactly as provided by the VM config.
    #[default]
    KeepConfigured,
    /// Adjust the kernel load address for boot protocols that require a
    /// reserved area inside the primary guest memory region.
    AdjustKernelForBootProtocol { protocol: VMBootProtocol },
}

/// A part of `AxVMConfig`, which represents a `VCpu`.
#[derive(Clone, Copy, Debug, Default)]
pub struct AxVCpuConfig {
    /// The entry address in GPA for the Bootstrap Processor (BSP).
    pub bsp_entry: GuestPhysAddr,
    /// The entry address in GPA for the Application Processor (AP).
    pub ap_entry: GuestPhysAddr,
}

/// Ramdisk image information.
#[derive(Debug, Default, Clone)]
pub struct RamdiskInfo {
    /// The load address in GPA for the ramdisk image.
    pub load_gpa: GuestPhysAddr,
    /// The size in bytes of the ramdisk image, `None` if not known yet.
    pub size: Option<usize>,
}

/// A part of `AxVMConfig`, which stores configuration attributes related to the load address of VM images.
#[derive(Debug, Default, Clone)]
pub struct VMImageConfig {
    /// The load address in GPA for the kernel image.
    pub kernel_load_gpa: GuestPhysAddr,
    /// Whether VM images are loaded from the host filesystem.
    pub loaded_from_filesystem: bool,
    /// The load address in GPA for the BIOS image, `None` if not used.
    pub bios_load_gpa: Option<GuestPhysAddr>,
    /// The load address in GPA for the device tree blob (DTB), `None` if not used.
    pub dtb_load_gpa: Option<GuestPhysAddr>,
    /// Ramdisk image info, `None` if not used.
    pub ramdisk: Option<RamdiskInfo>,
}

/// Physical interrupt source forwarded through a guest's virtual controller.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct PassthroughInterrupt {
    /// Architecture-local physical interrupt source number.
    pub source: u32,
    /// Trigger mode declared by firmware for the physical device.
    pub trigger: InterruptTriggerMode,
}

/// Runtime configuration for one VM.
#[derive(Debug)]
pub struct AxVMConfig {
    id: usize,
    name: String,
    pub(crate) phys_cpu_ls: PhysCpuList,
    /// vCPU configuration.
    pub cpu_config: AxVCpuConfig,
    /// VM image configuration.
    pub image_config: VMImageConfig,
    pass_through_devices: Vec<HostDeviceAssignment>,
    excluded_devices: Vec<Vec<String>>,
    pass_through_addresses: Vec<HostAddressAssignment>,
    reserved_address_ranges: Vec<ReservedAddressConfig>,
    pass_through_ports: Vec<HostPortAssignment>,
    address_space_policy: AddressSpacePolicy,
    memory_regions: Vec<VmMemConfig>,
    boot_policy: GuestBootPolicy,
    // Physical interrupt sources forwarded to the guest in passthrough mode.
    passthrough_irq_list: Vec<PassthroughInterrupt>,
    serial_profile: GuestSerialProfile,
    serial_firmware_identity: Option<GuestSerialFirmwareIdentity>,
    gic_profile: Option<GuestGicProfile>,
    plic_profile: Option<GuestPlicProfile>,
    timer_profile: Option<GuestTimerProfile>,
    serial_backend_factory: Arc<dyn SerialBackendFactory>,
    virtual_device_requests: Vec<VirtualDeviceRequest>,
    virtual_device_catalog: Arc<crate::ConfiguredDeviceCatalog>,
}

/// Parameters used to build an [`AxVMConfig`].
#[derive(Debug, Default)]
pub struct AxVMConfigParams {
    pub id: usize,
    pub name: String,
    pub phys_cpu_ls: PhysCpuList,
    pub cpu_config: AxVCpuConfig,
    pub image_config: VMImageConfig,
    pub pass_through_devices: Vec<HostDeviceAssignment>,
    pub excluded_devices: Vec<Vec<String>>,
    pub pass_through_addresses: Vec<HostAddressAssignment>,
    pub reserved_address_ranges: Vec<ReservedAddressConfig>,
    pub pass_through_ports: Vec<HostPortAssignment>,
    pub address_space_policy: AddressSpacePolicy,
    pub memory_regions: Vec<VmMemConfig>,
    pub boot_policy: GuestBootPolicy,
    /// Machine-owned virtual serial resources.
    pub serial_profile: Option<GuestSerialProfile>,
    /// App-owned backend factory for the mandatory virtual serial device.
    pub serial_backend_factory: Option<Arc<dyn SerialBackendFactory>>,
    /// Open-ended virtual-device requests parsed from guest configuration.
    pub virtual_device_requests: Vec<VirtualDeviceRequest>,
    /// Code-registered factories available to this VM.
    pub virtual_device_catalog: Option<Arc<crate::ConfiguredDeviceCatalog>>,
}

impl AxVMConfig {
    pub fn new(params: AxVMConfigParams) -> Self {
        let machine = crate::machine::current_machine_profile(params.phys_cpu_ls.cpu_num());
        let serial_profile = params.serial_profile.unwrap_or(machine.serial);
        Self {
            id: params.id,
            name: params.name,
            phys_cpu_ls: params.phys_cpu_ls,
            cpu_config: params.cpu_config,
            image_config: params.image_config,
            pass_through_devices: params.pass_through_devices,
            excluded_devices: params.excluded_devices,
            pass_through_addresses: params.pass_through_addresses,
            reserved_address_ranges: params.reserved_address_ranges,
            pass_through_ports: params.pass_through_ports,
            address_space_policy: params.address_space_policy,
            memory_regions: params.memory_regions,
            boot_policy: params.boot_policy,
            passthrough_irq_list: Vec::new(),
            serial_profile,
            serial_firmware_identity: None,
            gic_profile: machine.gic,
            plic_profile: machine.plic,
            timer_profile: machine.timer,
            serial_backend_factory: params
                .serial_backend_factory
                .unwrap_or_else(|| Arc::new(NullSerialBackendFactory)),
            virtual_device_requests: params.virtual_device_requests,
            virtual_device_catalog: params
                .virtual_device_catalog
                .unwrap_or_else(|| Arc::new(crate::ConfiguredDeviceCatalog::new())),
        }
    }

    #[cfg(test)]
    pub(crate) fn default_for_test(id: usize, name: &str) -> Self {
        Self::new(AxVMConfigParams {
            id,
            name: String::from(name),
            phys_cpu_ls: PhysCpuList::new(1, None, None),
            ..Default::default()
        })
    }

    /// Returns VM id.
    pub fn id(&self) -> usize {
        self.id
    }

    /// Returns VM name.
    pub fn name(&self) -> String {
        self.name.clone()
    }

    /// Returns configurations related to VM image load addresses.
    pub fn image_config(&self) -> &VMImageConfig {
        &self.image_config
    }

    /// Clears the configured DTB load address when no guest DTB is available.
    pub fn clear_dtb_load_gpa(&mut self) {
        self.image_config.dtb_load_gpa = None;
    }

    /// Sets the DTB load address used as an architecture boot argument.
    pub fn set_dtb_load_gpa(&mut self, dtb_load_gpa: GuestPhysAddr) {
        self.image_config.dtb_load_gpa = Some(dtb_load_gpa);
    }

    /// Returns whether VM images are loaded from the host filesystem.
    pub fn images_loaded_from_filesystem(&self) -> bool {
        self.image_config.loaded_from_filesystem
    }

    /// Returns the entry address in GPA for the Bootstrap Processor (BSP).
    pub fn bsp_entry(&self) -> GuestPhysAddr {
        // Retrieves BSP entry from the CPU configuration.
        self.cpu_config.bsp_entry
    }

    /// Returns the entry address in GPA for the Application Processor (AP).
    pub fn ap_entry(&self) -> GuestPhysAddr {
        // Retrieves AP entry from the CPU configuration.
        self.cpu_config.ap_entry
    }

    /// Returns a mutable reference to the physical CPU list.
    pub fn phys_cpu_ls_mut(&mut self) -> &mut PhysCpuList {
        &mut self.phys_cpu_ls
    }

    /// Returns the list of excluded devices.
    pub fn excluded_devices(&self) -> &[Vec<String>] {
        &self.excluded_devices
    }

    /// Adds one physical-device path to the passthrough exclusion set.
    pub fn exclude_device_path(&mut self, path: String) {
        if !self
            .excluded_devices
            .iter()
            .flatten()
            .any(|excluded| excluded == &path)
        {
            self.excluded_devices.push(std::vec![path]);
        }
    }

    /// Returns the list of passthrough address configurations.
    pub fn pass_through_addresses(&self) -> &[HostAddressAssignment] {
        &self.pass_through_addresses
    }

    /// Returns guest address ranges reserved from default passthrough mapping.
    pub fn reserved_address_ranges(&self) -> &[ReservedAddressConfig] {
        &self.reserved_address_ranges
    }

    /// Adds a guest address range reserved from default passthrough mapping.
    pub fn add_reserved_address_range(&mut self, range: ReservedAddressConfig) {
        self.reserved_address_ranges.push(range);
    }

    /// Returns the list of passthrough host I/O port configurations.
    pub fn pass_through_ports(&self) -> &[HostPortAssignment] {
        &self.pass_through_ports
    }

    /// Returns the guest physical address space population policy.
    pub fn address_space_policy(&self) -> AddressSpacePolicy {
        self.address_space_policy
    }

    /// Returns configurations related to VM memory regions.
    pub fn memory_regions(&self) -> &[VmMemConfig] {
        &self.memory_regions
    }

    /// Replaces configurations related to VM memory regions.
    pub fn set_memory_regions(&mut self, memory_regions: Vec<VmMemConfig>) {
        self.memory_regions = memory_regions;
    }

    /// Returns the policy used to adjust runtime boot image addresses.
    pub fn boot_policy(&self) -> GuestBootPolicy {
        self.boot_policy
    }

    /// Sets the policy used to adjust runtime boot image addresses.
    pub fn set_boot_policy(&mut self, boot_policy: GuestBootPolicy) {
        self.boot_policy = boot_policy;
    }

    /// Returns configurations related to VM passthrough devices.
    pub fn pass_through_devices(&self) -> &[HostDeviceAssignment] {
        &self.pass_through_devices
    }

    /// Adds a new passthrough device to the VM configuration.
    pub fn add_pass_through_device(&mut self, device: HostDeviceAssignment) {
        self.pass_through_devices.push(device);
    }

    /// Removes passthrough device from the VM configuration.
    pub fn remove_pass_through_device(&mut self, device: HostDeviceAssignment) {
        self.pass_through_devices.retain(|d| d != &device);
    }

    /// Clears all passthrough devices from the VM configuration.
    pub fn clear_pass_through_devices(&mut self) {
        self.pass_through_devices.clear();
    }

    /// Adds a physical interrupt source forwarded to the guest.
    pub fn add_pass_through_irq(&mut self, source: u32, trigger: InterruptTriggerMode) {
        let route = PassthroughInterrupt { source, trigger };
        if let Some(existing) = self
            .passthrough_irq_list
            .iter_mut()
            .find(|existing| existing.source == source)
        {
            *existing = route;
        } else {
            self.passthrough_irq_list.push(route);
        }
    }

    /// Returns the physical interrupt sources forwarded to the guest.
    pub fn pass_through_irqs(&self) -> &[PassthroughInterrupt] {
        &self.passthrough_irq_list
    }

    /// Returns whether the guest address space starts from host identity mappings.
    pub fn uses_passthrough_address_space(&self) -> bool {
        self.address_space_policy == AddressSpacePolicy::Passthrough
    }

    /// Returns the machine-owned virtual serial resources.
    pub(crate) const fn serial_profile(&self) -> GuestSerialProfile {
        self.serial_profile
    }

    /// Replaces the machine serial resources and firmware identity atomically.
    pub fn replace_machine_serial(
        &mut self,
        profile: GuestSerialProfile,
        identity: Option<GuestSerialFirmwareIdentity>,
    ) -> crate::AxVmResult {
        self.serial_profile = profile;
        self.serial_firmware_identity = identity;
        Ok(())
    }

    /// Returns firmware identity retained for the virtual serial node.
    pub fn serial_firmware_identity(&self) -> Option<&GuestSerialFirmwareIdentity> {
        self.serial_firmware_identity.as_ref()
    }

    /// Replaces the virtual GIC windows with host firmware resources.
    pub fn replace_machine_gic(&mut self, profile: GuestGicProfile) -> crate::AxVmResult {
        if self.gic_profile.is_none() {
            return Err(crate::AxVmError::invalid_config(
                "the selected machine has no AArch64 GIC",
            ));
        }
        let cpu_num = self.phys_cpu_ls.cpu_num().max(1);
        self.gic_profile = Some(profile.normalized_for_vcpus(cpu_num)?);
        Ok(())
    }

    /// Returns host firmware resources retained by the virtual GIC.
    pub fn gic_profile(&self) -> Option<&GuestGicProfile> {
        self.gic_profile.as_ref()
    }

    /// Replaces the AArch64 architectural timer resources with validated host firmware data.
    pub fn replace_machine_timer(&mut self, profile: GuestTimerProfile) -> crate::AxVmResult {
        if self.timer_profile.is_none() {
            return Err(crate::AxVmError::invalid_config(
                "the selected machine has no AArch64 architectural timer",
            ));
        }
        profile
            .validated_intids()
            .map_err(crate::AxVmError::invalid_config)?;
        self.timer_profile = Some(profile);
        Ok(())
    }

    /// Returns the machine-owned AArch64 architectural timer resources.
    pub fn timer_profile(&self) -> Option<&GuestTimerProfile> {
        self.timer_profile.as_ref()
    }

    /// Replaces the virtual PLIC window with host firmware resources.
    pub fn replace_machine_plic(&mut self, profile: GuestPlicProfile) -> crate::AxVmResult {
        if self.plic_profile.is_none() {
            return Err(crate::AxVmError::invalid_config(
                "the selected machine has no RISC-V PLIC",
            ));
        }
        profile.validate_for_vcpus(self.phys_cpu_ls.cpu_num())?;
        self.plic_profile = Some(profile);
        Ok(())
    }

    /// Returns host firmware resources retained by the virtual PLIC.
    pub fn plic_profile(&self) -> Option<&GuestPlicProfile> {
        self.plic_profile.as_ref()
    }

    /// Returns the factory that creates a backend for each virtual UART graph.
    pub fn serial_backend_factory(&self) -> Arc<dyn SerialBackendFactory> {
        self.serial_backend_factory.clone()
    }

    pub(crate) fn virtual_device_requests(&self) -> &[VirtualDeviceRequest] {
        &self.virtual_device_requests
    }

    pub(crate) fn virtual_device_catalog(&self) -> &crate::ConfiguredDeviceCatalog {
        &self.virtual_device_catalog
    }

    /// Relocate the guest kernel image while preserving the configured
    /// entry-point offsets relative to the load address.
    pub fn relocate_kernel_image(&mut self, kernel_load_gpa: GuestPhysAddr) {
        let old_load = self.image_config.kernel_load_gpa.as_usize();
        let new_load = kernel_load_gpa.as_usize();

        let bsp_offset = self
            .cpu_config
            .bsp_entry
            .as_usize()
            .checked_sub(old_load)
            .expect("BSP entry must not be below kernel load address");
        let ap_offset = self
            .cpu_config
            .ap_entry
            .as_usize()
            .checked_sub(old_load)
            .expect("AP entry must not be below kernel load address");

        self.image_config.kernel_load_gpa = kernel_load_gpa;
        self.cpu_config.bsp_entry = GuestPhysAddr::from(new_load + bsp_offset);
        self.cpu_config.ap_entry = GuestPhysAddr::from(new_load + ap_offset);
    }
}

impl Default for AxVMConfig {
    fn default() -> Self {
        Self::new(AxVMConfigParams::default())
    }
}

/// Represents the list of physical CPUs available for the VM.
#[derive(Debug, Default, Clone)]
pub struct PhysCpuList {
    cpu_num: usize,
    phys_cpu_ids: Option<Vec<usize>>,
    phys_cpu_sets: Option<Vec<usize>>,
}

impl PhysCpuList {
    /// Creates a physical CPU list.
    pub fn new(
        cpu_num: usize,
        phys_cpu_ids: Option<Vec<usize>>,
        phys_cpu_sets: Option<Vec<usize>>,
    ) -> Self {
        Self {
            cpu_num,
            phys_cpu_ids,
            phys_cpu_sets,
        }
    }

    /// Returns vCpu id list and its corresponding pCpu affinity list, as well as its physical id.
    /// If the pCpu affinity is None, it means the vCpu will be allocated to any available pCpu randomly.
    /// if the pCPU id is not provided, the vCpu's physical id will be set as vCpu id.
    ///
    /// Returns a vector of tuples, each tuple contains:
    /// - The vCpu id.
    /// - The pCpu affinity mask, `None` if not set.
    /// - The physical id of the vCpu, equal to vCpu id if not provided.
    pub fn get_vcpu_affinities_pcpu_ids(&self) -> Vec<(usize, Option<usize>, usize)> {
        if let Some(phys_cpu_ids) = &self.phys_cpu_ids
            && self.cpu_num != phys_cpu_ids.len()
        {
            error!(
                "ERROR!!!: cpu_num: {}, phys_cpu_ids: {:?}",
                self.cpu_num, self.phys_cpu_ids
            );
        }
        CurrentArch::vcpu_affinities(
            self.cpu_num,
            self.phys_cpu_ids.as_deref(),
            self.phys_cpu_sets.as_deref(),
        )
    }

    /// Returns the number of CPUs.
    pub fn cpu_num(&self) -> usize {
        self.cpu_num
    }

    /// Returns the physical CPU IDs.
    pub fn phys_cpu_ids(&self) -> &Option<Vec<usize>> {
        &self.phys_cpu_ids
    }

    /// Returns the physical CPU sets.
    pub fn phys_cpu_sets(&self) -> &Option<Vec<usize>> {
        &self.phys_cpu_sets
    }

    /// Sets the guest CPU sets.
    pub fn set_guest_cpu_sets(&mut self, phys_cpu_sets: Vec<usize>) {
        self.phys_cpu_sets = Some(phys_cpu_sets);
    }

    /// Sets the CPU IDs exposed to the guest.
    pub fn set_guest_phys_cpu_ids(&mut self, phys_cpu_ids: Vec<usize>) {
        self.phys_cpu_ids = Some(phys_cpu_ids);
    }
}

#[cfg(test)]
mod tests {
    use std::vec;

    use super::*;

    fn memory_region(gpa: usize, size: usize, map_type: VmMemMappingType) -> VmMemConfig {
        VmMemConfig {
            gpa,
            size,
            flags: 0x7,
            map_type,
        }
    }

    #[test]
    fn set_memory_regions_replaces_stale_snapshot_after_config_enrichment() {
        let main_memory = memory_region(0x8000_0000, 0x200000, VmMemMappingType::MapIdentical);
        let reserved_memory = memory_region(0x110000, 0x10000, VmMemMappingType::MapReserved);
        let mut config = AxVMConfig::default_for_test(1, "linux");

        config.set_memory_regions(vec![main_memory.clone()]);
        assert_eq!(config.memory_regions().len(), 1);

        config.set_memory_regions(vec![main_memory, reserved_memory]);

        let regions = config.memory_regions();
        assert_eq!(regions.len(), 2);
        assert_eq!(regions[1].gpa, 0x110000);
        assert_eq!(regions[1].size, 0x10000);
        assert_eq!(regions[1].map_type, VmMemMappingType::MapReserved);
    }

    #[cfg(target_arch = "x86_64")]
    #[test]
    fn controller_replacements_require_machine_capabilities() {
        let mut config = AxVMConfig::new(AxVMConfigParams {
            phys_cpu_ls: PhysCpuList::new(1, None, None),
            ..Default::default()
        });
        let gic = GuestGicProfile {
            compatible: "arm,gic-400".into(),
            node_path: "/interrupt-controller".into(),
            node_phandle: None,
            distributor: crate::machine::GuestMmioRegion {
                base: 0x1000,
                length: 0x1000,
            },
            cpu_region: crate::machine::GuestGicCpuRegion::CpuInterface(
                crate::machine::GuestMmioRegion {
                    base: 0x2000,
                    length: 0x2000,
                },
            ),
            its: Vec::new(),
        };
        let plic = GuestPlicProfile {
            node_path: "/plic".into(),
            node_phandle: None,
            base: 0x0c00_0000,
            length: 0x60_0000,
        };

        assert!(matches!(
            config.replace_machine_gic(gic),
            Err(crate::AxVmError::InvalidConfig { .. })
        ));
        assert!(matches!(
            config.replace_machine_plic(plic),
            Err(crate::AxVmError::InvalidConfig { .. })
        ));
    }
}