From 1fffbfb27ec6f4e5c1114e769da374d7b87c4abc Mon Sep 17 00:00:00 2001
From: Ludvig Liljenberg <4257730+ludfjig@users.noreply.github.com>
Date: Fri, 20 Mar 2026 14:15:18 -0700
Subject: [PATCH] Defer clearing tlb flush flag

Signed-off-by: Ludvig Liljenberg <4257730+ludfjig@users.noreply.github.com>
---
 .../src/hypervisor/hyperlight_vm/x86_64.rs    | 22 ++++++++++++-------
 1 file changed, 14 insertions(+), 8 deletions(-)

diff --git a/src/hyperlight_host/src/hypervisor/hyperlight_vm/x86_64.rs b/src/hyperlight_host/src/hypervisor/hyperlight_vm/x86_64.rs
index e8f21e119..d82f2c58b 100644
--- a/src/hyperlight_host/src/hypervisor/hyperlight_vm/x86_64.rs
+++ b/src/hyperlight_host/src/hypervisor/hyperlight_vm/x86_64.rs
@@ -292,7 +292,6 @@ impl HyperlightVm {
         let mut rflags = 1 << 1; // RFLAGS.1 is RES1
         if self.pending_tlb_flush {
             rflags |= 1 << 6; // set ZF if we need a tlb flush done before anything else executes
-            self.pending_tlb_flush = false;
         }
         // set RIP and RSP, reset others
         let regs = CommonRegisters {
@@ -319,13 +318,20 @@ impl HyperlightVm {
             .set_fpu(&CommonFpu::default())
             .map_err(DispatchGuestCallError::SetupRegs)?;
 
-        self.run(
-            mem_mgr,
-            host_funcs,
-            #[cfg(gdb)]
-            dbg_mem_access_fn,
-        )
-        .map_err(DispatchGuestCallError::Run)
+        let result = self
+            .run(
+                mem_mgr,
+                host_funcs,
+                #[cfg(gdb)]
+                dbg_mem_access_fn,
+            )
+            .map_err(DispatchGuestCallError::Run);
+
+        // Clear the TLB flush flag only after run() returns. The guest
+        // may have been cancelled before it executed the flush.
+        self.pending_tlb_flush = false;
+
+        result
     }
 
     /// Resets the following vCPU state: