Reset more vcpu state on snapshot restore

Signed-off-by: Ludvig Liljenberg <4257730+ludfjig@users.noreply.github.com>

Author: ludfjig

.github/workflows/ValidatePullRequest.yml

@@ -67,7 +67,7 @@ jobs:
       - docs-pr
       - build-guests
     strategy:
-      fail-fast: true
+      fail-fast: false
       matrix:
         hypervisor: [hyperv, 'hyperv-ws2025', mshv3, kvm]
         cpu: [amd, intel]

src/hyperlight_host/src/hypervisor/hyperlight_vm.rs

@@ -36,13 +36,11 @@ pub(crate) struct CommonFpu {
     pub fcw: u16,
     pub fsw: u16,
     pub ftwx: u8,
-    pub pad1: u8,
     pub last_opcode: u16,
     pub last_ip: u64,
     pub last_dp: u64,
     pub xmm: [[u8; 16]; 16],
     pub mxcsr: u32,
-    pub pad2: u32,
 }
 
 impl Default for CommonFpu {
@@ -52,13 +50,11 @@ impl Default for CommonFpu {
             fcw: FP_CONTROL_WORD_DEFAULT,
             fsw: 0,
             ftwx: 0,
-            pad1: 0,
             last_opcode: 0,
             last_ip: 0,
             last_dp: 0,
             xmm: [[0u8; 16]; 16],
             mxcsr: MXCSR_DEFAULT,
-            pad2: 0,
         }
     }
 }
@@ -71,13 +67,13 @@ impl From<&CommonFpu> for kvm_fpu {
             fcw: common_fpu.fcw,
             fsw: common_fpu.fsw,
             ftwx: common_fpu.ftwx,
-            pad1: common_fpu.pad1,
+            pad1: 0,
             last_opcode: common_fpu.last_opcode,
             last_ip: common_fpu.last_ip,
             last_dp: common_fpu.last_dp,
             xmm: common_fpu.xmm,
             mxcsr: common_fpu.mxcsr,
-            pad2: common_fpu.pad2,
+            pad2: 0,
         }
     }
 }
@@ -90,13 +86,13 @@ impl From<&CommonFpu> for FloatingPointUnit {
             fcw: common_fpu.fcw,
             fsw: common_fpu.fsw,
             ftwx: common_fpu.ftwx,
-            pad1: common_fpu.pad1,
+            pad1: 0,
             last_opcode: common_fpu.last_opcode,
             last_ip: common_fpu.last_ip,
             last_dp: common_fpu.last_dp,
             xmm: common_fpu.xmm,
             mxcsr: common_fpu.mxcsr,
-            pad2: common_fpu.pad2,
+            pad2: 0,
         }
     }
 }
@@ -109,13 +105,11 @@ impl From<&kvm_fpu> for CommonFpu {
             fcw: kvm_fpu.fcw,
             fsw: kvm_fpu.fsw,
             ftwx: kvm_fpu.ftwx,
-            pad1: kvm_fpu.pad1,
             last_opcode: kvm_fpu.last_opcode,
             last_ip: kvm_fpu.last_ip,
             last_dp: kvm_fpu.last_dp,
             xmm: kvm_fpu.xmm,
             mxcsr: kvm_fpu.mxcsr,
-            pad2: kvm_fpu.pad2,
         }
     }
 }
@@ -128,13 +122,11 @@ impl From<&FloatingPointUnit> for CommonFpu {
             fcw: mshv_fpu.fcw,
             fsw: mshv_fpu.fsw,
             ftwx: mshv_fpu.ftwx,
-            pad1: mshv_fpu.pad1,
             last_opcode: mshv_fpu.last_opcode,
             last_ip: mshv_fpu.last_ip,
             last_dp: mshv_fpu.last_dp,
             xmm: mshv_fpu.xmm,
             mxcsr: mshv_fpu.mxcsr,
-            pad2: mshv_fpu.pad2,
         }
     }
 }
@@ -174,7 +166,7 @@ impl From<&CommonFpu> for [(WHV_REGISTER_NAME, Align16<WHV_REGISTER_VALUE>); WHP
                 FpControl: fpu.fcw,
                 FpStatus: fpu.fsw,
                 FpTag: fpu.ftwx,
-                Reserved: fpu.pad1,
+                Reserved: 0,
                 LastFpOp: fpu.last_opcode,
                 Anonymous: WHV_X64_FP_CONTROL_STATUS_REGISTER_0_0 {
                     LastFpRip: fpu.last_ip,
@@ -293,7 +285,6 @@ impl TryFrom<&[(WHV_REGISTER_NAME, Align16<WHV_REGISTER_VALUE>)]> for CommonFpu
                     fpu.fcw = control.FpControl;
                     fpu.fsw = control.FpStatus;
                     fpu.ftwx = control.FpTag;
-                    fpu.pad1 = control.Reserved;
                     fpu.last_opcode = control.LastFpOp;
                     fpu.last_ip = unsafe { control.Anonymous.LastFpRip };
                 }
@@ -355,7 +346,6 @@ mod tests {
             fcw: 0x1234,
             fsw: 0x5678,
             ftwx: 0x9a,
-            pad1: 0xbc,
             last_opcode: 0xdef0,
             last_ip: 0xdeadbeefcafebabe,
             last_dp: 0xabad1deaf00dbabe,
@@ -365,7 +355,6 @@ mod tests {
                 [22u8; 16], [23u8; 16],
             ],
             mxcsr: 0x1f80,
-            pad2: 0,
         }
     }
 

