fix(wasmtime): correctly handle input-stream

crates/transport/src/frame/conn/mod.rs

@@ -23,6 +23,7 @@ use wasm_tokio::{
 };
 
 use crate::frame::MAX_INITIAL_PATH_CAPACITY;
+use crate::{copy_io_error, error_receiver_closed};
 
 mod client;
 mod server;
@@ -672,10 +673,6 @@ impl AsyncWrite for Outgoing {
     }
 }
 
-fn copy_io_error(err: &std::io::Error) -> std::io::Error {
-    std::io::Error::new(err.kind(), err.to_string())
-}
-
 #[instrument(level = "trace", skip_all, ret(level = "trace"))]
 async fn ingress(
     mut rx: impl AsyncRead + Unpin,
@@ -737,7 +734,7 @@ async fn ingress(
             trace!(?buf, "read data chunk");
             tx.send(Ok(buf.freeze()))
                 .await
-                .map_err(|_| receiver_closed())?;
+                .map_err(|_| error_receiver_closed())?;
         }
     }
 }
@@ -748,18 +745,13 @@ fn close_rx(state: &IngressState) -> std::io::Result<()> {
     trace!("all handles dropped, closing index trie receivers");
     let mut index = state.index.lock().unwrap_or_else(PoisonError::into_inner);
     if index.close_rx() || state.unprocessed.load(Ordering::Relaxed) {
-        return Err(unprocessed());
+        return Err(error_unprocessed());
     }
     Ok(())
 }
 
-/// Returns the error for data received for a closed receiver
-pub(crate) fn receiver_closed() -> std::io::Error {
-    std::io::Error::new(std::io::ErrorKind::BrokenPipe, "stream receiver closed")
-}
-
 /// Returns the error for received data, which was not processed
