diff --git a/src/neo/IO/Helper.cs b/src/neo/IO/Helper.cs index e165016139..e871bbaedf 100644 --- 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.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.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 index 1128270716..b126b7f485 100644 --- 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(() => Neo.IO.Helper.DecompressLz4(byteArray, byte.MaxValue - 1)); + Assert.ThrowsException(() => Neo.IO.Helper.DecompressLz4(byteArray, -1)); + + // Error length + + byteArray[0]++; + Assert.ThrowsException(() => Neo.IO.Helper.DecompressLz4(byteArray, byte.MaxValue)); + } + [TestMethod] public void TestAsSerializableArray() {