Fix unsoundness in generated Rust code

This commit fixes an soundness issue in the Rust code generator's generated
code. Previously unsound code was generated when:

* An import had a parameter
* That had either a list or a variant of an aggregate where:
* Where one field was a `list<T>` or `string`
* And another field was an owned resource

In this situation the Rust generator uses an "owned" type for the argument,
such as `MyStruct`. This is done to reflect how ownership of the resource in
the argument is lost when the function is called. The problem with this,
however, is that the rest of bindings generation assumes that imported
arguments are all "borrowed" meaning that raw pointers from lists/strings can
be passed through to the canonical ABI. This is not the case here because the
argument is owned, meaning that the generated code would record an argument to
be passed to the canonical ABI and then promptly deallocate the argument.

The fix here is preceded by a refactoring to how Rust manages owned types to
make the fix possible. The general idea for the fix though is that while `x:
MyStruct` is the argument to the function all references internally are through
`&x` to ensure that it remains rooted as an argument, preserving all pointers
to lists and such. This unfortunately means that ownership can no longer model
movement of resources and instead interior mutability must be used (since we
have to move out of `&Resource<T>` since everything is borrowed). Fixing that,
however, is probably not worth the complexity at this time.

Closes bytecodealliance#817
Closes bytecodealliance/wasmtime#7951

crates/guest-rust/src/lib.rs

@@ -5,8 +5,8 @@ extern crate alloc;
 use alloc::boxed::Box;
 use core::fmt;
 use core::marker;
-use core::mem::ManuallyDrop;
 use core::ops::{Deref, DerefMut};
+use core::sync::atomic::{AtomicU32, Ordering::Relaxed};
 
 #[cfg(feature = "macros")]
 pub use wit_bindgen_rust_macro::*;
@@ -188,7 +188,13 @@ type RawRep<T> = Option<T>;
 /// resources.
 #[repr(transparent)]
 pub struct Resource<T: WasmResource> {
-    handle: u32,
+    // NB: This would ideally be `u32` but it is not. The fact that this has
+    // interior mutability is not exposed in the API of this type except for the
+    // `take_handle` method which is supposed to in theory be private.
+    //
+    // This represents, almost all the time, a valid handle value. When it's
+    // invalid it's stored as `u32::MAX`.
+    handle: AtomicU32,
     _marker: marker::PhantomData<Box<T>>,
 }
 
