enforce resource borrow requirements for async calls

Signed-off-by: Joel Dice <joel.dice@fermyon.com>

Author: dicej

Cargo.lock

@@ -295,9 +295,10 @@ io-lifetimes = { version = "2.0.3", default-features = false }
 io-extras = "0.18.1"
 rustix = "0.38.43"
 # wit-bindgen:
-wit-bindgen = { version = "0.37.0", default-features = false }
-wit-bindgen-rt = { version = "0.37.0", default-features = false }
-wit-bindgen-rust-macro = { version = "0.37.0", default-features = false }
+# TODO: switch to crates.io once https://github.com/bytecodealliance/wit-bindgen/pull/1124 has been released
+wit-bindgen = { git = "https://github.com/bytecodealliance/wit-bindgen", default-features = false }
+wit-bindgen-rt = { git = "https://github.com/bytecodealliance/wit-bindgen", default-features = false }
+wit-bindgen-rust-macro = { git = "https://github.com/bytecodealliance/wit-bindgen", default-features = false }
 
 # wasm-tools family:
 wasmparser = { version = "0.223.0", default-features = false, features = ['simd'] }

Cargo.toml

@@ -94,6 +94,26 @@ mod test {
         }
     }
 
+    pub struct MyX;
+
+    impl borrowing_host::local::local::borrowing_types::HostX for Ctx {
+        fn new(&mut self) -> Result<Resource<MyX>> {
+            Ok(WasiView::table(self).push(MyX)?)
+        }
+
+        fn foo(&mut self, x: Resource<MyX>) -> Result<()> {
+            _ = WasiView::table(self).get(&x)?;
+            Ok(())
+        }
+
+        fn drop(&mut self, x: Resource<MyX>) -> Result<()> {
+            WasiView::table(self).delete(x)?;
+            Ok(())
+        }
+    }
+
+    impl borrowing_host::local::local::borrowing_types::Host for Ctx {}
+
     async fn test_round_trip(component: &[u8], input: &str, expected_output: &str) -> Result<()> {
         init_logger();
 
@@ -401,6 +421,22 @@ mod test {
         });
     }
 
