From 095a47844099f0a9e9b92de1629319895cd7b15d Mon Sep 17 00:00:00 2001
From: Andrew Steinborn <git@steinborn.me>
Date: Sun, 30 Dec 2018 04:53:47 -0500
Subject: [PATCH] Always ensure we don't pass a heap ByteBuf to the natives.

---
 .../natives/util/MoreByteBufUtils.java        |  2 +-
 .../netty/MinecraftCipherDecoder.java         |  8 ++++++-
 .../netty/MinecraftCipherEncoder.java         | 21 ++++++++++---------
 .../netty/MinecraftCompressDecoder.java       | 18 +++++++++-------
 .../netty/MinecraftCompressEncoder.java       |  8 ++++++-
 5 files changed, 37 insertions(+), 20 deletions(-)

diff --git a/native/src/main/java/com/velocitypowered/natives/util/MoreByteBufUtils.java b/native/src/main/java/com/velocitypowered/natives/util/MoreByteBufUtils.java
index e948da46d..460395e8a 100644
--- a/native/src/main/java/com/velocitypowered/natives/util/MoreByteBufUtils.java
+++ b/native/src/main/java/com/velocitypowered/natives/util/MoreByteBufUtils.java
@@ -25,7 +25,7 @@ public class MoreByteBufUtils {
       return buf.retain();
     }
 
-    // It's not, so we must make a memory copy.
+    // It's not, so we must make a direct copy.
     ByteBuf newBuf = alloc.directBuffer();
     newBuf.writeBytes(buf);
     return newBuf;
diff --git a/proxy/src/main/java/com/velocitypowered/proxy/protocol/netty/MinecraftCipherDecoder.java b/proxy/src/main/java/com/velocitypowered/proxy/protocol/netty/MinecraftCipherDecoder.java
index 373b88af5..4541fc902 100644
--- a/proxy/src/main/java/com/velocitypowered/proxy/protocol/netty/MinecraftCipherDecoder.java
+++ b/proxy/src/main/java/com/velocitypowered/proxy/protocol/netty/MinecraftCipherDecoder.java
@@ -2,6 +2,7 @@ package com.velocitypowered.proxy.protocol.netty;
 
 import com.google.common.base.Preconditions;
 import com.velocitypowered.natives.encryption.VelocityCipher;
+import com.velocitypowered.natives.util.MoreByteBufUtils;
 import io.netty.buffer.ByteBuf;
 import io.netty.channel.ChannelHandlerContext;
 import io.netty.handler.codec.ByteToMessageDecoder;
