Add length before compression data (neo-project#1768)

* Add length before compression * Optimize * Fix UT * Add UT * Add UT Co-authored-by: erikzhang <erik@neo.org>
neo-ngd · Jul 18, 2020 · 9d83a64 · 9d83a64
1 parent 85543c7
commit 9d83a64
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 7 deletions.
diff --git a/src/neo/IO/Helper.cs b/src/neo/IO/Helper.cs
@@ -1,6 +1,7 @@
 using K4os.Compression.LZ4;
 using System;
 using System.Buffers;
+using System.Buffers.Binary;
 using System.Collections.Generic;
 using System.IO;
 using System.Linq;
@@ -69,20 +70,19 @@ public static byte[] CompressLz4(this byte[] data)
             int maxLength = LZ4Codec.MaximumOutputSize(data.Length);
             using var buffer = MemoryPool<byte>.Shared.Rent(maxLength);
             int length = LZ4Codec.Encode(data, buffer.Memory.Span);
-            byte[] result = new byte[length];
-            buffer.Memory[..length].CopyTo(result);
+            byte[] result = new byte[sizeof(uint) + length];
+            BinaryPrimitives.WriteInt32LittleEndian(result, data.Length);
+            buffer.Memory[..length].CopyTo(result.AsMemory(4));
             return result;
         }
 
         public static byte[] DecompressLz4(this byte[] data, int maxOutput)
         {
-            var maxDecompressDataLength = data.Length * 255;
-            if (maxDecompressDataLength > 0) maxOutput = Math.Min(maxOutput, maxDecompressDataLength);
-            using var buffer = MemoryPool<byte>.Shared.Rent(maxOutput);
-            int length = LZ4Codec.Decode(data, buffer.Memory.Span);
+            int length = BinaryPrimitives.ReadInt32LittleEndian(data);
             if (length < 0 || length > maxOutput) throw new FormatException();
             byte[] result = new byte[length];
-            buffer.Memory[..length].CopyTo(result);
+            if (LZ4Codec.Decode(data.AsSpan(4), result) != length)
+                throw new FormatException();
             return result;
         }
 

diff --git a/tests/neo.UnitTests/IO/UT_IOHelper.cs b/tests/neo.UnitTests/IO/UT_IOHelper.cs
@@ -119,6 +119,37 @@ public void TestAsSerializable()
             }
         }
 
+        [TestMethod]
+        public void TestCompression()
+        {
+            var data = new byte[] { 1, 2, 3, 4 };
+            var byteArray = Neo.IO.Helper.CompressLz4(data);
+            var result = Neo.IO.Helper.DecompressLz4(byteArray, byte.MaxValue);
+
+            CollectionAssert.AreEqual(result, data);
+
+            // Compress
+
+            data = new byte[255];
+            for (int x = 0; x < data.Length; x++) data[x] = 1;
+
+            byteArray = Neo.IO.Helper.CompressLz4(data);
+            result = Neo.IO.Helper.DecompressLz4(byteArray, byte.MaxValue);
+
+            Assert.IsTrue(byteArray.Length < result.Length);
+            CollectionAssert.AreEqual(result, data);
+
+            // Error max length
+
+            Assert.ThrowsException<FormatException>(() => Neo.IO.Helper.DecompressLz4(byteArray, byte.MaxValue - 1));
+            Assert.ThrowsException<FormatException>(() => Neo.IO.Helper.DecompressLz4(byteArray, -1));
+
+            // Error length
+
+            byteArray[0]++;
+            Assert.ThrowsException<FormatException>(() => Neo.IO.Helper.DecompressLz4(byteArray, byte.MaxValue));
+        }
+
         [TestMethod]
         public void TestAsSerializableArray()
         {