Merge pull request #492 from Leymooo/feature/combine-varint-prefix-with-compression

Do not copy memory in case when packet needs to compress

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);
   }

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);
   }
 

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() {
   }

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);