From ad6830f1e90735407ae153a0ee9bf8e223e40d14 Mon Sep 17 00:00:00 2001
From: Rich DiCroce <Rich.DiCroce@scientificgames.com>
Date: Fri, 31 Jan 2020 06:11:06 -0500
Subject: [PATCH] Allow a limit to be set on the decompressed buffer size
for ZlibDecoders (#9924)

Motivation:
It is impossible to know in advance how much memory will be needed to
decompress a stream of bytes that was compressed using the DEFLATE
algorithm. In theory, up to 1032 times the compressed size could be
needed. For untrusted input, an attacker could exploit this to exhaust
the memory pool.

Modifications:
ZlibDecoder and its subclasses now support an optional limit on the size
of the decompressed buffer. By default, if the limit is reached,
decompression stops and a DecompressionException is thrown. Behavior
upon reaching the limit is modifiable by subclasses in case they desire
something else.

Result:
The decompressed buffer can now be limited to a configurable size, thus
mitigating the possibility of memory pool exhaustion.

---
 .../codec/compression/JZlibDecoder.java       | 57 +++++++++++++++-
 .../codec/compression/JdkZlibDecoder.java     | 58 +++++++++++++++--
 .../codec/compression/ZlibDecoder.java        | 65 +++++++++++++++++++
 .../handler/codec/compression/JZlibTest.java  |  4 +-
 .../codec/compression/JdkZlibTest.java        |  4 +-
 .../codec/compression/ZlibCrossTest1.java     |  4 +-
 .../codec/compression/ZlibCrossTest2.java     |  4 +-
 .../handler/codec/compression/ZlibTest.java   | 57 +++++++++++++++-
 8 files changed, 235 insertions(+), 18 deletions(-)

diff --git a/codec/src/main/java/io/netty/handler/codec/compression/JZlibDecoder.java b/codec/src/main/java/io/netty/handler/codec/compression/JZlibDecoder.java
index 5d23bb8..cfb104e 100644
--- a/codec/src/main/java/io/netty/handler/codec/compression/JZlibDecoder.java
+++ b/codec/src/main/java/io/netty/handler/codec/compression/JZlibDecoder.java
@@ -18,6 +18,7 @@ package io.netty.handler.codec.compression;
 import com.jcraft.jzlib.Inflater;
 import com.jcraft.jzlib.JZlib;
 import io.netty.buffer.ByteBuf;
+import io.netty.buffer.ByteBufAllocator;
 import io.netty.channel.ChannelHandlerContext;
 
 import java.util.List;
@@ -34,7 +35,21 @@ public class JZlibDecoder extends ZlibDecoder {
      * @throws DecompressionException if failed to initialize zlib
      */
     public JZlibDecoder() {
-        this(ZlibWrapper.ZLIB);
+        this(ZlibWrapper.ZLIB, 0);
+    }
+
+    /**
+     * Creates a new instance with the default wrapper ({@link ZlibWrapper#ZLIB})
+     * and specified maximum buffer allocation.
+     *
+     * @param maxAllocation
+     *          Maximum size of the decompression buffer. Must be &gt;= 0.
+     *          If zero, maximum size is decided by the {@link ByteBufAllocator}.
+     *
+     * @throws DecompressionException if failed to initialize zlib
+     */
+    public JZlibDecoder(int maxAllocation) {
+        this(ZlibWrapper.ZLIB, maxAllocation);
     }
 
     /**
@@ -43,6 +58,21 @@ public class JZlibDecoder extends ZlibDecoder {
      * @throws DecompressionException if failed to initialize zlib
      */
     public JZlibDecoder(ZlibWrapper wrapper) {
+        this(wrapper, 0);
+    }
+
+    /**
+     * Creates a new instance with the specified wrapper and maximum buffer allocation.
+     *
+     * @param maxAllocation
+     *          Maximum size of the decompression buffer. Must be &gt;= 0.
+     *          If zero, maximum size is decided by the {@link ByteBufAllocator}.
+     *
+     * @throws DecompressionException if failed to initialize zlib
+     */
+    public JZlibDecoder(ZlibWrapper wrapper, int maxAllocation) {
+        super(maxAllocation);
+
         if (wrapper == null) {
             throw new NullPointerException("wrapper");
         }
@@ -61,6 +91,22 @@ public class JZlibDecoder extends ZlibDecoder {
      * @throws DecompressionException if failed to initialize zlib
      */
     public JZlibDecoder(byte[] dictionary) {
+        this(dictionary, 0);
+    }
+
+    /**
+     * Creates a new instance with the specified preset dictionary and maximum buffer allocation.
+     * The wrapper is always {@link ZlibWrapper#ZLIB} because it is the only format that
+     * supports the preset dictionary.
+     *
+     * @param maxAllocation
+     *          Maximum size of the decompression buffer. Must be &gt;= 0.
+     *          If zero, maximum size is decided by the {@link ByteBufAllocator}.
+     *
+     * @throws DecompressionException if failed to initialize zlib
+     */
+    public JZlibDecoder(byte[] dictionary, int maxAllocation) {
+        super(maxAllocation);
         if (dictionary == null) {
             throw new NullPointerException("dictionary");
         }
@@ -110,11 +156,11 @@ public class JZlibDecoder extends ZlibDecoder {
             final int oldNextInIndex = z.next_in_index;
 
             // Configure output.
-            ByteBuf decompressed = ctx.alloc().heapBuffer(inputLength << 1);
+	    ByteBuf decompressed = prepareDecompressBuffer(ctx, null, inputLength << 1);
 
             try {
                 loop: for (;;) {
-                    decompressed.ensureWritable(z.avail_in << 1);
+		     decompressed = prepareDecompressBuffer(ctx, decompressed, z.avail_in << 1);
                     z.avail_out = decompressed.writableBytes();
                     z.next_out = decompressed.array();
                     z.next_out_index = decompressed.arrayOffset() + decompressed.writerIndex();
@@ -170,4 +216,9 @@ public class JZlibDecoder extends ZlibDecoder {
             z.next_out = null;
         }
     }
+
+    @Override
+    protected void decompressionBufferExhausted(ByteBuf buffer) {
+        finished = true;
+    }
 }
diff --git a/codec/src/main/java/io/netty/handler/codec/compression/JdkZlibDecoder.java b/codec/src/main/java/io/netty/handler/codec/compression/JdkZlibDecoder.java
index d9b4b91..74e6d5f 100644
--- a/codec/src/main/java/io/netty/handler/codec/compression/JdkZlibDecoder.java
+++ b/codec/src/main/java/io/netty/handler/codec/compression/JdkZlibDecoder.java
@@ -16,6 +16,7 @@
 package io.netty.handler.codec.compression;
 
 import io.netty.buffer.ByteBuf;
+import io.netty.buffer.ByteBufAllocator;
 import io.netty.channel.ChannelHandlerContext;
 
 import java.util.List;
@@ -63,7 +64,19 @@ public class JdkZlibDecoder extends ZlibDecoder {
      * Creates a new instance with the default wrapper ({@link ZlibWrapper#ZLIB}).
      */
     public JdkZlibDecoder() {
-        this(ZlibWrapper.ZLIB, null);
+        this(ZlibWrapper.ZLIB, null, 0);
+    }
+
+    /**
+     * Creates a new instance with the default wrapper ({@link ZlibWrapper#ZLIB})
+     * and the specified maximum buffer allocation.
+     *
+     * @param maxAllocation
+     *          Maximum size of the decompression buffer. Must be &gt;= 0.
+     *          If zero, maximum size is decided by the {@link ByteBufAllocator}.
+     */
+    public JdkZlibDecoder(int maxAllocation) {
+        this(ZlibWrapper.ZLIB, null, maxAllocation);
     }
 
     /**
@@ -72,7 +85,20 @@ public class JdkZlibDecoder extends ZlibDecoder {
      * supports the preset dictionary.
      */
     public JdkZlibDecoder(byte[] dictionary) {
-        this(ZlibWrapper.ZLIB, dictionary);
+        this(ZlibWrapper.ZLIB, dictionary, 0);
+    }
+
+    /**
+     * Creates a new instance with the specified preset dictionary and maximum buffer allocation.
+     * The wrapper is always {@link ZlibWrapper#ZLIB} because it is the only format that
+     * supports the preset dictionary.
+     *
+     * @param maxAllocation
+     *          Maximum size of the decompression buffer. Must be &gt;= 0.
+     *          If zero, maximum size is decided by the {@link ByteBufAllocator}.
+     */
+    public JdkZlibDecoder(byte[] dictionary, int maxAllocation) {
+        this(ZlibWrapper.ZLIB, dictionary, maxAllocation);
     }
 
     /**
@@ -81,10 +107,25 @@ public class JdkZlibDecoder extends ZlibDecoder {
      * supported atm.
      */
     public JdkZlibDecoder(ZlibWrapper wrapper) {
-        this(wrapper, null);
+        this(wrapper, null, 0);
     }
 
-    private JdkZlibDecoder(ZlibWrapper wrapper, byte[] dictionary) {
+    /**
+     * Creates a new instance with the specified wrapper and maximum buffer allocation.
+     * Be aware that only {@link ZlibWrapper#GZIP}, {@link ZlibWrapper#ZLIB} and {@link ZlibWrapper#NONE} are
+     * supported atm.
+     *
+     * @param maxAllocation
+     *          Maximum size of the decompression buffer. Must be &gt;= 0.
+     *          If zero, maximum size is decided by the {@link ByteBufAllocator}.
+     */
+    public JdkZlibDecoder(ZlibWrapper wrapper, int maxAllocation) {
+        this(wrapper, null, maxAllocation);
+    }
+
+    private JdkZlibDecoder(ZlibWrapper wrapper, byte[] dictionary, int maxAllocation) {
+        super(maxAllocation);
+
         if (wrapper == null) {
             throw new NullPointerException("wrapper");
         }
@@ -167,7 +208,7 @@ public class JdkZlibDecoder extends ZlibDecoder {
             inflater.setInput(array);
         }
 
-        ByteBuf decompressed = ctx.alloc().heapBuffer(inflater.getRemaining() << 1);
+	ByteBuf decompressed = prepareDecompressBuffer(ctx, null,inflater.getRemaining() << 1);
         try {
             boolean readFooter = false;
             while (!inflater.needsInput()) {
@@ -198,7 +239,7 @@ public class JdkZlibDecoder extends ZlibDecoder {
                     }
                     break;
                 } else {
-                    decompressed.ensureWritable(inflater.getRemaining() << 1);
+		    decompressed = prepareDecompressBuffer(ctx, decompressed,inflater.getRemaining() << 1);
                 }
             }
 
@@ -222,6 +263,11 @@ public class JdkZlibDecoder extends ZlibDecoder {
         }
     }
 
+    @Override
+    protected void decompressionBufferExhausted(ByteBuf buffer) {
+        finished = true;
+    }
+
     @Override
     protected void handlerRemoved0(ChannelHandlerContext ctx) throws Exception {
         super.handlerRemoved0(ctx);
diff --git a/codec/src/main/java/io/netty/handler/codec/compression/ZlibDecoder.java b/codec/src/main/java/io/netty/handler/codec/compression/ZlibDecoder.java
index d01bc6b..26fd3e7 100644
--- a/codec/src/main/java/io/netty/handler/codec/compression/ZlibDecoder.java
+++ b/codec/src/main/java/io/netty/handler/codec/compression/ZlibDecoder.java
@@ -16,6 +16,8 @@
 package io.netty.handler.codec.compression;
 
 import io.netty.buffer.ByteBuf;
+import io.netty.buffer.ByteBufAllocator;
+import io.netty.channel.ChannelHandlerContext;
 import io.netty.handler.codec.ByteToMessageDecoder;
 
 /**
@@ -23,9 +25,72 @@ import io.netty.handler.codec.ByteToMessageDecoder;
  */
 public abstract class ZlibDecoder extends ByteToMessageDecoder {
 
+    /**
+     * Maximum allowed size of the decompression buffer.
+     */
+    protected final int maxAllocation;
+
+    /**
+     * Same as {@link #ZlibDecoder(int)} with maxAllocation = 0.
+     */
+    public ZlibDecoder() {
+        this(0);
+    }
+
+    /**
+     * Construct a new ZlibDecoder.
+     * @param maxAllocation
+     *          Maximum size of the decompression buffer. Must be &gt;= 0.
+     *          If zero, maximum size is decided by the {@link ByteBufAllocator}.
+     */
+    public ZlibDecoder(int maxAllocation) {
+        if (maxAllocation < 0) {
+            throw new IllegalArgumentException("maxAllocation must be >= 0");
+        }
+        this.maxAllocation = maxAllocation;
+    }
+
     /**
      * Returns {@code true} if and only if the end of the compressed stream
      * has been reached.
      */
     public abstract boolean isClosed();
+
+    /**
+     * Allocate or expand the decompression buffer, without exceeding the maximum allocation.
+     * Calls {@link #decompressionBufferExhausted(ByteBuf)} if the buffer is full and cannot be expanded further.
+     */
+    protected ByteBuf prepareDecompressBuffer(ChannelHandlerContext ctx, ByteBuf buffer, int preferredSize) {
+        if (buffer == null) {
+            if (maxAllocation == 0) {
+                return ctx.alloc().heapBuffer(preferredSize);
+            }
+
+            return ctx.alloc().heapBuffer(Math.min(preferredSize, maxAllocation), maxAllocation);
+        }
+
+        // this always expands the buffer if possible, even if the expansion is less than preferredSize
+        // we throw the exception only if the buffer could not be expanded at all
+        // this means that one final attempt to deserialize will always be made with the buffer at maxAllocation
+        if (buffer.ensureWritable(preferredSize, true) == 1) {
+            // buffer must be consumed so subclasses don't add it to output
+            // we therefore duplicate it when calling decompressionBufferExhausted() to guarantee non-interference
+            // but wait until after to consume it so the subclass can tell how much output is really in the buffer
+            decompressionBufferExhausted(buffer.duplicate());
+            buffer.skipBytes(buffer.readableBytes());
+            throw new DecompressionException("Decompression buffer has reached maximum size: " + buffer.maxCapacity());
+        }
+
+        return buffer;
+    }
+
+    /**
+     * Called when the decompression buffer cannot be expanded further.
+     * Default implementation is a no-op, but subclasses can override in case they want to
+     * do something before the {@link DecompressionException} is thrown, such as log the
+     * data that was decompressed so far.
+     */
+    protected void decompressionBufferExhausted(ByteBuf buffer) {
+    }
+
 }
diff --git a/codec/src/test/java/io/netty/handler/codec/compression/JZlibTest.java b/codec/src/test/java/io/netty/handler/codec/compression/JZlibTest.java
index 28f3919..015559e 100644
--- a/codec/src/test/java/io/netty/handler/codec/compression/JZlibTest.java
+++ b/codec/src/test/java/io/netty/handler/codec/compression/JZlibTest.java
@@ -23,7 +23,7 @@ public class JZlibTest extends ZlibTest {
     }
 
     @Override
-    protected ZlibDecoder createDecoder(ZlibWrapper wrapper) {
-        return new JZlibDecoder(wrapper);
+    protected ZlibDecoder createDecoder(ZlibWrapper wrapper, int maxAllocation) {
+        return new JZlibDecoder(wrapper, maxAllocation);
     }
 }
diff --git a/codec/src/test/java/io/netty/handler/codec/compression/JdkZlibTest.java b/codec/src/test/java/io/netty/handler/codec/compression/JdkZlibTest.java
index 23f178d..fc53282 100644
--- a/codec/src/test/java/io/netty/handler/codec/compression/JdkZlibTest.java
+++ b/codec/src/test/java/io/netty/handler/codec/compression/JdkZlibTest.java
@@ -26,8 +26,8 @@ public class JdkZlibTest extends ZlibTest {
     }
 
     @Override
-    protected ZlibDecoder createDecoder(ZlibWrapper wrapper) {
-        return new JdkZlibDecoder(wrapper);
+    protected ZlibDecoder createDecoder(ZlibWrapper wrapper, int maxAllocation) {
+        return new JdkZlibDecoder(wrapper, maxAllocation);
     }
 
     @Test(expected = DecompressionException.class)
diff --git a/codec/src/test/java/io/netty/handler/codec/compression/ZlibCrossTest1.java b/codec/src/test/java/io/netty/handler/codec/compression/ZlibCrossTest1.java
index 9e16e1a..3c31274 100644
--- a/codec/src/test/java/io/netty/handler/codec/compression/ZlibCrossTest1.java
+++ b/codec/src/test/java/io/netty/handler/codec/compression/ZlibCrossTest1.java
@@ -23,7 +23,7 @@ public class ZlibCrossTest1 extends ZlibTest {
     }
 
     @Override
-    protected ZlibDecoder createDecoder(ZlibWrapper wrapper) {
-        return new JZlibDecoder(wrapper);
+    protected ZlibDecoder createDecoder(ZlibWrapper wrapper, int maxAllocation) {
+        return new JZlibDecoder(wrapper, maxAllocation);
     }
 }
diff --git a/codec/src/test/java/io/netty/handler/codec/compression/ZlibCrossTest2.java b/codec/src/test/java/io/netty/handler/codec/compression/ZlibCrossTest2.java
index 8717019..00c6e18 100644
--- a/codec/src/test/java/io/netty/handler/codec/compression/ZlibCrossTest2.java
+++ b/codec/src/test/java/io/netty/handler/codec/compression/ZlibCrossTest2.java
@@ -25,8 +25,8 @@ public class ZlibCrossTest2 extends ZlibTest {
     }
 
     @Override
-    protected ZlibDecoder createDecoder(ZlibWrapper wrapper) {
-        return new JdkZlibDecoder(wrapper);
+    protected ZlibDecoder createDecoder(ZlibWrapper wrapper, int maxAllocation) {
+        return new JdkZlibDecoder(wrapper, maxAllocation);
     }
 
     @Test(expected = DecompressionException.class)
diff --git a/codec/src/test/java/io/netty/handler/codec/compression/ZlibTest.java b/codec/src/test/java/io/netty/handler/codec/compression/ZlibTest.java
index 7c25ec4..9d79c81 100644
--- a/codec/src/test/java/io/netty/handler/codec/compression/ZlibTest.java
+++ b/codec/src/test/java/io/netty/handler/codec/compression/ZlibTest.java
@@ -15,7 +15,9 @@
  */
 package io.netty.handler.codec.compression;
 
+import io.netty.buffer.AbstractByteBufAllocator;
 import io.netty.buffer.ByteBuf;
+import io.netty.buffer.ByteBufAllocator;
 import io.netty.buffer.ByteBufInputStream;
 import io.netty.buffer.Unpooled;
 import io.netty.channel.embedded.EmbeddedChannel;
@@ -88,8 +90,12 @@ public abstract class ZlibTest {
         rand.nextBytes(BYTES_LARGE);
     }
 
+    protected ZlibDecoder createDecoder(ZlibWrapper wrapper) {
+        return createDecoder(wrapper, 0);
+    }
+
     protected abstract ZlibEncoder createEncoder(ZlibWrapper wrapper);
-    protected abstract ZlibDecoder createDecoder(ZlibWrapper wrapper);
+    protected abstract ZlibDecoder createDecoder(ZlibWrapper wrapper, int maxAllocation);
 
     @Test
     public void testGZIP2() throws Exception {
@@ -345,6 +351,25 @@ public abstract class ZlibTest {
         testCompressLarge(ZlibWrapper.GZIP, ZlibWrapper.ZLIB_OR_NONE);
     }
 
+    @Test
+    public void testMaxAllocation() throws Exception {
+        int maxAllocation = 1024;
+        ZlibDecoder decoder = createDecoder(ZlibWrapper.ZLIB, maxAllocation);
+        EmbeddedChannel chDecoder = new EmbeddedChannel(decoder);
+        TestByteBufAllocator alloc = new TestByteBufAllocator(chDecoder.alloc());
+        chDecoder.config().setAllocator(alloc);
+
+        try {
+            chDecoder.writeInbound(Unpooled.wrappedBuffer(deflate(BYTES_LARGE)));
+            fail("decompressed size > maxAllocation, so should have thrown exception");
+        } catch (DecompressionException e) {
+            assertTrue(e.getMessage().startsWith("Decompression buffer has reached maximum size"));
+            assertEquals(maxAllocation, alloc.getMaxAllocation());
+            assertTrue(decoder.isClosed());
+            assertFalse(chDecoder.finish());
+        }
+    }
+
     private static byte[] gzip(byte[] bytes) throws IOException {
         ByteArrayOutputStream out = new ByteArrayOutputStream();
         GZIPOutputStream stream = new GZIPOutputStream(out);
@@ -360,4 +385,34 @@ public abstract class ZlibTest {
         stream.close();
         return out.toByteArray();
     }
+
+    private static final class TestByteBufAllocator extends AbstractByteBufAllocator {
+        private ByteBufAllocator wrapped;
+        private int maxAllocation;
+
+        TestByteBufAllocator(ByteBufAllocator wrapped) {
+            this.wrapped = wrapped;
+        }
+
+        public int getMaxAllocation() {
+            return maxAllocation;
+        }
+
+        @Override
+        public boolean isDirectBufferPooled() {
+            return wrapped.isDirectBufferPooled();
+        }
+
+        @Override
+        protected ByteBuf newHeapBuffer(int initialCapacity, int maxCapacity) {
+            maxAllocation = Math.max(maxAllocation, maxCapacity);
+            return wrapped.heapBuffer(initialCapacity, maxCapacity);
+        }
+
+        @Override
+        protected ByteBuf newDirectBuffer(int initialCapacity, int maxCapacity) {
+            maxAllocation = Math.max(maxAllocation, maxCapacity);
+            return wrapped.directBuffer(initialCapacity, maxCapacity);
+        }
+    }
 }
-- 
2.23.0

