refactor(vmm): 2-step CpuConfiguration build for aarch64

As done in x86_64, build a custom `CpuConfiguration` with 2 steps:
1. Construct a base `CpuConfiguration`.
2. Apply CPU template to the default `CpuConfiguration`.

By providing the same interface between x86_64 and aarch64, it becomes
possible to unify some operations.

Signed-off-by: Takahiro Itazuri <[email protected]>
Co-authored-by: Egor Lazarchuk <[email protected]>

Author: zulinx86

src/vmm/src/arch/aarch64/regs.rs

@@ -51,6 +51,9 @@ pub enum Error {
     /// Failed to get a register value.
     #[error("Failed to get register {0}: {1}")]
     GetOneReg(u64, kvm_ioctls::Error),
+    /// Failed to set a register value.
+    #[error("Failed to set register {0}: {1}")]
+    SetOneReg(u64, kvm_ioctls::Error),
     /// Failed to get the register list.
     #[error("Failed to retrieve list of registers: {0}")]
     GetRegList(kvm_ioctls::Error),

src/vmm/src/builder.rs

@@ -840,9 +840,17 @@ pub fn configure_system_for_boot(
     }
     #[cfg(target_arch = "aarch64")]
     {
-        let cpu_config = crate::guest_config::aarch64::CpuConfiguration::new(
-            &vcpus[0].kvm_vcpu.fd,
-            &vm_config.cpu_template,
+        // Construct the base CpuConfiguration to apply CPU template onto.
+        let regs = vcpus[0]
+            .kvm_vcpu
+            .get_regs(&cpu_template.reg_list())
+            .map_err(GuestConfigError)?;
+        let cpu_config = crate::guest_config::aarch64::CpuConfiguration { regs };
+
+        // Apply CPU template to the base CpuConfiguration.
+        let cpu_config = crate::guest_config::aarch64::CpuConfiguration::apply_template(
+            cpu_config,
+            &cpu_template,
         )?;
 
         let vcpu_config = VcpuConfig {

src/vmm/src/guest_config/aarch64/mod.rs

@@ -4,16 +4,15 @@
 /// Module with CPU templates for aarch64
 pub mod static_cpu_templates;
 
-use kvm_ioctls::VcpuFd;
 pub use static_cpu_templates::*;
 
-use super::templates::{CpuTemplateType, CustomCpuTemplate};
+use super::templates::CustomCpuTemplate;
 use crate::arch::regs::{Aarch64Register, Error as ArchError};
 
 /// Errors thrown while configuring templates.
 #[derive(Debug, PartialEq, Eq, thiserror::Error)]
 #[error("Failed to create a guest cpu configuration: {0}")]
-pub struct Error(#[from] ArchError);
+pub struct Error(#[from] pub ArchError);
 
 /// CPU configuration for aarch64
 #[derive(Debug, Default, Clone, PartialEq, Eq)]
@@ -24,56 +23,16 @@ pub struct CpuConfiguration {
 
 impl CpuConfiguration {
     /// Creates new guest CPU config based on the provided template
-    pub fn new(vcpu: &VcpuFd, template: &Option<CpuTemplateType>) -> Result<Self, Error> {
-        match template {
-            Some(ref cpu_template) => Self::with_template(vcpu, cpu_template),
-            None => Ok(Self::default()),
+    pub fn apply_template(mut self, template: &CustomCpuTemplate) -> Result<Self, Error> {
+        for (modifier, reg) in template.reg_modifiers.iter().zip(self.regs.iter_mut()) {
+            reg.value = modifier.bitmap.apply(reg.value);
         }
-    }
-
-    /// Creates new guest CPU config based on the provided template
-    fn with_template(vcpu: &VcpuFd, template: &CpuTemplateType) -> Result<Self, Error> {
-        match template {
-            CpuTemplateType::Custom(template) => Self::with_applied_template(vcpu, template),
-            CpuTemplateType::Static(StaticCpuTemplate::V1N1) => {
-                let template = static_cpu_templates::v1n1::v1n1();
-                Self::with_applied_template(vcpu, &template)
-            }
-            _ => unreachable!("Options other than V1N1 are invalid"),
-        }
-    }
-
-    /// Creates new guest CPU config based on the provided template
-    fn with_applied_template(vcpu: &VcpuFd, template: &CustomCpuTemplate) -> Result<Self, Error> {
-        let regs = template
-            .reg_modifiers
-            .iter()
-            .map(|modifier| {
-                vcpu.get_one_reg(modifier.addr)
-                    .map(|value| Aarch64Register {
-                        id: modifier.addr,
-                        value: modifier.bitmap.apply(value),
-                    })
-                    .map_err(|e| Error(ArchError::GetSysRegister(modifier.addr, e)))
-            })
-            .collect::<Result<Vec<_>, Error>>()?;
-
-        Ok(Self { regs })
+        Ok(self)
     }
 
     /// Returns ids of registers that are changed
     /// by this template
     pub fn register_ids(&self) -> Vec<u64> {
         self.regs.iter().map(|reg| reg.id).collect()
     }
-
-    /// Applies cpu template to vcpu
-    /// Used inside Vcpu to configure it
-    pub fn apply(&self, vcpu: &VcpuFd) -> Result<(), ArchError> {
-        for Aarch64Register { id, value } in self.regs.iter() {
-            vcpu.set_one_reg(*id, *value)
-                .map_err(|error| ArchError::SetSysRegister(*id, error))?;
-        }
-        Ok(())
-    }
 }

src/vmm/src/guest_config/templates/aarch64.rs

@@ -168,6 +168,16 @@ where
     Ok(deserialized_number)
 }
 
+impl CustomCpuTemplate {
+    /// Get a list of register IDs that are modified by the CPU template.
+    pub fn reg_list(&self) -> Vec<u64> {
+        self.reg_modifiers
+            .iter()
+            .map(|modifier| modifier.addr)
+            .collect()
+    }
+}
+
 // TODO mark with #[cfg(test)] when we combine all crates into
 // one firecracker crate
 impl CustomCpuTemplate {

src/vmm/src/vstate/vcpu/aarch64.rs

@@ -14,10 +14,9 @@ use versionize_derive::Versionize;
 use vm_memory::{Address, GuestAddress, GuestMemoryMmap};
 
 use crate::arch::aarch64::regs::Aarch64Register;
-use crate::arch::regs::Error as ArchError;
 use crate::arch::regs::{
     get_mpstate, read_mpidr, restore_registers, save_core_registers, save_registers,
-    save_system_registers, set_mpstate,
+    save_system_registers, set_mpstate, Error as ArchError,
 };
 use crate::vcpu::VcpuConfig;
 use crate::vstate::vcpu::VcpuEmulation;
@@ -27,19 +26,19 @@ use crate::vstate::vm::Vm;
 #[derive(Debug, thiserror::Error)]
 pub enum Error {
     #[error("Error configuring the vcpu registers: {0}")]
-    ConfigureRegisters(crate::arch::aarch64::regs::Error),
+    ConfigureRegisters(ArchError),
     #[error("Error creating vcpu: {0}")]
     CreateVcpu(kvm_ioctls::Error),
     #[error("Error getting the vcpu preferred target: {0}")]
     GetPreferredTarget(kvm_ioctls::Error),
     #[error("Error initializing the vcpu: {0}")]
     Init(kvm_ioctls::Error),
-    #[error("Error applying template to the vcpu: {0}")]
-    ApplyCpuTemplate(crate::arch::aarch64::regs::Error),
+    #[error("Error applying template: {0}")]
+    ApplyCpuTemplate(ArchError),
     #[error("Failed to restore the state of the vcpu: {0}")]
-    RestoreState(crate::arch::aarch64::regs::Error),
+    RestoreState(ArchError),
     #[error("Failed to save the state of the vcpu: {0}")]
-    SaveState(crate::arch::aarch64::regs::Error),
+    SaveState(ArchError),
 }
 
 type Result<T> = result::Result<T, Error>;
@@ -97,6 +96,13 @@ impl KvmVcpu {
         kernel_load_addr: GuestAddress,
         vcpu_config: &VcpuConfig,
     ) -> std::result::Result<(), Error> {
+        for Aarch64Register { id, value } in vcpu_config.cpu_config.regs.iter() {
+            self.fd
+                .set_one_reg(*id, *value)
+                .map_err(|err| Error::ApplyCpuTemplate(ArchError::SetOneReg(*id, err)))?;
+        }
+        self.additional_register_ids = vcpu_config.cpu_config.register_ids();
+
         crate::arch::aarch64::regs::setup_boot_regs(
             &self.fd,
             self.index,
@@ -108,12 +114,6 @@ impl KvmVcpu {
         self.mpidr =
             crate::arch::aarch64::regs::read_mpidr(&self.fd).map_err(Error::ConfigureRegisters)?;
 
-        vcpu_config
-            .cpu_config
-            .apply(&self.fd)
-            .map_err(Error::ApplyCpuTemplate)?;
-        self.additional_register_ids = vcpu_config.cpu_config.register_ids();
-
         Ok(())
     }