From 37a4199d4338e859b23bec7aed0f54930d815dfe Mon Sep 17 00:00:00 2001
From: Andrew Steinborn <git@steinborn.me>
Date: Thu, 6 May 2021 19:56:45 -0400
Subject: [PATCH] Combine VarInt prefix encoding with compression

This saves us a memory copy in the common "there is no need to compress this packet" case.
---
 .../proxy/connection/MinecraftConnection.java | 10 ++++----
 ... MinecraftCompressorAndLengthEncoder.java} | 24 +++++++++++++++----
 2 files changed, 26 insertions(+), 8 deletions(-)
 rename proxy/src/main/java/com/velocitypowered/proxy/protocol/netty/{MinecraftCompressEncoder.java => MinecraftCompressorAndLengthEncoder.java} (76%)

diff --git a/proxy/src/main/java/com/velocitypowered/proxy/connection/MinecraftConnection.java b/proxy/src/main/java/com/velocitypowered/proxy/connection/MinecraftConnection.java
index 4d6318968..ba172db42 100644
--- a/proxy/src/main/java/com/velocitypowered/proxy/connection/MinecraftConnection.java
+++ b/proxy/src/main/java/com/velocitypowered/proxy/connection/MinecraftConnection.java
@@ -36,12 +36,13 @@ import com.velocitypowered.proxy.VelocityServer;
 import com.velocitypowered.proxy.connection.client.HandshakeSessionHandler;
 import com.velocitypowered.proxy.connection.client.LoginSessionHandler;
 import com.velocitypowered.proxy.connection.client.StatusSessionHandler;
+import com.velocitypowered.proxy.network.Connections;
 import com.velocitypowered.proxy.protocol.MinecraftPacket;
 import com.velocitypowered.proxy.protocol.StateRegistry;
 import com.velocitypowered.proxy.protocol.netty.MinecraftCipherDecoder;
 import com.velocitypowered.proxy.protocol.netty.MinecraftCipherEncoder;
 import com.velocitypowered.proxy.protocol.netty.MinecraftCompressDecoder;
-import com.velocitypowered.proxy.protocol.netty.MinecraftCompressEncoder;
+import com.velocitypowered.proxy.protocol.netty.MinecraftCompressorAndLengthEncoder;
 import com.velocitypowered.proxy.protocol.netty.MinecraftDecoder;
 import com.velocitypowered.proxy.protocol.netty.MinecraftEncoder;
 import com.velocitypowered.proxy.util.except.QuietDecoderException;
@@ -402,8 +403,8 @@ public class MinecraftConnection extends ChannelInboundHandlerAdapter {
     } else {
       MinecraftCompressDecoder decoder = (MinecraftCompressDecoder) channel.pipeline()
           .get(COMPRESSION_DECODER);
-      MinecraftCompressEncoder encoder = (MinecraftCompressEncoder) channel.pipeline()
-          .get(COMPRESSION_ENCODER);
+      MinecraftCompressorAndLengthEncoder encoder =
+          (MinecraftCompressorAndLengthEncoder) channel.pipeline().get(COMPRESSION_ENCODER);
       if (decoder != null && encoder != null) {
         decoder.setThreshold(threshold);
         encoder.setThreshold(threshold);
@@ -411,9 +412,10 @@ public class MinecraftConnection extends ChannelInboundHandlerAdapter {
         int level = server.getConfiguration().getCompressionLevel();
         VelocityCompressor compressor = Natives.compress.get().create(level);
 
-        encoder = new MinecraftCompressEncoder(threshold, compressor);
+        encoder = new MinecraftCompressorAndLengthEncoder(threshold, compressor);
         decoder = new MinecraftCompressDecoder(threshold, compressor);
 
+        channel.pipeline().remove(FRAME_ENCODER);
         channel.pipeline().addBefore(MINECRAFT_DECODER, COMPRESSION_DECODER, decoder);
         channel.pipeline().addBefore(MINECRAFT_ENCODER, COMPRESSION_ENCODER, encoder);
       }
diff --git a/proxy/src/main/java/com/velocitypowered/proxy/protocol/netty/MinecraftCompressEncoder.java b/proxy/src/main/java/com/velocitypowered/proxy/protocol/netty/MinecraftCompressorAndLengthEncoder.java
similarity index 76%
rename from proxy/src/main/java/com/velocitypowered/proxy/protocol/netty/MinecraftCompressEncoder.java
rename to proxy/src/main/java/com/velocitypowered/proxy/protocol/netty/MinecraftCompressorAndLengthEncoder.java
index f6cb495ca..6684ee35d 100644
--- a/proxy/src/main/java/com/velocitypowered/proxy/protocol/netty/MinecraftCompressEncoder.java
+++ b/proxy/src/main/java/com/velocitypowered/proxy/protocol/netty/MinecraftCompressorAndLengthEncoder.java
@@ -23,13 +23,14 @@ import com.velocitypowered.proxy.protocol.ProtocolUtils;
 import io.netty.buffer.ByteBuf;
 import io.netty.channel.ChannelHandlerContext;
 import io.netty.handler.codec.MessageToByteEncoder;
+import java.util.zip.DataFormatException;
 
-public class MinecraftCompressEncoder extends MessageToByteEncoder<ByteBuf> {
+public class MinecraftCompressorAndLengthEncoder extends MessageToByteEncoder<ByteBuf> {
 
   private int threshold;
   private final VelocityCompressor compressor;
 
-  public MinecraftCompressEncoder(int threshold, VelocityCompressor compressor) {
+  public MinecraftCompressorAndLengthEncoder(int threshold, VelocityCompressor compressor) {
     this.threshold = threshold;
     this.compressor = compressor;
   }
@@ -39,16 +40,31 @@ public class MinecraftCompressEncoder extends MessageToByteEncoder<ByteBuf> {
     int uncompressed = msg.readableBytes();
     if (uncompressed < threshold) {
       // Under the threshold, there is nothing to do.
+      ProtocolUtils.writeVarInt(out, uncompressed + 1);
       ProtocolUtils.writeVarInt(out, 0);
       out.writeBytes(msg);
     } else {
-      ProtocolUtils.writeVarInt(out, uncompressed);
+      handleCompressed(ctx, msg, out);
+    }
+  }
+
+  private void handleCompressed(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, out);
+        compressor.deflate(compatibleIn, tmpBuf);
       } finally {
         compatibleIn.release();
       }
+
+      ProtocolUtils.writeVarInt(out, tmpBuf.readableBytes());
+      out.writeBytes(tmpBuf);
+    } finally {
+      tmpBuf.release();
     }
   }