-pub(crate) fn unprocessed() -> std::io::Error {
+fn error_unprocessed() -> std::io::Error {
     std::io::Error::new(
         std::io::ErrorKind::BrokenPipe,
         "received data was not processed",

crates/transport/src/lib.rs

@@ -38,9 +38,11 @@ use core::mem;
 use core::pin::{Pin, pin};
 use core::task::{Context, Poll};
 
-use bytes::BytesMut;
+use bytes::{BufMut as _, BytesMut};
 use tokio::io::{AsyncRead, ReadBuf};
+use tokio_util::codec::Encoder as _;
 use tracing::trace;
+use wasm_tokio::Leb128Encoder;
 
 /// Buffered incoming stream used for decoding values
 ///
@@ -95,3 +97,20 @@ impl AsyncRead for BufferedIncoming {
         pin!(&mut self.inner).poll_read(cx, buf)
     }
 }
+
+fn encode_ready_bytes_header(n: usize, dst: &mut BytesMut) -> std::io::Result<()> {
+    let n = u32::try_from(n)
+        .map_err(|err| std::io::Error::new(std::io::ErrorKind::InvalidInput, err))?;
+    dst.reserve(1);
+    dst.put_u8(0x01);
+    Leb128Encoder.encode(n, dst)
+}
+
+fn copy_io_error(err: &std::io::Error) -> std::io::Error {
+    std::io::Error::new(err.kind(), err.to_string())
+}
+
+/// Returns the error for data received for a closed receiver
+fn error_receiver_closed() -> std::io::Error {
+    std::io::Error::new(std::io::ErrorKind::BrokenPipe, "stream receiver closed")
+}

crates/transport/src/pack.rs

@@ -8,8 +8,7 @@ use tokio::io::{AsyncRead, AsyncReadExt as _};
 use tokio::sync::mpsc;
 use tokio_util::codec::{Decoder, Encoder};
 
-use crate::value::encode_ready_bytes_header;
-use crate::{Decode, Deferred as _, DeferredFn, Encode, Index, Sink};
+use crate::{Decode, Deferred as _, DeferredFn, Encode, Index, Sink, encode_ready_bytes_header};
 
 /// [Sink] used by [pack], which collects the packed asynchronous values
 #[derive(Clone)]

crates/transport/src/value.rs

@@ -31,18 +31,11 @@ use wasm_tokio::{
     Leb128DecoderU64, Leb128DecoderU128, Leb128Encoder, Utf8Codec,
 };
 
-use crate::BufferedIncoming;
 use crate::frame::Outgoing;
+use crate::{BufferedIncoming, copy_io_error, encode_ready_bytes_header, error_receiver_closed};
 
-const READ_CHUNK_SIZE: usize = 8192;
-
-pub(crate) fn encode_ready_bytes_header(n: usize, dst: &mut BytesMut) -> std::io::Result<()> {
-    let n = u32::try_from(n)
-        .map_err(|err| std::io::Error::new(std::io::ErrorKind::InvalidInput, err))?;
-    dst.reserve(1);
-    dst.put_u8(0x01);
-    Leb128Encoder.encode(n, dst)
-}
+/// Size of chunks read from byte streams
+pub const READ_CHUNK_SIZE: usize = 8192;
 
 fn encode_ready_bytes<T: AsRef<[u8]>>(chunks: &[T], dst: &mut BytesMut) -> std::io::Result<()> {
     let n = chunks.iter().map(|chunk| chunk.as_ref().len()).sum();
@@ -1930,7 +1923,7 @@ where
                     };
                     let item = item?;
                     if tx.send(item).is_err() {
-                        return Err(crate::frame::receiver_closed());
+                        return Err(error_receiver_closed());
                     }
                     if let Some(rx) = dec.decoder_mut().take_deferred() {
                         let mut buf = mem::take(dec.read_buffer_mut());
@@ -2092,23 +2085,26 @@ impl<W: Sink> Encode<W> for Pin<Box<dyn Stream<Item = Bytes> + Send>> {
 }
 
 /// Encoder for `stream<list<u8>>` with [`AsyncRead`] support
-pub struct StreamEncoderRead<W = Outgoing> {
+///
+/// At most `MAX_READY` bytes are read from the stream on encode, any remaining
+/// bytes are transmitted asynchronously.
+pub struct StreamEncoderRead<W = Outgoing, const MAX_READY: usize = { usize::MAX }> {
     deferred: Option<DeferredFn<W>>,
 }
 
-impl<W> Default for StreamEncoderRead<W> {
+impl<W, const MAX_READY: usize> Default for StreamEncoderRead<W, MAX_READY> {
     fn default() -> Self {
         Self { deferred: None }
     }
 }
 
-impl<W> Deferred<W> for StreamEncoderRead<W> {
+impl<W, const MAX_READY: usize> Deferred<W> for StreamEncoderRead<W, MAX_READY> {
     fn take_deferred(&mut self) -> Option<DeferredFn<W>> {
         self.deferred.take()
     }
 }
 
-impl<W, S> tokio_util::codec::Encoder<S> for StreamEncoderRead<W>
+impl<W, S, const MAX_READY: usize> tokio_util::codec::Encoder<S> for StreamEncoderRead<W, MAX_READY>
 where
     W: Sink,
     S: AsyncRead + Send + Unpin + 'static,
@@ -2120,8 +2116,12 @@ where
         self.deferred = None;
         let mut ready = BytesMut::default();
         loop {
-            ready.reserve(READ_CHUNK_SIZE);
-            match items.read_buf(&mut ready).now_or_never() {
+            let n = MAX_READY.saturating_sub(ready.len());
+            if n == 0 {
+                break;
+            }
+            ready.reserve(n.min(READ_CHUNK_SIZE));
+            match items.read_buf(&mut (&mut ready).limit(n)).now_or_never() {
                 Some(Ok(0)) => return encode_ready_bytes(&[ready], dst),
                 Some(Ok(..)) => {}
                 Some(Err(err)) => return Err(err),
@@ -2219,7 +2219,10 @@ fn ensure_processed<D>(framed: &FramedRead<BufferedIncoming, D>) -> std::io::Res
     if framed.read_buffer().is_empty() && framed.get_ref().buffer.is_empty() {
         Ok(())
     } else {
-        Err(crate::frame::unprocessed())
+        Err(std::io::Error::new(
+            std::io::ErrorKind::InvalidData,
+            "received data after value end",
+        ))
     }
 }
 
@@ -2262,7 +2265,7 @@ where
                 })?;
                 trace!(i, end, "received stream chunk");
                 if tx.send(chunk).await.is_err() {
-                    return Err(crate::frame::receiver_closed());
+                    return Err(error_receiver_closed());
                 }
                 for (i, deferred) in zip(i.., mem::take(&mut framed.decoder_mut().deferred)) {
                     if let Some(deferred) = deferred {
@@ -2378,7 +2381,7 @@ impl tokio_util::codec::Decoder for StreamDecoderBytes {
                         }
                         trace!(?chunk, "received pending byte stream chunk");
                         if tx.send(chunk).await.is_err() {
-                            return Err(crate::frame::receiver_closed());
+                            return Err(error_receiver_closed());
                         }
                     }
                     Err(std::io::ErrorKind::UnexpectedEof.into())
@@ -2423,29 +2426,44 @@ impl tokio_util::codec::Decoder for StreamDecoderRead {
 
         // stream is pending
         let (tx, rx) = mpsc::channel(128);
+        let (done_tx, done_rx) = oneshot::channel();
         let dec = mem::take(&mut self.dec).into_inner();
         self.deferred = Some(Box::new(|mut r, path| {
             Box::pin(async move {
-                if !path.is_empty() {
-                    r = r.index(&path).map_err(std::io::Error::other)?;
-                }
-                let mut framed = FramedRead::new(r, dec);
-                trace!("receiving pending byte stream chunk");
-                while let Some(chunk) = framed.next().await {
-                    let chunk = chunk?;
-                    if chunk.is_empty() {
-                        trace!("received stream end");
-                        return ensure_processed(&framed);
+                let res = async {
+                    if !path.is_empty() {
+                        r = r.index(&path).map_err(std::io::Error::other)?;
                     }
-                    trace!(?chunk, "received byte stream chunk");
-                    if tx.send(std::io::Result::Ok(chunk)).await.is_err() {
-                        return Err(crate::frame::receiver_closed());
+                    let mut framed = FramedRead::new(r, dec);
+                    trace!("receiving pending byte stream chunk");
+                    while let Some(chunk) = framed.next().await {
+                        let chunk = chunk?;
+                        if chunk.is_empty() {
+                            trace!("received stream end");
+                            return ensure_processed(&framed);
+                        }
+                        trace!(?chunk, "received byte stream chunk");
+                        if tx.send(chunk).await.is_err() {
+                            return Err(error_receiver_closed());
+                        }
                     }
+                    Err(std::io::ErrorKind::UnexpectedEof.into())
                 }
-                Err(std::io::ErrorKind::UnexpectedEof.into())
+                .await;
+                _ = done_tx.send(res.as_ref().copied().map_err(copy_io_error));
+                res
             })
         }));
-        Ok(Some(Box::pin(StreamReader::new(ReceiverStream::new(rx)))))
+        let done = stream::once(done_rx).filter_map(|res| async move {
+            match res {
+                Ok(Ok(())) => None,
+                Ok(Err(err)) => Some(Err(err)),
+                Err(..) => Some(Err(std::io::ErrorKind::UnexpectedEof.into())),
+            }
+        });
+        Ok(Some(Box::pin(StreamReader::new(
+            ReceiverStream::new(rx).map(Ok).chain(done),
+        ))))
     }
 }
 
@@ -2566,6 +2584,11 @@ mod tests {
         enc.encode(tokio::io::empty(), &mut buf)?;
         assert!(enc.take_deferred().is_none());
 
+        let mut enc = StreamEncoderRead::<Outgoing, 2>::default();
+        enc.encode(b"fo".as_slice(), &mut buf)?;
+        assert!(enc.take_deferred().is_some());
+
+        let mut enc = StreamEncoderRead::<Outgoing>::default();
         let (_tx, rx) = tokio::io::duplex(1);
         enc.encode(rx, &mut buf)?;
         let err = tokio_util::io::StreamReader::new(stream::iter([Err::<Bytes, _>(
@@ -2621,7 +2644,7 @@ mod tests {
     async fn future_rejects_unprocessed() {
         let (_fut, f) = pending::<FutureBytes>();
         let err = f(buffered(b"\x03foo\x2a"), Vec::default()).await;
-        assert_eq!(err.unwrap_err().kind(), ErrorKind::BrokenPipe);
+        assert_eq!(err.unwrap_err().kind(), ErrorKind::InvalidData);
     }
 
     #[test_log::test(tokio::test)]
@@ -2632,7 +2655,7 @@ mod tests {
         buffer.extend_from_slice(&[0x2a; READ_CHUNK_SIZE]);
         let (_fut, f) = pending::<FutureBytes>();
         let err = f(buffered(&buffer), Vec::default()).await;
-        assert_eq!(err.unwrap_err().kind(), ErrorKind::BrokenPipe);
+        assert_eq!(err.unwrap_err().kind(), ErrorKind::InvalidData);
     }
 
     #[test_log::test(tokio::test)]
@@ -2647,14 +2670,27 @@ mod tests {
     async fn stream_rejects_unprocessed() {
         let (_stream, f) = pending::<StreamBytes>();
         let err = f(buffered(b"\x03foo\x00\x2a"), Vec::default()).await;
-        assert_eq!(err.unwrap_err().kind(), ErrorKind::BrokenPipe);
+        assert_eq!(err.unwrap_err().kind(), ErrorKind::InvalidData);
     }
 
     #[test_log::test(tokio::test)]
     async fn stream_rejects_truncated() {
         let (_stream, f) = pending::<StreamBytes>();
         let err = f(buffered(b"\x03foo"), Vec::default()).await;
         assert_eq!(err.unwrap_err().kind(), ErrorKind::UnexpectedEof);
+
+        let (mut r, f) = pending::<Pin<Box<dyn AsyncRead + Send>>>();
+        let err = f(buffered(b"\x03foo"), Vec::default()).await;
+        assert_eq!(err.unwrap_err().kind(), ErrorKind::UnexpectedEof);
+        let mut buf = Vec::default();
+        let err = r.read_to_end(&mut buf).await;
+        assert_eq!(err.unwrap_err().kind(), ErrorKind::UnexpectedEof);
+        assert_eq!(buf, b"foo");
+
+        let (mut r, f) = pending::<Pin<Box<dyn AsyncRead + Send>>>();
+        drop(f);
+        let err = r.read_to_end(&mut Vec::default()).await;
+        assert_eq!(err.unwrap_err().kind(), ErrorKind::UnexpectedEof);
     }
 
     #[test_log::test(tokio::test)]

crates/wasmtime/Cargo.toml

@@ -15,7 +15,7 @@ anyhow = { workspace = true, features = ["std"] }
 bytes = { workspace = true }
 futures = { workspace = true, features = ["alloc"] }
 tokio = { workspace = true, features = ["macros"] }
-tokio-util = { workspace = true, features = ["codec", "compat"] }
+tokio-util = { workspace = true, features = ["codec", "io"] }
 tracing = { workspace = true, features = ["attributes"] }
 uuid = { workspace = true, features = ["std", "v7"] }
 wasm-tokio = { workspace = true }