fix AsyncMessageReader to handle partial reads

Signed-off-by: Connor Tsui <[email protected]>

Author: connortsui20

vortex-ipc/src/messages/decoder.rs

@@ -46,9 +46,20 @@ enum State {
 
 #[derive(Debug)]
 pub enum PollRead {
+    /// A complete message was decoded.
     Some(DecoderMessage),
-    /// Returns the _total_ number of bytes needed to make progress.
-    /// Note this is _not_ the incremental number of bytes needed to make progress.
+    /// The decoder needs more data to make progress.
+    ///
+    /// The inner value is the *total* number of bytes the buffer should contain, not the
+    /// incremental amount needed. Callers should:
+    /// 1. Resize the buffer to this length.
+    /// 2. Fill the buffer completely (handling partial reads as needed).
+    /// 3. Only then call [`MessageDecoder::read_next`] again.
+    ///
+    /// The decoder checks [`bytes::Buf::remaining`] to determine available data, which for
+    /// [`bytes::BytesMut`] returns the buffer length regardless of how many bytes were actually
+    /// written. Calling `read_next` before the buffer is fully populated will cause the decoder
+    /// to read uninitialized or stale data.
     NeedMore(usize),
 }
 

vortex-ipc/src/messages/reader_async.rs

@@ -17,14 +17,31 @@ use crate::messages::DecoderMessage;
 use crate::messages::MessageDecoder;
 use crate::messages::PollRead;
 
+/// The state of an in-progress read operation.
+#[derive(Default)]
+enum ReadState {
+    /// Ready to consult the decoder for the next operation.
+    #[default]
+    AwaitingDecoder,
+    /// Filling the buffer with data from the underlying reader.
+    ///
+    /// Async readers may return fewer bytes than requested (partial reads), especially over network
+    /// connections. This state persists across multiple `poll_next` calls until the buffer is
+    /// completely filled, at which point we transition back to [`Self::AwaitingDecoder`].
+    Filling {
+        /// The number of bytes read into the buffer so far.
+        total_bytes_read: usize,
+    },
+}
+
 pin_project! {
     /// An IPC message reader backed by an `AsyncRead` stream.
     pub struct AsyncMessageReader<R> {
         #[pin]
         read: R,
         buffer: BytesMut,
         decoder: MessageDecoder,
-        bytes_read: usize,
+        state: ReadState,
     }
 }
 
@@ -34,7 +51,7 @@ impl<R> AsyncMessageReader<R> {
             read,
             buffer: BytesMut::new(),
             decoder: MessageDecoder::new(registry),
-            bytes_read: 0,
+            state: ReadState::default(),
         }
     }
 }
@@ -45,29 +62,35 @@ impl<R: AsyncRead> Stream for AsyncMessageReader<R> {
     fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
         let mut this = self.project();
         loop {
-            match this.decoder.read_next(this.buffer)? {
-                PollRead::Some(msg) => return Poll::Ready(Some(Ok(msg))),
-                PollRead::NeedMore(nbytes) => {
-                    this.buffer.resize(nbytes, 0x00);
+            match this.state {
+                ReadState::AwaitingDecoder => match this.decoder.read_next(this.buffer)? {
+                    PollRead::Some(msg) => return Poll::Ready(Some(Ok(msg))),
+                    // The decoder needs more data to be read into the buffer, so resize for more
+                    // space to read from `this.read`.
+                    PollRead::NeedMore(new_len) => {
+                        this.buffer.resize(new_len, 0x00);
+                        *this.state = ReadState::Filling {
+                            total_bytes_read: 0,
+                        };
+                    }
+                },
+                ReadState::Filling { total_bytes_read } => {
+                    let unfilled = &mut this.buffer.as_mut()[*total_bytes_read..];
 
-                    match ready!(
-                        this.read
-                            .as_mut()
-                            .poll_read(cx, &mut this.buffer.as_mut()[*this.bytes_read..])
-                    ) {
-                        Ok(0) => {
-                            // End of file
-                            return Poll::Ready(None);
-                        }
-                        Ok(nbytes) => {
-                            *this.bytes_read += nbytes;
-                            // If we've finished the read operation, then we continue the loop
-                            // and the decoder should present us with a new response.
-                            if *this.bytes_read == nbytes {
-                                *this.bytes_read = 0;
+                    match ready!(this.read.as_mut().poll_read(cx, unfilled)) {
+                        Err(e) => return Poll::Ready(Some(Err(e.into()))),
+                        Ok(0) => return Poll::Ready(None),
+                        // Continue to read until we fill the buffer. We need to do this in a loop
+                        // to handle partial reads.
+                        Ok(bytes_read) => {
+                            *total_bytes_read += bytes_read;
+
+                            if *total_bytes_read == this.buffer.len() {
+                                *this.state = ReadState::AwaitingDecoder;
+                            } else {
+                                debug_assert!(*total_bytes_read < this.buffer.len());
                             }
                         }
-                        Err(e) => return Poll::Ready(Some(Err(e.into()))),
                     }
                 }
             }