From 9f8c3cf74223ed0a8a834134be9c917b9f10ceb5 Mon Sep 17 00:00:00 2001
From: BD <tunnelshade@users.noreply.github.com>
Date: Sun, 24 Sep 2023 04:25:18 +0530
Subject: [PATCH] Merge pull request from GHSA-55g7-9cwv-5qfv

* Validate chunk size to be within a configured maximum

* Add constructors to have max size configurable

* Code cleanup

* Use 512MB for consistency

---------

Co-authored-by: Taro L. Saito <leo@xerial.org>
---
 .../SnappyHadoopCompatibleOutputStream.java   |  4 ++--
 .../org/xerial/snappy/SnappyInputStream.java  | 23 +++++++++++++++++++
 .../org/xerial/snappy/SnappyOutputStream.java |  6 ++++-
 .../xerial/snappy/SnappyOutputStreamTest.java | 12 ++++++++++
 .../java/org/xerial/snappy/SnappyTest.java    | 18 +++++++++++++++
 5 files changed, 60 insertions(+), 3 deletions(-)

diff --git a/src/main/java/org/xerial/snappy/SnappyHadoopCompatibleOutputStream.java b/src/main/java/org/xerial/snappy/SnappyHadoopCompatibleOutputStream.java
index 67d2d188..d9f69bf4 100644
--- a/src/main/java/org/xerial/snappy/SnappyHadoopCompatibleOutputStream.java
+++ b/src/main/java/org/xerial/snappy/SnappyHadoopCompatibleOutputStream.java
@@ -1,9 +1,9 @@
 package org.xerial.snappy;
 
-import java.io.OutputStream;
-
 import org.xerial.snappy.buffer.CachedBufferAllocator;
 
