From a8e0516d18b996b3a6aea80de97e017daa6e38a1 Mon Sep 17 00:00:00 2001
From: Leymooo <vjatseslav.maspanov@gmail.com>
Date: Fri, 7 May 2021 19:36:30 +0300
Subject: [PATCH] Also do not copy memory in case when packet needs to compress

---
 .../proxy/protocol/ProtocolUtils.java         | 10 +++++
 .../MinecraftCompressorAndLengthEncoder.java  | 45 ++++++++++---------
 .../netty/MinecraftVarintLengthEncoder.java   |  2 +-
 .../proxy/protocol/ProtocolUtilsTest.java     | 14 ++++++
 4 files changed, 49 insertions(+), 22 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 c7e25f528..a40598e43 100644
--- a/proxy/src/main/java/com/velocitypowered/proxy/protocol/ProtocolUtils.java
+++ b/proxy/src/main/java/com/velocitypowered/proxy/protocol/ProtocolUtils.java
@@ -147,6 +147,16 @@ public enum ProtocolUtils {
     }
   }
 
+  /**
+   * Writes all integers as 3 bytes Minecraft-style VarInt to the specified {@code buf}.
+   * @param buf the buffer to read from
+   * @param value the integer to write
+   */
+  public static void writeVarIntAs3Bytes(ByteBuf buf, int value) {
+    int w = (value & 0x7F | 0x80) << 16 | ((value >>> 7) & 0x7F | 0x80) << 8 | (value >>> 14);
+    buf.writeMedium(w);
+  }
+
   public static String readString(ByteBuf buf) {
     return readString(buf, DEFAULT_MAX_STRING_SIZE);
   }
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 6684ee35d..73cea902b 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
@@ -17,6 +17,8 @@
 
 package com.velocitypowered.proxy.protocol.netty;
 
+import static com.velocitypowered.proxy.protocol.netty.MinecraftVarintLengthEncoder.IS_JAVA_CIPHER;
+
 import com.velocitypowered.natives.compression.VelocityCompressor;
 import com.velocitypowered.natives.util.MoreByteBufUtils;
 import com.velocitypowered.proxy.protocol.ProtocolUtils;
@@ -50,36 +52,37 @@ public class MinecraftCompressorAndLengthEncoder extends MessageToByteEncoder<By
 
   private void handleCompressed(ChannelHandlerContext ctx, ByteBuf msg, ByteBuf out)
       throws DataFormatException {
+    ProtocolUtils.writeVarIntAs3Bytes(out, 0); //Dummy packet length
     int uncompressed = msg.readableBytes();
-    ByteBuf tmpBuf = MoreByteBufUtils.preferredBuffer(ctx.alloc(), compressor, uncompressed - 1);
+    ProtocolUtils.writeVarInt(out, uncompressed);
+    ByteBuf compatibleIn = MoreByteBufUtils.ensureCompatible(ctx.alloc(), compressor, msg);
     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);
+      compressor.deflate(compatibleIn, out);
     } finally {
-      tmpBuf.release();
+      compatibleIn.release();
     }
+
+    int writerIndex = out.writerIndex();
+    int packetLength = out.readableBytes() - 3;
+    out.writerIndex(0);
+    ProtocolUtils.writeVarIntAs3Bytes(out, packetLength); //Rewrite packet length
+    out.writerIndex(writerIndex);
   }
 
   @Override
   protected ByteBuf allocateBuffer(ChannelHandlerContext ctx, ByteBuf msg, boolean preferDirect)
       throws Exception {
-    // We allocate bytes to be compressed plus 1 byte. This covers two cases:
-    //
-    // - Compression
-    //    According to https://github.com/ebiggers/libdeflate/blob/master/libdeflate.h#L103,
-    //    if the data compresses well (and we do not have some pathological case) then the maximum
-    //    size the compressed size will ever be is the input size minus one.
-    // - Uncompressed
-    //    This is fairly obvious - we will then have one more than the uncompressed size.
-    int initialBufferSize = msg.readableBytes() + 1;
+    int uncompressed = msg.readableBytes();
+    if (uncompressed < threshold) {
+      int finalBufferSize = uncompressed + 1;
+      finalBufferSize += ProtocolUtils.varIntBytes(finalBufferSize);
+      return IS_JAVA_CIPHER
+          ? ctx.alloc().heapBuffer(finalBufferSize)
+          : ctx.alloc().directBuffer(finalBufferSize);
+    }
+
+    // (maximum data length after compression) + packet length varint + uncompressed data varint
+    int initialBufferSize = (uncompressed - 1) + 3 + ProtocolUtils.varIntBytes(uncompressed);
     return MoreByteBufUtils.preferredBuffer(ctx.alloc(), compressor, initialBufferSize);
   }
 
diff --git a/proxy/src/main/java/com/velocitypowered/proxy/protocol/netty/MinecraftVarintLengthEncoder.java b/proxy/src/main/java/com/velocitypowered/proxy/protocol/netty/MinecraftVarintLengthEncoder.java
index 94bfa84d0..f65764bf0 100644
--- a/proxy/src/main/java/com/velocitypowered/proxy/protocol/netty/MinecraftVarintLengthEncoder.java
+++ b/proxy/src/main/java/com/velocitypowered/proxy/protocol/netty/MinecraftVarintLengthEncoder.java
@@ -29,7 +29,7 @@ import io.netty.handler.codec.MessageToByteEncoder;
 public class MinecraftVarintLengthEncoder extends MessageToByteEncoder<ByteBuf> {
 
   public static final MinecraftVarintLengthEncoder INSTANCE = new MinecraftVarintLengthEncoder();
-  private static final boolean IS_JAVA_CIPHER = Natives.cipher.get() == JavaVelocityCipher.FACTORY;
+  public static final boolean IS_JAVA_CIPHER = Natives.cipher.get() == JavaVelocityCipher.FACTORY;
 
   private MinecraftVarintLengthEncoder() {
   }
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 57cd8f7da..58530a4a3 100644
--- a/proxy/src/test/java/com/velocitypowered/proxy/protocol/ProtocolUtilsTest.java
+++ b/proxy/src/test/java/com/velocitypowered/proxy/protocol/ProtocolUtilsTest.java
@@ -71,6 +71,20 @@ public class ProtocolUtilsTest {
     assertEquals(test, ProtocolUtils.readVarIntSafely(buf));
   }
 
+  @Test
+  void test3Bytes() {
+    ByteBuf buf = Unpooled.buffer(5);
+    for (int i = 0; i < 2097152; i += 31) {
+      writeReadTest3Bytes(buf, i);
+    }
+  }
+
+  private void writeReadTest3Bytes(ByteBuf buf, int test) {
+    buf.clear();
+    ProtocolUtils.writeVarIntAs3Bytes(buf, test);
+    assertEquals(test, ProtocolUtils.readVarInt(buf));
+  }
+
   @Test
   void testBytesWrittenAtBitBoundaries() {
     ByteBuf varintNew = Unpooled.buffer(5);