Use WHvResetPartition on windows

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

Author: ludfjig

src/hyperlight_host/src/hypervisor/hyperlight_vm/aarch64.rs

@@ -89,11 +89,15 @@ impl HyperlightVm {
         unimplemented!("get_snapshot_sregs")
     }
 
-    pub(crate) fn reset_vcpu(
+    pub(crate) fn reset_vm_state(&mut self) -> std::result::Result<(), RegisterError> {
+        unimplemented!("reset_vm_state")
+    }
+
+    pub(crate) fn restore_sregs(
         &mut self,
         _cr3: u64,
         _sregs: &CommonSpecialRegisters,
     ) -> std::result::Result<(), RegisterError> {
-        unimplemented!("reset_vcpu")
+        unimplemented!("restore_sregs")
     }
 }

src/hyperlight_host/src/hypervisor/hyperlight_vm/x86_64.rs

@@ -40,9 +40,9 @@ use crate::hypervisor::gdb::{
 };
 #[cfg(gdb)]
 use crate::hypervisor::gdb::{DebugError, DebugMemoryAccessError};
-use crate::hypervisor::regs::{
-    CommonDebugRegs, CommonFpu, CommonRegisters, CommonSpecialRegisters,
-};
+#[cfg(not(target_os = "windows"))]
+use crate::hypervisor::regs::CommonDebugRegs;
+use crate::hypervisor::regs::{CommonFpu, CommonRegisters, CommonSpecialRegisters};
 #[cfg(not(gdb))]
 use crate::hypervisor::virtual_machine::VirtualMachine;
 #[cfg(kvm)]
@@ -335,23 +335,37 @@ impl HyperlightVm {
     }
 
     /// Resets the following vCPU state:
-    /// - General purpose registers
-    /// - Debug registers
-    /// - XSAVE (includes FPU/SSE state with proper FCW and MXCSR defaults)
-    /// - Special registers (restored from snapshot, with CR3 updated to new page table location)
+    /// - On Windows: calls WHvResetPartition (resets all per-VP state including
+    ///   GP registers, debug registers, XSAVE, MSRs, APIC, etc.)
+    /// - On Linux: explicitly resets GP registers, debug registers, and XSAVE
+    ///
+    /// This does NOT restore special registers (except on windows). Call `restore_sregs` separately
+    /// after memory mappings are established.
     // TODO: check if other state needs to be reset
