Fix ZstandardStream truncating multi-frame zstd responses to the first frame

src/libraries/System.IO.Compression/src/System/IO/Compression/Zstandard/ZstandardDecoder.cs

@@ -156,6 +156,14 @@ private void InitializeDecoder()
         /// <param name="bytesConsumed">When this method returns, contains the number of bytes consumed from the source.</param>
         /// <param name="bytesWritten">When this method returns, contains the number of bytes written to the destination.</param>
         /// <returns>An <see cref="OperationStatus"/> indicating the result of the operation.</returns>
+        /// <remarks>
+        /// The method returns <see cref="OperationStatus.Done"/> after all
+        /// contents of a single Zstandard frame are decoded and returned. However,
+        /// Zstandard data streams may consist of multiple concatenated frames
+        /// (some of which may even be empty). To allow processing further
+        /// frames on the same <see cref="ZstandardDecoder" /> instance, call
+        /// <see cref="Reset" /> before calling <see cref="Decompress" /> on the rest of the data.
+        /// </remarks>
         /// <exception cref="ObjectDisposedException">The decoder has been disposed.</exception>
         /// <exception cref="IOException">An error occurred during decompression.</exception>
         public OperationStatus Decompress(ReadOnlySpan<byte> source, Span<byte> destination, out int bytesConsumed, out int bytesWritten)
@@ -344,7 +352,7 @@ public static bool TryDecompress(ReadOnlySpan<byte> source, Span<byte> destinati
             }
         }
 
-        /// <summary>Resets the decoder session, allowing reuse for the next decompression operation.</summary>
+        /// <summary>Resets the decoder session, allowing reuse for decompressing the next Zstandard frame.</summary>
         /// <exception cref="ObjectDisposedException">The decoder has been disposed.</exception>
         /// <exception cref="IOException">Failed to reset the decoder session.</exception>
         public void Reset()

src/libraries/System.IO.Compression/src/System/IO/Compression/Zstandard/ZstandardStream.Decompress.cs

@@ -13,6 +13,10 @@ public sealed partial class ZstandardStream
     {
         private ZstandardDecoder? _decoder;
         private bool _nonEmptyInput;
+        private bool _endOfStream;
+
+        // Length of a Zstandard frame magic number, in bytes.
+        private const int ZstdFrameMagicLength = 4;
 
         /// <summary>Initializes a new instance of the <see cref="ZstandardStream" /> class by using the specified stream and decoder instance.</summary>
         /// <param name="stream">The stream from which data to decompress is read.</param>
@@ -49,21 +53,33 @@ public ZstandardStream(Stream stream, ZstandardDecompressionOptions decompressio
             _decoder = new ZstandardDecoder(decompressionOptions);
         }
 
+        // Decompresses buffered input for the current frame only. Returns true when there is a result for
+        // the caller to act on (output was produced, the frame finished, or a zero-byte read should
+        // return), and false when more input is needed to make progress on the current frame. Crossing a
+        // frame boundary into the next concatenated frame is handled by Read/ReadAsync once this reports
+        // OperationStatus.Done.
         private bool TryDecompress(Span<byte> destination, out int bytesWritten, out OperationStatus lastResult)
         {
             Debug.Assert(_decoder != null);
 
-            // Decompress any data we may have in our buffer.
-            lastResult = _decoder.Decompress(_buffer.ActiveSpan, destination, out int bytesConsumed, out bytesWritten);
+            OperationStatus status = _decoder.Decompress(_buffer.ActiveSpan, destination, out int bytesConsumed, out bytesWritten);
             _buffer.Discard(bytesConsumed);
+            lastResult = status;
 
-            if (lastResult == OperationStatus.InvalidData)
+            if (status == OperationStatus.InvalidData)
             {
                 throw new InvalidDataException(SR.ZstandardStream_Decompress_InvalidData);
             }
 
-            // If we successfully decompressed any bytes, or if we've reached the end of the decompression, we're done.
-            if (bytesWritten != 0 || lastResult == OperationStatus.Done)
+            if (status == OperationStatus.Done)
+            {
+                // Reached the end of a frame. ZSTD_decompressStream reports end-of-frame, not end-of-stream:
+                // a zstd stream may be frames concatenated back-to-back (RFC 8878 section 3), so Read/ReadAsync
+                // decide whether another frame follows. This may be a zero-output frame (bytesWritten == 0).
+                return true;
+            }
+
+            if (bytesWritten != 0)
             {
                 return true;
             }
@@ -78,17 +94,16 @@ private bool TryDecompress(Span<byte> destination, out int bytesWritten, out Ope
                 // doesn't necessarily mean it'll decompress into at least one byte of output, but it's a reasonable approximation
                 // for the 99% case.  If it's wrong, it just means that a caller using zero-byte reads as a way to delay
                 // getting a buffer to use for a subsequent call may end up getting one earlier than otherwise preferred.
-                Debug.Assert(lastResult == OperationStatus.DestinationTooSmall);
+                Debug.Assert(status == OperationStatus.DestinationTooSmall);
                 if (_buffer.ActiveLength != 0)
                 {
-                    Debug.Assert(bytesWritten == 0);
                     return true;
                 }
             }
 
             Debug.Assert(
-                lastResult == OperationStatus.NeedMoreData ||
-                (lastResult == OperationStatus.DestinationTooSmall && destination.IsEmpty && _buffer.ActiveLength == 0), $"{nameof(lastResult)} == {lastResult}, {nameof(destination.Length)} == {destination.Length}");
+                status == OperationStatus.NeedMoreData ||
+                (status == OperationStatus.DestinationTooSmall && destination.IsEmpty && _buffer.ActiveLength == 0), $"{nameof(status)} == {status}, {nameof(destination.Length)} == {destination.Length}");
 
             return false;
         }
