Merge pull request #520 from CommunityToolkit/dev/fix-stream-read-of-t

Fix StreamExtensions.Read<T> when using buffered streams

Author: Sergio0694

src/CommunityToolkit.HighPerformance/Extensions/StreamExtensions.cs

@@ -207,27 +207,45 @@ public static unsafe T Read<T>(this Stream stream)
     {
 #if NETSTANDARD2_1_OR_GREATER
         T result = default;
-        int length = sizeof(T);
+        int bytesOffset = 0;
 
-        unsafe
+        // As per Stream.Read's documentation:
+        // "The total number of bytes read into the buffer. This can be less than the number of bytes allocated in the
+        // buffer if that many bytes are not currently available, or zero (0) if the end of the stream has been reached."
+        // Because of this, we have to loop until all requires bytes have been read, and only throw if the return is 0.
+        do
         {
-            if (stream.Read(new Span<byte>(&result, length)) != length)
+            int bytesRead = stream.Read(new Span<byte>((byte*)&result + bytesOffset, sizeof(T) - bytesOffset));
+
+            // A return value of 0 indicates that the end of the stream has been reached
+            if (bytesRead == 0)
             {
                 ThrowInvalidOperationExceptionForEndOfStream();
             }
+
+            bytesOffset += bytesRead;
         }
+        while (bytesOffset < sizeof(T));
 
         return result;
 #else
-        int length = sizeof(T);
-        byte[] buffer = ArrayPool<byte>.Shared.Rent(length);
+        int bytesOffset = 0;
+        byte[] buffer = ArrayPool<byte>.Shared.Rent(sizeof(T));
 
         try
         {
-            if (stream.Read(buffer, 0, length) != length)
+            do
             {
-                ThrowInvalidOperationExceptionForEndOfStream();
+                int bytesRead = stream.Read(buffer, bytesOffset, sizeof(T) - bytesOffset);
+
+                if (bytesRead == 0)
+                {
+                    ThrowInvalidOperationExceptionForEndOfStream();
+                }
+
+                bytesOffset += bytesRead;
             }
+            while (bytesOffset < sizeof(T));
 
             return Unsafe.ReadUnaligned<T>(ref buffer[0]);
         }

tests/CommunityToolkit.HighPerformance.UnitTests/Extensions/Test_StreamExtensions.cs

@@ -4,7 +4,6 @@
 
 using System;
 using System.IO;
-using CommunityToolkit.HighPerformance;
 using Microsoft.VisualStudio.TestTools.UnitTesting;
 
 namespace CommunityToolkit.HighPerformance.UnitTests.Extensions;
@@ -37,4 +36,83 @@ public void Test_StreamExtensions_ReadWrite()
 
         _ = Assert.ThrowsException<InvalidOperationException>(() => stream.Read<long>());
     }
+
+    // See https://github.com/CommunityToolkit/dotnet/issues/513
+    [TestMethod]
+    public void Test_StreamExtensions_ReadWrite_WithBufferedStream()
+    {
+        Stream stream = new BufferedStream();
+
+        stream.Write(true);
+        stream.Write(42);
+        stream.Write(3.14f);
+        stream.Write(unchecked(uint.MaxValue * 324823489204ul));
+
+        stream.Position = 0;
+
+        Assert.AreEqual(true, stream.Read<bool>());
+        Assert.AreEqual(42, stream.Read<int>());
+        Assert.AreEqual(3.14f, stream.Read<float>());
+        Assert.AreEqual(unchecked(uint.MaxValue * 324823489204ul), stream.Read<ulong>());
+
+        _ = Assert.ThrowsException<InvalidOperationException>(() => stream.Read<long>());
+    }
+
+    private sealed class BufferedStream : MemoryStream
+    {
+        private ReadOnlyMemory<byte> bufferedBytes;
+
+        public override int Read(byte[] buffer, int offset, int count)
+        {
+            if (this.bufferedBytes.IsEmpty)
+            {
+                this.bufferedBytes = ReadMoreBytes();
+            }
+
+            int bytesToCopy = Math.Min(this.bufferedBytes.Length, count);
+
+            this.bufferedBytes.Span.Slice(0, bytesToCopy).CopyTo(buffer.AsSpan(offset, count));
+            this.bufferedBytes = this.bufferedBytes.Slice(bytesToCopy);
+
+            return bytesToCopy;
+        }
+
+#if NET6_0_OR_GREATER
+        public override int Read(Span<byte> buffer)
+        {
+            if (this.bufferedBytes.IsEmpty)
+            {
+                this.bufferedBytes = ReadMoreBytes();
+            }
+
+            int bytesToCopy = Math.Min(this.bufferedBytes.Length, buffer.Length);
+
+            this.bufferedBytes.Span.Slice(0, bytesToCopy).CopyTo(buffer);
+            this.bufferedBytes = this.bufferedBytes.Slice(bytesToCopy);
+
+            return bytesToCopy;
+        }
+#endif
+
+        private byte[] ReadMoreBytes()
+        {
+            byte[] array = new byte[3];
+            int bytesOffset = 0;
+
+            do
+            {
+                int bytesRead = base.Read(array, bytesOffset, 3 - bytesOffset);
+
+                bytesOffset += bytesRead;
+
+                if (bytesRead == 0)
+                {
+                    return array.AsSpan(0, bytesOffset).ToArray();
+                }
+            }
+            while (bytesOffset < 3);
+
+            return array;
+        }
+    }
 }