fix aarch64 not able to restore snaps when SVE enabled

CHANGELOG.md

@@ -36,6 +36,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   Now cpu-template-helper will print warnings if it encounters SVE registers
   during the conversion process. This is because cpu templates are limited
   to only modify registers less than 128 bits.
+- [#4413](https://github.com/firecracker-microvm/firecracker/pull/4413):
+  Fixed a bug in the Firecracker that prevented it to restore snapshots of
+  VMs that had SVE enabled.
 - [#4414](https://github.com/firecracker-microvm/firecracker/pull/4360):
   Made `PATCH` requests to the `/machine-config` endpoint transactional, meaning
   Firecracker's configuration will be unchanged if the request returns an error.

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

@@ -109,6 +109,12 @@ arm64_sys_reg!(ID_AA64MMFR2_EL1, 3, 0, 0, 7, 2);
 // EL0 Virtual Timer Registers
 arm64_sys_reg!(KVM_REG_ARM_TIMER_CNT, 3, 3, 14, 3, 2);
 
+/// Vector lengths pseudo-register
+/// TODO: this can be removed after https://github.com/rust-vmm/kvm-bindings/pull/89
+/// is merged and new version is used in Firecracker.
+pub const KVM_REG_ARM64_SVE_VLS: u64 =
+    KVM_REG_ARM64 | KVM_REG_ARM64_SVE as u64 | KVM_REG_SIZE_U512 | 0xffff;
+
 /// Different aarch64 registers sizes
 #[derive(Debug)]
 pub enum RegSize {

src/vmm/src/arch/aarch64/vcpu.rs

@@ -179,17 +179,15 @@ pub fn get_all_registers_ids(vcpufd: &VcpuFd) -> Result<Vec<u64>, VcpuError> {
     }
 }
 
-/// Set the state of the system registers.
+/// Set the state of one system register.
 ///
 /// # Arguments
 ///
-/// * `regs` - Slice of registers to be set.
-pub fn set_registers(vcpufd: &VcpuFd, regs: &Aarch64RegisterVec) -> Result<(), VcpuError> {
-    for reg in regs.iter() {
-        vcpufd
-            .set_one_reg(reg.id, reg.as_slice())
-            .map_err(|e| VcpuError::SetOneReg(reg.id, e))?;
-    }
+/// * `reg` - Register to be set.
+pub fn set_register(vcpufd: &VcpuFd, reg: Aarch64RegisterRef) -> Result<(), VcpuError> {
+    vcpufd
+        .set_one_reg(reg.id, reg.as_slice())
+        .map_err(|e| VcpuError::SetOneReg(reg.id, e))?;
     Ok(())
 }
 
@@ -275,7 +273,9 @@ mod tests {
 
         vcpu.vcpu_init(&kvi).unwrap();
         get_all_registers(&vcpu, &mut regs).unwrap();
-        set_registers(&vcpu, &regs).unwrap();
+        for reg in regs.iter() {
+            set_register(&vcpu, reg).unwrap();
+        }
     }
 
     #[test]

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

@@ -12,11 +12,11 @@ use kvm_ioctls::*;
 use serde::{Deserialize, Serialize};
 
 use crate::arch::aarch64::regs::{
-    arm64_core_reg_id, offset__of, Aarch64RegisterVec, KVM_REG_ARM_TIMER_CNT,
+    arm64_core_reg_id, offset__of, Aarch64RegisterVec, KVM_REG_ARM64_SVE_VLS, KVM_REG_ARM_TIMER_CNT,
 };
 use crate::arch::aarch64::vcpu::{
     get_all_registers, get_all_registers_ids, get_mpidr, get_mpstate, get_registers, set_mpstate,
-    set_registers, setup_boot_regs, VcpuError as ArchError,
+    set_register, setup_boot_regs, VcpuError as ArchError,
 };
 use crate::cpu_config::aarch64::custom_cpu_template::VcpuFeatures;
 use crate::cpu_config::templates::CpuConfiguration;
@@ -139,7 +139,8 @@ impl KvmVcpu {
             kvi.features[index] = feature.bitmap.apply(kvi.features[index]);
         }
 
-        self.init_vcpu_fd(&kvi)?;
+        self.init_vcpu(&kvi)?;
+        self.finalize_vcpu(&kvi)?;
 
         self.kvi = if !vcpu_features.is_empty() {
             Some(kvi)
@@ -189,11 +190,31 @@ impl KvmVcpu {
             Some(kvi) => kvi,
             None => Self::default_kvi(vm_fd, self.index)?,
         };
+        self.kvi = state.kvi;
 
-        self.init_vcpu_fd(&kvi)?;
+        self.init_vcpu(&kvi)?;
 
-        self.kvi = state.kvi;
-        set_registers(&self.fd, &state.regs).map_err(KvmVcpuError::RestoreState)?;
+        // If KVM_REG_ARM64_SVE_VLS is present it needs to
+        // be set before vcpu is finalized.
+        if let Some(sve_vls_reg) = state
+            .regs
+            .iter()
+            .find(|reg| reg.id == KVM_REG_ARM64_SVE_VLS)
+        {
+            set_register(&self.fd, sve_vls_reg).map_err(KvmVcpuError::RestoreState)?;
+        }
+
+        self.finalize_vcpu(&kvi)?;
+
+        // KVM_REG_ARM64_SVE_VLS needs to be skipped after vcpu is finalized.
+        // If it is present it is handled in the code above.
+        for reg in state
+            .regs
+            .iter()
+            .filter(|reg| reg.id != KVM_REG_ARM64_SVE_VLS)
+        {
+            set_register(&self.fd, reg).map_err(KvmVcpuError::RestoreState)?;
+        }
         set_mpstate(&self.fd, state.mp_state).map_err(KvmVcpuError::RestoreState)?;
         Ok(())
     }
@@ -235,10 +256,14 @@ impl KvmVcpu {
     }
 
     /// Initializes internal vcpufd.
-    /// Does additional check for SVE and calls `vcpu_finalize` if
-    /// SVE is enabled.
-    fn init_vcpu_fd(&self, kvi: &kvm_bindings::kvm_vcpu_init) -> Result<(), KvmVcpuError> {
+    fn init_vcpu(&self, kvi: &kvm_bindings::kvm_vcpu_init) -> Result<(), KvmVcpuError> {
         self.fd.vcpu_init(kvi).map_err(KvmVcpuError::Init)?;
+        Ok(())
+    }
+
+    /// Checks for SVE feature and calls `vcpu_finalize` if
+    /// it is enabled.
+    fn finalize_vcpu(&self, kvi: &kvm_bindings::kvm_vcpu_init) -> Result<(), KvmVcpuError> {
         if (kvi.features[0] & (1 << kvm_bindings::KVM_ARM_VCPU_SVE)) != 0 {
             // KVM_ARM_VCPU_SVE has value 4 so casting to i32 is safe.
             #[allow(clippy::cast_possible_wrap)]

tests/integration_tests/functional/test_cpu_features.py

@@ -628,7 +628,7 @@ def test_cpu_cpuid_restore(microvm_factory, guest_kernel, msr_cpu_template):
     PLATFORM != "x86_64", reason="CPU features are masked only on x86_64."
 )
 @pytest.mark.parametrize("cpu_template", ["T2", "T2S", "C3"])
-def test_cpu_template(test_microvm_with_api, cpu_template):
+def test_cpu_template(test_microvm_with_api, cpu_template, microvm_factory):
     """
     Test masked and enabled cpu features against the expected template.
 
@@ -657,6 +657,15 @@ def test_cpu_template(test_microvm_with_api, cpu_template):
     check_masked_features(test_microvm, cpu_template)
     check_enabled_features(test_microvm, cpu_template)
 
+    # Check that cpu features are still correct
+    # after snap/restore cycle.
+    snapshot = test_microvm.snapshot_full()
+    restored_vm = microvm_factory.build()
+    restored_vm.spawn()
+    restored_vm.restore_from_snapshot(snapshot, resume=True)
+    check_masked_features(restored_vm, cpu_template)
+    check_enabled_features(restored_vm, cpu_template)
+
 
 def check_masked_features(test_microvm, cpu_template):
     """Verify the masked features of the given template."""

tests/integration_tests/functional/test_cpu_features_aarch64.py

@@ -108,6 +108,14 @@ def test_cpu_features_with_static_template(
     guest_kv = re.search(r"vmlinux-(\d+\.\d+)", guest_kernel.name).group(1)
     _check_cpu_features_arm(vm, guest_kv, "v1n1")
 
+    # Check that cpu features are still correct
+    # after snap/restore cycle.
+    snapshot = vm.snapshot_full()
+    restored_vm = microvm_factory.build()
+    restored_vm.spawn()
+    restored_vm.restore_from_snapshot(snapshot, resume=True)
+    _check_cpu_features_arm(restored_vm, guest_kv, "v1n1")
+
 
 @pytest.mark.skipif(
     PLATFORM != "aarch64",
@@ -128,3 +136,11 @@ def test_cpu_features_with_custom_template(
     vm.start()
     guest_kv = re.search(r"vmlinux-(\d+\.\d+)", guest_kernel.name).group(1)
     _check_cpu_features_arm(vm, guest_kv, custom_cpu_template["name"])
+
+    # Check that cpu features are still correct
+    # after snap/restore cycle.
+    snapshot = vm.snapshot_full()
+    restored_vm = microvm_factory.build()
+    restored_vm.spawn()
+    restored_vm.restore_from_snapshot(snapshot, resume=True)
+    _check_cpu_features_arm(restored_vm, guest_kv, custom_cpu_template["name"])