@@ -215,20 +221,33 @@ pub unsafe trait RustResource: WasmResource {
 impl<T: WasmResource> Resource<T> {
     #[doc(hidden)]
     pub unsafe fn from_handle(handle: u32) -> Self {
+        assert!(handle != u32::MAX);
         Self {
-            handle,
+            handle: AtomicU32::new(handle),
             _marker: marker::PhantomData,
         }
     }
 
+    /// Takes ownership of the handle owned by `resource`.
+    ///
+    /// Note that this ideally would be `into_handle` taking `Resource<T>` by
+    /// ownership. The code generator does not enable that in all situations,
+    /// unfortunately, so this is provided instead.
+    ///
+    /// Also note that `take_handle` is in theory only ever called on values
+    /// owned by a generated function. For example a generated function might
+    /// take `Resource<T>` as an argument but then call `take_handle` on a
+    /// reference to that argument. In that sense the dynamic nature of
+    /// `take_handle` should only be exposed internally to generated code, not
+    /// to user code.
     #[doc(hidden)]
-    pub fn into_handle(resource: Resource<T>) -> u32 {
-        ManuallyDrop::new(resource).handle
+    pub fn take_handle(resource: &Resource<T>) -> u32 {
+        resource.handle.swap(u32::MAX, Relaxed)
     }
 
     #[doc(hidden)]
     pub fn handle(resource: &Resource<T>) -> u32 {
-        resource.handle
+        resource.handle.load(Relaxed)
     }
 
     /// Creates a new Rust-defined resource from the underlying representation
@@ -261,7 +280,7 @@ impl<T: WasmResource> Resource<T> {
         T: RustResource,
     {
         unsafe {
-            let rep = T::rep(resource.handle);
+            let rep = T::rep(resource.handle.load(Relaxed));
             RawRep::take(&mut *(rep as *mut RawRep<T>)).unwrap()
         }
     }
@@ -280,7 +299,7 @@ impl<T: RustResource> Deref for Resource<T> {
 
     fn deref(&self) -> &T {
         unsafe {
-            let rep = T::rep(self.handle);
+            let rep = T::rep(self.handle.load(Relaxed));
             RawRep::as_ref(&*(rep as *const RawRep<T>)).unwrap()
         }
     }
@@ -289,7 +308,7 @@ impl<T: RustResource> Deref for Resource<T> {
 impl<T: RustResource> DerefMut for Resource<T> {
     fn deref_mut(&mut self) -> &mut T {
         unsafe {
-            let rep = T::rep(self.handle);
+            let rep = T::rep(self.handle.load(Relaxed));
             RawRep::as_mut(&mut *(rep as *mut RawRep<T>)).unwrap()
         }
     }
@@ -306,7 +325,15 @@ impl<T: WasmResource> fmt::Debug for Resource<T> {
 impl<T: WasmResource> Drop for Resource<T> {
     fn drop(&mut self) {
         unsafe {
-            T::drop(self.handle);
+            match self.handle.load(Relaxed) {
+                // If this handle was "taken" then don't do anything in the
+                // destructor.
+                u32::MAX => {}
+
+                // ... but otherwise do actually destroy it with the imported
+                // component model intrinsic as defined through `T`.
+                other => T::drop(other),
+            }
         }
     }
 }

crates/rust/src/bindgen.rs

@@ -413,8 +413,8 @@ impl Bindgen for FunctionBindgen<'_, '_> {
                 let rt = self.gen.gen.runtime_path();
                 let resource = dealias(self.gen.resolve, *resource);
                 results.push(match self.gen.gen.resources[&resource].direction {
-                    Direction::Import => format!("({op}).into_handle() as i32"),
-                    Direction::Export => format!("{rt}::Resource::into_handle({op}) as i32"),
+                    Direction::Import => format!("({op}).take_handle() as i32"),
+                    Direction::Export => format!("{rt}::Resource::take_handle(&{op}) as i32"),
                 });
             }
 

crates/rust/src/interface.rs

@@ -703,7 +703,6 @@ impl InterfaceGenerator<'_> {
             }
             let name = to_rust_ident(name);
             self.push_str(&name);
-            params.push(name);
             self.push_str(": ");
 
             // Select the "style" of mode that the parameter's type will be
@@ -731,6 +730,34 @@ impl InterfaceGenerator<'_> {
             let mode = self.type_mode_for(param, style, "'_");
             self.print_ty(param, mode);
             self.push_str(",");
+
+            // Depending on the style of this request vs what we got perhaps
+            // change how this argument is used.
+            //
+            // If the `mode` that was selected matches the requested style, then
+            // everything is as expected and the argument should be used as-is.
+            // If it differs though then that means that we requested a borrowed
+            // mode but a different mode ended up being selected. This situation
+            // indicates for example that an argument to an import should be
+            // borrowed but the argument's type means that it can't be borrowed.
+            // For example all arguments to imports are borrowed by default but
+            // owned resources cannot ever be borrowed, so they pop out here as
+            // owned instead.
+            //
+            // In such a situation the lower code still expects to be operating
+            // over borrows. For example raw pointers from lists are passed to
+            // the canonical ABI layer assuming that the lists are "rooted" by
+            // the caller. To uphold this invariant a borrow of the argument is
+            // recorded as the name of this parameter. That ensures that all
+            // access to the parameter is done indirectly which pretends, at
+            // least internally, that the argument was borrowed. The original
+            // motivation for this was #817.
+            if mode.style == style {
+                params.push(name);
+            } else {
+                assert!(style != TypeOwnershipStyle::Owned);
+                params.push(format!("&{name}"));
+            }
         }
         self.push_str(")");
         params
@@ -1772,8 +1799,8 @@ impl<'a> wit_bindgen_core::InterfaceGenerator<'a> for InterfaceGenerator<'a> {
                         }}
 
                         #[doc(hidden)]
-                        pub fn into_handle(self) -> u32 {{
-                            {rt}::Resource::into_handle(self.handle)
+                        pub fn take_handle(&self) -> u32 {{
+                            {rt}::Resource::take_handle(&self.handle)
                         }}
 
                         #[doc(hidden)]

tests/runtime/ownership.rs

@@ -1,15 +1,24 @@
 use anyhow::Result;
+use wasmtime::component::Resource;
 use wasmtime::Store;
 
-wasmtime::component::bindgen!(in "tests/runtime/ownership");
+wasmtime::component::bindgen!({
+    path: "tests/runtime/ownership",
+    with: {
+        "test:ownership/both-list-and-resource/the-resource": MyResource,
+    },
+});
 
 #[derive(Default)]
 pub struct MyImports {
     called_foo: bool,
     called_bar: bool,
     called_baz: bool,
+    last_resource_list: Option<Vec<String>>,
 }
 
+pub struct MyResource;
+
 impl lists::Host for MyImports {
     fn foo(&mut self, list: Vec<Vec<String>>) -> Result<Vec<Vec<String>>> {
         self.called_foo = true;
@@ -31,6 +40,31 @@ impl thing_in_and_out::Host for MyImports {
     }
 }
 
+impl test::ownership::both_list_and_resource::Host for MyImports {
+    fn list_and_resource(
+        &mut self,
+        value: test::ownership::both_list_and_resource::Thing,
+    ) -> Result<()> {
+        assert_eq!(value.b.rep(), 100);
+        assert!(value.b.owned());
+        let expected = self.last_resource_list.as_ref().unwrap();
+        assert_eq!(value.a, *expected);
+        Ok(())
+    }
+}
+
+impl test::ownership::both_list_and_resource::HostTheResource for MyImports {
+    fn new(&mut self, list: Vec<String>) -> Result<Resource<MyResource>> {
+        assert!(self.last_resource_list.is_none());
+        self.last_resource_list = Some(list);
+        Ok(Resource::new_own(100))
+    }
+
+    fn drop(&mut self, _val: Resource<MyResource>) -> Result<()> {
+        unreachable!()
+    }
+}
+
 #[test]
 fn run() -> Result<()> {
     for name in ["owning", "borrowing", "borrowing-duplicate-if-necessary"] {
@@ -52,6 +86,7 @@ fn run_test(exports: Ownership, store: &mut Store<crate::Wasi<MyImports>>) -> Re
     assert!(store.data().0.called_foo);
     assert!(store.data().0.called_bar);
     assert!(store.data().0.called_baz);
+    assert!(store.data().0.last_resource_list.is_some());
 
     Ok(())
 }

tests/runtime/ownership/borrowing-duplicate-if-necessary.rs

@@ -40,5 +40,14 @@ impl Guest for Exports {
             },
             thing_in_and_out::baz(value)
         );
+
+        let strings = vec!["foo", "bar", "baz"];
+        let resource = test::ownership::both_list_and_resource::TheResource::new(&strings);
+        test::ownership::both_list_and_resource::list_and_resource(
+            test::ownership::both_list_and_resource::Thing {
+                a: strings.iter().map(|s| s.to_string()).collect(),
+                b: resource,
+            },
+        );
     }
 }

tests/runtime/ownership/borrowing.rs

@@ -34,5 +34,14 @@ impl Guest for Exports {
             value: vec!["some value".to_owned(), "another value".to_owned()],
         };
         assert_eq!(value, thing_in_and_out::baz(&value));
+
+        let strings = vec!["foo", "bar", "baz"];
+        let resource = test::ownership::both_list_and_resource::TheResource::new(&strings);
+        test::ownership::both_list_and_resource::list_and_resource(
+            test::ownership::both_list_and_resource::Thing {
+                a: strings.iter().map(|s| s.to_string()).collect(),
+                b: resource,
+            },
+        );
     }
 }

tests/runtime/ownership/owning.rs

@@ -29,5 +29,14 @@ impl Guest for Exports {
             value: vec!["some value".to_owned(), "another value".to_owned()],
         };
         assert_eq!(value, thing_in_and_out::baz(&value));
+
+        let strings = vec!["foo".to_string(), "bar".to_string(), "baz".to_string()];
+        let resource = test::ownership::both_list_and_resource::TheResource::new(&strings);
+        test::ownership::both_list_and_resource::list_and_resource(
+            test::ownership::both_list_and_resource::Thing {
+                a: strings,
+                b: resource,
+            },
+        );
     }
 }

tests/runtime/ownership/world.wit

@@ -1,5 +1,17 @@
 package test:ownership;
 
+interface both-list-and-resource {
+    resource the-resource {
+        constructor(the-list: list<string>);
+    }
+    record thing {
+        a: list<string>,
+        b: the-resource
+    }
+
+    list-and-resource: func(a: thing);
+}
+
 world ownership {
     import lists: interface {
         foo: func(a: list<list<string>>) -> list<list<string>>;
@@ -23,5 +35,7 @@ world ownership {
         baz: func(a: thing) -> thing;
     }
 
+    import both-list-and-resource;
+
     export foo: func();
 }