feat(vmm): Change XSAVE type from kvm_xsave to Xsave in snapshot

Intel AMX is an XSTATE feature and TILEDATA is disabled by default
because it requires a larger area to save its state than the traditional
4096 bytes. Instead, Linux kernel allows VMMs to request the guest
permission via `arch_prctl()`. As such, the size of the XSTATE buffer
required to save XSTASTE is dynamic. To support dynamically-sized
buffer, `KVM_CAP_XSAVE2` was introduced with `KVM_GET_XSAVE2`.
Accordingly, kvm-bindings added `Xsave` that is an alias of
`FamStructWrapper` for the `kvm_xsave` struct with FAM in the end, and
kvm-ioctls added `get_xsave2()` for `KVM_GET_XSAVE2` and `set_xsave2()`
to take `Xsave` to call `KVM_SET_XSAVE`.

Change the type of `xsave` in `VcpuState` from `kvm_xsave` to `Xsave`.
Use `get_xsave2()` and `set_xsave2()`.

Signed-off-by: Takahiro Itazuri <[email protected]>

Author: zulinx86

Cargo.lock

@@ -22,8 +22,8 @@ displaydoc = "0.2.5"
 event-manager = "0.4.0"
 gdbstub = { version = "0.7.3", optional = true }
 gdbstub_arch = { version = "0.3.1", optional = true }
-kvm-bindings = { version = "0.11.0", features = ["fam-wrappers", "serde"] }
-kvm-ioctls = "0.20.0"
+kvm-bindings = { version = "0.11.1", features = ["fam-wrappers", "serde"] }
+kvm-ioctls = "0.21.0"
 libc = "0.2.170"
 linux-loader = "0.13.0"
 log = { version = "0.4.26", features = ["std", "serde"] }

src/vmm/Cargo.toml

@@ -10,12 +10,12 @@ use std::fmt::Debug;
 
 use kvm_bindings::{
     kvm_debugregs, kvm_lapic_state, kvm_mp_state, kvm_regs, kvm_sregs, kvm_vcpu_events, kvm_xcrs,
-    kvm_xsave, CpuId, Msrs, KVM_MAX_CPUID_ENTRIES, KVM_MAX_MSR_ENTRIES,
+    kvm_xsave, kvm_xsave2, CpuId, Msrs, Xsave, KVM_MAX_CPUID_ENTRIES, KVM_MAX_MSR_ENTRIES,
 };
 use kvm_ioctls::{VcpuExit, VcpuFd};
 use log::{error, warn};
 use serde::{Deserialize, Serialize};
-use vmm_sys_util::fam;
+use vmm_sys_util::fam::{self, FamStruct};
 
 use crate::arch::x86_64::gen::msr_index::{MSR_IA32_TSC, MSR_IA32_TSC_DEADLINE};
 use crate::arch::x86_64::interrupts;
@@ -74,8 +74,10 @@ pub enum KvmVcpuError {
     VcpuGetVcpuEvents(kvm_ioctls::Error),
     /// Failed to get KVM vcpu xcrs: {0}
     VcpuGetXcrs(kvm_ioctls::Error),
-    /// Failed to get KVM vcpu xsave: {0}
+    /// Failed to get KVM vcpu xsave via KVM_GET_XSAVE: {0}
     VcpuGetXsave(kvm_ioctls::Error),
+    /// Failed to get KVM vcpu xsave via KVM_GET_XSAVE2: {0}
+    VcpuGetXsave2(kvm_ioctls::Error),
     /// Failed to get KVM vcpu cpuid: {0}
     VcpuGetCpuid(kvm_ioctls::Error),
     /// Failed to get KVM TSC frequency: {0}
@@ -147,6 +149,10 @@ pub struct KvmVcpu {
     /// The list of MSRs to include in a VM snapshot, in the same order as KVM returned them
     /// from KVM_GET_MSR_INDEX_LIST
     msrs_to_save: Vec<u32>,
+    /// Byte size requiring to hold the dynamically-sized `kvm_xsave` struct.
+    ///
+    /// `None` if `KVM_CAP_XSAVE2` not supported.
+    xsave2_size: Option<usize>,
 }
 
 /// Vcpu peripherals
@@ -176,6 +182,7 @@ impl KvmVcpu {
             fd: kvm_vcpu,
             peripherals: Default::default(),
             msrs_to_save: vm.msrs_to_save().to_vec(),
+            xsave2_size: vm.xsave2_size(),
         })
     }
 
