From ec32def592fb05407722e7de87994b76007dc90f Mon Sep 17 00:00:00 2001
From: Andrew Steinborn <git@steinborn.me>
Date: Tue, 11 Jun 2019 02:04:38 -0400
Subject: [PATCH] Take advantage of pooled ByteBufs in Java cipher

ByteBufs are pooled in Velocity, so we can achieve lower memory usage
by using what Netty is going to provide to us for free.
---
 .../encryption/JavaVelocityCipher.java        | 58 ++++++++-----------
 1 file changed, 25 insertions(+), 33 deletions(-)

diff --git a/native/src/main/java/com/velocitypowered/natives/encryption/JavaVelocityCipher.java b/native/src/main/java/com/velocitypowered/natives/encryption/JavaVelocityCipher.java
index c7a14815e..88f7fa921 100644
--- a/native/src/main/java/com/velocitypowered/natives/encryption/JavaVelocityCipher.java
+++ b/native/src/main/java/com/velocitypowered/natives/encryption/JavaVelocityCipher.java
@@ -2,7 +2,7 @@ package com.velocitypowered.natives.encryption;
 
 import com.google.common.base.Preconditions;
 import io.netty.buffer.ByteBuf;
-import io.netty.buffer.Unpooled;
+import io.netty.buffer.ByteBufUtil;
 import io.netty.channel.ChannelHandlerContext;
 import java.security.GeneralSecurityException;
 import javax.crypto.Cipher;
@@ -23,9 +23,6 @@ public class JavaVelocityCipher implements VelocityCipher {
       return new JavaVelocityCipher(false, key);
     }
   };
-  private static final int INITIAL_BUFFER_SIZE = 1024 * 8;
-  private static final ThreadLocal<byte[]> inBufLocal = ThreadLocal.withInitial(
-      () -> new byte[INITIAL_BUFFER_SIZE]);
 
   private final Cipher cipher;
   private boolean disposed = false;
@@ -41,20 +38,12 @@ public class JavaVelocityCipher implements VelocityCipher {
     ensureNotDisposed();
 
     int inBytes = source.readableBytes();
-    ByteBuf asHeapBuf = asHeapBuf(source);
+    byte[] asBytes = ByteBufUtil.getBytes(source);
 
     int outputSize = cipher.getOutputSize(inBytes);
-    if (!destination.hasArray()) {
-      byte[] outBuf = new byte[outputSize];
-      cipher.update(asHeapBuf.array(), asHeapBuf.arrayOffset(), inBytes, outBuf);
-      destination.writeBytes(outBuf);
-    } else {
-      // If the destination we write to is an array, we can use the backing array directly.
-      destination.ensureWritable(outputSize);
-      int produced = cipher.update(asHeapBuf.array(), asHeapBuf.arrayOffset(), inBytes,
-          destination.array(), destination.arrayOffset());
-      destination.writerIndex(destination.writerIndex() + produced);
-    }
+    byte[] outBuf = new byte[outputSize];
+    cipher.update(asBytes, 0, inBytes, outBuf);
+    destination.writeBytes(outBuf);
   }
 
   @Override
@@ -62,28 +51,31 @@ public class JavaVelocityCipher implements VelocityCipher {
     ensureNotDisposed();
 
     int inBytes = source.readableBytes();
-    ByteBuf asHeapBuf = asHeapBuf(source);
-
+    ByteBuf asHeapBuf = toHeap(source);
     ByteBuf out = ctx.alloc().heapBuffer(cipher.getOutputSize(inBytes));
-    out.writerIndex(cipher.update(asHeapBuf.array(), asHeapBuf.arrayOffset(), inBytes, out.array(),
-        out.arrayOffset()));
-    return out;
+    try {
+      out.writerIndex(
+          cipher.update(asHeapBuf.array(), asHeapBuf.arrayOffset() + asHeapBuf.readerIndex(),
+              inBytes, out.array(), out.arrayOffset() + out.writerIndex()));
+      return out;
+    } catch (ShortBufferException e) {
+      out.release();
+      throw e;
+    } finally {
+      asHeapBuf.release();
+    }
   }
 
-  private static ByteBuf asHeapBuf(ByteBuf source) {
-    if (source.hasArray()) {
-      // If this byte buffer is backed by an array, we can just use this buffer directly.
-      return source;
+  private static ByteBuf toHeap(ByteBuf src) {
+    if (!src.isDirect()) {
+      return src.retain();
     }
 
-    int inBytes = source.readableBytes();
-    byte[] inBuf = inBufLocal.get();
-    if (inBuf.length <= inBytes) {
-      inBuf = new byte[inBytes];
-      inBufLocal.set(inBuf);
-    }
-    source.readBytes(inBuf, 0, inBytes);
-    return Unpooled.wrappedBuffer(inBuf, 0, inBytes);
+    // Copy into a temporary heap buffer. We could use a local buffer, but Netty pools all buffers,
+    // so we'd lose more than we gain.
+    ByteBuf asHeapBuf = src.alloc().heapBuffer(src.readableBytes());
+    asHeapBuf.writeBytes(src);
+    return asHeapBuf;
   }
 
   @Override