feat: only reset SYS_CNTPCT_EL0 if KVM_CAP_COUNTER_OFFSET is present

Now when KVM is a separate object, we can query optional capabilities
when configuring vcpus on aarch64 and only reset the SYS_CNTPCT_EL0
if KVM_CAP_COUNTER_OFFSET is present.

Signed-off-by: Egor Lazarchuk <yegorlz@amazon.co.uk>

Author: ShadowCurse

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

@@ -14,6 +14,7 @@ use log::warn;
 
 use super::get_fdt_addr;
 use super::regs::*;
+use crate::vstate::kvm::OptionalCapabilities;
 use crate::vstate::memory::GuestMemoryMmap;
 
 /// Errors thrown while setting aarch64 registers.
@@ -79,6 +80,7 @@ pub fn setup_boot_regs(
     cpu_id: u8,
     boot_ip: u64,
     mem: &GuestMemoryMmap,
+    optional_capabilities: &OptionalCapabilities,
 ) -> Result<(), VcpuError> {
     let kreg_off = offset_of!(kvm_regs, regs);
 
@@ -113,17 +115,16 @@ pub fn setup_boot_regs(
         // Resetting KVM_REG_ARM_PTIMER_CNT for single vcpu is enough because there is only
         // one timer struct with offsets per VM.
         // Because the access to KVM_REG_ARM_PTIMER_CNT is only present starting 6.4 kernel,
-        // we don't fail if ioctl returns an error.
+        // we only do the reset if KVM_CAP_COUNTER_OFFSET is present as it was added
+        // in the same patch series as the ability to set the KVM_REG_ARM_PTIMER_CNT register.
         // Path series which introduced the needed changes:
         // https://lore.kernel.org/all/20230330174800.2677007-1-maz@kernel.org/
         // Note: the value observed by the guest will still be above 0, because there is a delta
         // time between this resetting and first call to KVM_RUN.
-        let zero: u64 = 0;
-        if vcpufd
-            .set_one_reg(KVM_REG_ARM_PTIMER_CNT, &zero.to_le_bytes())
-            .is_err()
-        {
-            warn!("Unable to reset VM physical counter. VM will use host value instead.");
+        if optional_capabilities.counter_offset {
+            vcpufd
+                .set_one_reg(KVM_REG_ARM_PTIMER_CNT, &[0; 8])
+                .map_err(|err| VcpuError::SetOneReg(id, err))?;
         }
     }
     Ok(())
@@ -233,20 +234,21 @@ pub fn set_mpstate(vcpufd: &VcpuFd, state: kvm_mp_state) -> Result<(), VcpuError
 #[cfg(test)]
 mod tests {
     #![allow(clippy::undocumented_unsafe_blocks)]
-    use kvm_ioctls::Kvm;
 
     use super::*;
     use crate::arch::aarch64::layout;
     use crate::test_utils::arch_mem;
+    use crate::vstate::kvm::Kvm;
 
     #[test]
     fn test_setup_regs() {
-        let kvm = Kvm::new().unwrap();
-        let vm = kvm.create_vm().unwrap();
+        let kvm = Kvm::new(vec![]).unwrap();
+        let vm = kvm.fd.create_vm().unwrap();
         let vcpu = vm.create_vcpu(0).unwrap();
         let mem = arch_mem(layout::FDT_MAX_SIZE + 0x1000);
+        let optional_capabilities = kvm.optional_capabilities();
 
-        let res = setup_boot_regs(&vcpu, 0, 0x0, &mem);
+        let res = setup_boot_regs(&vcpu, 0, 0x0, &mem, &optional_capabilities);
         assert!(matches!(
             res.unwrap_err(),
             VcpuError::SetOneReg(0x6030000000100042, _)
@@ -256,14 +258,17 @@ mod tests {
         vm.get_preferred_target(&mut kvi).unwrap();
         vcpu.vcpu_init(&kvi).unwrap();
 
-        setup_boot_regs(&vcpu, 0, 0x0, &mem).unwrap();
+        setup_boot_regs(&vcpu, 0, 0x0, &mem, &optional_capabilities).unwrap();
 
         // Check that the register is reset on compatible kernels.
         // Because there is a delta in time between we reset the register and time we
         // read it, we cannot compare with 0. Instead we read CNTVNT_EL0 (Virtual counter)
         // from the host. The host value should be much bigger than the VM physical counter.
-        let mut reg_bytes = [0_u8; 8];
-        if vcpu.get_one_reg(SYS_CNTPCT_EL0, &mut reg_bytes).is_ok() {
+        if optional_capabilities.counter_offset {
+            let mut reg_bytes = [0_u8; 8];
+            vcpu.get_one_reg(SYS_CNTPCT_EL0, &mut reg_bytes).unwrap();
+            let counter_value = u64::from_le_bytes(reg_bytes);
+
             let virt_count = unsafe {
                 // Rust is so smart, it cannot figure out that `virt_count`
                 // is read later in the assert!, so it complains about unused assignment.
@@ -276,15 +281,15 @@ mod tests {
                 );
                 virt_count
             };
-            let counter_value = u64::from_le_bytes(reg_bytes);
+
             assert!(counter_value < virt_count);
         }
     }
 
     #[test]
     fn test_read_mpidr() {
-        let kvm = Kvm::new().unwrap();
-        let vm = kvm.create_vm().unwrap();
+        let kvm = Kvm::new(vec![]).unwrap();
+        let vm = kvm.fd.create_vm().unwrap();
         let vcpu = vm.create_vcpu(0).unwrap();
         let mut kvi: kvm_bindings::kvm_vcpu_init = kvm_bindings::kvm_vcpu_init::default();
         vm.get_preferred_target(&mut kvi).unwrap();
@@ -302,8 +307,8 @@ mod tests {
 
     #[test]
     fn test_get_set_regs() {
-        let kvm = Kvm::new().unwrap();
-        let vm = kvm.create_vm().unwrap();
+        let kvm = Kvm::new(vec![]).unwrap();
+        let vm = kvm.fd.create_vm().unwrap();
         let vcpu = vm.create_vcpu(0).unwrap();
         let mut kvi: kvm_bindings::kvm_vcpu_init = kvm_bindings::kvm_vcpu_init::default();
         vm.get_preferred_target(&mut kvi).unwrap();
@@ -324,8 +329,8 @@ mod tests {
     fn test_mpstate() {
         use std::os::unix::io::AsRawFd;
 
-        let kvm = Kvm::new().unwrap();
-        let vm = kvm.create_vm().unwrap();
+        let kvm = Kvm::new(vec![]).unwrap();
+        let vm = kvm.fd.create_vm().unwrap();
         let vcpu = vm.create_vcpu(0).unwrap();
         let mut kvi: kvm_bindings::kvm_vcpu_init = kvm_bindings::kvm_vcpu_init::default();
         vm.get_preferred_target(&mut kvi).unwrap();

src/vmm/src/builder.rs

@@ -811,16 +811,16 @@ pub fn configure_system_for_boot(
         cpu_config,
     };
 
-    // Configure vCPUs with normalizing and setting the generated CPU configuration.
-    for vcpu in vcpus.iter_mut() {
-        vcpu.kvm_vcpu
-            .configure(vmm.guest_memory(), entry_addr, &vcpu_config)
-            .map_err(VmmError::VcpuConfigure)
-            .map_err(Internal)?;
-    }
-
     #[cfg(target_arch = "x86_64")]
     {
+        // Configure vCPUs with normalizing and setting the generated CPU configuration.
+        for vcpu in vcpus.iter_mut() {
+            vcpu.kvm_vcpu
+                .configure(vmm.guest_memory(), entry_addr, &vcpu_config)
+                .map_err(VmmError::VcpuConfigure)
+                .map_err(Internal)?;
+        }
+
         // Write the kernel command line to guest memory. This is x86_64 specific, since on
         // aarch64 the command line will be specified through the FDT.
         let cmdline_size = boot_cmdline
@@ -855,6 +855,19 @@ pub fn configure_system_for_boot(
     }
     #[cfg(target_arch = "aarch64")]
     {
+        let optional_capabilities = vmm.kvm.optional_capabilities();
+        // Configure vCPUs with normalizing and setting the generated CPU configuration.
+        for vcpu in vcpus.iter_mut() {
+            vcpu.kvm_vcpu
+                .configure(
+                    vmm.guest_memory(),
+                    entry_addr,
+                    &vcpu_config,
+                    &optional_capabilities,
+                )
+                .map_err(VmmError::VcpuConfigure)
+                .map_err(Internal)?;
+        }
         let vcpu_mpidr = vcpus
             .iter_mut()
             .map(|cpu| cpu.kvm_vcpu.get_mpidr())

src/vmm/src/vstate/kvm.rs

@@ -1,4 +1,6 @@
-use kvm_bindings::{CpuId, MsrList, KVM_API_VERSION, KVM_MAX_CPUID_ENTRIES};
+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};
 
@@ -17,8 +19,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 MSR index list to save into snapshots: {0}
     GetMsrsToSave(crate::arch::x86_64::msr::MsrError),
+    #[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
@@ -127,7 +131,13 @@ impl Kvm {
         Ok(())
     }
 }
-
+#[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] = [
@@ -140,6 +150,16 @@ impl Kvm {
         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,
+        }
+    }
+
     /// Saves and returns the Kvm Vm state.
     pub fn save_state(&self) -> KvmState {
         KvmState {

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

@@ -22,6 +22,7 @@ use crate::cpu_config::aarch64::custom_cpu_template::VcpuFeatures;
 use crate::cpu_config::templates::CpuConfiguration;
 use crate::logger::{error, IncMetric, METRICS};
 use crate::vcpu::{VcpuConfig, VcpuError};
+use crate::vstate::kvm::{Kvm, OptionalCapabilities};
 use crate::vstate::memory::{Address, GuestAddress, GuestMemoryMmap};
 use crate::vstate::vcpu::VcpuEmulation;
 use crate::vstate::vm::Vm;
@@ -77,7 +78,7 @@ impl KvmVcpu {
     ///
     /// * `index` - Represents the 0-based CPU index between [0, max vcpus).
     /// * `vm` - The vm to which this vcpu will get attached.
-    pub fn new(index: u8, vm: &Vm) -> Result<Self, KvmVcpuError> {
+    pub fn new(index: u8, vm: &Vm, _: &Kvm) -> Result<Self, KvmVcpuError> {
         let kvm_vcpu = vm
             .fd()
             .create_vcpu(index.into())
@@ -115,6 +116,7 @@ impl KvmVcpu {
         guest_mem: &GuestMemoryMmap,
         kernel_load_addr: GuestAddress,
         vcpu_config: &VcpuConfig,
+        optional_capabilities: &OptionalCapabilities,
     ) -> Result<(), KvmVcpuError> {
         for reg in vcpu_config.cpu_config.regs.iter() {
             self.fd
@@ -127,6 +129,7 @@ impl KvmVcpu {
             self.index,
             kernel_load_addr.raw_value(),
             guest_mem,
+            optional_capabilities,
         )
         .map_err(KvmVcpuError::ConfigureRegisters)?;
 
@@ -312,7 +315,7 @@ mod tests {
 
     fn setup_vcpu(mem_size: usize) -> (Kvm, Vm, KvmVcpu, GuestMemoryMmap) {
         let (kvm, mut vm, vm_mem) = setup_vm_with_memory(mem_size);
-        let mut vcpu = KvmVcpu::new(0, &vm).unwrap();
+        let mut vcpu = KvmVcpu::new(0, &vm, &kvm).unwrap();
         vcpu.init(&[]).unwrap();
         vm.setup_irqchip(1).unwrap();
 
@@ -321,11 +324,11 @@ mod tests {
 
     #[test]
     fn test_create_vcpu() {
-        let (_, vm, _) = setup_vm_with_memory(0x1000);
+        let (kvm, vm, _) = setup_vm_with_memory(0x1000);
 
         unsafe { libc::close(vm.fd().as_raw_fd()) };
 
-        let err = KvmVcpu::new(0, &vm);
+        let err = KvmVcpu::new(0, &vm, &kvm);
         assert_eq!(
             err.err().unwrap().to_string(),
             "Error creating vcpu: Bad file descriptor (os error 9)".to_string()
@@ -337,7 +340,8 @@ mod tests {
 
     #[test]
     fn test_configure_vcpu() {
-        let (_, _, mut vcpu, vm_mem) = setup_vcpu(0x10000);
+        let (kvm, _, mut vcpu, vm_mem) = setup_vcpu(0x10000);
+        let optional_capabilities = kvm.optional_capabilities();
 
         let vcpu_config = VcpuConfig {
             vcpu_count: 1,
@@ -348,6 +352,7 @@ mod tests {
             &vm_mem,
             GuestAddress(crate::arch::get_kernel_start()),
             &vcpu_config,
+            &optional_capabilities,
         )
         .unwrap();
 
@@ -357,6 +362,7 @@ mod tests {
             &vm_mem,
             GuestAddress(crate::arch::get_kernel_start()),
             &vcpu_config,
+            &optional_capabilities,
         );
         assert_eq!(
             err.unwrap_err(),
@@ -372,8 +378,8 @@ mod tests {
 
     #[test]
     fn test_init_vcpu() {
-        let (_, mut vm, _) = setup_vm_with_memory(0x1000);
-        let mut vcpu = KvmVcpu::new(0, &vm).unwrap();
+        let (kvm, mut vm, _) = setup_vm_with_memory(0x1000);
+        let mut vcpu = KvmVcpu::new(0, &vm, &kvm).unwrap();
         vm.setup_irqchip(1).unwrap();
 
         // KVM_ARM_VCPU_PSCI_0_2 is set by default.
@@ -391,8 +397,8 @@ mod tests {
 
     #[test]
     fn test_vcpu_save_restore_state() {
-        let (_, mut vm, _) = setup_vm_with_memory(0x1000);
-        let mut vcpu = KvmVcpu::new(0, &vm).unwrap();
+        let (kvm, mut vm, _) = setup_vm_with_memory(0x1000);
+        let mut vcpu = KvmVcpu::new(0, &vm, &kvm).unwrap();
         vm.setup_irqchip(1).unwrap();
 
         // Calling KVM_GET_REGLIST before KVM_VCPU_INIT will result in error.
@@ -435,8 +441,8 @@ mod tests {
         //
         // This should fail with ENOEXEC.
         // https://elixir.bootlin.com/linux/v5.10.176/source/arch/arm64/kvm/arm.c#L1165
-        let (_, mut vm, _) = setup_vm_with_memory(0x1000);
-        let vcpu = KvmVcpu::new(0, &vm).unwrap();
+        let (kvm, mut vm, _) = setup_vm_with_memory(0x1000);
+        let vcpu = KvmVcpu::new(0, &vm, &kvm).unwrap();
         vm.setup_irqchip(1).unwrap();
 
         vcpu.dump_cpu_config().unwrap_err();
@@ -445,8 +451,8 @@ mod tests {
     #[test]
     fn test_dump_cpu_config_after_init() {
         // Test `dump_cpu_config()` after `KVM_VCPU_INIT`.
-        let (_, mut vm, _) = setup_vm_with_memory(0x1000);
-        let mut vcpu = KvmVcpu::new(0, &vm).unwrap();
+        let (kvm, mut vm, _) = setup_vm_with_memory(0x1000);
+        let mut vcpu = KvmVcpu::new(0, &vm, &kvm).unwrap();
         vm.setup_irqchip(1).unwrap();
         vcpu.init(&[]).unwrap();
 
@@ -455,10 +461,10 @@ mod tests {
 
     #[test]
     fn test_setup_non_boot_vcpu() {
-        let (_, vm, _) = setup_vm_with_memory(0x1000);
-        let mut vcpu1 = KvmVcpu::new(0, &vm).unwrap();
+        let (kvm, vm, _) = setup_vm_with_memory(0x1000);
+        let mut vcpu1 = KvmVcpu::new(0, &vm, &kvm).unwrap();
         vcpu1.init(&[]).unwrap();
-        let mut vcpu2 = KvmVcpu::new(1, &vm).unwrap();
+        let mut vcpu2 = KvmVcpu::new(1, &vm, &kvm).unwrap();
         vcpu2.init(&[]).unwrap();
     }
 

src/vmm/src/vstate/vcpu/mod.rs

@@ -1008,6 +1008,7 @@ pub(crate) mod tests {
                     smt: false,
                     cpu_config: crate::cpu_config::aarch64::CpuConfiguration::default(),
                 },
+                &kvm.optional_capabilities(),
             )
             .expect("failed to configure vcpu");
 

src/vmm/src/vstate/vm.rs

@@ -204,7 +204,6 @@ impl Vm {
                 .get_irqchip()
                 .save_device(mpidrs)
                 .map_err(VmError::SaveGic)?,
-            kvm_cap_modifiers: self.kvm_cap_modifiers.clone(),
         })
     }