From 6369a95ec99f4a1146ff3c38f1321770331dd4f5 Mon Sep 17 00:00:00 2001
From: Andrew Steinborn <git@steinborn.me>
Date: Sat, 8 May 2021 18:40:23 -0400
Subject: [PATCH] Readd safe and slow compression handling and hide it behind a
 system property

---
 .../MinecraftCompressorAndLengthEncoder.java  | 44 +++++++++++++++++--
 1 file changed, 40 insertions(+), 4 deletions(-)

diff --git a/proxy/src/main/java/com/velocitypowered/proxy/protocol/netty/MinecraftCompressorAndLengthEncoder.java b/proxy/src/main/java/com/velocitypowered/proxy/protocol/netty/MinecraftCompressorAndLengthEncoder.java
index 73cea902b..34d247bc7 100644
--- a/proxy/src/main/java/com/velocitypowered/proxy/protocol/netty/MinecraftCompressorAndLengthEncoder.java
+++ b/proxy/src/main/java/com/velocitypowered/proxy/protocol/netty/MinecraftCompressorAndLengthEncoder.java
@@ -29,6 +29,9 @@ import java.util.zip.DataFormatException;
 
 public class MinecraftCompressorAndLengthEncoder extends MessageToByteEncoder<ByteBuf> {
 
+  private static final boolean MUST_USE_SAFE_AND_SLOW_COMPRESSION_HANDLING =
+      Boolean.getBoolean("velocity.increased-compression-cap");
+
   private int threshold;
   private final VelocityCompressor compressor;
 
@@ -46,29 +49,62 @@ public class MinecraftCompressorAndLengthEncoder extends MessageToByteEncoder<By
       ProtocolUtils.writeVarInt(out, 0);
       out.writeBytes(msg);
     } else {
-      handleCompressed(ctx, msg, out);
+      if (MUST_USE_SAFE_AND_SLOW_COMPRESSION_HANDLING) {
+        handleCompressedSafe(ctx, msg, out);
+      } else {
+        handleCompressedFast(ctx, msg, out);
+      }
     }
   }
 
-  private void handleCompressed(ChannelHandlerContext ctx, ByteBuf msg, ByteBuf out)
+  private void handleCompressedFast(ChannelHandlerContext ctx, ByteBuf msg, ByteBuf out)
       throws DataFormatException {
-    ProtocolUtils.writeVarIntAs3Bytes(out, 0); //Dummy packet length
     int uncompressed = msg.readableBytes();
+
+    ProtocolUtils.writeVarIntAs3Bytes(out, 0); // Dummy packet length
     ProtocolUtils.writeVarInt(out, uncompressed);
     ByteBuf compatibleIn = MoreByteBufUtils.ensureCompatible(ctx.alloc(), compressor, msg);
+
+    int startCompressed = out.writerIndex();
     try {
       compressor.deflate(compatibleIn, out);
     } finally {
       compatibleIn.release();
     }
+    int compressedLength = out.writerIndex() - startCompressed;
+    if (compressedLength >= 1 << 21) {
+      throw new DataFormatException("The server sent a very large (over 2MiB compressed) packet. "
+          + "Please restart Velocity with the JVM flag -Dvelocity.increased-compression-cap=true "
+          + "to fix this issue.");
+    }
 
     int writerIndex = out.writerIndex();
     int packetLength = out.readableBytes() - 3;
     out.writerIndex(0);
-    ProtocolUtils.writeVarIntAs3Bytes(out, packetLength); //Rewrite packet length
+    ProtocolUtils.writeVarIntAs3Bytes(out, packetLength); // Rewrite packet length
     out.writerIndex(writerIndex);
   }
 
+  private void handleCompressedSafe(ChannelHandlerContext ctx, ByteBuf msg, ByteBuf out)
+      throws DataFormatException {
+    int uncompressed = msg.readableBytes();
+    ByteBuf tmpBuf = MoreByteBufUtils.preferredBuffer(ctx.alloc(), compressor, uncompressed - 1);
+    try {
+      ProtocolUtils.writeVarInt(tmpBuf, uncompressed);
+      ByteBuf compatibleIn = MoreByteBufUtils.ensureCompatible(ctx.alloc(), compressor, msg);
+      try {
+        compressor.deflate(compatibleIn, tmpBuf);
+      } finally {
+        compatibleIn.release();
+      }
+
+      ProtocolUtils.writeVarInt(out, tmpBuf.readableBytes());
+      out.writeBytes(tmpBuf);
+    } finally {
+      tmpBuf.release();
+    }
+  }
+
   @Override
   protected ByteBuf allocateBuffer(ChannelHandlerContext ctx, ByteBuf msg, boolean preferDirect)
       throws Exception {