@@ -130,42 +145,139 @@ public override int Read(Span<byte> buffer)
 
             try
             {
-                int bytesWritten;
-                OperationStatus lastResult;
-                while (!TryDecompress(buffer, out bytesWritten, out lastResult))
+                if (_endOfStream)
                 {
-                    _buffer.EnsureAvailableSpace(1);
+                    // A previous read reached the end of the final frame (or rejected trailing data). The
+                    // boundary probe may have left the decoder non-resumable, so never re-enter the decode
+                    // loop; report end-of-stream to all subsequent reads.
+                    return 0;
+                }
 
-                    int bytesRead = _stream.Read(_buffer.AvailableSpan);
-                    if (bytesRead <= 0)
+                while (true)
+                {
+                    int bytesWritten;
+                    OperationStatus lastResult;
+                    while (!TryDecompress(buffer, out bytesWritten, out lastResult))
                     {
-                        if (_nonEmptyInput && !buffer.IsEmpty)
-                            ThrowTruncatedInvalidData();
-                        break;
+                        _buffer.EnsureAvailableSpace(1);
+
+                        int bytesRead = _stream.Read(_buffer.AvailableSpan);
+                        if (bytesRead <= 0)
+                        {
+                            // The underlying stream ended in the middle of a frame, so the data is truncated.
+                            // A clean end after a completed frame is reported as Done by TryDecompress and is
+                            // resolved by the frame-boundary logic below, not here.
+                            if (_nonEmptyInput && !buffer.IsEmpty)
+                            {
+                                ThrowTruncatedInvalidData();
+                            }
+
+                            return 0;
+                        }
+
+                        _nonEmptyInput = true;
+
+                        if (bytesRead > _buffer.AvailableLength)
+                        {
+                            ThrowInvalidStream();
+                        }
+
+                        _buffer.Commit(bytesRead);
                     }
 
-                    _nonEmptyInput = true;
+                    if (lastResult != OperationStatus.Done || bytesWritten != 0)
+                    {
+                        // Output to hand back, or not at a finished-frame boundary: return to the caller.
+                        return bytesWritten;
+                    }
 
-                    if (bytesRead > _buffer.AvailableLength)
+                    // A frame finished with no pending output. A zstd stream may be frames concatenated
+                    // back-to-back (RFC 8878 section 3), so decide whether another frame follows before
+                    // reporting end-of-stream. async: false completes synchronously (it only ever takes the
+                    // synchronous read path).
+                    ValueTask<bool> advanceTask = AdvanceToNextFrame(async: false, cancellationToken: default);
+                    Debug.Assert(advanceTask.IsCompleted, "AdvanceToNextFrame should complete synchronously when async: false");
+                    if (!advanceTask.GetAwaiter().GetResult())
                     {
-                        ThrowInvalidStream();
+                        return bytesWritten;
                     }
+                }
+            }
+            finally
+            {
+                EndRWOperation();
+            }
+        }
 