src/hyperlight_host/src/hypervisor/regs/fpu.rs

@@ -18,15 +18,18 @@ use std::sync::LazyLock;
 
 #[cfg(gdb)]
 use kvm_bindings::kvm_guest_debug;
-use kvm_bindings::{kvm_debugregs, kvm_fpu, kvm_regs, kvm_sregs, kvm_userspace_memory_region};
+use kvm_bindings::{
+    kvm_debugregs, kvm_fpu, kvm_regs, kvm_sregs, kvm_userspace_memory_region, kvm_xsave,
+};
 use kvm_ioctls::Cap::UserMemory;
 use kvm_ioctls::{Kvm, VcpuExit, VcpuFd, VmFd};
 use tracing::{Span, instrument};
 
 #[cfg(gdb)]
 use crate::hypervisor::gdb::DebuggableVm;
 use crate::hypervisor::regs::{
-    CommonDebugRegs, CommonFpu, CommonRegisters, CommonSpecialRegisters,
+    CommonDebugRegs, CommonFpu, CommonRegisters, CommonSpecialRegisters, FP_CONTROL_WORD_DEFAULT,
+    MXCSR_DEFAULT,
 };
 use crate::hypervisor::virtual_machine::{VirtualMachine, VmExit};
 use crate::mem::memory_region::MemoryRegion;
@@ -142,12 +145,16 @@ impl VirtualMachine for KvmVm {
     }
 
     fn fpu(&self) -> Result<CommonFpu> {
+        // Note: On KVM this ignores MXCSR.
+        // See https://github.com/torvalds/linux/blob/d358e5254674b70f34c847715ca509e46eb81e6f/arch/x86/kvm/x86.c#L12554-L12599
         let kvm_fpu = self.vcpu_fd.get_fpu()?;
         Ok((&kvm_fpu).into())
     }
 
     fn set_fpu(&self, fpu: &CommonFpu) -> Result<()> {
         let kvm_fpu: kvm_fpu = fpu.into();
+        // Note: On KVM this ignores MXCSR.
+        // See https://github.com/torvalds/linux/blob/d358e5254674b70f34c847715ca509e46eb81e6f/arch/x86/kvm/x86.c#L12554-L12599
         self.vcpu_fd.set_fpu(&kvm_fpu)?;
         Ok(())
     }
@@ -174,7 +181,6 @@ impl VirtualMachine for KvmVm {
         Ok(())
     }
 
-    #[cfg(crashdump)]
     fn xsave(&self) -> Result<Vec<u8>> {
         let xsave = self.vcpu_fd.get_xsave()?;
         Ok(xsave
@@ -183,6 +189,51 @@ impl VirtualMachine for KvmVm {
             .flat_map(u32::to_le_bytes)
             .collect())
     }
+
+    fn reset_xsave(&self) -> Result<()> {
+        let mut xsave = kvm_xsave::default(); // default is zeroed 4KB buffer with no FAM
+
+        // XSAVE area layout from Intel SDM Vol. 1 Section 13.4.1:
+        // - Bytes 0-1: FCW (x87 FPU Control Word)
+        // - Bytes 24-27: MXCSR
+        // - Bytes 512-519: XSTATE_BV (bitmap of valid state components)
+        xsave.region[0] = FP_CONTROL_WORD_DEFAULT as u32;
+        xsave.region[6] = MXCSR_DEFAULT;
+        // XSTATE_BV = 0x3: bits 0,1 = x87 + SSE valid. This tells KVM to apply
+        // the legacy region from this buffer. Without this, some KVM versions
+        // may ignore set_xsave entirely when XSTATE_BV=0.
+        xsave.region[128] = 0x3;
+
+        // SAFETY: No dynamic features enabled, 4KB is sufficient
+        unsafe { self.vcpu_fd.set_xsave(&xsave)? };
+
+        Ok(())
+    }
+
+    #[cfg(test)]
+    #[cfg(feature = "init-paging")]
+    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
+                .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
+        unsafe { self.vcpu_fd.set_xsave(&xsave)? };
+
+        Ok(())
+    }
 }
 
 #[cfg(gdb)]

src/hyperlight_host/src/hypervisor/virtual_machine/kvm.rs

@@ -85,7 +85,7 @@ pub fn is_hypervisor_present() -> bool {
 }
 
 /// The hypervisor types available for the current platform
