From 8058157b060c321a43ff07367b4b2bbbc6100a7b Mon Sep 17 00:00:00 2001
From: Andrew Brown <andrew.brown@intel.com>
Date: Mon, 19 Sep 2022 15:46:15 -0700
Subject: [PATCH] wiggle: allow wiggle to use shared memory

`wiggle` looks for an exported `Memory` named `"memory"` to use for its
guest slices. This change allows it to use a `SharedMemory` if this is
the kind of memory used for the export.

It is `unsafe` to use shared memory in Wiggle because of broken Rust
guarantees: previously, Wiggle could hand out slices to WebAssembly
linear memory that could be concurrently modified by some other thread.
With the introduction of Wiggle's new `UnsafeGuestSlice` (#5225, #5229,
 #5264), Wiggle should now correctly communicate its guarantees through
its API.
---
 crates/wiggle/generate/src/wasmtime.rs | 20 +++++++++++++++-----
 1 file changed, 15 insertions(+), 5 deletions(-)

diff --git a/crates/wiggle/generate/src/wasmtime.rs b/crates/wiggle/generate/src/wasmtime.rs
index be3603754645..aaffe8989c0a 100644
--- a/crates/wiggle/generate/src/wasmtime.rs
+++ b/crates/wiggle/generate/src/wasmtime.rs
@@ -112,13 +112,23 @@ fn generate_func(
     };
 
     let body = quote! {
-        let mem = match caller.get_export("memory") {
-            Some(#rt::wasmtime_crate::Extern::Memory(m)) => m,
+        let (mem, ctx, shared) = match caller.get_export("memory") {
+            Some(#rt::wasmtime_crate::Extern::Memory(m)) => {
+                let (mem, ctx) = m.data_and_store_mut(&mut caller);
+                let ctx = get_cx(ctx);
+                (mem, ctx, false)
+            }
+            Some(#rt::wasmtime_crate::Extern::SharedMemory(m)) => {
+                // SAFETY: `WasmtimeGuestMemory` never hands out the mutable
+                // slice and the `unsafe`-ness is propagated by the `shared`
+                // flag.
+                let mem = unsafe { std::mem::transmute::<&[std::cell::UnsafeCell<u8>], &mut [u8]>(m.data()) };
+                let ctx = get_cx(caller.data_mut());
+                (mem, ctx, true)
+            }
             _ => #rt::anyhow::bail!("missing required memory export"),
         };
-        let (mem , ctx) = mem.data_and_store_mut(&mut caller);
-        let ctx = get_cx(ctx);
-        let mem = #rt::wasmtime::WasmtimeGuestMemory::new(mem, false);
+        let mem = #rt::wasmtime::WasmtimeGuestMemory::new(mem, shared);
         Ok(<#ret_ty>::from(#abi_func(ctx, &mem #(, #arg_names)*) #await_ ?))
     };