-                    _buffer.Commit(bytesRead);
+        // Called after TryDecompress reports OperationStatus.Done with no pending output (a finished frame).
+        // Decides whether another concatenated frame follows by reading up to a frame magic and feeding just
+        // those bytes to the decoder: a valid magic is accepted (NeedMoreData) so decoding continues, while
+        // bytes that are not a frame magic identify trailing data after the final frame. Returns true (decoder
+        // reset and ready) so the caller loops to decode the next frame, or false (stream complete; _endOfStream
+        // is set and a seekable base stream is rewound to the end of the compressed data) so the caller returns.
+        // Shared by Read (async: false) and ReadAsync (async: true).
+        private async ValueTask<bool> AdvanceToNextFrame(bool async, CancellationToken cancellationToken)
+        {
+            Debug.Assert(_decoder != null);
+
+            if (_buffer.ActiveLength < ZstdFrameMagicLength)
+            {
+                // Not enough buffered to tell a split next-frame magic from end-of-stream; read just enough
+                // to complete the magic. Limiting the read to exactly the missing bytes (rather than filling
+                // the available buffer) avoids consuming and hiding trailing bytes past the magic on a
+                // non-seekable stream, where they can't be rewound; any following frame's body is read by the
+                // Read/ReadAsync loop afterwards.
+                int needed = ZstdFrameMagicLength - _buffer.ActiveLength;
+                _buffer.EnsureAvailableSpace(needed);
+                int peeked = async
+                    ? await _stream.ReadAtLeastAsync(_buffer.AvailableMemory.Slice(0, needed), needed, throwOnEndOfStream: false, cancellationToken: cancellationToken).ConfigureAwait(false)
+                    : _stream.ReadAtLeast(_buffer.AvailableSpan.Slice(0, needed), needed, throwOnEndOfStream: false);
+                if (peeked > 0)
+                {
+                    _nonEmptyInput = true;
+                    _buffer.Commit(peeked);
                 }
+            }
 
-                // When decompression finishes, rewind the stream to the exact end of compressed data
-                if (lastResult == OperationStatus.Done && _stream.CanSeek)
+            if (_buffer.ActiveLength >= ZstdFrameMagicLength)
+            {
+                // Determine whether another concatenated frame follows by feeding the decoder exactly the next
+                // frame magic number. The decoder validates the magic against both standard and skippable
+                // frames: a valid magic leaves it asking for more input (NeedMoreData), so decoding continues,
+                // while bytes that are not a frame magic produce InvalidData, identifying trailing data after
+                // the final frame. Feeding only the magic (not the whole buffer) means a frame whose magic is
+                // valid but whose body is corrupt is not mistaken for trailing data here: the magic is accepted
+                // and the corrupt body is rejected by the subsequent decode in the Read/ReadAsync loop.
+                _decoder.Reset();
+
+                // The magic alone never decodes into output, so the decoder needs no real output space; a single
+                // scratch byte borrowed from the buffer's free region (which the decoder won't write to) satisfies
+                // the non-empty-destination requirement. A Span local / stackalloc can't be used here because this
+                // method is async.
+                _buffer.EnsureAvailableSpace(1);
+                if (_decoder.Decompress(_buffer.ActiveSpan.Slice(0, ZstdFrameMagicLength), _buffer.AvailableSpan.Slice(0, 1), out int bytesConsumed, out _) == OperationStatus.NeedMoreData)
                 {
-                    TryRewindStream(_stream);
+                    // A valid magic; the decoder has taken the magic bytes into its session to continue decoding
+                    // the rest of the frame. Drop them from the buffer and keep decoding.
+                    _buffer.Discard(bytesConsumed);
+                    return true;
                 }
 
-                return bytesWritten;
+                // Not a frame: leave the magic bytes buffered so they are included in the trailing-data rewind
+                // below (on a non-seekable stream they simply remain unconsumed).
             }
-            finally
+
+            // Trailing non-zstd data or end of input after the final frame: the stream is complete. Mark the
+            // stream ended so subsequent reads short-circuit to 0 without re-entering the (now non-resumable)
+            // decoder, and leave any trailing bytes on a seekable base stream by rewinding to the end of the
+            // compressed data, mirroring how DeflateStream handles data after the last gzip member.
+            _endOfStream = true;
+            if (_stream.CanSeek)
             {
-                EndRWOperation();
+                TryRewindStream(_stream);
             }
+
+            return false;
         }
 
         /// <summary>Asynchronously reads decompressed bytes from the underlying stream and places them in the specified array.</summary>