+import java.io.OutputStream;
+
 public class SnappyHadoopCompatibleOutputStream extends SnappyOutputStream
 {
     public SnappyHadoopCompatibleOutputStream(OutputStream out)
diff --git a/src/main/java/org/xerial/snappy/SnappyInputStream.java b/src/main/java/org/xerial/snappy/SnappyInputStream.java
index 9835cf90..7ccce1da 100755
--- a/src/main/java/org/xerial/snappy/SnappyInputStream.java
+++ b/src/main/java/org/xerial/snappy/SnappyInputStream.java
@@ -36,8 +36,11 @@
 public class SnappyInputStream
         extends InputStream
 {
+    public static final int MAX_CHUNK_SIZE = 512 * 1024 * 1024; // 512 MiB
+
     private boolean finishedReading = false;
     protected final InputStream in;
+    private final int maxChunkSize;
 
     private byte[] compressed;
     private byte[] uncompressed;
@@ -55,6 +58,21 @@ public class SnappyInputStream
     public SnappyInputStream(InputStream input)
             throws IOException
     {
+        this(input, MAX_CHUNK_SIZE);
+    }
+
+
+    /**
+     * Create a filter for reading compressed data as a uncompressed stream with provided maximum chunk size
+     *
+     * @param input
+     * @param maxChunkSize
+     * @throws IOException
+     */
+    public SnappyInputStream(InputStream input, int maxChunkSize)
+            throws IOException
+    {
+        this.maxChunkSize = maxChunkSize;
         this.in = input;
         readHeader();
     }
@@ -422,6 +440,11 @@ protected boolean hasNextChunk()
             throw new SnappyError(SnappyErrorCode.INVALID_CHUNK_SIZE, "chunkSize is too big or negative : " + chunkSize);
         }
 
+        // chunkSize is big
+        if (chunkSize > maxChunkSize) {
+            throw new SnappyError(SnappyErrorCode.FAILED_TO_UNCOMPRESS, String.format("Received chunkSize %,d is greater than max configured chunk size %,d", chunkSize, maxChunkSize));
+        }
+
         // extend the compressed data buffer size
         if (compressed == null || chunkSize > compressed.length) {
             // chunkSize exceeds limit
diff --git a/src/main/java/org/xerial/snappy/SnappyOutputStream.java b/src/main/java/org/xerial/snappy/SnappyOutputStream.java
index 0bab154f..290cc91b 100755
--- a/src/main/java/org/xerial/snappy/SnappyOutputStream.java
+++ b/src/main/java/org/xerial/snappy/SnappyOutputStream.java
@@ -59,6 +59,7 @@
 public class SnappyOutputStream
         extends OutputStream
 {
+    public static final int MAX_BLOCK_SIZE = 512 * 1024 * 1024; // 512 MiB
     static final int MIN_BLOCK_SIZE = 1 * 1024;
     static final int DEFAULT_BLOCK_SIZE = 32 * 1024; // Use 32kb for the default block size
 
@@ -84,7 +85,7 @@ public SnappyOutputStream(OutputStream out)
     /**
      * @param out
      * @param blockSize byte size of the internal buffer size
-     * @throws IOException
+     * @throws IllegalArgumentException when blockSize is larger than 512 MiB
      */
     public SnappyOutputStream(OutputStream out, int blockSize)
     {
@@ -95,6 +96,9 @@ public SnappyOutputStream(OutputStream out, int blockSize, BufferAllocatorFactor
     {
         this.out = out;
         this.blockSize = Math.max(MIN_BLOCK_SIZE, blockSize);
+        if (this.blockSize > MAX_BLOCK_SIZE){
+            throw new IllegalArgumentException(String.format("Provided chunk size %,d larger than max %,d", this.blockSize, MAX_BLOCK_SIZE));
+        }
         int inputSize = blockSize;
         int outputSize = SnappyCodec.HEADER_SIZE + 4 + Snappy.maxCompressedLength(blockSize);
 
diff --git a/src/test/java/org/xerial/snappy/SnappyOutputStreamTest.java b/src/test/java/org/xerial/snappy/SnappyOutputStreamTest.java
index f2b7774a..d7831cbe 100755
--- a/src/test/java/org/xerial/snappy/SnappyOutputStreamTest.java
+++ b/src/test/java/org/xerial/snappy/SnappyOutputStreamTest.java
@@ -34,6 +34,7 @@
 import java.nio.ByteOrder;
 
 import org.junit.Test;
+import org.junit.Assert;
 import org.xerial.snappy.buffer.BufferAllocatorFactory;
 import org.xerial.snappy.buffer.CachedBufferAllocator;
 import org.xerial.snappy.buffer.DefaultBufferAllocator;
@@ -106,6 +107,17 @@ public void bufferSize()
         is.close();
     }
 
+    @Test(expected = IllegalArgumentException.class)
+    public void invalidBlockSize()
+            throws Exception
+    {
+        // We rely on catch below, if there is no error this test will pass
+        // This can be done better with Assertions.assertThrows
+        Boolean exceptionThrown = false;
+        ByteArrayOutputStream b = new ByteArrayOutputStream();
+        SnappyOutputStream os = new SnappyOutputStream(b, 1024 * 1024 * 1024);
+    }
+
     @Test
     public void smallWrites()
             throws Exception
diff --git a/src/test/java/org/xerial/snappy/SnappyTest.java b/src/test/java/org/xerial/snappy/SnappyTest.java
index 30edf66c..52c3a005 100755
--- a/src/test/java/org/xerial/snappy/SnappyTest.java
+++ b/src/test/java/org/xerial/snappy/SnappyTest.java
@@ -379,6 +379,24 @@ public void isInvalidChunkLengthForSnappyInputStreamOutOfMemory()
         }
     }
 
+    /*
+    Tests sad cases for SnappyInputStream.read method
+    - Expects a failed to compress exception due to upper bounds chunk size
+    - {-126, 'S', 'N', 'A', 'P', 'P', 'Y', 0, 0, 0, 0, 0, 0, 0, 0, 0,(byte) 0x7f, (byte) 0xff, (byte) 0xff, (byte) 0xff}
+     */
+    @Test
+    public void isInvalidChunkLengthForSnappyInputStream()
+            throws Exception {
+        byte[] data = {-126, 'S', 'N', 'A', 'P', 'P', 'Y', 0, 0, 0, 0, 0, 0, 0, 0, 0, (byte) 0x7f, (byte) 0xff, (byte) 0xff, (byte) 0xff};
+        SnappyInputStream in = new SnappyInputStream(new ByteArrayInputStream(data));
+        byte[] out = new byte[50];
+        try {
+            in.read(out);
+        } catch (SnappyError error) {
+            Assert.assertEquals(error.errorCode, SnappyErrorCode.FAILED_TO_UNCOMPRESS);
+        }
+    }
+
     /*
     Tests happy cases for BitShuffle.shuffle method
     - double: 0, 10