Merge branch 'main' into cxymds/fix-1855-stack

crates/rio/src/compress_reader.rs

@@ -71,6 +71,7 @@ where
 
     /// Optional: allow users to customize block_size
     pub fn with_block_size(inner: R, block_size: usize, compression_algorithm: CompressionAlgorithm) -> Self {
+        debug_assert!(block_size > 0, "CompressReader block_size must be non-zero");
         Self {
             inner,
             buffer: Vec::new(),
@@ -183,11 +184,21 @@ pin_project! {
         buffer: Vec<u8>,
         buffer_pos: usize,
         finished: bool,
+        // A previously surfaced stream error is sticky: without this, a caller
+        // that polls again after an error would restart at the header phase and
+        // read a truncated tail as a clean EOF, converting the error into a
+        // silently short body.
+        poisoned: bool,
         // Fields for saving header read progress across polls
         header_buf: [u8; 8],
         header_read: usize,
-        header_done: bool,
-        // Fields for saving compressed block read progress across polls
+        // Fields for saving compressed block read progress across polls.
+        // `compressed_len > 0` means a block payload is in flight: the header has
+        // been fully parsed and `compressed_read` bytes of the payload are already
+        // consumed from the inner stream. The header phase must not run again (and
+        // must not reset `compressed_read`) until this block completes, or a
+        // `Poll::Pending` in the middle of a payload would silently drop the bytes
+        // read so far and desynchronize the block framing.
         compressed_buf: Vec<u8>,
         compressed_read: usize,
         compressed_len: usize,
@@ -205,9 +216,9 @@ where
             buffer: Vec::new(),
             buffer_pos: 0,
             finished: false,
+            poisoned: false,
             header_buf: [0u8; 8],
             header_read: 0,
-            header_done: false,
             compressed_buf: Vec::new(),
             compressed_read: 0,
             compressed_len: 0,
@@ -236,54 +247,74 @@ where
         if *this.finished {
             return Poll::Ready(Ok(()));
         }
-        // Read header
-        while !*this.header_done && *this.header_read < HEADER_LEN {
-            let mut temp = [0u8; HEADER_LEN];
-            let mut temp_buf = ReadBuf::new(&mut temp[0..HEADER_LEN - *this.header_read]);
-            match this.inner.as_mut().poll_read(cx, &mut temp_buf) {
-                Poll::Pending => return Poll::Pending,
-                Poll::Ready(Ok(())) => {
-                    let n = temp_buf.filled().len();
-                    if n == 0 {
-                        break;
+        if *this.poisoned {
+            return Poll::Ready(Err(io::Error::new(io::ErrorKind::InvalidData, "decompress reader previously failed")));
+        }
+
+        if *this.compressed_len == 0 {
+            // Read the 8-byte block header, resuming across polls via `header_read`.
+            while *this.header_read < HEADER_LEN {
+                let mut temp = [0u8; HEADER_LEN];
+                let mut temp_buf = ReadBuf::new(&mut temp[0..HEADER_LEN - *this.header_read]);
+                match this.inner.as_mut().poll_read(cx, &mut temp_buf) {
+                    Poll::Pending => return Poll::Pending,
+                    Poll::Ready(Ok(())) => {
+                        let n = temp_buf.filled().len();
+                        if n == 0 {
+                            if *this.header_read == 0 {
+                                // Clean EOF on a block boundary.
+                                *this.finished = true;
+                                return Poll::Ready(Ok(()));
+                            }
+                            *this.poisoned = true;
+                            return Poll::Ready(Err(io::Error::new(
+                                io::ErrorKind::UnexpectedEof,
+                                "unexpected EOF while reading compressed block header",
+                            )));
+                        }
+                        this.header_buf[*this.header_read..*this.header_read + n].copy_from_slice(&temp_buf.filled()[..n]);
+                        *this.header_read += n;
+                    }
+                    Poll::Ready(Err(e)) => {
+                        // error!("DecompressReader poll_read: read header error: {e}");
+                        *this.poisoned = true;
+                        return Poll::Ready(Err(e));
                     }
-                    this.header_buf[*this.header_read..*this.header_read + n].copy_from_slice(&temp_buf.filled()[..n]);
-                    *this.header_read += n;
-                }
-                Poll::Ready(Err(e)) => {
-                    // error!("DecompressReader poll_read: read header error: {e}");
-                    return Poll::Ready(Err(e));
                 }
             }
-            if *this.header_read < HEADER_LEN {
-                return Poll::Pending;
-            }
-        }
-        if !*this.header_done && *this.header_read == 0 {
-            return Poll::Ready(Ok(()));
-        }
-        let typ = this.header_buf[0];
-        let len = (this.header_buf[1] as usize) | ((this.header_buf[2] as usize) << 8) | ((this.header_buf[3] as usize) << 16);
-        let crc = (this.header_buf[4] as u32)
-            | ((this.header_buf[5] as u32) << 8)
-            | ((this.header_buf[6] as u32) << 16)
-            | ((this.header_buf[7] as u32) << 24);
-        *this.header_read = 0;
-        *this.header_done = true;
 
-        if typ == COMPRESS_TYPE_END {
+            let typ = this.header_buf[0];
+            let len =
+                (this.header_buf[1] as usize) | ((this.header_buf[2] as usize) << 8) | ((this.header_buf[3] as usize) << 16);
+            *this.header_read = 0;
+
+            // `CompressReader` never emits an end block — a stream terminates on
+            // inner EOF, which is what lets concatenated per-part streams decode as
+            // one. This branch is kept for streams that do carry the marker.
+            if typ == COMPRESS_TYPE_END {
+                *this.compressed_read = 0;
+                *this.compressed_len = 0;
+                *this.finished = true;
+                return Poll::Ready(Ok(()));
+            }
+            if typ != COMPRESS_TYPE_COMPRESSED && typ != COMPRESS_TYPE_UNCOMPRESSED {
+                // error!("DecompressReader unknown compression type: {typ}");
+                *this.poisoned = true;
+                return Poll::Ready(Err(io::Error::new(io::ErrorKind::InvalidData, "Unknown compression type")));
+            }
+            if len == 0 {
+                *this.poisoned = true;
+                return Poll::Ready(Err(io::Error::new(io::ErrorKind::InvalidData, "Invalid compressed block length")));
+            }
+
+            if this.compressed_buf.len() < len {
+                this.compressed_buf.resize(len, 0);
+            }
+            *this.compressed_len = len;
             *this.compressed_read = 0;
-            *this.compressed_len = 0;
-            *this.finished = true;
-            return Poll::Ready(Ok(()));
         }
 
-        if this.compressed_buf.len() < len {
-            this.compressed_buf.resize(len, 0);
-        }
-        *this.compressed_len = len;
-        *this.compressed_read = 0;
-
+        // Fill the in-flight block payload, resuming across polls via `compressed_read`.
         while *this.compressed_read < *this.compressed_len {
             let mut temp_buf = ReadBuf::new(&mut this.compressed_buf[*this.compressed_read..*this.compressed_len]);
             match this.inner.as_mut().poll_read(cx, &mut temp_buf) {
@@ -291,7 +322,13 @@ where
                 Poll::Ready(Ok(())) => {
                     let n = temp_buf.filled().len();
                     if n == 0 {
-                        break;
+                        *this.compressed_read = 0;
+                        *this.compressed_len = 0;
+                        *this.poisoned = true;
+                        return Poll::Ready(Err(io::Error::new(
+                            io::ErrorKind::UnexpectedEof,
+                            "unexpected EOF while reading compressed block payload",
+                        )));
                     }
                     *this.compressed_read += n;
                 }
@@ -299,10 +336,17 @@ where
                     // error!("DecompressReader poll_read: read compressed block error: {e}");
                     *this.compressed_read = 0;
                     *this.compressed_len = 0;
+                    *this.poisoned = true;
                     return Poll::Ready(Err(e));
                 }
             }
         }
+
+        let typ = this.header_buf[0];
+        let crc = (this.header_buf[4] as u32)
+            | ((this.header_buf[5] as u32) << 8)
+            | ((this.header_buf[6] as u32) << 16)
+            | ((this.header_buf[7] as u32) << 24);
         let compressed_buf = &this.compressed_buf[..*this.compressed_len];
         // `compressed_buf`'s length comes from the untrusted 24-bit header length field, so it
         // can be shorter than 16 bytes. `uvarint` is safe on any slice length (reads at most 10
@@ -316,6 +360,7 @@ where
         if uvarint <= 0 || uvarint as usize > compressed_buf.len() {
             *this.compressed_read = 0;
             *this.compressed_len = 0;
+            *this.poisoned = true;
             return Poll::Ready(Err(io::Error::new(io::ErrorKind::InvalidData, "Invalid compressed block length prefix")));
         }
         let compressed_data = &compressed_buf[uvarint as usize..];
@@ -326,21 +371,29 @@ where
                     // error!("DecompressReader decompress_block error: {e}");
                     *this.compressed_read = 0;
                     *this.compressed_len = 0;
+                    *this.poisoned = true;
                     return Poll::Ready(Err(e));
                 }
             }
-        } else if typ == COMPRESS_TYPE_UNCOMPRESSED {
-            compressed_data.to_vec()
         } else {
-            // error!("DecompressReader unknown compression type: {typ}");
+            // The header phase already rejected every type other than
+            // COMPRESS_TYPE_COMPRESSED / COMPRESS_TYPE_UNCOMPRESSED.
+            compressed_data.to_vec()
+        };
+        if decompressed.is_empty() {
+            // The writer never emits zero-length plaintext blocks; an empty
+            // decode surfacing as Ready(Ok) with no bytes would read as EOF and
+            // silently truncate the stream.
+            *this.poisoned = true;
             *this.compressed_read = 0;
             *this.compressed_len = 0;
-            return Poll::Ready(Err(io::Error::new(io::ErrorKind::InvalidData, "Unknown compression type")));
-        };
+            return Poll::Ready(Err(io::Error::new(io::ErrorKind::InvalidData, "Empty compressed block")));
+        }
         if decompressed.len() != uncompress_len as usize {
             // error!("DecompressReader decompressed length mismatch: {} != {}", decompressed.len(), uncompress_len);
             *this.compressed_read = 0;
             *this.compressed_len = 0;
+            *this.poisoned = true;
             return Poll::Ready(Err(io::Error::new(io::ErrorKind::InvalidData, "Decompressed length mismatch")));
         }
         let actual_crc = {
@@ -352,13 +405,13 @@ where
             // error!("DecompressReader CRC32 mismatch: actual {actual_crc} != expected {crc}");
             *this.compressed_read = 0;
             *this.compressed_len = 0;
+            *this.poisoned = true;
             return Poll::Ready(Err(io::Error::new(io::ErrorKind::InvalidData, "CRC32 mismatch")));
         }
         *this.buffer = decompressed;
         *this.buffer_pos = 0;
         *this.compressed_read = 0;
         *this.compressed_len = 0;
-        *this.header_done = false;
         let to_copy = min(buf.remaining(), this.buffer.len());
         buf.put_slice(&this.buffer[..to_copy]);
         *this.buffer_pos += to_copy;
@@ -493,6 +546,184 @@ mod tests {
         assert_eq!(&decompressed, &data);
     }
 
+    /// Wraps a reader so every other poll returns `Poll::Pending` and every
+    /// `Ready` poll serves at most `chunk` bytes. This is the shape a duplex
+    /// pipe produces when the erasure writer is slower than the decoder, which
+    /// is exactly what desynchronized the block framing before the resumable
+    /// payload state was added (rustfs/rustfs#5957 multipart GET truncation).
+    struct PendingChunkReader<R> {
+        inner: R,
+        chunk: usize,
+        pending_next: bool,
+    }
+
+    impl<R> PendingChunkReader<R> {
+        fn new(inner: R, chunk: usize) -> Self {
+            Self {
+                inner,
+                chunk,
+                pending_next: true,
+            }
+        }
+    }
+
+    impl<R: AsyncRead + Unpin> AsyncRead for PendingChunkReader<R> {
+        fn poll_read(
+            mut self: std::pin::Pin<&mut Self>,
+            cx: &mut std::task::Context<'_>,
+            buf: &mut tokio::io::ReadBuf<'_>,
+        ) -> std::task::Poll<std::io::Result<()>> {
+            if self.pending_next {
+                self.pending_next = false;
+                cx.waker().wake_by_ref();
+                return std::task::Poll::Pending;
+            }
+            self.pending_next = true;
+            let cap = self.chunk.min(buf.remaining());
+            let mut scratch = vec![0u8; cap];
+            let mut inner_buf = tokio::io::ReadBuf::new(&mut scratch);
+            match std::pin::Pin::new(&mut self.inner).poll_read(cx, &mut inner_buf) {
+                std::task::Poll::Ready(Ok(())) => {
+                    buf.put_slice(inner_buf.filled());
+                    std::task::Poll::Ready(Ok(()))
+                }
+                other => other,
+            }
+        }
+    }
+
+    fn patterned_payload(size: usize, seed: u8) -> Vec<u8> {
+        (0..size)
+            .map(|i| ((i as u64).wrapping_mul(2_654_435_761).wrapping_add(seed as u64) >> 3) as u8)
+            .collect()
+    }
+
+    /// Root-cause regression for the multipart compressed GET truncation: a
+    /// `Poll::Pending` in the middle of a block payload must not drop the bytes
+    /// already consumed. Before the resumable payload state, the decoder reset
+    /// `compressed_read` on every re-poll and surfaced
+    /// `LZ4 error: ERROR_frameType_unknown` mid-stream.
+    #[tokio::test]
+    async fn test_decompress_reader_survives_pending_mid_payload() {
+        let data = patterned_payload(100 * 1024, 7);
+        let mut compress_reader =
+            CompressReader::with_block_size(Cursor::new(data.clone()), 8192, CompressionAlgorithm::default());
+        let mut compressed = Vec::new();
+        compress_reader.read_to_end(&mut compressed).await.unwrap();
+
+        for chunk in [1usize, 3, 7, 8, 17, 1000, 8192] {
+            let inner = PendingChunkReader::new(Cursor::new(compressed.clone()), chunk);
+            let mut decompress_reader = DecompressReader::new(inner, CompressionAlgorithm::default());
+            let mut decompressed = Vec::new();
+            decompress_reader.read_to_end(&mut decompressed).await.unwrap();
+            assert_eq!(decompressed, data, "pending-chunked decode must be byte-exact for chunk={chunk}");
+        }
+    }
+
+    /// Two independently compressed streams concatenated back to back — the
+    /// on-disk shape of a compressed multipart object — must decode across the
+    /// stream boundary even when every poll can suspend mid-block.
+    #[tokio::test]
+    async fn test_decompress_reader_survives_pending_across_concatenated_streams() {
+        let part1 = patterned_payload(64 * 1024, 7);
+        let part2 = patterned_payload(24 * 1024, 61);
+
+        let mut stored = Vec::new();
+        for part in [&part1, &part2] {
+            let mut compress_reader =
+                CompressReader::with_block_size(Cursor::new(part.clone()), 8192, CompressionAlgorithm::default());
+            let mut compressed = Vec::new();
+            compress_reader.read_to_end(&mut compressed).await.unwrap();
+            stored.extend_from_slice(&compressed);
+        }
+
+        let mut expected = part1;
+        expected.extend_from_slice(&part2);
+
+        for chunk in [1usize, 5, 8, 13, 4096] {
+            let inner = PendingChunkReader::new(Cursor::new(stored.clone()), chunk);
+            let mut decompress_reader = DecompressReader::new(inner, CompressionAlgorithm::default());
+            let mut decompressed = Vec::new();
+            decompress_reader.read_to_end(&mut decompressed).await.unwrap();
+            assert_eq!(
+                decompressed, expected,
+                "concatenated part streams must decode byte-exact for chunk={chunk}"
+            );
+        }
+    }
+
+    /// After the first stream error, every further poll must keep failing.
+    /// Without the sticky poison a retrying caller would restart at the header
+    /// phase and read the truncated tail as a clean EOF — converting a hard
+    /// error into a silently short body.
+    #[tokio::test]
+    async fn test_decompress_reader_error_is_sticky() {
+        let data = patterned_payload(32 * 1024, 7);
+        let mut compress_reader = CompressReader::with_block_size(Cursor::new(data), 8192, CompressionAlgorithm::default());
+        let mut compressed = Vec::new();
+        compress_reader.read_to_end(&mut compressed).await.unwrap();
+        compressed.truncate(compressed.len() - 3);
+
+        let mut decompress_reader = DecompressReader::new(Cursor::new(compressed), CompressionAlgorithm::default());
+        let mut out = Vec::new();
+        let first = decompress_reader
+            .read_to_end(&mut out)
+            .await
+            .expect_err("truncated payload must error");
+        assert_eq!(first.kind(), std::io::ErrorKind::UnexpectedEof);
+
+        let mut retry = Vec::new();
+        let second = decompress_reader
+            .read_to_end(&mut retry)
+            .await
+            .expect_err("a poll after the first error must not turn into a clean EOF");
+        assert_eq!(second.kind(), std::io::ErrorKind::InvalidData);
+        assert!(retry.is_empty(), "no bytes may be produced after the stream failed");
+    }
+
+    /// A stream cut off in the middle of a block payload must fail with a clean
+    /// UnexpectedEof instead of decoding a short buffer.
+    #[tokio::test]
+    async fn test_decompress_reader_truncated_payload_is_unexpected_eof() {
+        let data = patterned_payload(32 * 1024, 7);
+        let mut compress_reader = CompressReader::with_block_size(Cursor::new(data), 8192, CompressionAlgorithm::default());
+        let mut compressed = Vec::new();
+        compress_reader.read_to_end(&mut compressed).await.unwrap();
+
+        compressed.truncate(compressed.len() - 3);
+        let mut decompress_reader = DecompressReader::new(Cursor::new(compressed), CompressionAlgorithm::default());
+        let mut out = Vec::new();
+        let err = decompress_reader
+            .read_to_end(&mut out)
+            .await
+            .expect_err("truncated payload must error");
+        assert_eq!(err.kind(), std::io::ErrorKind::UnexpectedEof);
+    }
+
+    /// A stream cut off in the middle of a block header must fail with a clean
+    /// UnexpectedEof instead of parsing a garbage header.
+    #[tokio::test]
+    async fn test_decompress_reader_truncated_header_is_unexpected_eof() {
+        let data = patterned_payload(12 * 1024, 7);
+        let mut compress_reader = CompressReader::with_block_size(Cursor::new(data), 8192, CompressionAlgorithm::default());
+        let mut compressed = Vec::new();
+        compress_reader.read_to_end(&mut compressed).await.unwrap();
+
+        // Keep the first full block plus 3 bytes of the next header.
+        let ln = (compressed[1] as usize) | ((compressed[2] as usize) << 8) | ((compressed[3] as usize) << 16);
+        let first_block_end = 8 + ln;
+        assert!(compressed.len() > first_block_end, "fixture must contain more than one block");
+        compressed.truncate(first_block_end + 3);
+
+        let mut decompress_reader = DecompressReader::new(Cursor::new(compressed), CompressionAlgorithm::default());
+        let mut out = Vec::new();
+        let err = decompress_reader
+            .read_to_end(&mut out)
+            .await
+            .expect_err("truncated header must error");
+        assert_eq!(err.kind(), std::io::ErrorKind::UnexpectedEof);
+    }
+
     // Regression: a corrupted block whose 24-bit length field is < 16 must not panic.
     // Header layout (HEADER_LEN = 8): [type, len_lo, len_mid, len_hi, crc0..crc3], then `len`
     // bytes of block body. Pre-fix, poll_read sliced `compressed_buf[0..16]` unconditionally,
@@ -518,6 +749,85 @@ mod tests {
         assert_eq!(res.unwrap_err().kind(), std::io::ErrorKind::InvalidData);
     }
 
+    // Header-level fail-closed matrix, built by hand so the decoder is exercised against bytes no
+    // encoder in this crate can produce. Header layout (HEADER_LEN = 8):
+    // [type, len_lo, len_mid, len_hi, crc0..crc3], then `len` body bytes = uvarint(plain_len) + data.
+    #[tokio::test]
+    async fn test_decompress_reader_header_validation_matrix() {
+        // Build a block whose body is `uvarint(plain.len()) + plain` (i.e. the
+        // COMPRESS_TYPE_UNCOMPRESSED shape), with the header CRC taken over the plaintext exactly
+        // like the production writer does.
+        fn build_raw_block(typ: u8, plain: &[u8], len_override: Option<usize>) -> Vec<u8> {
+            let crc = {
+                let mut hasher = crc_fast::Digest::new(crc_fast::CrcAlgorithm::Crc32IsoHdlc);
+                hasher.update(plain);
+                hasher.finalize() as u32
+            };
+            let mut uvarint_buf = [0u8; 10];
+            let int_len = put_uvarint(&mut uvarint_buf[..], plain.len() as u64);
+            let body_len = int_len + plain.len();
+            let len = len_override.unwrap_or(body_len);
+
+            let mut out = Vec::with_capacity(HEADER_LEN + body_len);
+            out.push(typ);
+            out.push((len & 0xFF) as u8);
+            out.push(((len >> 8) & 0xFF) as u8);
+            out.push(((len >> 16) & 0xFF) as u8);
+            out.extend_from_slice(&crc.to_le_bytes());
+            out.extend_from_slice(&uvarint_buf[..int_len]);
+            out.extend_from_slice(plain);
+            out
+        }
+
+        let plain = b"uncompressed passthrough payload";
+
+        // (a) A well-formed uncompressed block decodes to the plaintext verbatim.
+        let mut out = Vec::new();
+        DecompressReader::new(
+            Cursor::new(build_raw_block(COMPRESS_TYPE_UNCOMPRESSED, plain, None)),
+            CompressionAlgorithm::default(),
+        )
+        .read_to_end(&mut out)
+        .await
+        .expect("a well-formed uncompressed block must decode");
+        assert_eq!(out.as_slice(), plain.as_slice());
+
+        // (b) An unknown block type must be rejected instead of being treated as passthrough.
+        let mut out = Vec::new();
+        let err = DecompressReader::new(Cursor::new(build_raw_block(0x7E, plain, None)), CompressionAlgorithm::default())
+            .read_to_end(&mut out)
+            .await
+            .expect_err("unknown compression type must error");
+        assert_eq!(err.kind(), std::io::ErrorKind::InvalidData);
+        assert!(err.to_string().contains("Unknown compression type"), "got: {err}");
+
+        // (c) A zero-length block would stall the decoder, so it must be rejected up front.
+        let mut out = Vec::new();
+        let err = DecompressReader::new(
+            Cursor::new(build_raw_block(COMPRESS_TYPE_UNCOMPRESSED, plain, Some(0))),
+            CompressionAlgorithm::default(),
+        )
+        .read_to_end(&mut out)
+        .await
+        .expect_err("zero-length block must error");
+        assert_eq!(err.kind(), std::io::ErrorKind::InvalidData);
+        assert!(err.to_string().contains("Invalid compressed block length"), "got: {err}");
+
+        // (d) A block that decodes to zero plaintext bytes must be rejected: the
+        // writer never emits empty blocks, and an empty decode surfacing as
+        // Ready(Ok) with no bytes would read as EOF and silently truncate.
+        let mut out = Vec::new();
+        let err = DecompressReader::new(
+            Cursor::new(build_raw_block(COMPRESS_TYPE_UNCOMPRESSED, b"", None)),
+            CompressionAlgorithm::default(),
+        )
+        .read_to_end(&mut out)
+        .await
+        .expect_err("empty block must error");
+        assert_eq!(err.kind(), std::io::ErrorKind::InvalidData);
+        assert!(err.to_string().contains("Empty compressed block"), "got: {err}");
+    }
+
     // Directly exercises the length-prefix guard: an unterminated varint (all continuation bytes)
     // makes `uvarint` return 0, which must be rejected as an invalid length prefix.
     #[tokio::test]

crates/rio/src/hardlimit_reader.rs

@@ -24,12 +24,17 @@ pin_project! {
         #[pin]
         pub inner: R,
         remaining: i64,
+        scratch: Vec<u8>,
     }
 }
 
 impl<R> HardLimitReader<R> {
     pub fn new(inner: R, limit: i64) -> Self {
-        HardLimitReader { inner, remaining: limit }
+        HardLimitReader {
+            inner,
+            remaining: limit,
+            scratch: Vec::new(),
+        }
     }
 }
 
@@ -37,19 +42,21 @@ impl<R> AsyncRead for HardLimitReader<R>
 where
     R: AsyncRead,
 {
-    fn poll_read(mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<Result<()>> {
-        if self.remaining < 0 {
+    fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<Result<()>> {
+        let mut this = self.project();
+        if *this.remaining < 0 {
             return Poll::Ready(Err(Error::other("input provided more bytes than specified")));
         }
-        let original_filled = buf.filled().len();
-        if self.remaining == 0 {
+        if buf.remaining() == 0 {
+            return Poll::Ready(Ok(()));
+        }
+        if *this.remaining == 0 {
             let mut discard = [0u8; 8192];
             let mut discard_buf = ReadBuf::new(&mut discard);
-            return match self.as_mut().project().inner.poll_read(cx, &mut discard_buf) {
+            return match this.inner.as_mut().poll_read(cx, &mut discard_buf) {
                 Poll::Pending => Poll::Pending,
                 Poll::Ready(Ok(())) => {
                     if discard_buf.filled().is_empty() {
-                        debug_assert_eq!(buf.filled().len(), original_filled);
                         Poll::Ready(Ok(()))
                     } else {
                         Poll::Ready(Err(Error::other("input provided more bytes than specified")))
@@ -58,30 +65,46 @@ where
                 Poll::Ready(Err(err)) => Poll::Ready(Err(err)),
             };
         }
-        // Save the initial length
-        let before = original_filled;
 
-        // Poll the inner reader
-        let this = self.as_mut().project();
-        let poll = this.inner.poll_read(cx, buf);
-
-        if let Poll::Ready(Ok(())) = &poll {
-            let after = buf.filled().len();
-            let read = (after - before) as i64;
-            if read == 0 && *this.remaining > 0 {
-                return Poll::Ready(Err(Error::new(
-                    std::io::ErrorKind::UnexpectedEof,
-                    IncompleteBody {
-                        remaining: *this.remaining,
-                    },
-                )));
+        let remaining = match usize::try_from(*this.remaining) {
+            Ok(remaining) => remaining,
+            Err(_) => usize::MAX,
+        };
+        let allowed = remaining.min(buf.remaining());
+        let read = if allowed == buf.remaining() {
+            let before = buf.filled().len();
+            match this.inner.as_mut().poll_read(cx, buf) {
+                Poll::Pending => return Poll::Pending,
+                Poll::Ready(Err(err)) => return Poll::Ready(Err(err)),
+                Poll::Ready(Ok(())) => buf.filled().len() - before,
             }
-            *this.remaining -= read;
-            if *this.remaining < 0 {
-                return Poll::Ready(Err(Error::other("input provided more bytes than specified")));
+        } else {
+            this.scratch.resize(allowed, 0);
+            let mut scratch_buf = ReadBuf::new(&mut this.scratch[..allowed]);
+            match this.inner.as_mut().poll_read(cx, &mut scratch_buf) {
+                Poll::Pending => return Poll::Pending,
+                Poll::Ready(Err(err)) => return Poll::Ready(Err(err)),
+                Poll::Ready(Ok(())) => {
+                    let read = scratch_buf.filled().len();
+                    buf.put_slice(scratch_buf.filled());
+                    read
+                }
             }
+        };
+        if read == 0 {
+            return Poll::Ready(Err(Error::new(
+                std::io::ErrorKind::UnexpectedEof,
+                IncompleteBody {
+                    remaining: *this.remaining,
+                },
+            )));
         }
-        poll
+        let read = match i64::try_from(read) {
+            Ok(read) => read,
+            Err(_) => return Poll::Ready(Err(Error::other("read count exceeds i64::MAX"))),
+        };
+        *this.remaining -= read;
+        Poll::Ready(Ok(()))
     }
 }
 
@@ -140,7 +163,12 @@ mod tests {
         assert!(err.is_some());
 
         let err = err.unwrap();
-        assert_eq!(err.kind(), std::io::ErrorKind::Other);
+        assert_eq!(err.kind(), std::io::ErrorKind::UnexpectedEof);
+        assert!(
+            err.get_ref()
+                .and_then(|source| source.downcast_ref::<std::io::Error>())
+                .is_some_and(|source| source.to_string().contains("more bytes than specified"))
+        );
     }
 
     #[tokio::test]
@@ -155,6 +183,17 @@ mod tests {
         assert_eq!(&buf, data);
     }
 
+    #[tokio::test]
+    async fn test_hardlimit_reader_zero_capacity_read_does_not_consume_input() {
+        let mut reader = HardLimitReader::new(BufReader::new(&b"abc"[..]), 3);
+        let mut empty = [];
+
+        assert_eq!(reader.read(&mut empty).await.expect("zero-capacity read should succeed"), 0);
+        let mut out = Vec::new();
+        reader.read_to_end(&mut out).await.expect("input should remain readable");
+        assert_eq!(out, b"abc");
+    }
+
     #[tokio::test]
     async fn test_hardlimit_reader_short_input_returns_unexpected_eof() {
         let data = b"abc";
@@ -195,4 +234,18 @@ mod tests {
         assert_eq!(err.kind(), std::io::ErrorKind::Other);
         assert!(err.to_string().contains("more bytes than specified"));
     }
+
+    #[tokio::test]
+    async fn test_hardlimit_reader_caps_each_read_before_reporting_extra_bytes() {
+        let mut reader = HardLimitReader::new(BufReader::new(&b"abcdef"[..]), 3);
+        let mut out = Vec::new();
+
+        let err = reader
+            .read_to_end(&mut out)
+            .await
+            .expect_err("bytes beyond the declared limit must be rejected");
+
+        assert_eq!(out, b"abc");
+        assert!(err.to_string().contains("more bytes than specified"));
+    }
 }