amd64: add missing barrier for certain page-table updates

Certain page-table updates that do not require TLB invalidation were
previously missing any kind of fence at all to keep table walks
coherent. This commit introduces a new barrier for this scenario.

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

Author: syntactically

src/hyperlight_guest_bin/src/arch/amd64/exception/handle.rs

@@ -76,6 +76,12 @@ fn handle_stack_pagefault(gva: u64) {
                 executable: false,
             }),
         );
+        // This is mapping an entry for the first time, so we don't
+        // need to worry about TLB maintenance.  We do (probably) need
+        // to worry about coherence between these stores and the page
+        // table walker, but this page won't be accessed again until
+        // after [`iret`], which is a serializing instruction, so
+        // that's already handled as well.
     }
 }
 
@@ -113,6 +119,8 @@ fn handle_cow_pagefault(_phys: PhysAddr, virt: VirtAddr, perms: CowMapping) {
                 executable: perms.executable,
             }),
         );
+        // This is updating an entry that was already valid, changing
+        // its OA, so we need to actually invalidate the TLB for it.
         core::arch::asm!("invlpg [{}]", in(reg) target_virt, options(readonly, nostack, preserves_flags));
     }
 }

src/hyperlight_guest_bin/src/arch/amd64/init.rs

@@ -121,6 +121,7 @@ unsafe fn init_stack() -> u64 {
                 executable: false,
             }),
         );
+        crate::paging::barrier::first_valid_same_ctx();
     }
     MAIN_STACK_TOP_GVA
 }

src/hyperlight_guest_bin/src/arch/amd64/machine.rs

@@ -229,6 +229,7 @@ impl ProcCtrl {
                     executable: false,
                 }),
             );
+            crate::paging::barrier::first_valid_same_ctx();
             let ptr = ptr as *mut Self;
             (&raw mut (*ptr).gdt).write_bytes(0u8, 1);
             (&raw mut (*ptr).tss).write_bytes(0u8, 1);

src/hyperlight_guest_bin/src/paging.rs

@@ -144,3 +144,58 @@ 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().try_phys_to_virt(gpa)
 }
+
+/// Barriers that other code may need to use when updating page tables
+pub mod barrier {
+    /// Call this function when a virtual address has just been made
+    /// valid for the first time after the last tlb invalidate that
+    /// affected it, and it will be used for the first time in the
+    /// same execution context as has made the modification.
+    ///
+    /// On most architectures, TLBs will not cache invalid entries, so
+    /// this does not need to issue a TLB. However, it does need to
+    /// ensure coherency between the previous writes and any future
+    /// uses by a page table walker.
+    ///
+    /// # Architecture-specific (amd64) notes
+    ///
+    /// The exact details around page walk coherency on amd64 seem a
+    /// bit fuzzy. The Intel manual notes that a serialising
+    /// instruction is necessary specifically to synchronise table
+    /// walks performed during instruction fetch [1], but is
+    /// relatively quiet about other page walks. The AMD manual notes
+    /// [2] that "a table entry is allowed to be upgraded (by marking
+    /// it as present, or by removing its write, execute or supervisor
+    /// restrictions) without explicitly maintaining TLB coherency",
+    /// but only states that TLB any upper-level TLB cache entries
+    /// will be flushed before re-walking to confirm the fault, which
+    /// does not clearly seem strong enough.
+    ///
+    /// In some limited testing, `mfence` typically seems to be
+    /// enough, but as it is not a serializing instruction on Intel
+    /// platforms, we assume it may not be quite good enough.  `cpuid`
+    /// is likely to be very slow, since we are definitely running
+    /// under a hypervisor (and often even nested). Currently, for
+    /// simplicity's sake, this just copies cr0 to itself, but other
+    /// options (including the `serialize` instruction where
+    /// available) could be worth exploring.
+    ///
+    /// [1] Intel 64 and IA-32 Architectures Software Developer's Manual, Volume 3: System Programming Guide
+    ///         Chapter 5: Paging
+    ///             §5.10: Caching Translation Information
+    ///                 §5.10.4: Invalidation of TLBs and Paging-Structure Caches
+    ///                     §5.10.4.3: Optional Invalidation
+    /// [2] AMD64 Architecture Programmer's Manual, Volume 2: System Programming
+    ///         Section 5: Page Translation and Protection
+    ///             §5.5: Translation-Lookaside Buffer
+    ///                 §5.5.3: TLB Management
+    #[inline(always)]
+    pub fn first_valid_same_ctx() {
+        unsafe {
+            core::arch::asm!("
+                mov rax, cr0
+                mov cr0, rax
+            ", out("rax") _);
+        }
+    }
+}

src/tests/rust_guests/simpleguest/src/main.rs

@@ -593,6 +593,7 @@ fn read_mapped_buffer(base: u64, len: u64, do_map: bool) -> Vec<u8> {
                     executable: true,
                 }),
             );
+            hyperlight_guest_bin::paging::barrier::first_valid_same_ctx();
         }
     }
 
@@ -623,7 +624,8 @@ fn write_mapped_buffer(base: u64, len: u64) -> bool {
                 writable: true,
                 executable: true,
             }),
-        )
+        );
+        hyperlight_guest_bin::paging::barrier::first_valid_same_ctx();
     };
 
     let data = unsafe { core::slice::from_raw_parts_mut(base, len) };
@@ -650,7 +652,8 @@ fn exec_mapped_buffer(base: u64, len: u64) -> bool {
                 writable: true,
                 executable: true,
             }),
-        )
+        );
+        hyperlight_guest_bin::paging::barrier::first_valid_same_ctx();
     };
 
     let data = unsafe { core::slice::from_raw_parts(base, len) };