wasmtime: Ensure that Func::wrap'd return values are compatible with the current store

crates/wasmtime/src/func.rs

@@ -202,7 +202,18 @@ macro_rules! getters {
                     let weak_store = instance.store.weak();
                     let weak_store = WeakStore(&weak_store);
 
-                    $( let $args = $args.into_abi_for_arg(weak_store); )*
+                    $(
+                        // Because this returned closure is not marked `unsafe`,
+                        // we have to check that incoming values are compatible
+                        // with our store.
+                        if !$args.compatible_with_store(weak_store) {
+                            return Err(Trap::new(
+                                "attempt to pass cross-`Store` value to Wasm as function argument"
+                            ));
+                        }
+
+                        let $args = $args.into_abi_for_arg(weak_store);
+                    )*
 
                     invoke_wasm_and_catch_traps(anyfunc.as_ref().vmctx, &instance.store, || {
                         ret = Some(fnptr(
@@ -827,6 +838,10 @@ pub unsafe trait WasmTy {
     #[doc(hidden)]
     type Abi: Copy;
 
+    // Is this value compatible with the given store?
+    #[doc(hidden)]
+    fn compatible_with_store<'a>(&self, store: WeakStore<'a>) -> bool;
+
     // Convert this value into its ABI representation, when passing a value into
     // Wasm as an argument.
     #[doc(hidden)]
@@ -871,6 +886,10 @@ pub unsafe trait WasmRet {
     #[doc(hidden)]
     type Abi: Copy;
 
+    // Same as `WasmTy::compatible_with_store`.
+    #[doc(hidden)]
+    fn compatible_with_store<'a>(&self, store: WeakStore<'a>) -> bool;
+
     // Similar to `WasmTy::into_abi_for_arg` but used when host code is
     // returning a value into Wasm, rather than host code passing an argument to
     // a Wasm call. Unlike `into_abi_for_arg`, implementors of this method can
@@ -904,6 +923,11 @@ pub unsafe trait WasmRet {
 unsafe impl WasmTy for () {
     type Abi = Self;
 
+    #[inline]
+    fn compatible_with_store<'a>(&self, _store: WeakStore<'a>) -> bool {
+        true
+    }
+
     #[inline]
     fn into_abi_for_arg<'a>(self, _store: WeakStore<'a>) -> Self::Abi {}
 
@@ -926,6 +950,11 @@ unsafe impl WasmTy for () {
 unsafe impl WasmTy for i32 {
     type Abi = Self;
 
+    #[inline]
+    fn compatible_with_store<'a>(&self, _store: WeakStore<'a>) -> bool {
+        true
+    }
+
     #[inline]
     fn into_abi_for_arg<'a>(self, _store: WeakStore<'a>) -> Self::Abi {
         self
@@ -966,6 +995,11 @@ unsafe impl WasmTy for i32 {
 unsafe impl WasmTy for u32 {
     type Abi = <i32 as WasmTy>::Abi;
 
+    #[inline]
+    fn compatible_with_store<'a>(&self, _store: WeakStore<'a>) -> bool {
+        true
+    }
+
     #[inline]
     fn into_abi_for_arg<'a>(self, _store: WeakStore<'a>) -> Self::Abi {
         self as i32
@@ -998,6 +1032,11 @@ unsafe impl WasmTy for u32 {
 unsafe impl WasmTy for i64 {
     type Abi = Self;
 
+    #[inline]
+    fn compatible_with_store<'a>(&self, _store: WeakStore<'a>) -> bool {
+        true
+    }
+
     #[inline]
     fn into_abi_for_arg<'a>(self, _store: WeakStore<'a>) -> Self::Abi {
         self
@@ -1038,6 +1077,11 @@ unsafe impl WasmTy for i64 {
 unsafe impl WasmTy for u64 {
     type Abi = <i64 as WasmTy>::Abi;
 
+    #[inline]
+    fn compatible_with_store<'a>(&self, _store: WeakStore<'a>) -> bool {
+        true
+    }
+
     #[inline]
     fn into_abi_for_arg<'a>(self, _store: WeakStore<'a>) -> Self::Abi {
         self as i64
@@ -1070,6 +1114,11 @@ unsafe impl WasmTy for u64 {
 unsafe impl WasmTy for f32 {
     type Abi = Self;
 
+    #[inline]
+    fn compatible_with_store<'a>(&self, _store: WeakStore<'a>) -> bool {
+        true
+    }
+
     #[inline]
     fn into_abi_for_arg<'a>(self, _store: WeakStore<'a>) -> Self::Abi {
         self
@@ -1110,6 +1159,11 @@ unsafe impl WasmTy for f32 {
 unsafe impl WasmTy for f64 {
     type Abi = Self;
 
+    #[inline]
+    fn compatible_with_store<'a>(&self, _store: WeakStore<'a>) -> bool {
+        true
+    }
+
     #[inline]
     fn into_abi_for_arg<'a>(self, _store: WeakStore<'a>) -> Self::Abi {
         self
@@ -1150,6 +1204,11 @@ unsafe impl WasmTy for f64 {
 unsafe impl WasmTy for Option<ExternRef> {
     type Abi = *mut u8;
 
+    #[inline]
+    fn compatible_with_store<'a>(&self, _store: WeakStore<'a>) -> bool {
+        true
+    }
+
     #[inline]
     fn into_abi_for_arg<'a>(self, store: WeakStore<'a>) -> Self::Abi {
         if let Some(x) = self {
@@ -1205,6 +1264,16 @@ unsafe impl WasmTy for Option<ExternRef> {
 unsafe impl WasmTy for Option<Func> {
     type Abi = *mut wasmtime_runtime::VMCallerCheckedAnyfunc;
 
+    #[inline]
+    fn compatible_with_store<'a>(&self, store: WeakStore<'a>) -> bool {
+        if let Some(f) = self {
+            let store = Store::upgrade(store.0).unwrap();
+            Store::same(&store, f.store())
+        } else {
+            true
+        }
+    }
+
     #[inline]
     fn into_abi_for_arg<'a>(self, _store: WeakStore<'a>) -> Self::Abi {
         if let Some(f) = self {
@@ -1251,6 +1320,11 @@ where
 {
     type Abi = <T as WasmTy>::Abi;
 
+    #[inline]
+    fn compatible_with_store<'a>(&self, store: WeakStore<'a>) -> bool {
+        <Self as WasmTy>::compatible_with_store(self, store)
+    }
+
     #[inline]
     unsafe fn into_abi_for_ret<'a>(self, store: WeakStore<'a>) -> Self::Abi {
         <Self as WasmTy>::into_abi_for_arg(self, store)
@@ -1288,6 +1362,14 @@ where
 {
     type Abi = <T as WasmTy>::Abi;
 
+    #[inline]
+    fn compatible_with_store<'a>(&self, store: WeakStore<'a>) -> bool {
+        match self {
+            Ok(x) => <T as WasmTy>::compatible_with_store(x, store),
+            Err(_) => true,
+        }
+    }
+
     #[inline]
     unsafe fn into_abi_for_ret<'a>(self, store: WeakStore<'a>) -> Self::Abi {
         match self {
@@ -1419,6 +1501,27 @@ impl Caller<'_> {
     }
 }
 
+#[inline(never)]
+#[cold]
+unsafe fn raise_cross_store_trap() -> ! {
+    #[derive(Debug)]
+    struct CrossStoreError;
+
+    impl std::error::Error for CrossStoreError {}
+
+    impl fmt::Display for CrossStoreError {
+        fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+            write!(
+                f,
+                "host function attempted to return cross-`Store` \
+                 value to Wasm",
+            )
+        }
+    }
+
+    raise_user_trap(Box::new(CrossStoreError));
+}
+
 macro_rules! impl_into_func {
     ($(
         ($($args:ident)*)
@@ -1482,8 +1585,17 @@ macro_rules! impl_into_func {
                         }))
                     };
                     match ret {
-                        Ok(ret) => ret.into_abi_for_ret(weak_store),
                         Err(panic) => wasmtime_runtime::resume_panic(panic),
+                        Ok(ret) => {
+                            // Because the wrapped function is not `unsafe`, we
+                            // can't assume it returned a value that is
+                            // compatible with this store.
+                            if !ret.compatible_with_store(weak_store) {
+                                raise_cross_store_trap();
+                            }
+
+                            ret.into_abi_for_ret(weak_store)
+                        }
                     }
                 }
 

tests/all/func.rs

@@ -442,3 +442,68 @@ fn func_write_nothing() -> anyhow::Result<()> {
         .contains("function attempted to return an incompatible value"));
     Ok(())
 }
+
+#[test]
+// Note: Cranelift only supports refrerence types (used in the wasm in this
+// test) on x64.
+#[cfg(target_arch = "x86_64")]
+fn return_cross_store_value() -> anyhow::Result<()> {
+    let wasm = wat::parse_str(
+        r#"
+            (import "" "" (func (result funcref)))
+
+            (func (export "run") (result funcref)
+                call 0
+            )
+        "#,
+    )?;
+    let mut config = Config::new();
+    config.wasm_reference_types(true);
+    let engine = Engine::new(&config);
+    let module = Module::new(&engine, &wasm)?;
+
+    let store1 = Store::new(&engine);
+    let store2 = Store::new(&engine);
+
+    let store2_func = Func::wrap(&store2, || {});
+    let return_cross_store_func = Func::wrap(&store1, move || Some(store2_func.clone()));
+
+    let instance = Instance::new(&store1, &module, &[return_cross_store_func.into()])?;
+
+    let run = instance.get_func("run").unwrap();
+    let result = run.call(&[]);
+    assert!(result.is_err());
+    assert!(result.unwrap_err().to_string().contains("cross-`Store`"));
+
+    Ok(())
+}
+
+#[test]
+// Note: Cranelift only supports refrerence types (used in the wasm in this
+// test) on x64.
+#[cfg(target_arch = "x86_64")]
+fn pass_cross_store_arg() -> anyhow::Result<()> {
+    let mut config = Config::new();
+    config.wasm_reference_types(true);
+    let engine = Engine::new(&config);
+
+    let store1 = Store::new(&engine);
+    let store2 = Store::new(&engine);
+
+    let store1_func = Func::wrap(&store1, |_: Option<Func>| {});
+    let store2_func = Func::wrap(&store2, || {});
+
+    // Using regular `.call` fails with cross-Store arguments.
+    assert!(store1_func
+        .call(&[Val::FuncRef(Some(store2_func.clone()))])
+        .is_err());
+
+    // And using `.get` followed by a function call also fails with cross-Store
+    // arguments.
+    let f = store1_func.get1::<Option<Func>, ()>()?;
+    let result = f(Some(store2_func));
+    assert!(result.is_err());
+    assert!(result.unwrap_err().to_string().contains("cross-`Store`"));
+
+    Ok(())
+}