diff --git a/src/java/net/jpountz/lz4/LZ4BlockInputStream.java b/src/java/net/jpountz/lz4/LZ4BlockInputStream.java index 3467cf10..db1036eb 100644 --- a/src/java/net/jpountz/lz4/LZ4BlockInputStream.java +++ b/src/java/net/jpountz/lz4/LZ4BlockInputStream.java @@ -23,6 +23,7 @@ import static net.jpountz.lz4.LZ4BlockOutputStream.HEADER_LENGTH; import static net.jpountz.lz4.LZ4BlockOutputStream.MAGIC; import static net.jpountz.lz4.LZ4BlockOutputStream.MAGIC_LENGTH; +import static net.jpountz.lz4.LZ4Constants.MAX_COMPRESSION_RATIO; import java.io.EOFException; import java.io.FilterInputStream; @@ -41,12 +42,11 @@ * support {@link #mark(int)}/{@link #reset()}. * Once a read fails with an {@link IOException}, every later read or skip on * this stream throws an {@link IOException} as well. - *
Use {@link Builder#withAcceptOversizedBlocks(boolean)} only for trusted - * inputs that may contain noncanonical legacy LZ4 blocks. Enabling it restores + *
Use {@link Builder#withAcceptOversizedBlocks(boolean)} only for inputs + * that may contain noncanonical legacy LZ4 blocks. Enabling it restores * acceptance of blocks whose compressed length is greater than or equal to the - * original length, which may permit large attacker-controlled allocations. - * Even then, blocks whose compressed length exceeds the LZ4 compress bound of - * their original length are rejected.
+ * original length. Even then, blocks whose compressed length exceeds the LZ4 + * compress bound of their original length are rejected. * @see LZ4BlockOutputStream */ public class LZ4BlockInputStream extends FilterInputStream { @@ -273,6 +273,10 @@ private void refill() throws IOException { } catch (RuntimeException e) { failure = new IOException("Stream is corrupted", e); throw failure; + } catch (Error e) { + // e.g. OutOfMemoryError: the block was partially consumed, so the stream cannot continue + failure = new IOException(e); + throw e; } } @@ -314,7 +318,9 @@ private boolean readBlock() throws IOException { || (compressionMethod == COMPRESSION_METHOD_RAW && originalLen != compressedLen) || (compressionMethod == COMPRESSION_METHOD_LZ4 && !acceptOversizedBlocks && compressedLen >= originalLen) // even an oversized LZ4 block cannot be larger than the LZ4 compress bound of its original length - || (compressionMethod == COMPRESSION_METHOD_LZ4 && compressedLen > LZ4Utils.maxCompressedLength(originalLen))) { + || (compressionMethod == COMPRESSION_METHOD_LZ4 && compressedLen > LZ4Utils.maxCompressedLength(originalLen)) + // nor can it decompress to more than the maximum compression ratio allows + || (compressionMethod == COMPRESSION_METHOD_LZ4 && originalLen > (long) compressedLen * MAX_COMPRESSION_RATIO)) { throw new IOException("Stream is corrupted"); } if (originalLen == 0 && compressedLen == 0) { @@ -328,18 +334,19 @@ private boolean readBlock() throws IOException { finished = true; return true; } - if (buffer.length < originalLen) { - buffer = new byte[Math.max(originalLen, buffer.length * 3 / 2)]; - } + // The lengths come from the untrusted header, so only allocate as the block data actually arrives + final int maxBlockSize = 1 << compressionLevel; switch (compressionMethod) { case COMPRESSION_METHOD_RAW: - readFully(buffer, originalLen); + buffer = LZ4Utils.readFullyGrowing(in, buffer, originalLen, maxBlockSize); break; case COMPRESSION_METHOD_LZ4: - if (compressedBuffer.length < compressedLen) { - compressedBuffer = new byte[Math.max(compressedLen, compressedBuffer.length * 3 / 2)]; + compressedBuffer = LZ4Utils.readFullyGrowing(in, compressedBuffer, compressedLen, + LZ4Utils.maxCompressedLength(maxBlockSize)); + // originalLen is at most MAX_COMPRESSION_RATIO times the compressed data that was read + if (buffer.length < originalLen) { + buffer = new byte[LZ4Utils.grownLength(buffer.length, originalLen, maxBlockSize)]; } - readFully(compressedBuffer, compressedLen); try { if (fastDecompressor == null) { final int decompressedLen = safeDecompressor.decompress(compressedBuffer, 0, compressedLen, buffer, 0, originalLen); @@ -368,8 +375,8 @@ private boolean readBlock() throws IOException { return true; } - // Like readFully(), except it signals incomplete reads by returning - // the number of bytes read (less than len) instead of throwing EOFException. + // Reads up to len bytes, and signals incomplete reads by returning + // the number of bytes read (less than len). private int tryReadFully(byte[] b, int len) throws IOException { int read = 0; while (read < len) { @@ -382,12 +389,6 @@ private int tryReadFully(byte[] b, int len) throws IOException { return read; } - private void readFully(byte[] b, int len) throws IOException { - if (tryReadFully(b, len) != len) { - throw new EOFException("Stream ended prematurely"); - } - } - @Override public boolean markSupported() { return false; @@ -439,10 +440,9 @@ public Builder withStopOnEmptyBlock(boolean stopOnEmptyBlock) { /** * Registers whether the builder should accept oversized blocks whose * compressed length is greater than or equal to the original length. - *Only enable this for trusted inputs. Accepting such blocks may permit - * large attacker-controlled allocations. Blocks whose compressed length - * exceeds the LZ4 compress bound of their original length are rejected - * even when this is enabled.
+ *Only enable this for inputs that may contain such noncanonical legacy + * blocks. Blocks whose compressed length exceeds the LZ4 compress bound of + * their original length are rejected even when this is enabled.
* * @param acceptOversizedBlocks whether oversized blocks should be accepted * @return current builder instance diff --git a/src/java/net/jpountz/lz4/LZ4Constants.java b/src/java/net/jpountz/lz4/LZ4Constants.java index 6bab187d..d0116a2c 100644 --- a/src/java/net/jpountz/lz4/LZ4Constants.java +++ b/src/java/net/jpountz/lz4/LZ4Constants.java @@ -31,6 +31,9 @@ enum LZ4Constants { static final int MIN_MATCH = 4; + // Each LZ4 match-length extension byte adds at most 255 decompressed bytes. + static final int MAX_COMPRESSION_RATIO = 255; + static final int HASH_LOG = MEMORY_USAGE - 2; static final int HASH_TABLE_SIZE = 1 << HASH_LOG; diff --git a/src/java/net/jpountz/lz4/LZ4DecompressorWithLength.java b/src/java/net/jpountz/lz4/LZ4DecompressorWithLength.java index eb0c2770..9e039aa7 100644 --- a/src/java/net/jpountz/lz4/LZ4DecompressorWithLength.java +++ b/src/java/net/jpountz/lz4/LZ4DecompressorWithLength.java @@ -21,6 +21,8 @@ import net.jpountz.util.ByteBufferUtils; import net.jpountz.util.SafeUtils; +import static net.jpountz.lz4.LZ4Constants.MAX_COMPRESSION_RATIO; + /** * Convenience class to decompress data compressed by {@link LZ4CompressorWithLength}. * This decompressor is NOT compatible with any other compressors in lz4-java @@ -39,8 +41,6 @@ public class LZ4DecompressorWithLength { - // Each LZ4 match-length extension byte adds at most 255 decompressed bytes. - private static final int MAX_COMPRESSION_RATIO = 255; private static final int DEFAULT_MAX_DECOMPRESSED_LENGTH = Integer.getInteger( "net.jpountz.lz4.LZ4DecompressorWithLength.maxDecompressedLength", 64 * 1024 * 1024); diff --git a/src/java/net/jpountz/lz4/LZ4FrameInputStream.java b/src/java/net/jpountz/lz4/LZ4FrameInputStream.java index 3bd0fec8..fccfdb93 100644 --- a/src/java/net/jpountz/lz4/LZ4FrameInputStream.java +++ b/src/java/net/jpountz/lz4/LZ4FrameInputStream.java @@ -20,6 +20,7 @@ import net.jpountz.xxhash.XXHash32; import net.jpountz.xxhash.XXHashFactory; +import java.io.EOFException; import java.io.FilterInputStream; import java.io.IOException; import java.io.InputStream; @@ -27,6 +28,7 @@ import java.nio.ByteOrder; import java.util.Locale; +import static net.jpountz.lz4.LZ4Constants.MAX_COMPRESSION_RATIO; import static net.jpountz.lz4.LZ4Utils.notEnoughSpace; /** @@ -51,6 +53,7 @@ public class LZ4FrameInputStream extends FilterInputStream { static final String DESCRIPTOR_HASH_MISMATCH = "Stream frame descriptor corrupted"; static final String INVALID_DESCRIPTOR = "Invalid or unsupported frame descriptor"; static final int MAGIC_SKIPPABLE_BASE = 0x184D2A50; + private static final byte[] EMPTY = new byte[0]; private final LZ4SafeDecompressor decompressor; private final XXHash32 checksum; @@ -306,6 +309,29 @@ private void readBlock() throws IOException { throw fail(e); } catch (RuntimeException e) { throw fail(new IOException("Stream is corrupted", e)); + } catch (Error e) { + // e.g. OutOfMemoryError: the block was partially consumed, so the stream cannot continue + fail(new IOException(e)); + throw e; + } + } + + // Reads the data of a block of the given length, growing buf as the data arrives + private byte[] readBlockData(byte[] buf, int len) throws IOException { + try { + return LZ4Utils.readFullyGrowing(in, buf == null ? EMPTY : buf, len, maxBlockSize); + } catch (EOFException e) { + // report truncation like the other parts of the frame do, not as an EOFException that callers may take for a + // clean end of data + throw new IOException(PREMATURE_EOS); + } + } + + private void setRawBuffer(byte[] rawBuffer) { + if (rawBuffer != this.rawBuffer) { + this.rawBuffer = rawBuffer; + buffer = ByteBuffer.wrap(rawBuffer); + buffer.limit(0); // don't expose buffer contents if this block fails } } @@ -348,31 +374,17 @@ private void readBlock0() throws IOException { throw new IOException(String.format(Locale.ROOT, "Block size %s exceeded max: %s", blockSize, maxBlockSize)); } - // The buffers may be larger than maxBlockSize if a previous frame had a larger block size. Never shrink them. - if (rawBuffer == null || rawBuffer.length < maxBlockSize) { - rawBuffer = new byte[maxBlockSize]; - buffer = ByteBuffer.wrap(rawBuffer); - buffer.limit(0); // don't expose buffer contents if this block fails - } - final byte[] tmpBuffer; // Use a temporary buffer, potentially one used for compression + // The block size comes from the untrusted input, so the buffers are only allocated as the block data actually + // arrives. They may be larger than maxBlockSize if a previous frame had a larger block size. Never shrink them. + final byte[] tmpBuffer; if (compressed) { - if (compressedBuffer == null || compressedBuffer.length < maxBlockSize) { - compressedBuffer = new byte[maxBlockSize]; // Reused during different compressions - } + compressedBuffer = readBlockData(compressedBuffer, blockSize); tmpBuffer = compressedBuffer; } else { + setRawBuffer(readBlockData(rawBuffer, blockSize)); tmpBuffer = rawBuffer; } - int offset = 0; - while (offset < blockSize) { - final int lastRead = in.read(tmpBuffer, offset, blockSize - offset); - if (lastRead < 0) { - throw new IOException(PREMATURE_EOS); - } - offset += lastRead; - } - // verify block checksum if (frameInfo.isEnabled(LZ4FrameOutputStream.FLG.Bits.BLOCK_CHECKSUM)) { final int hashCheck = readInt(in); @@ -383,8 +395,14 @@ private void readBlock0() throws IOException { final int currentBufferSize; if (compressed) { + // the block cannot decompress to more than the maximum compression ratio allows + final int maxDecompressedSize = (int) Math.min(maxBlockSize, (long) blockSize * MAX_COMPRESSION_RATIO); + if (rawBuffer == null || rawBuffer.length < maxDecompressedSize) { + final int oldLength = rawBuffer == null ? 0 : rawBuffer.length; + setRawBuffer(new byte[LZ4Utils.grownLength(oldLength, maxDecompressedSize, maxBlockSize)]); + } try { - currentBufferSize = decompressor.decompress(tmpBuffer, 0, blockSize, rawBuffer, 0, maxBlockSize); + currentBufferSize = decompressor.decompress(tmpBuffer, 0, blockSize, rawBuffer, 0, maxDecompressedSize); } catch (LZ4Exception e) { throw new IOException(e); } diff --git a/src/java/net/jpountz/lz4/LZ4Utils.java b/src/java/net/jpountz/lz4/LZ4Utils.java index 08f85903..8b8f3e56 100644 --- a/src/java/net/jpountz/lz4/LZ4Utils.java +++ b/src/java/net/jpountz/lz4/LZ4Utils.java @@ -16,6 +16,9 @@ * limitations under the License. */ +import java.io.EOFException; +import java.io.IOException; +import java.io.InputStream; import java.nio.ByteBuffer; import java.util.Arrays; @@ -30,6 +33,7 @@ enum LZ4Utils { ; private static final int MAX_INPUT_SIZE = 0x7E000000; + private static final int MIN_GROWN_BUFFER_SIZE = 64 * 1024; static int maxCompressedLength(int length) { if (length < 0) { @@ -40,6 +44,40 @@ static int maxCompressedLength(int length) { return length + length / 255 + 16; } + /** + * Reads exactly {@code len} bytes from {@code in} into the start of {@code buf} and returns the buffer holding + * them, which is {@code buf} if it is large enough. Otherwise, the buffer is grown only as data actually arrives, + * to at most twice the data read so far (but at least {@code min(len, 64 KiB)}), so that a length read from + * untrusted input cannot cause a large allocation before the data is present. The buffer is grown geometrically + * beyond {@code len}, up to {@code maxLen}, so that a sequence of slightly larger reads does not reallocate every + * time. + * + * @param maxLen the largest length any read into this buffer can have, at least {@code len} + * @throws EOFException if the stream ends before {@code len} bytes were read + */ + static byte[] readFullyGrowing(InputStream in, byte[] buf, int len, int maxLen) throws IOException { + int read = 0; + while (read < len) { + if (read == buf.length) { + buf = Arrays.copyOf(buf, grownLength(buf.length, Math.min(len, MIN_GROWN_BUFFER_SIZE), maxLen)); + } + final int r = in.read(buf, read, Math.min(len, buf.length) - read); + if (r < 0) { + throw new EOFException("Stream ended prematurely"); + } + read += r; + } + return buf; + } + + /** + * Returns the length to grow a buffer of length {@code length} to: twice its length, but at least + * {@code minLength} and at most {@code maxLength}. + */ + static int grownLength(int length, int minLength, int maxLength) { + return (int) Math.min(maxLength, Math.max(minLength, 2L * length)); + } + /** * Returns {@code true} if {@code available < required}. *diff --git a/src/test/net/jpountz/lz4/LZ4BlockStreamingTest.java b/src/test/net/jpountz/lz4/LZ4BlockStreamingTest.java index 8648d7b9..e6c0ef75 100644 --- a/src/test/net/jpountz/lz4/LZ4BlockStreamingTest.java +++ b/src/test/net/jpountz/lz4/LZ4BlockStreamingTest.java @@ -24,10 +24,12 @@ import java.io.IOException; import java.io.InputStream; import java.io.OutputStream; +import java.lang.management.ManagementFactory; import java.nio.ByteBuffer; import java.nio.ByteOrder; import java.nio.charset.Charset; import java.util.Arrays; +import java.util.Random; import java.util.function.Consumer; import java.util.zip.Adler32; import java.util.zip.CRC32; @@ -36,7 +38,9 @@ import com.carrotsearch.randomizedtesting.annotations.ParametersFactory; import net.jpountz.xxhash.XXHashFactory; +import org.junit.Assume; import org.junit.Test; +import org.junit.function.ThrowingRunnable; import static org.junit.Assert.*; import com.carrotsearch.randomizedtesting.annotations.Repeat; @@ -677,4 +681,123 @@ public void testFailureIsStickyAfterInvalidOriginalLength() throws IOException { assertFailedStream(in); } } + + private static byte[] blockHeader(int compressionMethod, int compressedLen, int originalLen) { + return ByteBuffer.allocate(LZ4BlockOutputStream.HEADER_LENGTH) + .order(ByteOrder.LITTLE_ENDIAN) + .put(LZ4BlockOutputStream.MAGIC) + .put((byte) (compressionMethod | 0x0F)) + .putInt(compressedLen) + .putInt(originalLen) + .array(); + } + + /** + * Runs the given action and returns the number of bytes allocated by the current thread while doing so. + */ + private static long allocatedBytes(ThrowingRunnable action) throws Throwable { + Assume.assumeTrue(ManagementFactory.getThreadMXBean() instanceof com.sun.management.ThreadMXBean); + final com.sun.management.ThreadMXBean threadMXBean = (com.sun.management.ThreadMXBean) ManagementFactory.getThreadMXBean(); + Assume.assumeTrue(threadMXBean.isThreadAllocatedMemorySupported() && threadMXBean.isThreadAllocatedMemoryEnabled()); + + final long threadId = Thread.currentThread().getId(); + long allocated = 0; + // the first iteration is a warm-up, so that one-time allocations (e.g. class initialization) are not counted + for (int i = 0; i < 2; i++) { + final long allocatedBefore = threadMXBean.getThreadAllocatedBytes(threadId); + action.run(); + allocated = threadMXBean.getThreadAllocatedBytes(threadId) - allocatedBefore; + } + return allocated; + } + + @Test + public void testTruncatedBlockDoesNotAllocateBlockSize() throws Throwable { + final int maxBlockSize = 1 << 25; + final byte[][] headers = { + blockHeader(LZ4BlockOutputStream.COMPRESSION_METHOD_LZ4, maxBlockSize - 1, maxBlockSize), + blockHeader(LZ4BlockOutputStream.COMPRESSION_METHOD_RAW, maxBlockSize, maxBlockSize), + }; + for (byte[] header : headers) { + // a header claiming a 32 MiB block, followed by a little payload and then the end of the stream + final byte[] truncated = Arrays.copyOf(header, header.length + 1000); + final long allocated = allocatedBytes(() -> assertThrows(EOFException.class, + () -> lz4BlockInputStreamBuilder().build(new ByteArrayInputStream(truncated)).read())); + assertTrue("Allocated " + allocated + " bytes", allocated < (1 << 20)); + } + } + + @Test + public void testRejectsOriginalLengthAboveMaxCompressionRatio() { + final byte[] header = blockHeader(LZ4BlockOutputStream.COMPRESSION_METHOD_LZ4, 4, 4 * 255 + 1); + final LZ4BlockInputStream in = lz4BlockInputStreamBuilder().build(new ByteArrayInputStream(header)); + final IOException e = assertThrows(IOException.class, in::read); + assertEquals("Stream is corrupted", e.getMessage()); + + // at the maximum compression ratio, the header is accepted, and reading fails only because the payload is missing + final byte[] maxRatioHeader = blockHeader(LZ4BlockOutputStream.COMPRESSION_METHOD_LZ4, 4, 4 * 255); + final LZ4BlockInputStream maxRatioIn = lz4BlockInputStreamBuilder().build(new ByteArrayInputStream(maxRatioHeader)); + assertThrows(EOFException.class, maxRatioIn::read); + } + + @Test + public void testMaximallyCompressibleBlocks() throws IOException { + final int blockSize = 1 << 25; + final byte[] data = new byte[blockSize + 12345]; + for (LZ4Compressor compressor : new LZ4Compressor[] { + LZ4Factory.fastestJavaInstance().fastCompressor(), + LZ4Factory.fastestJavaInstance().highCompressor(17), + }) { + final ByteArrayOutputStream out = new ByteArrayOutputStream(); + try (LZ4BlockOutputStream os = new LZ4BlockOutputStream(out, blockSize, compressor)) { + os.write(data); + } + try (LZ4BlockInputStream in = lz4BlockInputStreamBuilder().build(new ByteArrayInputStream(out.toByteArray()))) { + assertArrayEquals(data, in.readAllBytes()); + } + } + } + + @Test + public void testGrowingBlocksAllocation() throws Throwable { + // blocks that are each slightly larger than the previous one, as written by flushing growing messages + final int blockCount = 256; + final int baseSize = 256 << 10; + final byte[] data = new byte[baseSize + blockCount]; + new Random(42).nextBytes(data); + for (boolean compressible : new boolean[] {true, false}) { + final byte[] blockData = compressible ? new byte[data.length] : data; + if (compressible) { + // compressible, but with a compressed length that grows with the block + for (int i = 0; i < blockData.length; i++) { + blockData[i] = (byte) (data[i] & 1); + } + } + final ByteArrayOutputStream out = new ByteArrayOutputStream(); + long expectedLength = 0; + try (LZ4BlockOutputStream os = new LZ4BlockOutputStream(out, 1 << 20, + LZ4Factory.fastestJavaInstance().fastCompressor(), XXHashFactory.fastestInstance().newStreamingHash32(LZ4BlockOutputStream.DEFAULT_SEED).asChecksum(), true)) { + for (int i = 0; i < blockCount; i++) { + os.write(blockData, 0, baseSize + i); + os.flush(); + expectedLength += baseSize + i; + } + } + final byte[] compressed = out.toByteArray(); + final byte[] buf = new byte[1 << 10]; + final long finalExpectedLength = expectedLength; + final long allocated = allocatedBytes(() -> { + long length = 0; + try (LZ4BlockInputStream in = lz4BlockInputStreamBuilder().build(new ByteArrayInputStream(compressed))) { + int n; + while ((n = in.read(buf)) != -1) { + length += n; + } + } + assertEquals(finalExpectedLength, length); + }); + // the buffers grow geometrically, rather than being reallocated for every larger block + assertTrue("Allocated " + allocated + " bytes", allocated < 3 * (1L << 20)); + } + } } diff --git a/src/test/net/jpountz/lz4/LZ4FrameIOStreamTest.java b/src/test/net/jpountz/lz4/LZ4FrameIOStreamTest.java index 2cc18ffc..de2b242b 100644 --- a/src/test/net/jpountz/lz4/LZ4FrameIOStreamTest.java +++ b/src/test/net/jpountz/lz4/LZ4FrameIOStreamTest.java @@ -649,8 +649,10 @@ private static long allocatedBytesForDecode(byte[] input, long expectedLength) t @Test public void testBlockBuffersReusedAcrossFrames() throws IOException { - final byte[] data = new byte[] {1, 2, 3}; - final int frameCount = 2000; + // incompressible, so that the compressed block is large enough to need the full 4 MiB output buffer + final byte[] data = new byte[20000]; + new Random(42).nextBytes(data); + final int frameCount = 200; final ByteArrayOutputStream baos = new ByteArrayOutputStream(); for (int i = 0; i < frameCount; i++) { writeFrame(baos, BD_4MB, new byte[0]); @@ -658,8 +660,40 @@ public void testBlockBuffersReusedAcrossFrames() throws IOException { writeFrame(baos, FLG_BLOCK_INDEPENDENCE | FLG_CONTENT_CHECKSUM, BD_4MB, data, true); } final long allocated = allocatedBytesForDecode(baos.toByteArray(), (long) frameCount * data.length); - // one raw and one compressed buffer of 4 MiB each, plus some slack - Assert.assertTrue("Allocated " + allocated + " bytes", allocated < 3 * (8L << 20)); + // one 4 MiB output buffer and one compressed buffer, plus some slack per frame. Without reuse, this would be + // frameCount * 4 MiB. + Assert.assertTrue("Allocated " + allocated + " bytes", allocated < (4L << 20) + frameCount * 8192L); + } + + @Test + public void testGrowingBlocksAllocation() throws IOException { + Assume.assumeTrue(testSize == 0); // does not depend on the test size + // blocks that are each slightly larger than the previous one, as written by flushing growing messages + final int blockCount = 256; + final int baseSize = 256 << 10; + final byte[] data = new byte[baseSize + blockCount]; + new Random(42).nextBytes(data); + for (boolean compress : new boolean[] {true, false}) { + final ByteArrayOutputStream baos = new ByteArrayOutputStream(); + baos.write(frameHeader(FLG_BLOCK_INDEPENDENCE, BD_4MB), 0, 7); + final ByteBuffer intBuffer = ByteBuffer.allocate(4).order(ByteOrder.LITTLE_ENDIAN); + long expectedLength = 0; + for (int i = 0; i < blockCount; i++) { + final int size = baseSize + i; + final byte[] block = compress + ? LZ4Factory.fastestJavaInstance().fastCompressor().compress(data, 0, size) + : Arrays.copyOf(data, size); + intBuffer.putInt(0, compress ? block.length : block.length | LZ4FrameOutputStream.LZ4_FRAME_INCOMPRESSIBLE_MASK); + baos.write(intBuffer.array(), 0, 4); + baos.write(block, 0, block.length); + expectedLength += size; + } + intBuffer.putInt(0, 0); // EndMark + baos.write(intBuffer.array(), 0, 4); + final long allocated = allocatedBytesForDecode(baos.toByteArray(), expectedLength); + // the buffers grow geometrically, rather than being reallocated for every larger block + Assert.assertTrue("Allocated " + allocated + " bytes", allocated < 3 * (4L << 20)); + } } @Test @@ -677,6 +711,67 @@ public void testFramesWithoutBlocksAllocation() throws IOException { Assert.assertTrue("Allocated " + skippableAllocated + " bytes", skippableAllocated < frameCount * 64L); } + @Test + public void testTruncatedBlockDoesNotAllocateBlockSize() throws IOException { + Assume.assumeTrue(testSize == 0); // does not depend on the test size + Assume.assumeTrue(ManagementFactory.getThreadMXBean() instanceof com.sun.management.ThreadMXBean); + final com.sun.management.ThreadMXBean threadMXBean = (com.sun.management.ThreadMXBean) ManagementFactory.getThreadMXBean(); + Assume.assumeTrue(threadMXBean.isThreadAllocatedMemorySupported() && threadMXBean.isThreadAllocatedMemoryEnabled()); + + final int maxBlockSize = 4 << 20; + for (int blockHeader : new int[] {maxBlockSize, maxBlockSize | LZ4FrameOutputStream.LZ4_FRAME_INCOMPRESSIBLE_MASK}) { + // a frame header and a block header claiming a 4 MiB block, followed by a little payload and then the end of + // the stream + final byte[] frameHeader = frameHeader(FLG_BLOCK_INDEPENDENCE, BD_4MB); + final ByteBuffer input = ByteBuffer.allocate(frameHeader.length + 4 + 1000).order(ByteOrder.LITTLE_ENDIAN); + input.put(frameHeader).putInt(blockHeader); + + final long threadId = Thread.currentThread().getId(); + long allocated = 0; + // the first iteration is a warm-up, so that one-time allocations (e.g. class initialization) are not counted + for (int i = 0; i < 2; i++) { + final long allocatedBefore = threadMXBean.getThreadAllocatedBytes(threadId); + try (InputStream is = new LZ4FrameInputStream(new ByteArrayInputStream(input.array()))) { + final IOException e = Assert.assertThrows(IOException.class, is::read); + // a plain IOException, not an EOFException that callers may take for a clean end of data + Assert.assertEquals(IOException.class, e.getClass()); + Assert.assertEquals(LZ4FrameInputStream.PREMATURE_EOS, e.getMessage()); + } + allocated = threadMXBean.getThreadAllocatedBytes(threadId) - allocatedBefore; + } + Assert.assertTrue("Allocated " + allocated + " bytes", allocated < (1 << 20)); + } + } + + @Test + public void testSmallCompressedBlockDoesNotAllocateBlockSize() throws IOException { + Assume.assumeTrue(testSize == 0); // does not depend on the test size + final byte[] data = new byte[] {1, 2, 3}; + final ByteArrayOutputStream baos = new ByteArrayOutputStream(); + writeFrame(baos, BD_4MB, data); + final long allocated = allocatedBytesForDecode(baos.toByteArray(), data.length); + Assert.assertTrue("Allocated " + allocated + " bytes", allocated < (1 << 20)); + } + + @Test + public void testMaximallyCompressibleBlocks() throws IOException { + Assume.assumeTrue(testSize == 0); // does not depend on the test size + final byte[] data = new byte[(4 << 20) + 12345]; + for (LZ4Compressor compressor : new LZ4Compressor[] { + LZ4Factory.fastestJavaInstance().fastCompressor(), + LZ4Factory.fastestJavaInstance().highCompressor(17), + }) { + final ByteArrayOutputStream baos = new ByteArrayOutputStream(); + try (OutputStream os = new LZ4FrameOutputStream(baos, LZ4FrameOutputStream.BLOCKSIZE.SIZE_4MB, -1L, + compressor, XXHashFactory.fastestInstance().hash32(), LZ4FrameOutputStream.FLG.Bits.BLOCK_INDEPENDENCE)) { + os.write(data); + } + try (InputStream is = new LZ4FrameInputStream(new ByteArrayInputStream(baos.toByteArray()))) { + Assert.assertArrayEquals(data, is.readAllBytes()); + } + } + } + @Test public void testBlockLargerThanFrameBlockSizeAfterLargerFrame() throws IOException { final byte[] first = new byte[] {1, 2, 3};