@@ -17,7 +18,12 @@ public class MinecraftCipherDecoder extends ByteToMessageDecoder {
 
   @Override
   protected void decode(ChannelHandlerContext ctx, ByteBuf in, List<Object> out) throws Exception {
-    out.add(cipher.process(ctx, in));
+    ByteBuf compatible = MoreByteBufUtils.ensureCompatible(ctx.alloc(), cipher, in);
+    try {
+      out.add(cipher.process(ctx, compatible));
+    } finally {
+      compatible.release();
+    }
   }
 
   @Override
diff --git a/proxy/src/main/java/com/velocitypowered/proxy/protocol/netty/MinecraftCipherEncoder.java b/proxy/src/main/java/com/velocitypowered/proxy/protocol/netty/MinecraftCipherEncoder.java
index 3df7e8287..dd944a19f 100644
--- a/proxy/src/main/java/com/velocitypowered/proxy/protocol/netty/MinecraftCipherEncoder.java
+++ b/proxy/src/main/java/com/velocitypowered/proxy/protocol/netty/MinecraftCipherEncoder.java
@@ -2,11 +2,13 @@ package com.velocitypowered.proxy.protocol.netty;
 
 import com.google.common.base.Preconditions;
 import com.velocitypowered.natives.encryption.VelocityCipher;
+import com.velocitypowered.natives.util.MoreByteBufUtils;
 import io.netty.buffer.ByteBuf;
 import io.netty.channel.ChannelHandlerContext;
-import io.netty.handler.codec.MessageToByteEncoder;
+import io.netty.handler.codec.MessageToMessageEncoder;
+import java.util.List;
 
-public class MinecraftCipherEncoder extends MessageToByteEncoder<ByteBuf> {
+public class MinecraftCipherEncoder extends MessageToMessageEncoder<ByteBuf> {
 
   private final VelocityCipher cipher;
 
@@ -15,14 +17,13 @@ public class MinecraftCipherEncoder extends MessageToByteEncoder<ByteBuf> {
   }
 
   @Override
-  protected void encode(ChannelHandlerContext ctx, ByteBuf msg, ByteBuf out) throws Exception {
-    cipher.process(msg, out);
-  }
-
-  @Override
-  protected ByteBuf allocateBuffer(ChannelHandlerContext ctx, ByteBuf msg, boolean preferDirect)
-      throws Exception {
-    return ctx.alloc().directBuffer(msg.readableBytes());
+  protected void encode(ChannelHandlerContext ctx, ByteBuf msg, List<Object> out) throws Exception {
+    ByteBuf compatible = MoreByteBufUtils.ensureCompatible(ctx.alloc(), cipher, msg);
+    try {
+      out.add(cipher.process(ctx, compatible));
+    } finally {
+      compatible.release();
+    }
   }
 
   @Override
diff --git a/proxy/src/main/java/com/velocitypowered/proxy/protocol/netty/MinecraftCompressDecoder.java b/proxy/src/main/java/com/velocitypowered/proxy/protocol/netty/MinecraftCompressDecoder.java
index f8afa0d48..2c4aab5fc 100644
--- a/proxy/src/main/java/com/velocitypowered/proxy/protocol/netty/MinecraftCompressDecoder.java
+++ b/proxy/src/main/java/com/velocitypowered/proxy/protocol/netty/MinecraftCompressDecoder.java
@@ -3,6 +3,7 @@ package com.velocitypowered.proxy.protocol.netty;
 import static com.velocitypowered.proxy.protocol.util.NettyPreconditions.checkFrame;
 
 import com.velocitypowered.natives.compression.VelocityCompressor;
+import com.velocitypowered.natives.util.MoreByteBufUtils;
 import com.velocitypowered.proxy.protocol.ProtocolUtils;
 import io.netty.buffer.ByteBuf;
 import io.netty.channel.ChannelHandlerContext;
@@ -22,22 +23,23 @@ public class MinecraftCompressDecoder extends MessageToMessageDecoder<ByteBuf> {
   }
 
   @Override
-  protected void decode(ChannelHandlerContext ctx, ByteBuf msg, List<Object> out) throws Exception {
-    int expectedUncompressedSize = ProtocolUtils.readVarInt(msg);
+  protected void decode(ChannelHandlerContext ctx, ByteBuf in, List<Object> out) throws Exception {
+    int expectedUncompressedSize = ProtocolUtils.readVarInt(in);
     if (expectedUncompressedSize == 0) {
       // Strip the now-useless uncompressed size, this message is already uncompressed.
-      out.add(msg.retainedSlice());
-      msg.skipBytes(msg.readableBytes());
+      out.add(in.retainedSlice());
+      in.skipBytes(in.readableBytes());
       return;
     }
 
     checkFrame(expectedUncompressedSize >= threshold,
         "Uncompressed size %s is greater than threshold %s",
         expectedUncompressedSize, threshold);
-    ByteBuf uncompressed = ctx.alloc()
-        .buffer(Math.min(expectedUncompressedSize, MAXIMUM_INITIAL_BUFFER_SIZE));
+    ByteBuf compatibleIn = MoreByteBufUtils.ensureCompatible(ctx.alloc(), compressor, in);
+    ByteBuf uncompressed = ctx.alloc().directBuffer(Math.min(expectedUncompressedSize,
+        MAXIMUM_INITIAL_BUFFER_SIZE));
     try {
-      compressor.inflate(msg, uncompressed);
+      compressor.inflate(compatibleIn, uncompressed);
       checkFrame(expectedUncompressedSize == uncompressed.readableBytes(),
           "Mismatched compression sizes (got %s, expected %s)",
           uncompressed.readableBytes(), expectedUncompressedSize);
@@ -45,6 +47,8 @@ public class MinecraftCompressDecoder extends MessageToMessageDecoder<ByteBuf> {
     } catch (Exception e) {
       uncompressed.release();
       throw e;
+    } finally {
+      compatibleIn.release();
     }
   }
 
diff --git a/proxy/src/main/java/com/velocitypowered/proxy/protocol/netty/MinecraftCompressEncoder.java b/proxy/src/main/java/com/velocitypowered/proxy/protocol/netty/MinecraftCompressEncoder.java
index 60d7c5d19..63f75d472 100644
--- a/proxy/src/main/java/com/velocitypowered/proxy/protocol/netty/MinecraftCompressEncoder.java
+++ b/proxy/src/main/java/com/velocitypowered/proxy/protocol/netty/MinecraftCompressEncoder.java
@@ -1,6 +1,7 @@
 package com.velocitypowered.proxy.protocol.netty;
 
 import com.velocitypowered.natives.compression.VelocityCompressor;
+import com.velocitypowered.natives.util.MoreByteBufUtils;
 import com.velocitypowered.proxy.protocol.ProtocolUtils;
 import io.netty.buffer.ByteBuf;
 import io.netty.channel.ChannelHandlerContext;
@@ -25,7 +26,12 @@ public class MinecraftCompressEncoder extends MessageToByteEncoder<ByteBuf> {
       out.writeBytes(msg);
     } else {
       ProtocolUtils.writeVarInt(out, uncompressed);
-      compressor.deflate(msg, out);
+      ByteBuf compatibleIn = MoreByteBufUtils.ensureCompatible(ctx.alloc(), compressor, msg);
+      try {
+        compressor.deflate(compatibleIn, out);
+      } finally {
+        compatibleIn.release();
+      }
     }
   }