hypervisor: Fix live migration when AMX is configured

AMX requires dynamically enabling certain large state components which
leads to an increase in the size of kvm_xsave. This was not taken into
account by Cloud hypervisor until now.

We solve this by refactoring `XSaveState` to directly wrap `kvm::Xsave`
and always ensure (via a OnceLocked static variable) that all operations
on the wrapped xsave state obtain an instance of the right size.

Signed-Off-by: Oliver Anderson <[email protected]>
On-behalf-of: SAP [email protected]

Author: olivereanderson

hypervisor/Cargo.toml

@@ -21,7 +21,10 @@ cfg-if = { workspace = true }
 concat-idents = "1.1.5"
 igvm = { workspace = true, optional = true }
 igvm_defs = { workspace = true, optional = true }
-kvm-bindings = { workspace = true, optional = true, features = ["serde"] }
+kvm-bindings = { workspace = true, optional = true, features = [
+  "fam-wrappers",
+  "serde",
+] }
 kvm-ioctls = { workspace = true, optional = true }
 libc = { workspace = true }
 log = { workspace = true }

hypervisor/src/arch/x86/mod.rs

@@ -12,11 +12,13 @@
 //
 
 use core::fmt;
-
-use crate::{CpuVendor, Hypervisor};
+#[cfg(feature = "kvm")]
+use std::sync::OnceLock;
 
 use thiserror::Error;
 
+use crate::{CpuVendor, Hypervisor};
+
 #[cfg(all(feature = "mshv_emulator", target_arch = "x86_64"))]
 pub mod emulator;
 pub mod gdt;
@@ -330,19 +332,43 @@ pub enum AmxGuestSupportError {
     AmxGuestTileRequest { errno: i64 },
 }
 
-#[serde_with::serde_as]
+/// The length of the XSAVE flexible array member (FAM).
+/// This length increases when arch_prctl is utilized to dynamically add state components.
+///
+/// IMPORTANT: This static should only be updated via methods on [`XsaveState`].
+#[cfg(feature = "kvm")]
+static XSAVE_FAM_LENGTH: OnceLock<usize> = OnceLock::new();
+
 #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
-pub struct XsaveState {
-    #[serde_as(as = "[_; 1024usize]")]
-    pub region: [u32; 1024usize],
-}
+pub struct XsaveState(#[cfg(feature = "kvm")] pub(crate) kvm_bindings::Xsave);
 
 impl XsaveState {
     const ARCH_GET_XCOMP_SUPP: usize = 0x1021;
     const ARCH_REQ_XCOMP_GUEST_PERM: usize = 0x1025;
     const ARCH_XCOMP_TILECFG: usize = 17;
     const ARCH_XCOMP_TILEDATA: usize = 18;
 
+    /// Construct an instance via the given initializer.
+    ///
+    /// As long as dynamically enabled state components have only been enabled
+    /// through static methods on this struct it is guaranteed that the
+    /// initialization routine is given an Xsave struct of the expected size.
+    #[cfg(feature = "kvm")]
+    pub(crate) fn with_initializer<F, E>(
+        mut init: F,
+    ) -> Result<Self, Box<dyn std::error::Error + Send + Sync + 'static>>
+    where
+        F: FnMut(&mut kvm_bindings::Xsave) -> Result<(), E>,
+        E: Into<Box<dyn std::error::Error + Send + Sync + 'static>>,
+    {
+        let fam_length = XSAVE_FAM_LENGTH.get().unwrap_or(&0);
+
+        let mut xsave = kvm_bindings::Xsave::new(*fam_length)?;
+
+        init(&mut xsave).map_err(Into::into)?;
+        Ok(Self(xsave))
+    }
+
     /// This function enables the AMX related TILECFG and TILEDATA state components for guests.
     ///
     /// # Background
@@ -353,9 +379,28 @@ impl XsaveState {
         hypervisor: &dyn Hypervisor,
     ) -> Result<(), AmxGuestSupportError> {
         Self::amx_supported(hypervisor)?;
-
         Self::request_guest_amx_support()?;
 
+        // If we are using the KVM hypervisor we meed to query for the new xsave2 size and update
+        // `XSAVE_FAM_LENGTH` accordingly.
+        #[cfg(feature = "kvm")]
+        {
+            // Obtain the number of bytes the kvm_xsave struct requires.
+            // This number is documented to always be at least 4096 bytes, but
+            let size = hypervisor.check_extension_int(kvm_ioctls::Cap::Xsave2);
+            // Reality check: We should at least have this number of bytes and probably more as we have enabled
+            // AMX tiles. If this is not the case, it is probably best to panic.
+            assert!(size >= 4096);
+            let fam_length = {
+                // Computation is documented in `[kvm_bindings::kvm_xsave2::len]`
+                ((size as usize) - size_of::<kvm_bindings::kvm_xsave>())
+                    .div_ceil(size_of::<kvm_bindings::__u32>())
+            };
+            XSAVE_FAM_LENGTH
+                .set(fam_length)
+                .expect("This should only be set once");
+        }
+
         Ok(())
     }
 
@@ -420,10 +465,3 @@ impl XsaveState {
         }
     }
 }