@@ -207,37 +319,60 @@ public override async ValueTask<int> ReadAsync(Memory<byte> buffer, Cancellation
 
             try
             {
-                int bytesWritten;
-                OperationStatus lastResult;
-                while (!TryDecompress(buffer.Span, out bytesWritten, out lastResult))
+                if (_endOfStream)
                 {
-                    _buffer.EnsureAvailableSpace(1);
+                    // A previous read reached the end of the final frame (or rejected trailing data). The
+                    // boundary probe may have left the decoder non-resumable, so never re-enter the decode
+                    // loop; report end-of-stream to all subsequent reads.
+                    return 0;
+                }
 
-                    int bytesRead = await _stream.ReadAsync(_buffer.AvailableMemory, cancellationToken).ConfigureAwait(false);
-                    if (bytesRead <= 0)
+                while (true)
+                {
+                    int bytesWritten;
+                    OperationStatus lastResult;
+                    while (!TryDecompress(buffer.Span, out bytesWritten, out lastResult))
                     {
-                        if (_nonEmptyInput && !buffer.IsEmpty)
-                            ThrowTruncatedInvalidData();
-                        break;
+                        _buffer.EnsureAvailableSpace(1);
+
+                        int bytesRead = await _stream.ReadAsync(_buffer.AvailableMemory, cancellationToken).ConfigureAwait(false);
+                        if (bytesRead <= 0)
+                        {
+                            // The underlying stream ended in the middle of a frame, so the data is truncated.
+                            // A clean end after a completed frame is reported as Done by TryDecompress and is
+                            // resolved by the frame-boundary logic below, not here.
+                            if (_nonEmptyInput && !buffer.IsEmpty)
+                            {
+                                ThrowTruncatedInvalidData();
+                            }
+
+                            return 0;
+                        }
+
+                        _nonEmptyInput = true;
+
+                        if (bytesRead > _buffer.AvailableLength)
+                        {
+                            ThrowInvalidStream();
+                        }
+
+                        _buffer.Commit(bytesRead);
                     }
 
-                    _nonEmptyInput = true;
-
-                    if (bytesRead > _buffer.AvailableLength)
+                    if (lastResult != OperationStatus.Done || bytesWritten != 0)
                     {
-                        ThrowInvalidStream();
+                        // Output to hand back, or not at a finished-frame boundary: return to the caller.
+                        return bytesWritten;
                     }
 
-                    _buffer.Commit(bytesRead);
-                }
-
-                // When decompression finishes, rewind the stream to the exact end of compressed data
-                if (lastResult == OperationStatus.Done && _stream.CanSeek)
-                {
-                    TryRewindStream(_stream);
+                    // A frame finished with no pending output. A zstd stream may be frames concatenated
+                    // back-to-back (RFC 8878 section 3), so decide whether another frame follows before
+                    // reporting end-of-stream.
+                    if (!await AdvanceToNextFrame(async: true, cancellationToken: cancellationToken).ConfigureAwait(false))
+                    {
+                        return bytesWritten;
+                    }
                 }
-
-                return bytesWritten;
             }
             finally
             {