@@ -496,7 +503,27 @@ impl KvmVcpu {
             .map_err(KvmVcpuError::VcpuGetMpState)?;
         let regs = self.fd.get_regs().map_err(KvmVcpuError::VcpuGetRegs)?;
         let sregs = self.fd.get_sregs().map_err(KvmVcpuError::VcpuGetSregs)?;
-        let xsave = self.fd.get_xsave().map_err(KvmVcpuError::VcpuGetXsave)?;
+        let xsave = match self.xsave2_size {
+            Some(size) => {
+                // SAFETY: The subtraction never underflows because
+                // `KVM_CHECK_EXTENSION(KVM_CAP_XSAVE2)` returns at least 4096 bytes which equals
+                // the size of `kvm_xsave`.
+                // https://docs.kernel.org/virt/kvm/api.html#kvm-get-xsave2
+                //
+                // Note that the FAM entry type is `__u32` although `xsave2_xize` is in bytes.
+                let fam_len = (size - std::mem::size_of::<kvm_xsave>())
+                    .div_ceil(std::mem::size_of::<<kvm_xsave2 as FamStruct>::Entry>());
+                let mut xsave = Xsave::new(fam_len).map_err(KvmVcpuError::Fam)?;
+                // SAFETY: Safe because `xsave` is allocated above with enough size to save XSTATE.
+                unsafe { self.fd.get_xsave2(&mut xsave) }.map_err(KvmVcpuError::VcpuGetXsave2)?;
+                xsave
+            }
+            // `KVM_CAP_XSAVE2` not supported
+            None => Xsave::from(vec![kvm_xsave2 {
+                len: 0,
+                xsave: self.fd.get_xsave().map_err(KvmVcpuError::VcpuGetXsave)?,
+            }]),
+        };
         let xcrs = self.fd.get_xcrs().map_err(KvmVcpuError::VcpuGetXcrs)?;
         let debug_regs = self
             .fd
@@ -601,9 +628,17 @@ impl KvmVcpu {
         self.fd
             .set_sregs(&state.sregs)
             .map_err(KvmVcpuError::VcpuSetSregs)?;
-        self.fd
-            .set_xsave(&state.xsave)
-            .map_err(KvmVcpuError::VcpuSetXsave)?;
+        // SAFETY: Safe unless the snapshot is corrupted.
+        unsafe {
+            // kvm-ioctl's `set_xsave2()` can be called even on kernel versions not supporting
+            // `KVM_CAP_XSAVE2`, because it internally calls `KVM_SET_XSAVE` API that was extended
+            // by Linux kernel. Thus, `KVM_SET_XSAVE2` API does not exist as a KVM interface.
+            // However, kvm-ioctl added `set_xsave2()` to allow users to pass `Xsave` instead of the
+            // older `kvm_xsave`.
+            self.fd
+                .set_xsave2(&state.xsave)
+                .map_err(KvmVcpuError::VcpuSetXsave)?;
+        }
         self.fd
             .set_xcrs(&state.xcrs)
             .map_err(KvmVcpuError::VcpuSetXcrs)?;
@@ -684,7 +719,7 @@ pub struct VcpuState {
     /// Xcrs.
     pub xcrs: kvm_xcrs,
     /// Xsave.
-    pub xsave: kvm_xsave,
+    pub xsave: Xsave,
     /// Tsc khz.
     pub tsc_khz: Option<u32>,
 }
@@ -744,7 +779,7 @@ mod tests {
                 sregs: Default::default(),
                 vcpu_events: Default::default(),
                 xcrs: Default::default(),
-                xsave: Default::default(),
+                xsave: Xsave::new(0).unwrap(),
                 tsc_khz: Some(0),
             }
         }

src/vmm/src/vstate/vcpu/x86_64.rs

@@ -7,7 +7,7 @@ use kvm_bindings::{
     kvm_clock_data, kvm_irqchip, kvm_pit_config, kvm_pit_state2, MsrList, KVM_CLOCK_TSC_STABLE,
     KVM_IRQCHIP_IOAPIC, KVM_IRQCHIP_PIC_MASTER, KVM_IRQCHIP_PIC_SLAVE, KVM_PIT_SPEAKER_DUMMY,
 };
-use kvm_ioctls::VmFd;
+use kvm_ioctls::{Cap, VmFd};
 use serde::{Deserialize, Serialize};
 
 use crate::arch::x86_64::msr::MsrError;
@@ -19,6 +19,8 @@ use crate::vstate::vm::VmError;
 #[cfg(target_arch = "x86_64")]
 #[derive(Debug, PartialEq, Eq, thiserror::Error, displaydoc::Display)]
 pub enum ArchVmError {
+    /// Failed to check KVM capability (0): {1}
+    CheckCapability(Cap, kvm_ioctls::Error),
     /// Set PIT2 error: {0}
     SetPit2(kvm_ioctls::Error),
     /// Set clock error: {0}
@@ -48,6 +50,10 @@ pub enum ArchVmError {
 pub struct ArchVm {
     pub(super) fd: VmFd,
     msrs_to_save: MsrList,
+    /// Cache of byte size requiring to hold the dynamically-sized `kvm_xsave` struct.
+    ///
+    /// `None` if `KVM_CAP_XSAVE2` not supported.
+    xsave2_size: Option<usize>,
 }
 
 impl ArchVm {
@@ -57,10 +63,31 @@ impl ArchVm {
 
         let msrs_to_save = kvm.msrs_to_save().map_err(ArchVmError::GetMsrsToSave)?;
 
+        // `KVM_CAP_XSAVE2` was introduced to support dynamically-sized XSTATE buffer in kernel
+        // v5.17. `KVM_GET_EXTENSION(KVM_CAP_XSAVE2)` returns the required size in byte if
+        // supported; otherwise returns 0.
+        // https://github.com/torvalds/linux/commit/be50b2065dfa3d88428fdfdc340d154d96bf6848
+        // Cache the value in order not to call it at each vCPU creation.
+        let xsave2_size = match fd.check_extension_int(Cap::Xsave2) {
+            -1 => {
+                return Err(VmError::Arch(ArchVmError::CheckCapability(
+                    Cap::Xsave2,
+                    vmm_sys_util::errno::Error::last(),
+                )));
+            }
+            0 => None,
+            // SAFETY: Safe because the possible negative value is only -1 and is handled above.
+            ret => Some(ret.try_into().unwrap()),
+        };
+
         fd.set_tss_address(u64_to_usize(crate::arch::x86_64::layout::KVM_TSS_ADDRESS))
             .map_err(ArchVmError::SetTssAddress)?;
 
-        Ok(ArchVm { fd, msrs_to_save })
+        Ok(ArchVm {
+            fd,
+            msrs_to_save,
+            xsave2_size,
+        })
     }
 
     pub(super) fn arch_pre_create_vcpus(&mut self, _: u8) -> Result<(), ArchVmError> {
@@ -162,6 +189,11 @@ impl ArchVm {
     pub fn msrs_to_save(&self) -> &[u32] {
         self.msrs_to_save.as_slice()
     }
+
+    /// Gets the byte size of the `kvm_xsave` struct.
+    pub fn xsave2_size(&self) -> Option<usize> {
+        self.xsave2_size
+    }
 }
 
 #[derive(Default, Deserialize, Serialize)]