-
-impl Default for XsaveState {
-    fn default() -> Self {
-        // SAFETY: this is plain old data structure
-        unsafe { ::std::mem::zeroed() }
-    }
-}

hypervisor/src/kvm/mod.rs

@@ -2856,26 +2856,23 @@ impl KvmVcpu {
     /// X86 specific call that returns the vcpu's current "xsave struct".
     ///
     fn get_xsave(&self) -> cpu::Result<XsaveState> {
-        Ok(self
-            .fd
-            .get_xsave()
-            .map_err(|e| cpu::HypervisorCpuError::GetXsaveState(e.into()))?
-            .into())
+        XsaveState::with_initializer(|state|
+            // SAFETY: Any configured dynamically enabled state components are always enabled via
+            // static methods on `XsaveState` hence we know that `state` has the expected size.
+            unsafe { self.fd.get_xsave2(state) })
+        .map_err(|e| cpu::HypervisorCpuError::GetXsaveState(anyhow::Error::from_boxed(e)))
     }
 
     #[cfg(target_arch = "x86_64")]
     ///
     /// X86 specific call that sets the vcpu's current "xsave struct".
     ///
     fn set_xsave(&self, xsave: &XsaveState) -> cpu::Result<()> {
-        let xsave: kvm_bindings::kvm_xsave = (*xsave).clone().into();
-        // SAFETY: Here we trust the kernel not to read past the end of the kvm_xsave struct
-        // when calling the kvm-ioctl library function.
-        unsafe {
-            self.fd
-                .set_xsave(&xsave)
-                .map_err(|e| cpu::HypervisorCpuError::SetXsaveState(e.into()))
-        }
+        // SAFETY: Any configured dynamically enabled state components are always enabled via
+        // static methods on `XsaveState` hence we know that the wrapped instance has the
+        // expected size.
+        unsafe { self.fd.set_xsave2(&xsave.0) }
+            .map_err(|e| cpu::HypervisorCpuError::SetXsaveState(e.into()))
     }
 
     #[cfg(target_arch = "x86_64")]

hypervisor/src/kvm/x86_64/mod.rs

@@ -291,18 +291,3 @@ impl From<MsrEntry> for kvm_msr_entry {
         }
     }
 }
-
-impl From<kvm_xsave> for XsaveState {
-    fn from(s: kvm_xsave) -> Self {
-        Self { region: s.region }
-    }
-}
-
-impl From<XsaveState> for kvm_xsave {
-    fn from(s: XsaveState) -> Self {
-        Self {
-            region: s.region,
-            extra: Default::default(),
-        }
-    }
-}

hypervisor/src/mshv/mod.rs

@@ -398,7 +398,6 @@ impl hypervisor::Hypervisor for MshvHypervisor {
         }
     }
 
-
     #[cfg(feature = "kvm")]
     fn check_extension_int(&self, _capability: kvm_ioctls::Cap) -> i32 {
         unimplemented!()