Allocate stream block buffers only as block data arrives

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.
- * <p>Use {@link Builder#withAcceptOversizedBlocks(boolean)} only for trusted
- * inputs that may contain noncanonical legacy LZ4 blocks. Enabling it restores
+ * <p>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.</p>
+ * original length. Even then, blocks whose compressed length exceeds the LZ4
+ * compress bound of their original length are rejected.</p>
  * @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.
-     * <p>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.</p>
+     * <p>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.</p>
      *
      * @param acceptOversizedBlocks whether oversized blocks should be accepted
      * @return current builder instance

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;
 

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);
 

src/java/net/jpountz/lz4/LZ4FrameInputStream.java

@@ -20,13 +20,15 @@
 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;
 import java.nio.ByteBuffer;
 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);
       }

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}.
    * <p>