+    mod borrowing_host {
+        wasmtime::component::bindgen!({
+            path: "wit",
+            world: "borrowing-host",
+            trappable_imports: true,
+            concurrent_imports: true,
+            concurrent_exports: true,
+            async: {
+                only_imports: []
+            },
+            with: {
+                "local:local/borrowing-types/x": super::MyX,
+            }
+        });
+    }
+
     impl yield_host::local::local::continue_::Host for Ctx {
         fn set_continue(&mut self, v: bool) {
             self.continue_ = v;
@@ -555,6 +591,100 @@ mod test {
         test_run(&compose(caller, callee).await?).await
     }
 
+    async fn test_run_bool(component: &[u8], v: bool) -> Result<()> {
+        init_logger();
+
+        let mut config = Config::new();
+        config.debug_info(true);
+        config.cranelift_debug_verifier(true);
+        config.wasm_component_model(true);
+        config.wasm_component_model_async(true);
+        config.async_support(true);
+        config.epoch_interruption(true);
+
+        let engine = Engine::new(&config)?;
+
+        let component = Component::new(&engine, component)?;
+
+        let mut linker = Linker::new(&engine);
+
+        wasmtime_wasi::add_to_linker_async(&mut linker)?;
+        borrowing_host::BorrowingHost::add_to_linker(&mut linker, |ctx| ctx)?;
+
+        let mut store = Store::new(
+            &engine,
+            Ctx {
+                wasi: WasiCtxBuilder::new().inherit_stdio().build(),
+                table: ResourceTable::default(),
+                drop_count: 0,
+                continue_: false,
+                wakers: Arc::new(Mutex::new(None)),
+            },
+        );
+        store.set_epoch_deadline(1);
+
+        std::thread::spawn(move || {
+            std::thread::sleep(Duration::from_secs(10));
+            engine.increment_epoch();
+        });
+
+        let borrowing_host =
+            borrowing_host::BorrowingHost::instantiate_async(&mut store, &component, &linker)
+                .await?;
+
+        // Start three concurrent calls and then join them all:
+        let mut promises = PromisesUnordered::new();
+        for _ in 0..3 {
+            promises.push(
+                borrowing_host
+                    .local_local_run_bool()
+                    .call_run(&mut store, v)
+                    .await?,
+            );
+        }
+
+        while let Some(()) = promises.next(&mut store).await? {
+            // continue
+        }
+
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn async_borrowing_caller() -> Result<()> {
+        let caller = &fs::read(test_programs_artifacts::ASYNC_BORROWING_CALLER_COMPONENT).await?;
+        let callee = &fs::read(test_programs_artifacts::ASYNC_BORROWING_CALLEE_COMPONENT).await?;
+        test_run_bool(&compose(caller, callee).await?, false).await
+    }
+
+    #[tokio::test]
+    async fn async_borrowing_caller_misbehave() -> Result<()> {
+        let caller = &fs::read(test_programs_artifacts::ASYNC_BORROWING_CALLER_COMPONENT).await?;
+        let callee = &fs::read(test_programs_artifacts::ASYNC_BORROWING_CALLEE_COMPONENT).await?;
+        let error = format!(
+            "{:?}",
+            test_run_bool(&compose(caller, callee).await?, true)
+                .await
+                .unwrap_err()
+        );
+        assert!(error.contains("unknown handle index"), "{error}");
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn async_borrowing_callee() -> Result<()> {
+        let callee = &fs::read(test_programs_artifacts::ASYNC_BORROWING_CALLEE_COMPONENT).await?;
+        test_run_bool(callee, false).await
+    }
+
+    #[tokio::test]
+    async fn async_borrowing_callee_misbehave() -> Result<()> {
+        let callee = &fs::read(test_programs_artifacts::ASYNC_BORROWING_CALLEE_COMPONENT).await?;
+        let error = format!("{:?}", test_run_bool(callee, true).await.unwrap_err());
+        assert!(error.contains("unknown handle index"), "{error}");
+        Ok(())
+    }
+
     mod transmit {
         wasmtime::component::bindgen!({
             path: "wit",

crates/misc/component-async-tests/src/lib.rs

@@ -46,6 +46,23 @@ interface post-return {
   get-post-return-value: func() -> string;
 }
 
+interface borrowing-types {
+  resource x {
+    constructor();
+    foo: func();
+  }
+}
+
+interface borrowing {
+  use borrowing-types.{x};
+
+  foo: func(x: borrow<x>, misbehave: bool);
+}
+
+interface run-bool {
+  run: func(v: bool);
+}
+
 world yield-caller {
   import continue;
   import ready;
@@ -97,3 +114,18 @@ world post-return-caller {
 world post-return-callee {
   export post-return;
 }
+
+world borrowing-caller {
+  import borrowing;
+  export run-bool;
+}
+
+world borrowing-callee {
+  export borrowing;
+  export run-bool;
+}
+
+world borrowing-host {
+  import borrowing-types;
+  export run-bool;
+}

crates/misc/component-async-tests/wit/test.wit

@@ -0,0 +1,49 @@
+mod bindings {
+    wit_bindgen::generate!({
+        path: "../misc/component-async-tests/wit",
+        world: "borrowing-callee",
+        async: {
+            exports: [
+                "local:local/borrowing#foo",
+                "local:local/run-bool#run"
+            ]
+        }
+    });
+
+    use super::Component;
+    export!(Component);
+}
+
+use {
+    bindings::{
+        exports::local::local::{borrowing::Guest as Borrowing, run_bool::Guest as RunBool},
+        local::local::borrowing_types::X,
+    },
+    futures::future,
+    std::future::Future,
+    wit_bindgen_rt::async_support,
+};
+
+struct Component;
+
+impl Borrowing for Component {
+    fn foo(x: &X, misbehave: bool) -> impl Future<Output = ()> + 'static {
+        let handle = x.handle();
+        async_support::spawn(async move {
+            if misbehave {
+                unsafe { X::from_handle(handle) }.foo();
+            }
+        });
+        x.foo();
+        future::ready(())
+    }
+}
+
+impl RunBool for Component {
+    async fn run(misbehave: bool) {
+        Self::foo(&X::new(), misbehave).await
+    }
+}
+
+// Unused function; required since this file is built as a `bin`:
+fn main() {}

crates/test-programs/src/bin/async_borrowing_callee.rs

@@ -0,0 +1,33 @@
+mod bindings {
+    wit_bindgen::generate!({
+        path: "../misc/component-async-tests/wit",
+        world: "borrowing-caller",
+        async: {
+            imports: [
+                "local:local/borrowing#foo"
+            ],
+            exports: [
+                "local:local/run-bool#run"
+            ]
+        }
+    });
+
+    use super::Component;
+    export!(Component);
+}
+
+use bindings::{
+    exports::local::local::run_bool::Guest,
+    local::local::{borrowing::foo, borrowing_types::X},
+};
+
+struct Component;
+
+impl Guest for Component {
+    async fn run(misbehave: bool) {
+        foo(&X::new(), misbehave).await
+    }
+}
+
+// Unused function; required since this file is built as a `bin`:
+fn main() {}