Merge pull request #178 from bytecodealliance/dicej/fix-172

trap if closing future to which no value has been written

Author: dicej

crates/wasmtime/src/runtime/component/concurrent/futures_and_streams.rs

@@ -1416,6 +1416,8 @@ struct TransmitState {
     writer_watcher: Option<oneshot::Sender<()>>,
     /// Like `writer_watcher`, but for the reverse direction.
     reader_watcher: Option<oneshot::Sender<()>>,
+    /// Whether the write end may be closed or not.
+    may_close_writer: bool,
 }
 
 impl Default for TransmitState {
@@ -1427,6 +1429,7 @@ impl Default for TransmitState {
             write: WriteState::Open,
             reader_watcher: None,
             writer_watcher: None,
+            may_close_writer: true,
         }
     }
 }
@@ -1550,7 +1553,7 @@ impl Instance {
         store
             .as_context_mut()
             .with_detached_instance(self, |mut store, instance| {
-                let (write, read) = instance.new_transmit()?;
+                let (write, read) = instance.new_transmit(TransmitKind::Future)?;
 
                 Ok((
                     FutureWriter::new(
@@ -1591,7 +1594,7 @@ impl Instance {
         store
             .as_context_mut()
             .with_detached_instance(self, |mut store, instance| {
-                let (write, read) = instance.new_transmit()?;
+                let (write, read) = instance.new_transmit(TransmitKind::Stream)?;
 
                 Ok((
                     StreamWriter::new(
@@ -1899,7 +1902,10 @@ impl ComponentInstance {
 
     /// Allocate a new future or stream, including the `TransmitState` and the
     /// `TransmitHandle`s corresponding to the read and write ends.
-    fn new_transmit(&mut self) -> Result<(TableId<TransmitHandle>, TableId<TransmitHandle>)> {
+    fn new_transmit(
+        &mut self,
+        kind: TransmitKind,
+    ) -> Result<(TableId<TransmitHandle>, TableId<TransmitHandle>)> {
         let state_id = self.push(TransmitState::default())?;
 
         let write = self.push(TransmitHandle::new(state_id))?;
@@ -1909,6 +1915,10 @@ impl ComponentInstance {
         state.write_handle = write;
         state.read_handle = read;
 
+        if let TransmitKind::Future = kind {
+            state.may_close_writer = false;
+        }
+
         log::trace!("new transmit: state {state_id:?}; write {write:?}; read {read:?}",);
 
         Ok((write, read))
@@ -1941,7 +1951,10 @@ impl ComponentInstance {
     /// write ends to the (sub-)component instance to which the specified
     /// `TableIndex` belongs.
     fn guest_new(&mut self, ty: TableIndex) -> Result<ResourcePair> {
-        let (write, read) = self.new_transmit()?;
+        let (write, read) = self.new_transmit(match ty {
+            TableIndex::Future(_) => TransmitKind::Future,
+            TableIndex::Stream(_) => TransmitKind::Stream,
+        })?;
         let read = self
             .state_table(ty)
             .insert(read.rep(), waitable_state(ty, StreamFutureState::Read))?;
@@ -1975,6 +1988,7 @@ impl ComponentInstance {
         let transmit = self
             .get_mut(transmit_id)
             .with_context(|| format!("retrieving state for transmit [{transmit_rep}]"))?;
+        transmit.may_close_writer = true;
 
         let new_state = if let ReadState::Closed = &transmit.read {
             ReadState::Closed
@@ -2176,8 +2190,11 @@ impl ComponentInstance {
 
     /// Cancel a pending stream or future write from the host.
     ///
-    /// `rep` is the `TransmitState` rep for the stream or future.
-    fn host_cancel_write(&mut self, rep: u32) -> Result<ReturnCode> {
+    /// # Arguments
+    ///
+    /// * `rep` - The `TransmitState` rep for the stream or future.
+    /// * `kind` - Whether `rep` is for a stream or a future.
+    fn host_cancel_write(&mut self, rep: u32, kind: TransmitKind) -> Result<ReturnCode> {
         let transmit_id = TableId::<TransmitState>::new(rep);
         let transmit = self.get_mut(transmit_id)?;
 
@@ -2208,6 +2225,10 @@ impl ComponentInstance {
 
         log::trace!("cancelled write {transmit_id:?}");
 
+        if let (TransmitKind::Future, ReturnCode::Cancelled(0)) = (kind, code) {
+            transmit.may_close_writer = false;
+        }
+
         Ok(code)
     }
 
@@ -2259,6 +2280,10 @@ impl ComponentInstance {
             .get_mut(transmit_id)
             .with_context(|| format!("error closing writer {transmit_rep}"))?;
 
+        if !transmit.may_close_writer {
+            bail!("cannot close future write end without first writing a value")
+        }
+
         transmit.writer_watcher = None;
 
         // Existing queued transmits must be updated with information for the impending writer closure
@@ -2652,6 +2677,7 @@ impl ComponentInstance {
         let transmit_id = self.get(transmit_handle)?.state;
         log::trace!("guest_write {transmit_handle:?} (handle {handle}; state {transmit_id:?})",);
         let transmit = self.get_mut(transmit_id)?;
+        transmit.may_close_writer = true;
         let new_state = if let ReadState::Closed = &transmit.read {
             ReadState::Closed
         } else {
@@ -3018,10 +3044,11 @@ impl ComponentInstance {
         writer: u32,
         _async_: bool,
     ) -> Result<ReturnCode> {
-        let (rep, WaitableState::Stream(_, state) | WaitableState::Future(_, state)) =
-            self.state_table(ty).get_mut_by_index(writer)?
-        else {
-            bail!("invalid stream or future handle");
+        let (rep, state) = self.state_table(ty).get_mut_by_index(writer)?;
+        let (state, kind) = match state {
+            WaitableState::Stream(_, state) => (state, TransmitKind::Stream),
+            WaitableState::Future(_, state) => (state, TransmitKind::Future),
+            _ => bail!("invalid stream or future handle"),
         };
         let id = TableId::<TransmitHandle>::new(rep);
         log::trace!("guest cancel write {id:?} (handle {writer})");
@@ -3037,7 +3064,7 @@ impl ComponentInstance {
             }
         }
         let rep = self.get(id)?.state.rep();
-        self.host_cancel_write(rep)
+        self.host_cancel_write(rep, kind)
     }
 
     /// Cancel a pending read for the specified stream or future from the guest.

tests/misc_testsuite/component-model-async/future-cancel-read-closed.wast

@@ -8,11 +8,13 @@
   (core instance $libc (instantiate $libc))
   (core func $read (canon future.read $f async (memory $libc "mem")))
   (core func $cancel (canon future.cancel-read $f))
+  (core func $write (canon future.write $f async (memory $libc "mem")))
   (core func $close-write (canon future.close-writable $f))
   (core module $m
     (import "" "new" (func $new (result i64)))
     (import "" "read" (func $read (param i32 i32) (result i32)))
     (import "" "cancel" (func $cancel (param i32) (result i32)))
+    (import "" "write" (func $write (param i32 i32) (result i32)))
     (import "" "close-write" (func $close-write (param i32)))
 
     (func (export "f") (result i32)
@@ -34,6 +36,10 @@
 
       ;; close the write end
       local.get $write
+      i32.const 0
+      call $write
+      drop
+      local.get $write
       call $close-write
 
       ;; cancel the read, returning the result
@@ -47,11 +53,12 @@
       (export "new" (func $new))
       (export "read" (func $read))
       (export "cancel" (func $cancel))
+      (export "write" (func $write))
       (export "close-write" (func $close-write))
     ))
   ))
 
   (func (export "f") (result u32) (canon lift (core func $i "f")))
 )
 
-(assert_return (invoke "f") (u32.const 1)) ;; expect CLOSED status (not CANCELLED)
+(assert_return (invoke "f") (u32.const 0x11)) ;; expect CLOSED status (not CANCELLED)

tests/misc_testsuite/component-model-async/futures-must-write.wast

@@ -0,0 +1,128 @@
+;;! component_model_async = true
+
+;; This test contains two components $C and $D that test that a trap occurs
+;; when closing the writable end of a future (in $C) before having written
+;; a value while closing the readable end of a future (in $D) before reading
+;; a value is fine.
+;;
+;; (Copied from
+;; https://github.com/WebAssembly/component-model/blob/future-trap/test/async/futures-must-write.wast)
+(component
+  (component $C
+    (core module $Memory (memory (export "mem") 1))
+    (core instance $memory (instantiate $Memory))
+    (core module $CM
+      (import "" "mem" (memory 1))
+      (import "" "future.new" (func $future.new (result i64)))
+      (import "" "future.write" (func $future.write (param i32 i32) (result i32)))
+      (import "" "future.close-writable" (func $future.close-writable (param i32)))
+
+      (global $fw (mut i32) (i32.const 0))
+
+      (func $start-future (export "start-future") (result i32)
+        ;; create a new future, return the readable end to the caller
+        (local $ret64 i64)
+        (local.set $ret64 (call $future.new))
+        (global.set $fw (i32.wrap_i64 (i64.shr_u (local.get $ret64) (i64.const 32))))
+        (i32.wrap_i64 (local.get $ret64))
+      )
+      (func $attempt-write (export "attempt-write") (result i32)
+        ;; because the caller already closed the readable end, this write will eagerly
+        ;; return CLOSED having written no values.
+        (local $ret i32)
+        (local.set $ret (call $future.write (global.get $fw) (i32.const 42)))
+        (if (i32.ne (i32.const 0x01 (; CLOSED=1 | (0<<4) ;)) (local.get $ret))
+          (then
+            (i32.load (i32.add (local.get $ret) (i32.const 0x8000_0000)))
+          unreachable))
+
+        ;; return without trapping
+        (i32.const 42)
+      )
+      (func $close-writable (export "close-writable")
+        ;; maybe boom
+        (call $future.close-writable (global.get $fw))
+      )
+    )
+    (type $FT (future u8))
+    (canon future.new $FT (core func $future.new))
+    (canon future.write $FT async (memory $memory "mem") (core func $future.write))
+    (canon future.close-writable $FT (core func $future.close-writable))
+    (core instance $cm (instantiate $CM (with "" (instance
+      (export "mem" (memory $memory "mem"))
+      (export "future.new" (func $future.new))
+      (export "future.write" (func $future.write))
+      (export "future.close-writable" (func $future.close-writable))
+    ))))
+    (func (export "start-future") (result (future u8)) (canon lift (core func $cm "start-future")))
+    (func (export "attempt-write") (result u32) (canon lift (core func $cm "attempt-write")))
+    (func (export "close-writable") (canon lift (core func $cm "close-writable")))
+  )
+  (component $D
+    (import "c" (instance $c
+      (export "start-future" (func (result (future u8))))
+      (export "attempt-write" (func (result u32)))
+      (export "close-writable" (func))
+    ))
+
+    (core module $Memory (memory (export "mem") 1))
+    (core instance $memory (instantiate $Memory))
+    (core module $Core
+      (import "" "mem" (memory 1))
+      (import "" "future.read" (func $future.read (param i32 i32) (result i32)))
+      (import "" "future.close-readable" (func $future.close-readable (param i32)))
+      (import "" "start-future" (func $start-future (result i32)))
+      (import "" "attempt-write" (func $attempt-write (result i32)))
+      (import "" "close-writable" (func $close-writable))
+
+      (func $close-readable-future-before-read (export "close-readable-future-before-read") (result i32)
+        ;; call 'start-future' to get the future we'll be working with
+        (local $fr i32)
+        (local.set $fr (call $start-future))
+        (if (i32.ne (i32.const 1) (local.get $fr))
+          (then unreachable))
+
+        ;; ok to immediately close the readable end
+        (call $future.close-readable (local.get $fr))
+
+        ;; the callee will see that we closed the readable end when it tries to write
+        (call $attempt-write)
+      )
+      (func $close-writable-future-before-write (export "close-writable-future-before-write")
+        ;; call 'start-future' to get the future we'll be working with
+        (local $fr i32)
+        (local.set $fr (call $start-future))
+        (if (i32.ne (i32.const 1) (local.get $fr))
+          (then unreachable))
+
+        ;; boom
+        (call $close-writable)
+      )
+    )
+    (type $FT (future u8))
+    (canon future.new $FT (core func $future.new))
+    (canon future.read $FT async (memory $memory "mem") (core func $future.read))
+    (canon future.close-readable $FT (core func $future.close-readable))
+    (canon lower (func $c "start-future") (core func $start-future'))
+    (canon lower (func $c "attempt-write") (core func $attempt-write'))
+    (canon lower (func $c "close-writable") (core func $close-writable'))
+    (core instance $core (instantiate $Core (with "" (instance
+      (export "mem" (memory $memory "mem"))
+      (export "future.new" (func $future.new))
+      (export "future.read" (func $future.read))
+      (export "future.close-readable" (func $future.close-readable))
+      (export "start-future" (func $start-future'))
+      (export "attempt-write" (func $attempt-write'))
+      (export "close-writable" (func $close-writable'))
+    ))))
+    (func (export "close-readable-future-before-read") (result u32) (canon lift (core func $core "close-readable-future-before-read")))
+    (func (export "close-writable-future-before-write") (canon lift (core func $core "close-writable-future-before-write")))
+  )
+  (instance $c (instantiate $C))
+  (instance $d (instantiate $D (with "c" (instance $c))))
+  (func (export "close-writable-future-before-write") (alias export $d "close-writable-future-before-write"))
+  (func (export "close-readable-future-before-read") (alias export $d "close-readable-future-before-read"))
+)
+
+(assert_return (invoke "close-readable-future-before-read") (u32.const 42))
+(assert_trap (invoke "close-writable-future-before-write") "cannot close future write end without first writing a value")