Merge pull request #254 from Ichbinjoe/aes-nocopy

Make AES crypto operations no-copy

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

@@ -35,50 +35,25 @@ public class JavaVelocityCipher implements VelocityCipher {
   }
 
   @Override
-  public void process(ByteBuf source, ByteBuf destination) throws ShortBufferException {
+  public void process(ByteBuf source) {
     ensureNotDisposed();
+    Preconditions.checkArgument(source.hasArray(), "No source array");
 
     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(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();
+      cipher.update(source.array(), source.arrayOffset(), inBytes, source.array(),
+          source.arrayOffset());
+    } catch (ShortBufferException ex) {
+      /* This _really_ shouldn't happen - AES CFB8 will work in place.
+         If you run into this, that means that for whatever reason the Java Runtime has determined
+         that the output buffer needs more bytes than the input buffer. When we are working with
+         AES-CFB8, the output size is equal to the input size. See the problem? */
+      throw new AssertionError("Cipher update did not operate in place and requested a larger "
+              + "buffer than the source buffer");
     }
   }
 
-  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;
@@ -90,6 +65,6 @@ public class JavaVelocityCipher implements VelocityCipher {
 
   @Override
   public BufferPreference preferredBufferType() {
-    return BufferPreference.HEAP_PREFERRED;
+    return BufferPreference.HEAP_REQUIRED;
   }
 }

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,
-        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;
-    }
+    impl.process(ctx, base, len, base, encrypt);
   }
 
   @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, ByteBuf destination) throws ShortBufferException;
-
-  ByteBuf process(ChannelHandlerContext ctx, ByteBuf source) throws ShortBufferException;
+  void process(ByteBuf source);
 }

native/src/main/java/com/velocitypowered/natives/util/BufferPreference.java

@@ -1,6 +1,10 @@
 package com.velocitypowered.natives.util;
 
 public enum BufferPreference {
+  /**
+   * A heap buffer is required.
+   */
+  HEAP_REQUIRED,
   /**
    * A heap buffer is preferred (but not required).
    */

native/src/main/java/com/velocitypowered/natives/util/MoreByteBufUtils.java

@@ -20,28 +20,30 @@ public class MoreByteBufUtils {
    * @return a buffer compatible with the native
    */
   public static ByteBuf ensureCompatible(ByteBufAllocator alloc, Native nativeStuff, ByteBuf buf) {
-    if (nativeStuff.preferredBufferType() != BufferPreference.DIRECT_REQUIRED
-        || buf.hasMemoryAddress()) {
+    if (isCompatible(nativeStuff, buf)) {
       return buf.retain();
     }
 
     // It's not, so we must make a direct copy.
-    ByteBuf newBuf = alloc.directBuffer(buf.readableBytes());
+    ByteBuf newBuf = preferredBuffer(alloc, nativeStuff, buf.readableBytes());
     newBuf.writeBytes(buf);
     return newBuf;
   }
 
-  /**
-   * Creates a {@link ByteBuf} that will have the best performance with the specified
-   * {@code nativeStuff}.
-   *
-   * @param alloc the {@link ByteBufAllocator} to use
-   * @param nativeStuff the native we are working with
-   * @return a buffer compatible with the native
-   */
-  public static ByteBuf preferredBuffer(ByteBufAllocator alloc, Native nativeStuff) {
-    return nativeStuff.preferredBufferType() != BufferPreference.HEAP_PREFERRED
-        ? alloc.directBuffer() : alloc.heapBuffer();
+  private static boolean isCompatible(Native nativeStuff, ByteBuf buf) {
+    BufferPreference preferred = nativeStuff.preferredBufferType();
+    switch (preferred) {
+      case DIRECT_PREFERRED:
+      case HEAP_PREFERRED:
+        // The native prefers this type, but doesn't strictly require we provide it.
+        return true;
+      case DIRECT_REQUIRED:
+        return buf.hasMemoryAddress();
+      case HEAP_REQUIRED:
+        return buf.hasArray();
+      default:
+        throw new AssertionError("Preferred buffer type unknown");
+    }
   }
 
   /**
@@ -55,7 +57,15 @@ public class MoreByteBufUtils {
    */
   public static ByteBuf preferredBuffer(ByteBufAllocator alloc, Native nativeStuff,
       int initialCapacity) {
-    return nativeStuff.preferredBufferType() != BufferPreference.HEAP_PREFERRED
-        ? alloc.directBuffer(initialCapacity) : alloc.heapBuffer(initialCapacity);
+    switch (nativeStuff.preferredBufferType()) {
+      case HEAP_REQUIRED:
+      case HEAP_PREFERRED:
+        return alloc.heapBuffer(initialCapacity);
+      case DIRECT_PREFERRED:
+      case DIRECT_REQUIRED:
+        return alloc.directBuffer(initialCapacity);
+      default:
+        throw new AssertionError("Preferred buffer type unknown");
+    }
   }
 }

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);
-      decrypt.process(dest, decryptionBuf);
-      source.readerIndex(0);
-      assertTrue(ByteBufUtil.equals(source, decryptionBuf));
+      encrypt.process(workingBuf);
+      decrypt.process(workingBuf);
+      assertTrue(ByteBufUtil.equals(source, workingBuf));
     } finally {
       source.release();
-      dest.release();
-      decryptionBuf.release();
+      workingBuf.release();
       decrypt.dispose();
       encrypt.dispose();
     }

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

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