amd64: slightly improve TLB flushing

This is still somewhat broken (see
#721), since it
relies on the flush_tlb function ending up in the big area around the
code that is still identity mapped. So, this code should probably move
somewhere else. However, this fixes a very significant bug in the
prior code, where the flush_tlb() function called on guest function
dispatch could return to the wrong address, due to the correct return
address having been pushed onto the wrong page due to a stale TLB
entry.

Signed-off-by: Lucy Menon <168595099+syntactically@users.noreply.github.com>

Author: syntactically

src/hyperlight_guest_bin/src/arch/amd64/dispatch.rs

@@ -0,0 +1,70 @@
+/*
+Copyright 2025  The Hyperlight Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+ */
+
+//! This module contains the architecture-specific dispatch function
+//! entry sequence
+
+// The extern "C" here is a lie until #![feature(abi_custom)] is stabilised
+unsafe extern "C" {
+    /// There are two reasons that we need an architecture-specific
+    /// assembly stub currently: (1) to ensure that the fact that the
+    /// dispatch function never returns but can be executed again is
+    /// OK (i.e. we must ensure that there is no stack frame
+    /// teardown); (2) at least presently, to ensure that the TLB is
+    /// flushed in certain cases.
+    ///
+    /// # TLB flushing
+    ///
+    /// The hyperlight host likes to use one partition and reset it in
+    /// various ways; if that has happened, there might stale TLB
+    /// entries hanging around from the former user of the
+    /// partition. Flushing the TLB here is not quite the right thing
+    /// to do, since incorrectly cached entries could make even this
+    /// code not exist, but regrettably there is not a simple way for
+    /// the host to trigger flushing when it ought to happen, so for
+    /// now this works in practice, since the text segment is always
+    /// part of the big identity-mapped region at the base of the
+    /// guest.  The stack, however, is not part of the snapshot region
+    /// which is (in practice) idmapped for the relevant area, so this
+    /// cannot touch the stack until the flush is done.
+    ///
+    /// Currently this just always flips CR4.PGE back and forth to
+    /// trigger a tlb flush. We should use a faster approach where
+    /// available
+    ///
+    /// # ABI
+    ///
+    /// The ZF should be set if a TLB flush is required on this call
+    /// (e.g. the first call after a snapshot restore)
+    ///
+    /// The stack pointer should, unusually for amd64, by 16-byte
+    /// aligned at the beginning of the function---no return address
+    /// should be pushed.
+    pub(crate) unsafe fn dispatch_function(must_flush_tlb: u64);
+}
+core::arch::global_asm!("
+    .global dispatch_function
+    dispatch_function:
+    jnz flush_done
+    mov rdi, cr4
+    xor rdi, 0x80
+    mov cr4, rdi
+    xor rdi, 0x80
+    mov cr4, rdi
+    flush_done:
+    call {internal_dispatch_function}\n
+    hlt\n
+", internal_dispatch_function = sym crate::guest_function::call::internal_dispatch_function);

src/hyperlight_guest_bin/src/arch/amd64/mod.rs

@@ -15,6 +15,7 @@ limitations under the License.
  */
 
 pub(crate) mod context;
+pub(crate) mod dispatch;
 pub mod exception;
 mod init;
 mod layout;

src/hyperlight_guest_bin/src/guest_function/call.rs

@@ -22,7 +22,6 @@ use hyperlight_common::flatbuffer_wrappers::function_call::{FunctionCall, Functi
 use hyperlight_common::flatbuffer_wrappers::function_types::{FunctionCallResult, ParameterType};
 use hyperlight_common::flatbuffer_wrappers::guest_error::{ErrorCode, GuestError};
 use hyperlight_guest::error::{HyperlightGuestError, Result};
-use hyperlight_guest::exit::halt;
 use tracing::{Span, instrument};
 
 use crate::{GUEST_HANDLE, REGISTERED_GUEST_FUNCTIONS};
@@ -79,14 +78,8 @@ pub(crate) fn call_guest_function(function_call: FunctionCall) -> Result<Vec<u8>
     }
 }
 
-// This function is marked as no_mangle/inline to prevent the compiler from inlining it , if its inlined the epilogue will not be called
-// and we will leak memory as the epilogue will not be called as halt() is not going to return.
-//
-// This function may panic, as we have no other ways of dealing with errors at this level
-#[unsafe(no_mangle)]
-#[inline(never)]
 #[instrument(skip_all, parent = Span::current(), level= "Trace")]
-fn internal_dispatch_function() {
+pub(crate) fn internal_dispatch_function() {
     let handle = unsafe { GUEST_HANDLE };
 
     #[cfg(debug_assertions)]
@@ -115,34 +108,3 @@ fn internal_dispatch_function() {
         }
     }
 }
-
-// This is implemented as a separate function to make sure that epilogue in the internal_dispatch_function is called before the halt()
-// which if it were included in the internal_dispatch_function cause the epilogue to not be called because the halt() would not return
-// when running in the hypervisor.
-#[instrument(skip_all, parent = Span::current(), level= "Trace")]
-pub(crate) extern "C" fn dispatch_function() {
-    // The hyperlight host likes to use one partition and reset it in
-    // various ways; if that has happened, there might stale TLB
-    // entries hanging around from the former user of the
-    // partition. Flushing the TLB here is not quite the right thing
-    // to do, since incorrectly cached entries could make even this
-    // code not exist, but regrettably there is not a simple way for
-    // the host to trigger flushing when it ought to happen, so for
-    // now this works in practice, since the text segment is always
-    // part of the big identity-mapped region at the base of the
-    // guest.
-    crate::paging::flush_tlb();
-
-    // Read the current TSC to report it to the host with the spans/events
-    // This helps calculating the timestamps relative to the guest call
-    #[cfg(feature = "trace_guest")]
-    {
-        let guest_start_tsc = hyperlight_guest_tracing::invariant_tsc::read_tsc();
-        // Reset the trace state for the new guest function call with the new start TSC
-        // This clears any existing spans/events from previous calls ensuring a clean state
-        hyperlight_guest_tracing::new_call(guest_start_tsc);
-    }
-
-    internal_dispatch_function();
-    halt();
-}

src/hyperlight_guest_bin/src/lib.rs

@@ -20,8 +20,8 @@ extern crate alloc;
 
 use core::fmt::Write;
 
+use arch::dispatch::dispatch_function;
 use buddy_system_allocator::LockedHeap;
-use guest_function::call::dispatch_function;
 use guest_function::register::GuestFunctionRegister;
 use guest_logger::init_logger;
 use hyperlight_common::flatbuffer_wrappers::guest_error::ErrorCode;

src/hyperlight_guest_bin/src/paging.rs

@@ -144,22 +144,3 @@ pub fn virt_to_phys(gva: vmem::VirtAddr) -> impl Iterator<Item = vmem::Mapping>
 pub fn phys_to_virt(gpa: vmem::PhysAddr) -> Option<*mut u8> {
     GuestMappingOperations::new().fallible_phys_to_virt(gpa)
 }
-
-pub fn flush_tlb() {
-    // Currently this just always flips CR4.PGE back and forth to
-    // trigger a tlb flush. We should use a faster approach where
-    // available
-    let mut orig_cr4: u64;
-    unsafe {
-        asm!("mov {}, cr4", out(reg) orig_cr4);
-    }
-    let tmp_cr4: u64 = orig_cr4 ^ (1 << 7); // CR4.PGE
-    unsafe {
-        asm!(
-            "mov cr4, {}",
-            "mov cr4, {}",
-            in(reg) tmp_cr4,
-            in(reg) orig_cr4
-        );
-    }
-}

src/hyperlight_host/src/hypervisor/hyperlight_vm.rs

@@ -104,6 +104,8 @@ pub(crate) struct HyperlightVm {
 
     mmap_regions: Vec<(u32, MemoryRegion)>, // Later mapped regions (slot number, region)
 
+    pending_tlb_flush: bool,
+
     #[cfg(gdb)]
     gdb_conn: Option<DebugCommChannel<DebugResponse, DebugMsg>>,
     #[cfg(gdb)]
@@ -421,6 +423,8 @@ impl HyperlightVm {
 
             mmap_regions: Vec::new(),
 
+            pending_tlb_flush: false,
+
             #[cfg(gdb)]
             gdb_conn,
             #[cfg(gdb)]
@@ -625,6 +629,15 @@ impl HyperlightVm {
         Ok(self.vm.sregs()?)
     }
 
+    /// Set the current base page table physical address
+    pub(crate) fn set_root_pt(&mut self, addr: u64) -> Result<(), AccessPageTableError> {
+        let mut sregs = self.vm.sregs()?;
+        sregs.cr3 = addr;
+        self.vm.set_sregs(&sregs)?;
+        self.pending_tlb_flush = true;
+        Ok(())
+    }
+
     /// Get the current stack top virtual address
     pub(crate) fn get_stack_top(&mut self) -> u64 {
         self.rsp_gva
@@ -662,17 +675,25 @@ impl HyperlightVm {
         let NextAction::Call(dispatch_func_addr) = self.entrypoint else {
             return Err(DispatchGuestCallError::Uninitialized);
         };
+        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 {
             rip: dispatch_func_addr,
             // We usually keep the top of the stack 16-byte
-            // aligned. However, the ABI requirement is that the stack
-            // be aligned _before a call instruction_, which means
-            // that the stack needs to actually be ≡ 8 mod 16 at the
-            // first instruction (since, on x64, a call instruction
-            // automatically pushes a return address).
-            rsp: self.rsp_gva - 8,
-            rflags: 1 << 1,
+            // aligned. Since the usual ABI requirement is that the
+            // stack be aligned _before a call instruction_, one might
+            // expect that the stack pointer here needs to actually be
+            // ≡ 8 mod 16 at the first instruction (since, on x64, a
+            // call instruction automatically pushes a return
+            // address).  However, the x64 entry stub in
+            // hyperlight_guest::arch::dispatch handles this itself,
+            // so we do use the aligned address here.
+            rsp: self.rsp_gva,
+            rflags,
             ..Default::default()
         };
         self.vm

src/hyperlight_host/tests/integration_test.rs

@@ -704,13 +704,13 @@ fn log_message() {
     //  - internal_dispatch_function does a log::trace! in debug mode
     //  - logs from trace level tracing spans created as logs because of the tracing `log` feature
     //    - 4 from evolve call (hyperlight_main + halt)
-    //    - 16 from guest call (internal_dispatch_function + others)
+    //    - 12 from guest call (internal_dispatch_function + others)
     // and are multiplied because we make 6 calls to `log_test_messages`
     // NOTE: These numbers need to be updated if log messages or spans are added/removed
     let num_fixed_trace_log = if cfg!(debug_assertions) {
-        (1 + 20) * 6
+        (1 + 16) * 6
     } else {
-        20 * 6
+        16 * 6
     };
 
     let tests = vec![