Ensure to always fill the buffer in the exact length

CollapseLauncher · Sep 24, 2023 · aae89d8 · aae89d8
1 parent de287be
commit aae89d8
Showing 1 changed file with 26 additions and 4 deletions.
diff --git a/Hi3Helper.SharpHDiffPatch/PatchCore/PatchCore.cs b/Hi3Helper.SharpHDiffPatch/PatchCore/PatchCore.cs
@@ -37,6 +37,28 @@ internal struct CoverHeader
         internal int nextCoverIndex;
     }
 
+#if !NET7_0_OR_GREATER
+    internal static class StreamExtension
+    {
+        public static int ReadExactly(this Stream stream, Span<byte> buffer)
+        {
+            int totalRead = 0;
+            while (totalRead < buffer.Length)
+            {
+                int read = stream.Read(buffer.Slice(totalRead));
+                if (read == 0)
+                {
+                    return totalRead;
+                }
+
+                totalRead += read;
+            }
+
+            return totalRead;
+        }
+    }
+#endif
+
     internal class PatchCore
     {
         private const int _kSignTagBit = 1;
@@ -230,7 +252,7 @@ private static void _patch_add_old_with_rle(Stream outCache, ref RLERefClipStruc
             rleLoader.rleInputClip.BaseStream.Position = oldPos;
 
             Span<byte> tempBuffer = new byte[decodeStep];
-            rleLoader.rleInputClip.BaseStream.Read(tempBuffer);
+            rleLoader.rleInputClip.BaseStream.ReadExactly(tempBuffer);
             outCache.Write(tempBuffer);
             outCache.Position = lastPos;
             _TBytesRle_load_stream_decode_add(ref rleLoader, outCache, decodeStep);
@@ -239,7 +261,7 @@ private static void _patch_add_old_with_rle(Stream outCache, ref RLERefClipStruc
         private static void _TOutStreamCache_copyFromClip(Stream outCache, BinaryReader copyReader, long copyLength)
         {
             Span<byte> buffer = new byte[copyLength];
-            copyReader.BaseStream.Read(buffer);
+            copyReader.BaseStream.ReadExactly(buffer);
             long lastPos = outCache.Position;
             outCache.Write(buffer);
             outCache.Position = lastPos;
@@ -315,10 +337,10 @@ private static unsafe void _TBytesRle_load_stream_mem_add(ref RLERefClipStruct r
 
             Span<byte> rleData = stackalloc byte[decodeStep];
             Span<byte> oldData = stackalloc byte[decodeStep];
-            rleLoader.rleCodeClip.BaseStream.Read(rleData);
+            rleLoader.rleCodeClip.BaseStream.ReadExactly(rleData);
 
             long lastPosCopy = outCache.Position;
-            outCache.Read(oldData);
+            outCache.ReadExactly(oldData);
             outCache.Position = lastPosCopy;
 
             fixed (byte* rlePtr = rleData)