Make set_xsave take slice

Author: ludfjig

src/hyperlight_host/src/hypervisor/hyperlight_vm.rs

@@ -1305,8 +1305,9 @@ mod tests {
     ///   Bytes 160-415: XMM0-XMM15 (16 regs × 16 bytes)
     ///   Bytes 416-511: Reserved
     ///   Bytes 512-575: XSAVE header (preserve XCOMP_BV at 520-527)
-    fn dirty_xsave(current_xsave: &[u8]) -> [u32; 1024] {
-        let mut xsave = [0u32; 1024];
+    #[allow(clippy::needless_range_loop)]
+    fn dirty_xsave(current_xsave: &[u8]) -> Vec<u32> {
+        let mut xsave = vec![0u32; current_xsave.len() / 4];
 
         // FCW (bytes 0-1) + FSW (bytes 2-3) - pack into xsave[0]
         // FCW = 0x0F7F (different from default 0x037F), FSW = 0x1234
@@ -1328,28 +1329,36 @@ mod tests {
         }
 
         // ST0-ST7/MM0-MM7 (bytes 32-159, indices 8-39)
-        for item in xsave.iter_mut().take(40).skip(8) {
-            *item = 0xCAFEBABE;
+        for i in 8..40 {
+            xsave[i] = 0xCAFEBABE;
         }
         // XMM0-XMM15 (bytes 160-415, indices 40-103)
-        for item in xsave.iter_mut().take(104).skip(40) {
-            *item = 0xDEADBEEF;
+        for i in 40..104 {
+            xsave[i] = 0xDEADBEEF;
+        }
+
+        // Reserved area (bytes 416-511, indices 104-127) - dirty it
+        for i in 104..128 {
+            xsave[i] = 0xABCDEF12;
         }
 
         // Preserve entire XSAVE header (bytes 512-575, indices 128-143) from current state
         // This includes XSTATE_BV (512-519) and XCOMP_BV (520-527) which
         // MSHV/WHP require to be set correctly for compacted format.
         // Failure to do this will cause set_xsave to fail on MSHV.
-        if current_xsave.len() >= 576 {
-            #[allow(clippy::needless_range_loop)]
-            for i in 128..144 {
-                let byte_offset = i * 4;
-                xsave[i] = u32::from_le_bytes(
-                    current_xsave[byte_offset..byte_offset + 4]
-                        .try_into()
-                        .unwrap(),
-                );
-            }
+        for i in 128..144 {
+            let byte_offset = i * 4;
+            xsave[i] = u32::from_le_bytes(
+                current_xsave[byte_offset..byte_offset + 4]
+                    .try_into()
+                    .unwrap(),
+            );
+        }
+
+        // Dirty extended state areas (bytes 576+, indices 144+)
+        // These include XSAVE areas for SSE, AVX, MPX, etc.
+        for i in 144..xsave.len() {
+            xsave[i] = 0x12345678;
         }
 
         xsave

src/hyperlight_host/src/hypervisor/hyperv_linux.rs

@@ -252,14 +252,26 @@ impl Hypervisor for MshvVm {
 
     #[cfg(test)]
     #[cfg(feature = "init-paging")]
-    fn set_xsave(&self, xsave: &[u32; 1024]) -> Result<()> {
-        let mut buf = XSave::default();
+    fn set_xsave(&self, xsave: &[u32]) -> Result<()> {
+        const MSHV_XSAVE_SIZE: usize = 4096;
+        if std::mem::size_of_val(xsave) != MSHV_XSAVE_SIZE {
+            return Err(new_error!(
+                "Provided xsave size {} does not match MSHV supported size {}",
+                std::mem::size_of_val(xsave),
+                MSHV_XSAVE_SIZE
+            ));
+        }
+
         // Safety: all valid u32 values are 4 valid u8 values
         let (prefix, bytes, suffix) = unsafe { xsave.align_to() };
         if !prefix.is_empty() || !suffix.is_empty() {
             return Err(new_error!("Invalid xsave buffer alignment"));
         }
-        buf.buffer.copy_from_slice(bytes);
+        let buf = XSave {
+            buffer: bytes
+                .try_into()
+                .map_err(|_| new_error!("mshv xsave must be 4096 u8s"))?,
+        };
         self.vcpu_fd.set_xsave(&buf)?;
         Ok(())
     }

src/hyperlight_host/src/hypervisor/hyperv_windows.rs

@@ -556,7 +556,7 @@ impl Hypervisor for WhpVm {
 
     #[cfg(test)]
     #[cfg(feature = "init-paging")]
-    fn set_xsave(&self, xsave: &[u32; 1024]) -> Result<()> {
+    fn set_xsave(&self, xsave: &[u8]) -> Result<()> {
         use crate::HyperlightError;
 
         // Get the required buffer size by calling with NULL buffer.
@@ -581,7 +581,7 @@ impl Hypervisor for WhpVm {
             return Err(HyperlightError::WindowsAPIError(e));
         }
 
-        let provided_size = std::mem::size_of_val(xsave) as u32;
+        let provided_size = xsave.len() as u32;
         if buffer_size_needed > provided_size {
             return Err(new_error!(
                 "Xsave buffer too small: needed {} bytes, provided {} bytes",

src/hyperlight_host/src/hypervisor/kvm.rs

@@ -207,10 +207,21 @@ impl Hypervisor for KvmVm {
 
     #[cfg(test)]
     #[cfg(feature = "init-paging")]
-    fn set_xsave(&self, xsave: &[u32; 1024]) -> Result<()> {
+    fn set_xsave(&self, xsave: &[u32]) -> Result<()> {
+        const KVM_XSAVE_SIZE: usize = 4096;
+
+        if std::mem::size_of_val(xsave) != KVM_XSAVE_SIZE {
+            return Err(new_error!(
+                "Provided xsave size {} does not match KVM supported size {}",
+                std::mem::size_of_val(xsave),
+                KVM_XSAVE_SIZE
+            ));
+        }
         let xsave = kvm_xsave {
-            region: *xsave,
-            ..Default::default() // zeroed 4KB buffer with no FAM
+            region: xsave
+                .try_into()
+                .map_err(|_| new_error!("kvm xsave must be 1024 u32s"))?,
+            ..Default::default()
         };
         // Safety: Safe because we only copy 4096 bytes
         // and have not enabled any dynamic xsave features

src/hyperlight_host/src/hypervisor/mod.rs

@@ -134,7 +134,7 @@ pub(crate) trait Hypervisor: Debug + Send {
     /// Set xsave - only used for tests
     #[cfg(test)]
     #[cfg(feature = "init-paging")]
-    fn set_xsave(&self, xsave: &[u32; 1024]) -> Result<()>;
+    fn set_xsave(&self, xsave: &[u32]) -> Result<()>;
 
     /// Get the debug registers of the vCPU
     #[allow(dead_code)]