Make AES crypto operations use one buffer

All AES implementations being used are 'copy safe', where the source and
destination arrays may be the same. Lets save ourself a copy and reap
the performance wins!

native/src/main/java/com/velocitypowered/natives/encryption/JavaVelocityCipher.java

@@ -35,50 +35,24 @@ public class JavaVelocityCipher implements VelocityCipher {
   }
 
   @Override
-  public void process(ByteBuf source, ByteBuf destination) throws ShortBufferException {
+  public void process(ByteBuf source) {
     ensureNotDisposed();
 
     int inBytes = source.readableBytes();
     byte[] asBytes = ByteBufUtil.getBytes(source);
 
-    int outputSize = cipher.getOutputSize(inBytes);
-    byte[] outBuf = new byte[outputSize];
-    cipher.update(asBytes, 0, inBytes, outBuf);
-    destination.writeBytes(outBuf);
-  }
-
-  @Override
-  public ByteBuf process(ChannelHandlerContext ctx, ByteBuf source) throws ShortBufferException {
-    ensureNotDisposed();
-
-    int inBytes = source.readableBytes();
-    ByteBuf asHeapBuf = toHeap(source);
-    ByteBuf out = ctx.alloc().heapBuffer(cipher.getOutputSize(inBytes));
     try {
-      out.writerIndex(
+      cipher.update(asBytes, 0, inBytes, asBytes);
-          cipher.update(asHeapBuf.array(), asHeapBuf.arrayOffset() + asHeapBuf.readerIndex(),
+    } catch (ShortBufferException ex) {
-              inBytes, out.array(), out.arrayOffset() + out.writerIndex()));
+      /* This _really_ shouldn't happen - AES CFB8 will work in place.
-      return out;
+         If you run into this, that means that for whatever reason the Java Runtime has determined
-    } catch (ShortBufferException e) {
+         that the output buffer needs more bytes than the input buffer. When we are working with
-      out.release();
+         AES-CFB8, the output size is equal to the input size. See the problem? */
-      throw e;
+      throw new AssertionError("Cipher update did not operate in place and requested a larger "
-    } finally {
+              + "buffer than the source buffer");
-      asHeapBuf.release();
     }
   }
 
-  private static ByteBuf toHeap(ByteBuf src) {
-    if (src.hasArray()) {
-      return src.retain();
-    }
-
-    // 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
   public void dispose() {
     disposed = true;

native/src/main/java/com/velocitypowered/natives/encryption/NativeVelocityCipher.java

@@ -33,40 +33,14 @@ public class NativeVelocityCipher implements VelocityCipher {
   }
 
   @Override
-  public void process(ByteBuf source, ByteBuf destination) throws ShortBufferException {
+  public void process(ByteBuf source) {
     ensureNotDisposed();
     source.memoryAddress();
-    destination.memoryAddress();
 
-    // The exact amount we read in is also the amount we write out.
+    long base = source.memoryAddress() + source.readerIndex();
     int len = source.readableBytes();
-    destination.ensureWritable(len);
 
-    impl.process(ctx, source.memoryAddress() + source.readerIndex(), len,
+    impl.process(ctx, base, len, base, encrypt);
-        destination.memoryAddress() + destination.writerIndex(), encrypt);
-
-    source.skipBytes(len);
-    destination.writerIndex(destination.writerIndex() + len);
-  }
-
-  @Override
-  public ByteBuf process(ChannelHandlerContext ctx, ByteBuf source) throws ShortBufferException {
-    ensureNotDisposed();
-    source.memoryAddress(); // sanity check
-
-    int len = source.readableBytes();
-    ByteBuf out = ctx.alloc().directBuffer(len);
-
-    try {
-      impl.process(this.ctx, source.memoryAddress() + source.readerIndex(), len,
-          out.memoryAddress(), encrypt);
-      source.skipBytes(len);
-      out.writerIndex(len);
-      return out;
-    } catch (Exception e) {
-      out.release();
-      throw e;
-    }
   }
 
   @Override

native/src/main/java/com/velocitypowered/natives/encryption/VelocityCipher.java

@@ -7,8 +7,5 @@ import io.netty.channel.ChannelHandlerContext;
 import javax.crypto.ShortBufferException;
 
 public interface VelocityCipher extends Disposable, Native {
-
+  void process(ByteBuf source);
-  void process(ByteBuf source, ByteBuf destination) throws ShortBufferException;
-
-  ByteBuf process(ChannelHandlerContext ctx, ByteBuf source) throws ShortBufferException;
 }

native/src/test/java/com/velocitypowered/natives/encryption/VelocityCipherTest.java

@@ -56,20 +56,18 @@ class VelocityCipherTest {
     VelocityCipher encrypt = factory.forEncryption(new SecretKeySpec(AES_KEY, "AES"));
 
     ByteBuf source = bufSupplier.get();
-    ByteBuf dest = bufSupplier.get();
-    ByteBuf decryptionBuf = bufSupplier.get();
 
     source.writeBytes(TEST_DATA);
 
+    ByteBuf workingBuf = source.copy();
+
     try {
-      encrypt.process(source, dest);
+      encrypt.process(workingBuf);
-      decrypt.process(dest, decryptionBuf);
+      decrypt.process(workingBuf);
-      source.readerIndex(0);
+      assertTrue(ByteBufUtil.equals(source, workingBuf));
-      assertTrue(ByteBufUtil.equals(source, decryptionBuf));
     } finally {
       source.release();
-      dest.release();
+      workingBuf.release();
-      decryptionBuf.release();
       decrypt.dispose();
       encrypt.dispose();
     }

proxy/src/main/java/com/velocitypowered/proxy/protocol/netty/MinecraftCipherDecoder.java

@@ -20,9 +20,11 @@ public class MinecraftCipherDecoder extends ByteToMessageDecoder {
   protected void decode(ChannelHandlerContext ctx, ByteBuf in, List<Object> out) throws Exception {
     ByteBuf compatible = MoreByteBufUtils.ensureCompatible(ctx.alloc(), cipher, in);
     try {
-      out.add(cipher.process(ctx, compatible));
+      cipher.process(compatible);
-    } finally {
+      out.add(compatible);
-      compatible.release();
+    } catch (Exception e) {
+      compatible.release(); // compatible will never be used if we throw an exception
+      throw e;
     }
   }
 

proxy/src/main/java/com/velocitypowered/proxy/protocol/netty/MinecraftCipherEncoder.java

@@ -20,9 +20,11 @@ public class MinecraftCipherEncoder extends MessageToMessageEncoder<ByteBuf> {
   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));
+      cipher.process(compatible);
-    } finally {
+      out.add(compatible);
-      compatible.release();
+    } catch (Exception e) {
+      compatible.release(); // compatible will never be used if we throw an exception
+      throw e;
     }
   }