From 60a3bf082f9fa0a09a2f4181e3abe9af0d341b5a Mon Sep 17 00:00:00 2001
From: md_5 <git@md-5.net>
Date: Mon, 27 Jan 2025 20:08:58 +1100
Subject: [PATCH] Preallocate compression output buffer to remove unnecessary
 resizing

---
 .../main/java/net/md_5/bungee/jni/zlib/NativeZlib.java    | 8 +++++++-
 native/src/test/java/net/md_5/bungee/NativeZlibTest.java  | 2 +-
 .../java/net/md_5/bungee/compress/PacketDecompressor.java | 8 +++++++-
 3 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/native/src/main/java/net/md_5/bungee/jni/zlib/NativeZlib.java b/native/src/main/java/net/md_5/bungee/jni/zlib/NativeZlib.java
index 4dc558d53..6fa9c6fd3 100644
--- a/native/src/main/java/net/md_5/bungee/jni/zlib/NativeZlib.java
+++ b/native/src/main/java/net/md_5/bungee/jni/zlib/NativeZlib.java
@@ -52,7 +52,13 @@ public class NativeZlib implements BungeeZlib
 
         while ( !nativeCompress.finished && ( compress || in.isReadable() ) )
         {
-            out.ensureWritable( OUTPUT_BUFFER_SIZE );
+            if ( compress )
+            {
+                out.ensureWritable( OUTPUT_BUFFER_SIZE );
+            } else
+            {
+                Preconditions.checkArgument( out.isWritable(), "Output buffer overrun" );
+            }
 
             int processed;
             try
diff --git a/native/src/test/java/net/md_5/bungee/NativeZlibTest.java b/native/src/test/java/net/md_5/bungee/NativeZlibTest.java
index 112a3b83b..b40484a26 100644
--- a/native/src/test/java/net/md_5/bungee/NativeZlibTest.java
+++ b/native/src/test/java/net/md_5/bungee/NativeZlibTest.java
@@ -64,7 +64,7 @@ public class NativeZlibTest
 
         zlib.process( originalBuf, compressed );
 
-        ByteBuf uncompressed = Unpooled.directBuffer();
+        ByteBuf uncompressed = Unpooled.directBuffer( dataBuf.length, dataBuf.length );
 
         zlib.init( false, 0 );
         zlib.process( compressed, uncompressed );
diff --git a/proxy/src/main/java/net/md_5/bungee/compress/PacketDecompressor.java b/proxy/src/main/java/net/md_5/bungee/compress/PacketDecompressor.java
index 445ee947a..8e89d4b65 100644
--- a/proxy/src/main/java/net/md_5/bungee/compress/PacketDecompressor.java
+++ b/proxy/src/main/java/net/md_5/bungee/compress/PacketDecompressor.java
@@ -7,10 +7,12 @@ import io.netty.handler.codec.MessageToMessageDecoder;
 import java.util.List;
 import net.md_5.bungee.jni.zlib.BungeeZlib;
 import net.md_5.bungee.protocol.DefinedPacket;
+import net.md_5.bungee.protocol.OverflowPacketException;
 
 public class PacketDecompressor extends MessageToMessageDecoder<ByteBuf>
 {
 
+    private static final int MAX_DECOMPRESSED_LEN = 1 << 23;
     private final BungeeZlib zlib = CompressFactory.zlib.newInstance();
 
     @Override
@@ -34,8 +36,12 @@ public class PacketDecompressor extends MessageToMessageDecoder<ByteBuf>
             out.add( in.retain() );
         } else
         {
-            ByteBuf decompressed = ctx.alloc().directBuffer();
+            if ( size > MAX_DECOMPRESSED_LEN )
+            {
+                throw new OverflowPacketException( "Packet may not be larger than " + MAX_DECOMPRESSED_LEN + " bytes" );
+            }
 
+            ByteBuf decompressed = ctx.alloc().directBuffer( size, size );
             try
             {
                 zlib.process( in, decompressed );