refactor: move arch specific KVM code into modules

Follow up of the refactoring we did in the past around vstate/vm.rs,
where we separate architecture specific code in modules, keeping only
the arch-independent code in 'kvm/mod.rs'.

No functional changes intended.

Signed-off-by: Babis Chalios <[email protected]>

Author: bchalios

src/vmm/src/vstate/kvm/aarch64.rs

@@ -0,0 +1,66 @@
+// Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+use std::convert::Infallible;
+
+use kvm_ioctls::Kvm as KvmFd;
+
+use crate::cpu_config::templates::KvmCapability;
+
+/// ['Kvm'] initialization can't fail for Aarch64
+pub type KvmArchError = Infallible;
+
+/// Errors associated with the wrappers over KVM ioctls.
+/// Needs `rustfmt::skip` to make multiline comments work
+/// Optional capabilities.
+#[derive(Debug, Default)]
+pub struct OptionalCapabilities {
+    /// KVM_CAP_COUNTER_OFFSET
+    pub counter_offset: bool,
+}
+
+/// Struct with kvm fd and kvm associated parameters.
+#[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>,
+}
+
+impl Kvm {
+    pub(crate) 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,
+    ];
+
+    /// Initialize [`Kvm`] type for Aarch64 architecture
+    pub fn init_arch(
+        fd: KvmFd,
+        max_memslots: usize,
+        kvm_cap_modifiers: Vec<KvmCapability>,
+    ) -> Result<Self, KvmArchError> {
+        Ok(Self {
+            fd,
+            max_memslots,
+            kvm_cap_modifiers,
+        })
+    }
+
+    /// 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,
+        }
+    }
+}

src/vmm/src/vstate/kvm.rs renamed to src/vmm/src/vstate/kvm/mod.rs

@@ -2,16 +2,22 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use kvm_bindings::KVM_API_VERSION;
-#[cfg(target_arch = "x86_64")]
-use kvm_bindings::{CpuId, KVM_MAX_CPUID_ENTRIES, MsrList};
 use kvm_ioctls::Kvm as KvmFd;
 use serde::{Deserialize, Serialize};
 
-#[cfg(target_arch = "x86_64")]
-use crate::arch::x86_64::xstate::{XstateError, request_dynamic_xstate_features};
 use crate::cpu_config::templates::KvmCapability;
 use crate::vstate::memory::{GuestMemory, GuestMemoryMmap};
 
+#[cfg(target_arch = "aarch64")]
+mod aarch64;
+#[cfg(target_arch = "x86_64")]
+mod x86_64;
+
+#[cfg(target_arch = "aarch64")]
+pub use aarch64::{Kvm, KvmArchError, OptionalCapabilities};
+#[cfg(target_arch = "x86_64")]
+pub use x86_64::{Kvm, KvmArchError};
+
 /// Errors associated with the wrappers over KVM ioctls.
 /// Needs `rustfmt::skip` to make multiline comments work
 #[rustfmt::skip]
@@ -24,29 +30,10 @@ pub enum KvmError {
     /**  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,
+    /// Architecture specific error: {0}
+    ArchError(#[from] KvmArchError)
 }
 
 impl Kvm {
@@ -67,36 +54,7 @@ impl Kvm {
 
         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)
+        Ok(Kvm::init_arch(kvm_fd, max_memslots, kvm_cap_modifiers)?)
     }
 
     /// Check guest memory does not have more regions than kvm allows.
@@ -144,55 +102,6 @@ impl Kvm {
         }
     }
 }
-#[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)]

src/vmm/src/vstate/kvm/x86_64.rs

@@ -0,0 +1,74 @@
+// Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+use kvm_bindings::{CpuId, KVM_MAX_CPUID_ENTRIES, MsrList};
+use kvm_ioctls::Kvm as KvmFd;
+
+use crate::arch::x86_64::xstate::{XstateError, request_dynamic_xstate_features};
+use crate::cpu_config::templates::KvmCapability;
+
+/// Architecture specific error for KVM initialization
+#[derive(Debug, thiserror::Error, displaydoc::Display)]
+pub enum KvmArchError {
+    /// Failed to get supported cpuid: {0}
+    GetSupportedCpuId(kvm_ioctls::Error),
+    /// Failed to request permission for dynamic XSTATE features: {0}
+    XstateFeatures(XstateError),
+}
+
+/// Struct with kvm fd and kvm associated parameters.
+#[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>,
+    /// Supported CpuIds.
+    pub supported_cpuid: CpuId,
+}
+
+impl Kvm {
+    pub(crate) 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,
+    ];
+
+    /// Initialize [`Kvm`] type for x86_64 architecture
+    pub fn init_arch(
+        fd: KvmFd,
+        max_memslots: usize,
+        kvm_cap_modifiers: Vec<KvmCapability>,
+    ) -> Result<Self, KvmArchError> {
+        request_dynamic_xstate_features().map_err(KvmArchError::XstateFeatures)?;
+
+        let supported_cpuid = fd
+            .get_supported_cpuid(KVM_MAX_CPUID_ENTRIES)
+            .map_err(KvmArchError::GetSupportedCpuId)?;
+
+        Ok(Kvm {
+            fd,
+            max_memslots,
+            kvm_cap_modifiers,
+            supported_cpuid,
+        })
+    }
+
+    /// Msrs needed to be saved on snapshot creation.
+    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)
+    }
+}