From 7e932eaad4ca0d0c21833cb4f65368fbe5f87d4c Mon Sep 17 00:00:00 2001
From: Andrew Steinborn <git@steinborn.me>
Date: Fri, 27 Oct 2023 15:35:04 -0400
Subject: [PATCH] Modify `MinecraftCompressorAndLengthEncoder` to handle 8MiB
 outgoing packets like vanilla supports now

---
 .../velocitypowered/proxy/protocol/ProtocolUtils.java | 11 ++++++-----
 .../netty/MinecraftCompressorAndLengthEncoder.java    | 10 +++++-----
 .../proxy/protocol/ProtocolUtilsTest.java             | 10 +++++-----
 3 files changed, 16 insertions(+), 15 deletions(-)

diff --git a/proxy/src/main/java/com/velocitypowered/proxy/protocol/ProtocolUtils.java b/proxy/src/main/java/com/velocitypowered/proxy/protocol/ProtocolUtils.java
index e65ee056b..23f9de87c 100644
--- a/proxy/src/main/java/com/velocitypowered/proxy/protocol/ProtocolUtils.java
+++ b/proxy/src/main/java/com/velocitypowered/proxy/protocol/ProtocolUtils.java
@@ -163,16 +163,17 @@ public enum ProtocolUtils {
   }
 
   /**
-   * Writes the specified {@code value} as a 21-bit Minecraft VarInt to the specified {@code buf}.
-   * The upper 11 bits will be discarded.
+   * Writes the specified {@code value} as a 28-bit Minecraft VarInt to the specified {@code buf}.
+   * The upper 4 bits will be discarded.
    *
    * @param buf   the buffer to read from
    * @param value the integer to write
    */
-  public static void write21BitVarInt(ByteBuf buf, int value) {
+  public static void write28BitVarInt(ByteBuf buf, int value) {
     // See https://steinborn.me/posts/performance/how-fast-can-you-write-a-varint/
-    int w = (value & 0x7F | 0x80) << 16 | ((value >>> 7) & 0x7F | 0x80) << 8 | (value >>> 14);
-    buf.writeMedium(w);
+    int w = (value & 0x7F | 0x80) << 24 | (((value >>> 7) & 0x7F | 0x80) << 16)
+        | ((value >>> 14) & 0x7F | 0x80) << 8 | (value >>> 21);
+    buf.writeInt(w);
   }
 
   public static String readString(ByteBuf buf) {
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 90952a729..d020d93c6 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
@@ -57,7 +57,7 @@ public class MinecraftCompressorAndLengthEncoder extends MessageToByteEncoder<By
       throws DataFormatException {
     int uncompressed = msg.readableBytes();
 
-    ProtocolUtils.write21BitVarInt(out, 0); // Dummy packet length
+    ProtocolUtils.write28BitVarInt(out, 0); // Dummy packet length
     ProtocolUtils.writeVarInt(out, uncompressed);
     ByteBuf compatibleIn = MoreByteBufUtils.ensureCompatible(ctx.alloc(), compressor, msg);
 
@@ -68,14 +68,14 @@ public class MinecraftCompressorAndLengthEncoder extends MessageToByteEncoder<By
       compatibleIn.release();
     }
     int compressedLength = out.writerIndex() - startCompressed;
-    if (compressedLength >= 1 << 21) {
-      throw new DataFormatException("The server sent a very large (over 2MiB compressed) packet.");
+    if (compressedLength >= 1 << 23) {
+      throw new DataFormatException("The server sent a very large (over 8MiB compressed) packet.");
     }
 
     int writerIndex = out.writerIndex();
     int packetLength = out.readableBytes() - 3;
     out.writerIndex(0);
-    ProtocolUtils.write21BitVarInt(out, packetLength); // Rewrite packet length
+    ProtocolUtils.write28BitVarInt(out, packetLength); // Rewrite packet length
     out.writerIndex(writerIndex);
   }
 
@@ -92,7 +92,7 @@ public class MinecraftCompressorAndLengthEncoder extends MessageToByteEncoder<By
     }
 
     // (maximum data length after compression) + packet length varint + uncompressed data varint
-    int initialBufferSize = (uncompressed - 1) + 3 + ProtocolUtils.varIntBytes(uncompressed);
+    int initialBufferSize = (uncompressed - 1) + 4 + ProtocolUtils.varIntBytes(uncompressed);
     return MoreByteBufUtils.preferredBuffer(ctx.alloc(), compressor, initialBufferSize);
   }
 
diff --git a/proxy/src/test/java/com/velocitypowered/proxy/protocol/ProtocolUtilsTest.java b/proxy/src/test/java/com/velocitypowered/proxy/protocol/ProtocolUtilsTest.java
index f476de67c..62e8274ca 100644
--- a/proxy/src/test/java/com/velocitypowered/proxy/protocol/ProtocolUtilsTest.java
+++ b/proxy/src/test/java/com/velocitypowered/proxy/protocol/ProtocolUtilsTest.java
@@ -74,16 +74,16 @@ public class ProtocolUtilsTest {
   }
 
   @Test
-  void test3Bytes() {
+  void test4Bytes() {
     ByteBuf buf = Unpooled.buffer(5);
-    for (int i = 0; i < 2097152; i += 31) {
-      writeReadTest3Bytes(buf, i);
+    for (int i = 0; i < (1 << 28) - 1; i += 31) {
+      writeReadTest4Bytes(buf, i);
     }
   }
 
-  private void writeReadTest3Bytes(ByteBuf buf, int test) {
+  private void writeReadTest4Bytes(ByteBuf buf, int test) {
     buf.clear();
-    ProtocolUtils.write21BitVarInt(buf, test);
+    ProtocolUtils.write28BitVarInt(buf, test);
     assertEquals(test, ProtocolUtils.readVarInt(buf));
   }