-#[derive(PartialEq, Eq, Debug)]
+#[derive(PartialEq, Eq, Debug, Copy, Clone)]
 pub(crate) enum HypervisorType {
     #[cfg(kvm)]
     Kvm,
@@ -170,12 +170,17 @@ pub(crate) trait VirtualMachine: Debug + Send {
     #[allow(dead_code)]
     fn debug_regs(&self) -> Result<CommonDebugRegs>;
     /// Set the debug registers of the vCPU
-    #[allow(dead_code)]
     fn set_debug_regs(&self, drs: &CommonDebugRegs) -> Result<()>;
 
-    /// xsave
-    #[cfg(crashdump)]
+    /// Get xsave
+    #[allow(dead_code)]
     fn xsave(&self) -> Result<Vec<u8>>;
+    /// Reset xsave to default state
+    fn reset_xsave(&self) -> Result<()>;
+    /// Set xsave - only used for tests
+    #[cfg(test)]
+    #[cfg(feature = "init-paging")]
+    fn set_xsave(&self, xsave: &[u32]) -> Result<()>;
 
     /// Get partition handle
     #[cfg(target_os = "windows")]

src/hyperlight_host/src/hypervisor/virtual_machine/mod.rs

@@ -21,7 +21,7 @@ use std::sync::LazyLock;
 #[cfg(gdb)]
 use mshv_bindings::{DebugRegisters, hv_message_type_HVMSG_X64_EXCEPTION_INTERCEPT};
 use mshv_bindings::{
-    FloatingPointUnit, SpecialRegisters, StandardRegisters, hv_message_type,
+    FloatingPointUnit, SpecialRegisters, StandardRegisters, XSave, hv_message_type,
     hv_message_type_HVMSG_GPA_INTERCEPT, hv_message_type_HVMSG_UNMAPPED_GPA,
     hv_message_type_HVMSG_X64_HALT, hv_message_type_HVMSG_X64_IO_PORT_INTERCEPT,
     hv_partition_property_code_HV_PARTITION_PROPERTY_SYNTHETIC_PROC_FEATURES,
@@ -34,7 +34,8 @@ use tracing::{Span, instrument};
 #[cfg(gdb)]
 use crate::hypervisor::gdb::DebuggableVm;
 use crate::hypervisor::regs::{
-    CommonDebugRegs, CommonFpu, CommonRegisters, CommonSpecialRegisters,
+    CommonDebugRegs, CommonFpu, CommonRegisters, CommonSpecialRegisters, FP_CONTROL_WORD_DEFAULT,
+    MXCSR_DEFAULT,
 };
 use crate::hypervisor::virtual_machine::{VirtualMachine, VmExit};
 use crate::mem::memory_region::{MemoryRegion, MemoryRegionFlags};
@@ -222,11 +223,65 @@ impl VirtualMachine for MshvVm {
         Ok(())
     }
 
-    #[cfg(crashdump)]
     fn xsave(&self) -> Result<Vec<u8>> {
         let xsave = self.vcpu_fd.get_xsave()?;
         Ok(xsave.buffer.to_vec())
     }
+
+    fn reset_xsave(&self) -> Result<()> {
+        let current_xsave = self.vcpu_fd.get_xsave()?;
+        if current_xsave.buffer.len() < 576 {
+            // Minimum: 512 legacy + 64 header
+            return Err(new_error!(
+                "Unexpected xsave length {}",
+                current_xsave.buffer.len()
+            ));
+        }
+
+        let mut buf = XSave::default(); // default is zeroed 4KB buffer
+
+        // Copy XCOMP_BV (offset 520-527) - preserves feature mask + compacted bit
+        buf.buffer[520..528].copy_from_slice(&current_xsave.buffer[520..528]);
+
+        // XSAVE area layout from Intel SDM Vol. 1 Section 13.4.1:
+        // - Bytes 0-1: FCW (x87 FPU Control Word)
+        // - Bytes 24-27: MXCSR
+        // - Bytes 512-519: XSTATE_BV (bitmap of valid state components)
+        buf.buffer[0..2].copy_from_slice(&FP_CONTROL_WORD_DEFAULT.to_le_bytes());
+        buf.buffer[24..28].copy_from_slice(&MXCSR_DEFAULT.to_le_bytes());
+        // XSTATE_BV = 0x3: bits 0,1 = x87 + SSE valid. Explicitly tell hypervisor
+        // to apply the legacy region from this buffer for consistent behavior.
+        buf.buffer[512..520].copy_from_slice(&0x3u64.to_le_bytes());
+
+        self.vcpu_fd.set_xsave(&buf)?;
+        Ok(())
+    }
+
+    #[cfg(test)]
+    #[cfg(feature = "init-paging")]
+    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"));
+        }
+        let buf = XSave {
+            buffer: bytes
+                .try_into()
+                .map_err(|_| new_error!("mshv xsave must be 4096 u8s"))?,
+        };
+        self.vcpu_fd.set_xsave(&buf)?;
+        Ok(())
+    }
 }
 
 #[cfg(gdb)]