From 3245e424835463cc69ceec0482344e0a84e7ab4e Mon Sep 17 00:00:00 2001
From: Ludvig Liljenberg <4257730+ludfjig@users.noreply.github.com>
Date: Tue, 3 Nov 2026 16:26:16 -0800
Subject: [PATCH] Restore the guest TSC exactly on WHP

On WHP hosts with TscAdjustSupport, restore does not rewind the guest
TSC. Writing TSC shifts TSC_ADJUST by the same amount, and writing
TSC_ADJUST shifts TSC back, so restoring both keeps the time since the
snapshot.

Clear TscAdjustSupport in the partition's processor feature banks, as
the MSHV driver does. The probe then drops TSC_ADJUST from the reset
set, and restore_rewinds_tsc runs on WHP too.

Signed-off-by: Ludvig Liljenberg <4257730+ludfjig@users.noreply.github.com>
---
 CHANGELOG.md                                  |  3 ++
 docs/msr.md                                   | 10 ++----
 .../src/hypervisor/regs/x86_64/msrs.rs        |  1 -
 .../hypervisor/virtual_machine/whp/x86_64.rs  | 35 +++++++++++++++++++
 .../src/sandbox/initialized.rs                |  4 +--
 5 files changed, 42 insertions(+), 11 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 3d430b0b5..118c3a861 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -4,6 +4,9 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/).
 
 ## [Prerelease] - Unreleased
 
+### Fixed
+* Fix restore on WHP not rewinding the guest TSC on hosts with `TSC_ADJUST` support.
+
 ## [v0.19.0] - 2026-11-04
 
 ### Added
diff --git a/docs/msr.md b/docs/msr.md
index aa5bf3794..689a738d7 100644
--- a/docs/msr.md
+++ b/docs/msr.md
@@ -180,13 +180,9 @@ MTRRs cannot be read.
 
 ### TSC
 
-Hyper-V stores `TSC` and `TSC_ADJUST` independently. While time runs, it
-preserves `TSC - TSC_ADJUST`: writing `TSC` also changes `TSC_ADJUST`, while
-writing `TSC_ADJUST` changes the internal TSC offset.
-
-Restore writes `TSC` before `TSC_ADJUST`. The two writes remove the guest's
-delta without freezing partition time. KVM also restores `TSC` so all backends
-use the same guest-time semantics.
+Restore sets the guest TSC back to its snapshot value. MSHV and WHP partitions
+disable `TscAdjustSupport`. That feature makes a write to `TSC` or `TSC_ADJUST`
+move the other by the same amount, so restoring both would undo the rewind.
 
 ## Retained-state inventory
 
diff --git a/src/hyperlight_host/src/hypervisor/regs/x86_64/msrs.rs b/src/hyperlight_host/src/hypervisor/regs/x86_64/msrs.rs
index 775c26a18..119c4076f 100644
--- a/src/hyperlight_host/src/hypervisor/regs/x86_64/msrs.rs
+++ b/src/hyperlight_host/src/hypervisor/regs/x86_64/msrs.rs
@@ -183,7 +183,6 @@ const NON_MTRR_RESETTABLE_MSRS: &[u32] = &[
     // MSHV maps XSS through its U_XSS alias.
     MSR_XSS,
     MSR_TSC,
-    // Hyper-V stores TSC_ADJUST independently from TSC.
     MSR_TSC_ADJUST,
     // Host probing omits TSC_AUX when RDTSCP is unavailable.
     MSR_TSC_AUX,
diff --git a/src/hyperlight_host/src/hypervisor/virtual_machine/whp/x86_64.rs b/src/hyperlight_host/src/hypervisor/virtual_machine/whp/x86_64.rs
index 57342a532..9c0d27247 100644
--- a/src/hyperlight_host/src/hypervisor/virtual_machine/whp/x86_64.rs
+++ b/src/hyperlight_host/src/hypervisor/virtual_machine/whp/x86_64.rs
@@ -286,6 +286,8 @@ impl WhpVm {
             )
             .map_err(|e| CreateVmError::SetPartitionProperty(e.into()))?;
 
+            Self::disable_tsc_adjust(p)?;
+
             #[cfg(feature = "hw-interrupts")]
             Self::enable_lapic_emulation(p)?;
 
@@ -320,6 +322,39 @@ impl WhpVm {
         })
     }
 
+    /// `TscAdjustSupport` makes a write to `TSC` or `TSC_ADJUST` move the other
+    /// by the same amount, so restoring both would undo the TSC rewind.
+    fn disable_tsc_adjust(partition: WHV_PARTITION_HANDLE) -> Result<(), CreateVmError> {
+        // Bit 18 of `WHV_PROCESSOR_FEATURES1`, the second bank.
+        const TSC_ADJUST_SUPPORT: u64 = 1 << 18;
+
+        let mut banks = WHV_PROCESSOR_FEATURES_BANKS {
+            BanksCount: 2,
+            ..Default::default()
+        };
+        let size = std::mem::size_of_val(&banks) as u32;
+        // SAFETY: `banks` is a buffer of `size` bytes, and every view of its
+        // union is plain `u64` data.
+        unsafe {
+            WHvGetPartitionProperty(
+                partition,
+                WHvPartitionPropertyCodeProcessorFeaturesBanks,
+                &mut banks as *mut _ as *mut c_void,
+                size,
+                None,
+            )
+            .map_err(|e| CreateVmError::GetPartitionProperty(e.into()))?;
+            banks.Anonymous.AsUINT64[1] &= !TSC_ADJUST_SUPPORT;
+            WHvSetPartitionProperty(
+                partition,
+                WHvPartitionPropertyCodeProcessorFeaturesBanks,
+                &banks as *const _ as *const c_void,
+                size,
+            )
+            .map_err(|e| CreateVmError::SetPartitionProperty(e.into()))
+        }
+    }
+
     fn get_registers(
         &self,
         names: &[WHV_REGISTER_NAME],
diff --git a/src/hyperlight_host/src/sandbox/initialized.rs b/src/hyperlight_host/src/sandbox/initialized.rs
index db3b7bcec..cc2c6830b 100644
--- a/src/hyperlight_host/src/sandbox/initialized.rs
+++ b/src/hyperlight_host/src/sandbox/initialized.rs
@@ -4832,10 +4832,8 @@ mod tests {
 
         /// KVM ignores a host TSC write within one second of the value it
         /// expects, so restores closer together than that exercise the KVM path.
-        /// The 500 ms gaps tolerate host stalls shorter than that. WHP is
-        /// excluded because it keeps elapsed time on hosts with `TSC_ADJUST`.
+        /// The 500 ms gaps tolerate host stalls shorter than that.
         #[test]
-        #[cfg(any(kvm, mshv3))]
         fn restore_rewinds_tsc() {
             let mut sbox = SandboxBuilder::from_file(simple_guest_as_pathbuf())
                 .build()