From 2cf67e242e65e6391787de561ad0f3da6ad42d99 Mon Sep 17 00:00:00 2001 From: Jonas Konrad Date: Thu, 22 Oct 2026 10:27:59 +0200 Subject: [PATCH] Allocate stream block buffers only as block data arrives LZ4BlockInputStream sized its output buffer from the originalLen and its compressed buffer from the compressedLen of a block header before reading any of the block data. A 21 byte header could therefore make it allocate about 64 MiB, which stays referenced while the underlying stream stalls. LZ4FrameInputStream did the same with two buffers of the frame's maximum block size, up to 8 MiB, for a block header of any size. Read block data into buffers that only grow as the data arrives, to at most twice the data read so far, and grow them geometrically up to the maximum block size, so that a sequence of slightly larger blocks does not reallocate every time. Allocate the output buffer of a compressed block only once the compressed data has been read, and size it for the maximum LZ4 compression ratio of 255, which LZ4BlockInputStream now also checks in the block header. The memory a peer can make a stream hold is now proportional to the data it actually sent. The block stream's output buffer growth is also capped at the maximum block size now. A truncated block in LZ4FrameInputStream is still reported as a plain IOException, not as an EOFException that callers may take for a clean end of data. An Error such as OutOfMemoryError while reading a block now also makes the stream fail, since the block was already partially consumed. Co-Authored-By: Claude Opus 5.5 --- .../net/jpountz/lz4/LZ4BlockInputStream.java | 50 +++---- src/java/net/jpountz/lz4/LZ4Constants.java | 3 + .../lz4/LZ4DecompressorWithLength.java | 4 +- .../net/jpountz/lz4/LZ4FrameInputStream.java | 58 ++++++--- src/java/net/jpountz/lz4/LZ4Utils.java | 38 ++++++ .../jpountz/lz4/LZ4BlockStreamingTest.java | 123 ++++++++++++++++++ .../net/jpountz/lz4/LZ4FrameIOStreamTest.java | 103 ++++++++++++++- 7 files changed, 328 insertions(+), 51 deletions(-) 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};