fix(vmm): enable KVM_CAP_ARM_WRITABLE_IMP_ID_REGS on aarch64

CHANGELOG.md

@@ -43,6 +43,13 @@ and this project adheres to
   host page cache, so same-host reads still see the full contents. Block device
   backing files are always `fsync`'d regardless. See
   [snapshot documentation](docs/snapshotting/snapshot-support.md).
+- [#6116](https://github.com/firecracker-microvm/firecracker/pull/6116): On
+  aarch64, enable `KVM_CAP_ARM_WRITABLE_IMP_ID_REGS` when the host kernel offers
+  it (Linux 6.15 and later), adding support for custom CPU templates that modify
+  the implementation ID registers (`MIDR_EL1`, `REVIDR_EL1`, `AIDR_EL1`). Such
+  templates previously failed at boot with
+  `Failed to set register ... Invalid argument` because KVM rejects the write
+  unless the capability is enabled on the VM.
 
 ### Changed
 

docs/cpu_templates/cpu-templates.md

@@ -50,7 +50,11 @@ Firecracker supports two types of CPU templates:
 > CPU templates for ARM (both static and custom) require the following patch to
 > be available in the host kernel:
 > [Support writable CPU ID registers from userspace](https://lore.kernel.org/kvm/20230212215830.2975485-1-jingzhangos@google.com/#t).
-> Otherwise KVM will fail to write to the ARM registers.
+> Otherwise KVM will fail to write to the ARM registers. Additionally, modifying
+> the implementation ID registers (`MIDR_EL1`, `REVIDR_EL1`, `AIDR_EL1`)
+> requires a host kernel with `KVM_CAP_ARM_WRITABLE_IMP_ID_REGS` (Linux 6.15 or
+> later):
+> [KVM: arm64: writable MIDR/REVIDR](https://lore.kernel.org/lkml/20250210154953.27002-1-sebott@redhat.com/).
 
 ## Static CPU templates
 

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

@@ -577,6 +577,36 @@ mod tests {
         (vm, vcpu)
     }
 
+    #[test]
+    fn test_writable_imp_id_regs() {
+        // MIDR_EL1: op0=3, op1=0, CRn=0, CRm=0, op2=0.
+        const MIDR_EL1: u64 = 0x6030_0000_0013_c000;
+        // An arbitrary valid MIDR value (implementer Arm, part Neoverse N1).
+        const FAKE_MIDR: u64 = 0x410f_d0c0;
+
+        // `KvmVm::new` enables KVM_CAP_ARM_WRITABLE_IMP_ID_REGS when the host
+        // kernel offers it, before any vCPU is created.
+        let vm = setup_vm_with_memory(0x1000);
+        if vm
+            .fd()
+            .check_extension_raw(u64::from(kvm_bindings::KVM_CAP_ARM_WRITABLE_IMP_ID_REGS))
+            != 1
+        {
+            // Host kernel predates writable implementation ID registers
+            // (Linux 6.15); there is nothing further to verify.
+            return;
+        }
+        let mut vcpu = KvmVcpu::new(0, &vm).unwrap();
+        vcpu.init(&[]).unwrap();
+
+        let mut val = [0u8; 8];
+        vcpu.fd
+            .set_one_reg(MIDR_EL1, &FAKE_MIDR.to_le_bytes())
+            .unwrap();
+        vcpu.fd.get_one_reg(MIDR_EL1, &mut val).unwrap();
+        assert_eq!(u64::from_le_bytes(val), FAKE_MIDR);
+    }
+
     #[test]
     fn test_create_vcpu() {
         let vm = setup_vm_with_memory(0x1000);

src/vmm/src/arch/aarch64/vm.rs

@@ -3,15 +3,30 @@
 
 use std::sync::Mutex;
 
+use kvm_bindings::{KVM_CAP_ARM_WRITABLE_IMP_ID_REGS, KVMIO, kvm_enable_cap};
 use serde::{Deserialize, Serialize};
+use vmm_sys_util::errno;
+use vmm_sys_util::ioctl::ioctl_with_ref;
+use vmm_sys_util::ioctl_iow_nr;
 
 use crate::Kvm;
 use crate::arch::aarch64::gic::GicState;
+use crate::logger::warn;
 use crate::snapshot::Persist;
 use crate::vstate::memory::{GuestMemoryExtension, GuestMemoryState};
 use crate::vstate::resources::{ResourceAllocator, ResourceAllocatorState};
 use crate::vstate::vm::{VmCommon, VmError};
 
+// TODO(https://github.com/rust-vmm/kvm/pull/382): kvm-ioctls does not expose
+// `VmFd::enable_cap` on aarch64 yet; this is the same definition it uses
+// internally. Replace the direct ioctl with `enable_cap` once a release
+// containing that PR is available.
+#[allow(missing_docs)]
+mod ioctls {
+    use super::*;
+    ioctl_iow_nr!(KVM_ENABLE_CAP, KVMIO, 0xa3, kvm_enable_cap);
+}
+
 /// Structure representing the current architecture's understand of what a "virtual machine" is.
 #[derive(Debug)]
 pub struct KvmVm {
@@ -38,6 +53,39 @@ impl KvmVm {
     /// Create a new `KvmVm` struct.
     pub fn new(kvm: Kvm) -> Result<KvmVm, VmError> {
         let common = Self::create_common(kvm)?;
+
+        // KVM gates writes to the implementation ID registers (MIDR_EL1,
+        // REVIDR_EL1, AIDR_EL1) behind KVM_CAP_ARM_WRITABLE_IMP_ID_REGS,
+        // which must be enabled before any vCPU is created. Without it, a
+        // custom CPU template that modifies these registers fails at boot
+        // with EINVAL when the template is applied. Enabling the capability
+        // on its own does not change guest-visible state: the registers keep
+        // their host values unless a template rewrites them, and writes of
+        // unchanged values (e.g. on snapshot restore) were already accepted
+        // before this capability existed.
+        if common
+            .fd
+            .check_extension_raw(u64::from(KVM_CAP_ARM_WRITABLE_IMP_ID_REGS))
+            == 1
+        {
+            let cap = kvm_enable_cap {
+                cap: KVM_CAP_ARM_WRITABLE_IMP_ID_REGS,
+                ..Default::default()
+            };
+            // SAFETY: The ioctl is safe because we allocated the struct and
+            // the kernel will only read the size of the struct.
+            let ret = unsafe { ioctl_with_ref(&common.fd, ioctls::KVM_ENABLE_CAP(), &cap) };
+            if ret != 0 {
+                // Not fatal: a VM whose CPU template does not touch the
+                // implementation ID registers is unaffected, and one that
+                // does will fail loudly when the template is applied.
+                warn!(
+                    "Failed to enable KVM_CAP_ARM_WRITABLE_IMP_ID_REGS: {}",
+                    errno::Error::last()
+                );
+            }
+        }
+
         Ok(KvmVm {
             common,
             irqchip_handle: None,