Fix a panic in Bytes{,Mut} StreamProducer impls (#12878)

This commit fixes the logic of these `impl`s to match the `Vec`-style
blocks to relinquish the entire buffer to Wasmtime immediately. This
fixes an issue where `split_off` is called with too large a value which
can panic.

Author: alexcrichton

Cargo.toml

@@ -96,7 +96,7 @@ rustix = { workspace = true, features = ["mm", "process"] }
 
 [dev-dependencies]
 # depend again on wasmtime to activate its default features for tests
-wasmtime = { workspace = true, features = ['default', 'anyhow', 'winch', 'pulley', 'all-arch', 'call-hook', 'memory-protection-keys', 'component-model-async'] }
+wasmtime = { workspace = true, features = ['default', 'anyhow', 'winch', 'pulley', 'all-arch', 'call-hook', 'memory-protection-keys', 'component-model-async', 'component-model-async-bytes'] }
 env_logger = { workspace = true }
 log = { workspace = true }
 filecheck = { workspace = true }

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

@@ -732,24 +732,13 @@ impl<D> StreamProducer<D> for bytes::Bytes {
     type Buffer = Cursor<Self>;
 
     fn poll_produce<'a>(
-        mut self: Pin<&mut Self>,
+        self: Pin<&mut Self>,
         _: &mut Context<'_>,
-        mut store: StoreContextMut<'a, D>,
+        _store: StoreContextMut<'a, D>,
         mut dst: Destination<'a, Self::Item, Self::Buffer>,
         _: bool,
     ) -> Poll<Result<StreamResult>> {
-        let cap = dst.remaining(&mut store);
-        let Some(cap) = cap.and_then(core::num::NonZeroUsize::new) else {
-            // on 0-length or host reads, buffer the bytes
-            dst.set_buffer(Cursor::new(mem::take(self.get_mut())));
-            return Poll::Ready(Ok(StreamResult::Dropped));
-        };
-        let cap = cap.into();
-        // data does not fit in destination, fill it and buffer the rest
-        dst.set_buffer(Cursor::new(self.split_off(cap)));
-        let mut dst = dst.as_direct(store, cap);
-        dst.remaining().copy_from_slice(&self);
-        dst.mark_written(cap);
+        dst.set_buffer(Cursor::new(mem::take(self.get_mut())));
         Poll::Ready(Ok(StreamResult::Dropped))
     }
 }
@@ -760,24 +749,13 @@ impl<D> StreamProducer<D> for bytes::BytesMut {
     type Buffer = Cursor<Self>;
 
     fn poll_produce<'a>(
-        mut self: Pin<&mut Self>,
+        self: Pin<&mut Self>,
         _: &mut Context<'_>,
-        mut store: StoreContextMut<'a, D>,
+        _store: StoreContextMut<'a, D>,
         mut dst: Destination<'a, Self::Item, Self::Buffer>,
         _: bool,
     ) -> Poll<Result<StreamResult>> {
-        let cap = dst.remaining(&mut store);
-        let Some(cap) = cap.and_then(core::num::NonZeroUsize::new) else {
-            // on 0-length or host reads, buffer the bytes
-            dst.set_buffer(Cursor::new(mem::take(self.get_mut())));
-            return Poll::Ready(Ok(StreamResult::Dropped));
-        };
-        let cap = cap.into();
-        // data does not fit in destination, fill it and buffer the rest
-        dst.set_buffer(Cursor::new(self.split_off(cap)));
-        let mut dst = dst.as_direct(store, cap);
-        dst.remaining().copy_from_slice(&self);
-        dst.mark_written(cap);
+        dst.set_buffer(Cursor::new(mem::take(self.get_mut())));
         Poll::Ready(Ok(StreamResult::Dropped))
     }
 }

tests/all/component_model/async.rs

@@ -1267,3 +1267,60 @@ async fn concurrent_sync_calls_to_async_host() -> Result<()> {
 
     Ok(())
 }
+
+#[tokio::test]
+#[cfg_attr(miri, ignore)]
+async fn bytes_stream_producer() -> Result<()> {
+    let mut config = Config::new();
+    config.wasm_component_model_async(true);
+    let engine = Engine::new(&config)?;
+
+    let component = Component::new(
+        &engine,
+        r#"
+        (component
+            (core module $libc (memory (export "mem") 1))
+            (core instance $libc (instantiate $libc))
+            (core module $m
+                (import "" "mem" (memory 1))
+                (import "" "stream.read" (func $stream.read (param i32 i32 i32) (result i32)))
+
+                (func (export "read") (param i32 i32) (result i32)
+                    (call $stream.read (local.get 0) (i32.const 0) (local.get 1))
+                )
+            )
+            (type $s (stream u8))
+            (core func $stream.read (canon stream.read $s async (memory $libc "mem")))
+            (core instance $i (instantiate $m
+                (with "" (instance
+                    (export "mem" (memory $libc "mem"))
+                    (export "stream.read" (func $stream.read))
+                ))
+            ))
+            (func (export "read") (param "s" (stream u8)) (param "l" u32) (result u32)
+                (canon lift (core func $i "read")))
+        )
+    "#,
+    )?;
+
+    let linker = Linker::new(&engine);
+    let mut store = Store::new(&engine, ());
+    let instance = linker.instantiate_async(&mut store, &component).await?;
+    let func = instance.get_typed_func::<(StreamReader<u8>, u32), (u32,)>(&mut store, "read")?;
+
+    // read less than the capacity
+    let reader = StreamReader::new(&mut store, bytes::Bytes::from_static(b"hello"))?;
+    assert_eq!(
+        func.call_async(&mut store, (reader, 1)).await?,
+        ((1 << 4) | 0,),
+    );
+
+    // read more than the capacity
+    let reader = StreamReader::new(&mut store, bytes::Bytes::from_static(b"hello"))?;
+    assert_eq!(
+        func.call_async(&mut store, (reader, 100)).await?,
+        ((5 << 4) | 1,),
+    );
+
+    Ok(())
+}