Merge pull request #263 from Cysharp/hadashiA/fix-brotli-infinite-loop

Fix infinite loop in BrotliEncoder
Cysharp · Mar 21, 2024 · f9aaa81 · f9aaa81
2 parents d1440ca + e0a7ae2
commit f9aaa81
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 8 deletions.
diff --git a/src/MemoryPack.Core/Compression/BrotliCompressor.cs b/src/MemoryPack.Core/Compression/BrotliCompressor.cs
@@ -206,15 +206,17 @@ public void CopyTo<TBufferWriter>(ref MemoryPackWriter<TBufferWriter> memoryPack
         var encoder = new BrotliEncoder(quality, window);
         try
         {
-            var writtenNotAdvanced = 0;
+            var bytesWritten = 0;
             foreach (var item in bufferWriter)
             {
-                writtenNotAdvanced = CompressCore(ref encoder, item.Span, ref memoryPackWriter, initialLength: null, isFinalBlock: false);
+                var span = item.Span;
+                if (span.Length <= 0) continue;
+                bytesWritten += CompressCore(ref encoder, span, ref memoryPackWriter, initialLength: null, isFinalBlock: false);
             }
 
             // call BrotliEncoderOperation.Finish
-            var finalBlockLength = (writtenNotAdvanced == 0) ? null : (int?)(writtenNotAdvanced + 10);
-            CompressCore(ref encoder, ReadOnlySpan<byte>.Empty, ref memoryPackWriter, initialLength: finalBlockLength, isFinalBlock: true);
+            var finalBlockMaxLength = BrotliUtils.BrotliEncoderMaxCompressedSize(bytesWritten) - bytesWritten;
+            CompressCore(ref encoder, ReadOnlySpan<byte>.Empty, ref memoryPackWriter, initialLength: finalBlockMaxLength, isFinalBlock: true);
         }
         finally
         {
@@ -257,7 +259,7 @@ static int CompressCore<TBufferWriter>(ref BrotliEncoder encoder, ReadOnlySpan<b
         where TBufferWriter : class, IBufferWriter<byte>
 #endif
     {
-        var writtenNotAdvanced = 0;
+        var totalWritten = 0;
 
         var lastResult = OperationStatus.DestinationTooSmall;
         while (lastResult == OperationStatus.DestinationTooSmall)
@@ -266,21 +268,20 @@ static int CompressCore<TBufferWriter>(ref BrotliEncoder encoder, ReadOnlySpan<b
             var dest = MemoryMarshal.CreateSpan(ref spanRef, destBufferWriter.BufferLength);
 
             lastResult = encoder.Compress(source, dest, out int bytesConsumed, out int bytesWritten, isFinalBlock: isFinalBlock);
-            writtenNotAdvanced += bytesConsumed;
+            totalWritten += bytesWritten;
 
             if (lastResult == OperationStatus.InvalidData) MemoryPackSerializationException.ThrowCompressionFailed();
             if (bytesWritten > 0)
             {
                 destBufferWriter.Advance(bytesWritten);
-                writtenNotAdvanced = 0;
             }
             if (bytesConsumed > 0)
             {
                 source = source.Slice(bytesConsumed);
             }
         }
 
-        return writtenNotAdvanced;
+        return totalWritten;
     }
 
     public void Dispose()

diff --git a/tests/MemoryPack.Tests/BrotliTest.cs b/tests/MemoryPack.Tests/BrotliTest.cs
@@ -1,5 +1,7 @@
 using MemoryPack.Compression;
 using System;
+using System.Buffers;
+using System.IO.Compression;
 
 namespace MemoryPack.Tests;
 
@@ -14,6 +16,44 @@ public void LargeByteArray()
         data.MemDecmpDeserialize(bin);
     }
 
+    [Fact]
+    public void EncodeEmptyCntent()
+    {
+        var buffer = new ArrayBufferWriter<byte>();
+        using var state = MemoryPackWriterOptionalStatePool.Rent(null);
+        var writer = new MemoryPackWriter<ArrayBufferWriter<byte>>(ref buffer, state);
+
+        using var compressor = new BrotliCompressor(CompressionLevel.Fastest);
+        compressor.CopyTo(ref writer);
+
+        using var decompressor = new BrotliDecompressor();
+        decompressor.Decompress(compressor.ToArray()).ToArray().Should().BeEmpty();
+    }
+
+    [Fact]
+    public void EncodeEmptyFinalBlock()
+    {
+        using var state = MemoryPackWriterOptionalStatePool.Rent(null);
+
+        var compressor = new BrotliCompressor(CompressionLevel.Fastest);
+        var coWriter = new MemoryPackWriter<BrotliCompressor>(ref compressor, state);
+
+        var bytes = new byte[248];
+        Random.Shared.NextBytes(bytes);
+        coWriter.WriteUnmanagedArray(bytes);
+        coWriter.Flush();
+
+        var buffer = new ArrayBufferWriter<byte>();
+        compressor.CopyTo(buffer);
+
+        using var readerState = MemoryPackReaderOptionalStatePool.Rent(null);
+        using var decompressor = new BrotliDecompressor();
+        var decompressed = decompressor.Decompress(compressor.ToArray());
+        var reader = new MemoryPackReader(in decompressed, readerState);
+
+        reader.ReadArray<byte>().Should().BeEquivalentTo(bytes);
+        compressor.Dispose();
+    }
 }
 
 [MemoryPackable]