fixup! Implement wasmtime_memory_image APIs

Fixes to some barrier logic, along with some minor formatting nits

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

Author: syntactically

src/hyperlight_wasm_runtime/src/platform.rs

@@ -175,6 +175,15 @@ pub extern "C" fn wasmtime_init_traps(handler: wasmtime_trap_handler_t) -> i32 {
 // Copy a VA range to a new VA. Old and new VA, and len, must be
 // page-aligned.
 fn copy_va_mapping(base: *const u8, len: usize, to_va: *mut u8, remap_original: bool) {
+    // TODO: all this barrier code is amd64 specific. It should be
+    // refactored to use some better architecture-independent APIs.
+    //
+    // On amd64, "upgrades" including the first time that a a valid
+    // translation exists for a VA, only need a light (serialising
+    // instruction) barrier.  Since invlpg is also a barrier, we don't
+    // even need that, if we did do a downgrade remap just before.
+    let mut needs_first_valid_exposure_barrier = false;
+
     // TODO: make this more efficient by directly exposing the ability
     // to traverse an entire VA range in
     // hyperlight_guest_bin::paging::virt_to_phys, and coalescing
@@ -188,18 +197,25 @@ fn copy_va_mapping(base: *const u8, len: usize, to_va: *mut u8, remap_original:
             // Skip unmapped pages, since they will be unmapped in
             // both the original and the new copy
             vmem::MappingKind::Unmapped => continue,
-            vmem::MappingKind::Basic(bm) if bm.writable => (vmem::MappingKind::Cow(vmem::CowMapping {
-                readable: bm.readable,
-                executable: bm.executable,
-            }), true),
-            vmem::MappingKind::Basic(bm) => (vmem::MappingKind::Basic(vmem::BasicMapping {
-                readable: bm.readable,
-                writable: false,
-                executable: bm.executable,
-            }), false),
+            vmem::MappingKind::Basic(bm) if bm.writable => (
+                vmem::MappingKind::Cow(vmem::CowMapping {
+                    readable: bm.readable,
+                    executable: bm.executable,
+                }),
+                true,
+            ),
+            vmem::MappingKind::Basic(bm) => (
+                vmem::MappingKind::Basic(vmem::BasicMapping {
+                    readable: bm.readable,
+                    writable: false,
+                    executable: bm.executable,
+                }),
+                false,
+            ),
             vmem::MappingKind::Cow(cm) => (vmem::MappingKind::Cow(cm), false),
         };
-        if remap_original && was_writable {
+        let do_downgrade = remap_original && was_writable;
+        if do_downgrade {
             // If necessary, remap the original page as Cow, instead
             // of whatever it is now, to ensure that any more writes to
             // that region do not change the image base.
@@ -225,6 +241,21 @@ fn copy_va_mapping(base: *const u8, len: usize, to_va: *mut u8, remap_original:
                 new_kind,
             );
         }
+        if do_downgrade {
+            // Since we have downgraded a page from writable to CoW we
+            // need to do an invlpg on it. Because invlpg is a
+            // serialising instruction, we don't need need the other
+            // barrier for the new mapping.
+            unsafe {
+                core::arch::asm!("invlpg [{}]", in(reg) mapping.virt_base, options(readonly, nostack, preserves_flags));
+            }
+            needs_first_valid_exposure_barrier = false;
+        } else {
+            needs_first_valid_exposure_barrier = true;
+        }
+    }
+    if needs_first_valid_exposure_barrier {
+        paging::barrier::first_valid_same_ctx();
     }
 }
 
@@ -259,7 +290,7 @@ pub extern "C" fn wasmtime_memory_image_map_at(
 
 #[no_mangle]
 pub extern "C" fn wasmtime_memory_image_free(_image: *mut c_void) {
-    /* This should never be called in practice, because we simple
+    /* This should never be called in practice, because we simply
      * restore the snapshot rather than actually unload/destroy instances */
     panic!("wasmtime_memory_image_free");
 }