Fix unsoundness in generated Rust code (#846)

* Refactor how Rust handles ownership

This commit is the next in a long line of refactors to try to model how
the Rust generator handles ownership. The Rust generator has an
`ownership` knob for deciding how to render types. The Rust generator
additionally models ownership of resources/lists in params/results of
imports/exports. Putting all of this together has led to many attempts
to wrangle this in a form that's understandable but every time a bug
comes up I feel like I don't understand what's going on. I've made
another attempt in this commit.

Here the goal is to centralize knowledge about types as much as
possible. Instead of being spread out over a few methods everything is
hopefully now centralized in a single type and a single "main method".
All other little pieces stem from these two and are modeled to be
relatively simple compared to the "main method". Whether or not this
stands the test of time we'll see.

This does change generated code in some situations as can be seen by the
test that was updated. Nothing major should be changed, but a few more
borrows are now required in places which probably should have been
borrows before.

More comments are found in the code about specific changes made here.

* 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 #817
Closes bytecodealliance/wasmtime#7951

Author: alexcrichton

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"),
                 });
             }
 
@@ -435,41 +435,39 @@ impl Bindgen for FunctionBindgen<'_, '_> {
 
                 let dealiased_resource = dealias(resolve, *resource);
 
-                results.push(
-                    if let Direction::Export = self.gen.gen.resources[&dealiased_resource].direction
-                    {
-                        match handle {
-                            Handle::Borrow(_) => {
-                                let name = resolve.types[*resource]
-                                    .name
-                                    .as_deref()
-                                    .unwrap()
-                                    .to_upper_camel_case();
-                                let rt = self.gen.gen.runtime_path();
-                                format!(
-                                    "{rt}::Resource::<{name}>::lift_borrow({op} as u32 as usize)"
-                                )
-                            }
-                            Handle::Own(_) => {
-                                let name = self.gen.type_path(dealiased_resource, true);
-                                format!("{name}::from_handle({op} as u32)")
-                            }
+                let result = if let Direction::Export =
+                    self.gen.gen.resources[&dealiased_resource].direction
+                {
+                    match handle {
+                        Handle::Borrow(_) => {
+                            let name = resolve.types[*resource]
+                                .name
+                                .as_deref()
+                                .unwrap()
+                                .to_upper_camel_case();
+                            let rt = self.gen.gen.runtime_path();
+                            format!("{rt}::Resource::<{name}>::lift_borrow({op} as u32 as usize)")
                         }
-                    } else if prefix == "" {
-                        let name = self.gen.type_path(dealiased_resource, true);
-                        format!("{name}::from_handle({op} as u32)")
-                    } else {
-                        let tmp = format!("handle{}", self.tmp());
-                        self.handle_decls.push(format!("let {tmp};"));
-                        let name = self.gen.type_path(dealiased_resource, true);
-                        format!(
-                            "{{\n
-                                {tmp} = {name}::from_handle({op} as u32);
-                                {prefix}{tmp}
-                            }}"
-                        )
-                    },
-                );
+                        Handle::Own(_) => {
+                            let name = self.gen.type_path(dealiased_resource, true);
+                            format!("{name}::from_handle({op} as u32)")
+                        }
+                    }
+                } else if prefix == "" {
+                    let name = self.gen.type_path(dealiased_resource, true);
+                    format!("{name}::from_handle({op} as u32)")
+                } else {
+                    let tmp = format!("handle{}", self.tmp());
+                    self.handle_decls.push(format!("let {tmp};"));
+                    let name = self.gen.type_path(dealiased_resource, true);
+                    format!(
+                        "{{\n
+                            {tmp} = {name}::from_handle({op} as u32);
+                            {prefix}{tmp}
+                        }}"
+                    )
+                };
+                results.push(result);
             }
 
             Instruction::RecordLower { ty, record, .. } => {