Repository: firecracker-microvm/firecracker
File: src/vmm/src/vstate/kvm.rs (Rust)

Code before:

#[cfg(test)]
pub(crate) mod tests {
    use super::*;

    #[test]
    fn test_combine_capabilities() {
        // Default caps for x86_64 and aarch64 both have KVM_CAP_IOEVENTFD and don't have
        // KVM_CAP_IOMMU caps.
        let additional_capabilities = vec![
            KvmCapability::Add(kvm_bindings::KVM_CAP_IOMMU),
            KvmCapability::Remove(kvm_bindings::KVM_CAP_IOEVENTFD),
        ];

        let combined_caps = Kvm::combine_capabilities(&additional_capabilities);
        assert!(combined_caps
            .iter()
            .any(|c| *c == kvm_bindings::KVM_CAP_IOMMU));
        assert!(!combined_caps
            .iter()
            .any(|c| *c == kvm_bindings::KVM_CAP_IOEVENTFD));
    }

    #[cfg(target_arch = "x86_64")]
    #[test]
    fn test_request_xstate_feature_permission() {
        // Test request_xstate_feature_permission() for Intel AMX.
        Kvm::request_xstate_feature_permission(arch_prctl::ARCH_XCOMP_TILEDATA).unwrap();

        let mut supported_xfeatures: libc::c_ulong = 0;
        // SAFETY: Safe because the third input (`addr`) is a valid `c_ulong` pointer.
        match SyscallReturnCode(unsafe {
            libc::syscall(
                libc::SYS_arch_prctl,
                arch_prctl::ARCH_GET_XCOMP_SUPP,
                &mut supported_xfeatures as *mut libc::c_ulong,
            )
        })
        .into_empty_result()
        {
            Ok(_) => {} // Continue this test
            Err(err) if err.raw_os_error() == Some(libc::EINVAL) => {
                // Dynamic XSTATE feature enabling is not supported in the first place, so nothing
                // to test on this kernel version.
                return;
            }
            Err(err) => panic!("Unexpected error: {}", err),
        };

        // If Intel AMX is not supported, nothing to test on this CPU.
        let intel_amx_feature_mask: libc::c_ulong = 1u64 << arch_prctl::ARCH_XCOMP_TILEDATA;

Review comment (security):
Sorry for second review, I had this written out and then GitHub apparently deleted it when I pressed "submit review" >.>

There's a lot of duplication between this test and the `request_xstate_feature_permission` function, and admittedly, I don't quite follow how this test is actually testing that function beyond the first line (which I can see as a sort of test about "this function never returns an error even on hardware/kernels that doesnt support AMX). But everything after that is just a copy-paste of the function itself :thinking:. I honestly think we can either remove everything after ln 301, or try to make it somehow more connected to the behavior of `request_xstate_feature_permission`, e.g. by having it return `Result<bool>` (with `Ok(true)` indicating that the feature was successfully enabled, and `Ok(false)` indicating that the feature wasn't supported, e.g. because of running on 5.10 or on hardware without AMX). Then the code below could assert that the return value of  `request_xstate_feature_permission` actually matches what the prctls return (and for the prctrls themselves, it'd probably be nice to add some helper functions, such as `fn supported_xsave_features() -> Result<c_ulong>` and `enabled_xsave_feature(feature: c_ulong) -> Result<()>` to avoid repeating all the EINVAL handling between function and test)

Code after:
// Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use kvm_bindings::KVM_API_VERSION;
#[cfg(target_arch = "x86_64")]
use kvm_bindings::{CpuId, MsrList, KVM_MAX_CPUID_ENTRIES};
use kvm_ioctls::Kvm as KvmFd;
use serde::{Deserialize, Serialize};

#[cfg(target_arch = "x86_64")]
use crate::arch::x86_64::xstate::{request_dynamic_xstate_features, XstateError};
use crate::cpu_config::templates::KvmCapability;
use crate::vstate::memory::{GuestMemory, GuestMemoryMmap};

/// Errors associated with the wrappers over KVM ioctls.
/// Needs `rustfmt::skip` to make multiline comments work
#[rustfmt::skip]
#[derive(Debug, thiserror::Error, displaydoc::Display)]
pub enum KvmError {
    /// The host kernel reports an invalid KVM API version: {0}
    ApiVersion(i32),
    /// Missing KVM capabilities: {0:#x?}
    Capabilities(u32),
    /**  Error creating KVM object: {0} Make sure the user launching the firecracker process is \
    configured on the /dev/kvm file's ACL. */
    Kvm(kvm_ioctls::Error),
    #[cfg(target_arch = "x86_64")]
    /// Failed to get supported cpuid: {0}
    GetSupportedCpuId(kvm_ioctls::Error),
    /// The number of configured slots is bigger than the maximum reported by KVM
    NotEnoughMemorySlots,
    #[cfg(target_arch = "x86_64")]
    /// Failed to request permission for dynamic XSTATE features: {0}
    XstateFeatures(XstateError),
}

/// Struct with kvm fd and kvm associated paramenters.
#[derive(Debug)]
pub struct Kvm {
    /// KVM fd.
    pub fd: KvmFd,
    /// Maximum number of memory slots allowed by KVM.
    pub max_memslots: usize,
    /// Additional capabilities that were specified in cpu template.
    pub kvm_cap_modifiers: Vec<KvmCapability>,

    #[cfg(target_arch = "x86_64")]
    /// Supported CpuIds.
    pub supported_cpuid: CpuId,
}

impl Kvm {
    /// Create `Kvm` struct.
    pub fn new(kvm_cap_modifiers: Vec<KvmCapability>) -> Result<Self, KvmError> {
        let kvm_fd = KvmFd::new().map_err(KvmError::Kvm)?;

        // Check that KVM has the correct version.
        // Safe to cast because this is a constant.
        #[allow(clippy::cast_possible_wrap)]
        if kvm_fd.get_api_version() != KVM_API_VERSION as i32 {
            return Err(KvmError::ApiVersion(kvm_fd.get_api_version()));
        }

        let total_caps = Self::combine_capabilities(&kvm_cap_modifiers);
        // Check that all desired capabilities are supported.
        Self::check_capabilities(&kvm_fd, &total_caps).map_err(KvmError::Capabilities)?;

        let max_memslots = kvm_fd.get_nr_memslots();

        #[cfg(target_arch = "aarch64")]
        {
            Ok(Self {
                fd: kvm_fd,
                max_memslots,
                kvm_cap_modifiers,
            })
        }

        #[cfg(target_arch = "x86_64")]
        {
            request_dynamic_xstate_features().map_err(KvmError::XstateFeatures)?;

            let supported_cpuid = kvm_fd
                .get_supported_cpuid(KVM_MAX_CPUID_ENTRIES)
                .map_err(KvmError::GetSupportedCpuId)?;

            Ok(Kvm {
                fd: kvm_fd,
                max_memslots,
                kvm_cap_modifiers,
                supported_cpuid,
            })
        }
    }

    /// Msrs needed to be saved on snapshot creation.
    #[cfg(target_arch = "x86_64")]
    pub fn msrs_to_save(&self) -> Result<MsrList, crate::arch::x86_64::msr::MsrError> {
        crate::arch::x86_64::msr::get_msrs_to_save(&self.fd)
    }

    /// Check guest memory does not have more regions than kvm allows.
    pub fn check_memory(&self, guest_mem: &GuestMemoryMmap) -> Result<(), KvmError> {
        if guest_mem.num_regions() > self.max_memslots {
            Err(KvmError::NotEnoughMemorySlots)
        } else {
            Ok(())
        }
    }

    fn combine_capabilities(kvm_cap_modifiers: &[KvmCapability]) -> Vec<u32> {
        let mut total_caps = Self::DEFAULT_CAPABILITIES.to_vec();
        for modifier in kvm_cap_modifiers.iter() {
            match modifier {
                KvmCapability::Add(cap) => {
                    if !total_caps.contains(cap) {
                        total_caps.push(*cap);
                    }
                }
                KvmCapability::Remove(cap) => {
                    if let Some(pos) = total_caps.iter().position(|c| c == cap) {
                        total_caps.swap_remove(pos);
                    }
                }
            }
        }
        total_caps
    }

    fn check_capabilities(kvm_fd: &KvmFd, capabilities: &[u32]) -> Result<(), u32> {
        for cap in capabilities {
            // If capability is not supported kernel will return 0.
            if kvm_fd.check_extension_raw(u64::from(*cap)) == 0 {
                return Err(*cap);
            }
        }
        Ok(())
    }

    /// Saves and returns the Kvm state.
    pub fn save_state(&self) -> KvmState {
        KvmState {
            kvm_cap_modifiers: self.kvm_cap_modifiers.clone(),
        }
    }
}
#[cfg(target_arch = "aarch64")]
/// Optional capabilities.
#[derive(Debug, Default)]
pub struct OptionalCapabilities {
    /// KVM_CAP_COUNTER_OFFSET
    pub counter_offset: bool,
}
#[cfg(target_arch = "aarch64")]
impl Kvm {
    const DEFAULT_CAPABILITIES: [u32; 7] = [
        kvm_bindings::KVM_CAP_IOEVENTFD,
        kvm_bindings::KVM_CAP_IRQFD,
        kvm_bindings::KVM_CAP_USER_MEMORY,
        kvm_bindings::KVM_CAP_ARM_PSCI_0_2,
        kvm_bindings::KVM_CAP_DEVICE_CTRL,
        kvm_bindings::KVM_CAP_MP_STATE,
        kvm_bindings::KVM_CAP_ONE_REG,
    ];

    /// Returns struct with optional capabilities statuses.
    pub fn optional_capabilities(&self) -> OptionalCapabilities {
        OptionalCapabilities {
            counter_offset: self
                .fd
                .check_extension_raw(kvm_bindings::KVM_CAP_COUNTER_OFFSET.into())
                != 0,
        }
    }
}

#[cfg(target_arch = "x86_64")]
impl Kvm {
    const DEFAULT_CAPABILITIES: [u32; 14] = [
        kvm_bindings::KVM_CAP_IRQCHIP,
        kvm_bindings::KVM_CAP_IOEVENTFD,
        kvm_bindings::KVM_CAP_IRQFD,
        kvm_bindings::KVM_CAP_USER_MEMORY,
        kvm_bindings::KVM_CAP_SET_TSS_ADDR,
        kvm_bindings::KVM_CAP_PIT2,
        kvm_bindings::KVM_CAP_PIT_STATE2,
        kvm_bindings::KVM_CAP_ADJUST_CLOCK,
        kvm_bindings::KVM_CAP_DEBUGREGS,
        kvm_bindings::KVM_CAP_MP_STATE,
        kvm_bindings::KVM_CAP_VCPU_EVENTS,
        kvm_bindings::KVM_CAP_XCRS,
        kvm_bindings::KVM_CAP_XSAVE,
        kvm_bindings::KVM_CAP_EXT_CPUID,
    ];
}

/// Structure holding an general specific VM state.
#[derive(Debug, Default, Serialize, Deserialize)]
pub struct KvmState {
    /// Additional capabilities that were specified in cpu template.
    pub kvm_cap_modifiers: Vec<KvmCapability>,
}

#[cfg(test)]
pub(crate) mod tests {
    use super::*;

    #[test]
    fn test_combine_capabilities() {
        // Default caps for x86_64 and aarch64 both have KVM_CAP_IOEVENTFD and don't have
        // KVM_CAP_IOMMU caps.
        let additional_capabilities = vec![
            KvmCapability::Add(kvm_bindings::KVM_CAP_IOMMU),
            KvmCapability::Remove(kvm_bindings::KVM_CAP_IOEVENTFD),
        ];

        let combined_caps = Kvm::combine_capabilities(&additional_capabilities);
        assert!(combined_caps
            .iter()
            .any(|c| *c == kvm_bindings::KVM_CAP_IOMMU));
        assert!(!combined_caps
            .iter()
            .any(|c| *c == kvm_bindings::KVM_CAP_IOEVENTFD));
    }
}