-    pub(crate) fn reset_vcpu(
+    pub(crate) fn reset_vm_state(&mut self) -> std::result::Result<(), RegisterError> {
+        #[cfg(target_os = "windows")]
+        self.vm.reset_partition()?;
+
+        #[cfg(not(target_os = "windows"))]
+        {
+            self.vm.set_regs(&CommonRegisters {
+                rflags: 1 << 1, // Reserved bit always set
+                ..Default::default()
+            })?;
+            self.vm.set_debug_regs(&CommonDebugRegs::default())?;
+            self.vm.reset_xsave()?;
+        }
+
+        Ok(())
+    }
+
+    /// Restores special registers from snapshot with CR3 updated to the
+    /// new page table location.
+    pub(crate) fn restore_sregs(
         &mut self,
         cr3: u64,
         sregs: &CommonSpecialRegisters,
     ) -> std::result::Result<(), RegisterError> {
-        self.vm.set_regs(&CommonRegisters {
-            rflags: 1 << 1, // Reserved bit always set
-            ..Default::default()
-        })?;
-        self.vm.set_debug_regs(&CommonDebugRegs::default())?;
-        self.vm.reset_xsave()?;
-
         #[cfg(not(feature = "nanvix-unstable"))]
         {
             // Restore the full special registers from snapshot, but update CR3
@@ -885,7 +899,9 @@ mod tests {
     use super::*;
     #[cfg(kvm)]
     use crate::hypervisor::regs::FP_CONTROL_WORD_DEFAULT;
-    use crate::hypervisor::regs::{CommonSegmentRegister, CommonTableRegister, MXCSR_DEFAULT};
+    use crate::hypervisor::regs::{
+        CommonDebugRegs, CommonSegmentRegister, CommonTableRegister, MXCSR_DEFAULT,
+    };
     use crate::hypervisor::virtual_machine::VirtualMachine;
     use crate::mem::layout::SandboxMemoryLayout;
     use crate::mem::memory_region::{GuestMemoryRegion, MemoryRegionFlags};
@@ -1629,7 +1645,8 @@ mod tests {
         assert_eq!(got_sregs, expected_sregs);
 
         // Reset the vCPU
-        hyperlight_vm.reset_vcpu(0, &default_sregs()).unwrap();
+        hyperlight_vm.reset_vm_state().unwrap();
+        hyperlight_vm.restore_sregs(0, &default_sregs()).unwrap();
 
         // Verify registers are reset to defaults
         assert_regs_reset(hyperlight_vm.vm.as_ref());
@@ -1758,7 +1775,8 @@ mod tests {
             assert_eq!(regs, expected_dirty);
 
             // Reset vcpu
-            hyperlight_vm.reset_vcpu(0, &default_sregs()).unwrap();
+            hyperlight_vm.reset_vm_state().unwrap();
+            hyperlight_vm.restore_sregs(0, &default_sregs()).unwrap();
 
             // Check registers are reset to defaults
             assert_regs_reset(hyperlight_vm.vm.as_ref());
@@ -1882,7 +1900,8 @@ mod tests {
             }
 
             // Reset vcpu
-            hyperlight_vm.reset_vcpu(0, &default_sregs()).unwrap();
+            hyperlight_vm.reset_vm_state().unwrap();
+            hyperlight_vm.restore_sregs(0, &default_sregs()).unwrap();
 
             // Check FPU is reset to defaults
             assert_fpu_reset(hyperlight_vm.vm.as_ref());
@@ -1933,7 +1952,8 @@ mod tests {
             assert_eq!(debug_regs, expected_dirty);
 
             // Reset vcpu
-            hyperlight_vm.reset_vcpu(0, &default_sregs()).unwrap();
+            hyperlight_vm.reset_vm_state().unwrap();
+            hyperlight_vm.restore_sregs(0, &default_sregs()).unwrap();
 
             // Check debug registers are reset to default values
             assert_debug_regs_reset(hyperlight_vm.vm.as_ref());
@@ -1982,9 +2002,10 @@ mod tests {
             assert_eq!(sregs, expected_dirty);
 
             // Reset vcpu
-            hyperlight_vm.reset_vcpu(0, &default_sregs()).unwrap();
+            hyperlight_vm.reset_vm_state().unwrap();
+            hyperlight_vm.restore_sregs(0, &default_sregs()).unwrap();
 
-            // Check registers are reset to defaults (CR3 is 0 as passed to reset_vcpu)
+            // Check registers are reset to defaults
             let sregs = hyperlight_vm.vm.sregs().unwrap();
             let mut expected_reset = CommonSpecialRegisters::standard_64bit_defaults(0);
             normalize_sregs_for_run_tests(&mut expected_reset, &sregs);
@@ -2020,7 +2041,11 @@ mod tests {
             let root_pt_addr = ctx.ctx.vm.get_root_pt().unwrap();
             let segment_state = ctx.ctx.vm.get_snapshot_sregs().unwrap();
 
-            ctx.ctx.vm.reset_vcpu(root_pt_addr, &segment_state).unwrap();
+            ctx.ctx.vm.reset_vm_state().unwrap();
+            ctx.ctx
+                .vm
+                .restore_sregs(root_pt_addr, &segment_state)
+                .unwrap();
 
             // Re-run from entrypoint (flag=1 means guest skips dirty phase, just does FXSAVE)
             // Use stack_top - 8 to match initialise()'s behavior (simulates call pushing return addr)

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

@@ -107,7 +107,7 @@ pub(crate) enum HypervisorType {
 /// Minimum XSAVE buffer size: 512 bytes legacy region + 64 bytes header.
 /// Only used by MSHV and WHP which use compacted XSAVE format and need to
 /// validate buffer size before accessing XCOMP_BV.
-#[cfg(any(mshv3, target_os = "windows"))]
+#[cfg(mshv3)]
 pub(crate) const XSAVE_MIN_SIZE: usize = 576;
 
 /// Standard XSAVE buffer size (4KB) used by KVM and MSHV.
@@ -244,6 +244,9 @@ pub enum RegisterError {
     #[cfg(target_os = "windows")]
     #[error("Failed to convert WHP registers: {0}")]
     ConversionFailed(String),
+    #[cfg(target_os = "windows")]
+    #[error("Failed to reset partition: {0}")]
+    ResetPartition(HypervisorError),
 }
 
 /// Map memory error
@@ -341,18 +344,32 @@ pub(crate) trait VirtualMachine: Debug + Send {
     #[allow(dead_code)]
     fn debug_regs(&self) -> std::result::Result<CommonDebugRegs, RegisterError>;
     /// Set the debug registers of the vCPU
+    #[allow(dead_code)] // Called on Linux (reset_vm_state) and via gdb, but dead on Windows release builds without gdb
     fn set_debug_regs(&self, drs: &CommonDebugRegs) -> std::result::Result<(), RegisterError>;
 
     /// Get xsave
     #[allow(dead_code)]
     fn xsave(&self) -> std::result::Result<Vec<u8>, RegisterError>;
     /// Reset xsave to default state
+    #[cfg(any(kvm, mshv3))]
     fn reset_xsave(&self) -> std::result::Result<(), RegisterError>;
     /// Set xsave - only used for tests
     #[cfg(test)]
     #[cfg(not(feature = "nanvix-unstable"))]
     fn set_xsave(&self, xsave: &[u32]) -> std::result::Result<(), RegisterError>;
 
+    /// Reset the partition using WHvResetPartition.
+    ///
+    /// Resets all per-VP state to architectural defaults: GP registers, segment
+    /// registers, control registers, EFER, MSRs, debug registers, full XSAVE
+    /// state (x87/SSE/AVX/AVX-512/AMX), XCR0, APIC/LAPIC, pending interrupts,
+    /// and VMCS/VMCB internals.
+    ///
+    /// Does NOT reset: GPA memory mappings or contents, partition configuration,
+    /// notification ports, or VPCI device state.
+    #[cfg(target_os = "windows")]
+    fn reset_partition(&mut self) -> std::result::Result<(), RegisterError>;
+
     /// Get partition handle
     #[cfg(target_os = "windows")]
     fn partition_handle(&self) -> windows::Win32::System::Hypervisor::WHV_PARTITION_HANDLE;

src/hyperlight_host/src/hypervisor/virtual_machine/whp.rs

@@ -32,17 +32,16 @@ use windows_result::HRESULT;
 use crate::hypervisor::gdb::{DebugError, DebuggableVm};
 use crate::hypervisor::regs::{
     Align16, CommonDebugRegs, CommonFpu, CommonRegisters, CommonSpecialRegisters,
-    FP_CONTROL_WORD_DEFAULT, MXCSR_DEFAULT, WHP_DEBUG_REGS_NAMES, WHP_DEBUG_REGS_NAMES_LEN,
-    WHP_FPU_NAMES, WHP_FPU_NAMES_LEN, WHP_REGS_NAMES, WHP_REGS_NAMES_LEN, WHP_SREGS_NAMES,
-    WHP_SREGS_NAMES_LEN,
+    WHP_DEBUG_REGS_NAMES, WHP_DEBUG_REGS_NAMES_LEN, WHP_FPU_NAMES, WHP_FPU_NAMES_LEN,
+    WHP_REGS_NAMES, WHP_REGS_NAMES_LEN, WHP_SREGS_NAMES, WHP_SREGS_NAMES_LEN,
 };
 use crate::hypervisor::surrogate_process::SurrogateProcess;
 use crate::hypervisor::surrogate_process_manager::get_surrogate_process_manager;
 #[cfg(feature = "hw-interrupts")]
 use crate::hypervisor::virtual_machine::x86_64::hw_interrupts::TimerThread;
 use crate::hypervisor::virtual_machine::{
     CreateVmError, HypervisorError, MapMemoryError, RegisterError, RunVcpuError, UnmapMemoryError,
-    VirtualMachine, VmExit, XSAVE_MIN_SIZE,
+    VirtualMachine, VmExit,
 };
 use crate::hypervisor::wrappers::HandleWrapper;
 use crate::mem::memory_region::{MemoryRegion, MemoryRegionFlags, MemoryRegionType};
@@ -607,6 +606,7 @@ impl VirtualMachine for WhpVm {
         })
     }
 
+    #[allow(dead_code)]
     fn set_debug_regs(&self, drs: &CommonDebugRegs) -> std::result::Result<(), RegisterError> {
         let whp_regs: [(WHV_REGISTER_NAME, Align16<WHV_REGISTER_VALUE>); WHP_DEBUG_REGS_NAMES_LEN] =
             drs.into();
@@ -666,85 +666,6 @@ impl VirtualMachine for WhpVm {
         Ok(xsave_buffer)
     }
 
-    fn reset_xsave(&self) -> std::result::Result<(), RegisterError> {
-        // WHP uses compacted XSAVE format (bit 63 of XCOMP_BV set).
-        // We cannot just zero out the xsave area, we need to preserve the XCOMP_BV.
-
-        // Get the required buffer size by calling with NULL buffer.
-        let mut buffer_size_needed: u32 = 0;
-
-        let result = unsafe {
-            WHvGetVirtualProcessorXsaveState(
-                self.partition,
-                0,
-                std::ptr::null_mut(),
-                0,
-                &mut buffer_size_needed,
-            )
-        };
-
-        // Expect insufficient buffer error; any other error is unexpected
-        if let Err(e) = result
-            && e.code() != windows::Win32::Foundation::WHV_E_INSUFFICIENT_BUFFER
-        {
-            return Err(RegisterError::GetXsaveSize(e.into()));
-        }
-
-        if buffer_size_needed < XSAVE_MIN_SIZE as u32 {
-            return Err(RegisterError::XsaveSizeMismatch {
-                expected: XSAVE_MIN_SIZE as u32,
-                actual: buffer_size_needed,
-            });
-        }
-
-        // Create a buffer to hold the current state (to get the correct XCOMP_BV)
-        let mut current_state = vec![0u8; buffer_size_needed as usize];
-        let mut written_bytes = 0;
-        unsafe {
-            WHvGetVirtualProcessorXsaveState(
-                self.partition,
-                0,
-                current_state.as_mut_ptr() as *mut std::ffi::c_void,
-                buffer_size_needed,
-                &mut written_bytes,
-            )
-            .map_err(|e| RegisterError::GetXsave(e.into()))?;
-        };
-
-        // Zero out most of the buffer, preserving only XCOMP_BV (520-528).
-        // Extended components with XSTATE_BV bit=0 will use their init values.
-        //
-        // - Legacy region (0-512): x87 FPU + SSE state
-        // - XSTATE_BV (512-520): Feature bitmap
-        // - XCOMP_BV (520-528): Compaction bitmap + format bit (KEEP)
-        // - Reserved (528-576): Header padding
-        // - Extended (576+): AVX, AVX-512, MPX, PKRU, AMX, etc.
-        current_state[0..520].fill(0);
-        current_state[528..].fill(0);
-
-        // 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)
-        current_state[0..2].copy_from_slice(&FP_CONTROL_WORD_DEFAULT.to_le_bytes());
-        current_state[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.
-        current_state[512..520].copy_from_slice(&0x3u64.to_le_bytes());
-
-        unsafe {
-            WHvSetVirtualProcessorXsaveState(
-                self.partition,
-                0,
-                current_state.as_ptr() as *const std::ffi::c_void,
-                buffer_size_needed,
-            )
-            .map_err(|e| RegisterError::SetXsave(e.into()))?;
-        }
-
-        Ok(())
-    }
-
     #[cfg(test)]
     #[cfg(not(feature = "nanvix-unstable"))]
     fn set_xsave(&self, xsave: &[u32]) -> std::result::Result<(), RegisterError> {
@@ -791,6 +712,36 @@ impl VirtualMachine for WhpVm {
         Ok(())
     }
 
+    fn reset_partition(&mut self) -> std::result::Result<(), RegisterError> {
+        unsafe {
+            WHvResetPartition(self.partition)
+                .map_err(|e| RegisterError::ResetPartition(e.into()))?;
+
+            // WHvResetPartition resets LAPIC to power-on defaults.
+            // Re-initialize it when LAPIC emulation is active.
+            //
+            // set_sregs filters out APIC_BASE (the snapshot sregs
+            // default it to 0, which would disable the LAPIC). This
+            // means APIC_BASE is never written after the VMCB rebuild,
+            // so the AVIC clean bit stays set and the CPU may use stale
+            // cached AVIC state. Re-writing the post-reset value forces
+            // the hypervisor to dirty SVM_CLEAN_FIELD_AVIC.
+            #[cfg(feature = "hw-interrupts")]
+            {
+                Self::init_lapic_bulk(self.partition)
+                    .map_err(|e| RegisterError::ResetPartition(e.into()))?;
+
+                let apic_base = self.sregs()?.apic_base;
+                self.set_registers(&[(
+                    WHvX64RegisterApicBase,
+                    Align16(WHV_REGISTER_VALUE { Reg64: apic_base }),
+                )])
+                .map_err(|e| RegisterError::ResetPartition(e.into()))?;
+            }
+        }
+        Ok(())
+    }
+
     /// Get the partition handle for this VM
     fn partition_handle(&self) -> WHV_PARTITION_HANDLE {
         self.partition