Native cryptography support using mbed TLS

native/README.md

@@ -10,7 +10,10 @@ traditional Java fallbacks.
 
 ## Encryption
 
-* No natives available yet, this will use the support inside your Java install.
+* **Supported platforms**: macOS 10.13, Linux amd64
+* **Rationale**: Using a C library for encryption means we can limit memory copies. Prior to Java 7, this was the only
+  way to use AES-NI extensions on modern processors, but this is less important since JDK 8 has native support.
+* **Note**: Due to U.S. restrictions on cryptography export, this native is provided in source code form only for now.
 
 ## OS support
 

native/compile-linux.sh

@@ -1,4 +1,10 @@
 #!/bin/bash
 
 # Modify as you need.
-gcc -O3 -I$JAVA_HOME/include/ -I$JAVA_HOME/include/linux/ -shared -lz src/main/c/*.c -o src/main/resources/linux_x64/velocity-compress.so
+MBEDTLS_ROOT=mbedtls-2.12.0
+CFLAGS="-O3 -I$JAVA_HOME/include/ -I$JAVA_HOME/include/linux/ -fPIC -shared"
+gcc $CFLAGS -lz src/main/c/jni_util.c src/main/c/jni_zlib_deflate.c src/main/c/jni_zlib_inflate.c \
+    -o src/main/resources/linux_x64/velocity-compress.so
+gcc $CFLAGS  -I $MBEDTLS_ROOT/include -shared $MBEDTLS_ROOT/library/aes.c $MBEDTLS_ROOT/library/aesni.c \
+    $MBEDTLS_ROOT/library/platform.c $MBEDTLS_ROOT/library/platform_util.c src/main/c/jni_util.c src/main/c/jni_cipher.c \
+    -o src/main/resources/linux_x64/velocity-cipher.so

native/compile-osx.sh

@@ -1,5 +1,12 @@
 #!/bin/bash
 
 # Modify as you need.
+MBEDTLS_ROOT=mbedtls-2.12.0
 export JAVA_HOME=/Library/Java/JavaVirtualMachines/jdk1.8.0_144.jdk/Contents/Home
-clang -O3 -I$JAVA_HOME/include/ -I$JAVA_HOME/include/darwin/ -shared -lz src/main/c/*.c -o src/main/resources/macosx/velocity-compress.dylib
+CFLAGS="-O3 -I$JAVA_HOME/include/ -I$JAVA_HOME/include/darwin/ -fPIC -shared"
+
+clang $CFLAGS -lz src/main/c/jni_util.c src/main/c/jni_zlib_deflate.c src/main/c/jni_zlib_inflate.c \
+    -o src/main/resources/macosx/velocity-compress.dylib
+clang $CFLAGS -I $MBEDTLS_ROOT/include -shared $MBEDTLS_ROOT/library/aes.c $MBEDTLS_ROOT/library/aesni.c \
+    $MBEDTLS_ROOT/library/platform.c $MBEDTLS_ROOT/library/platform_util.c src/main/c/jni_util.c src/main/c/jni_cipher.c \
+    -o src/main/resources/macosx/velocity-cipher.dylib

native/src/main/c/jni_cipher.c

@@ -0,0 +1,81 @@
+#include <jni.h>
+#include <stdlib.h>
+#include <string.h>
+#include <mbedtls/aes.h>
+#include "jni_util.h"
+
+typedef unsigned char byte;
+
+typedef struct {
+    mbedtls_aes_context cipher;
+    byte *key;
+} velocity_cipher_context;
+
+JNIEXPORT jlong JNICALL
+Java_com_velocitypowered_natives_encryption_MbedtlsAesImpl_init(JNIEnv *env,
+    jobject obj,
+    jbyteArray key)
+{
+    velocity_cipher_context *ctx = malloc(sizeof(velocity_cipher_context));
+    if (ctx == NULL) {
+        throwException(env, "java/lang/OutOfMemoryError", "cipher allocate context");
+        return 0;
+    }
+
+    jsize keyLen = (*env)->GetArrayLength(env, key);
+    jbyte* keyBytes = (*env)->GetPrimitiveArrayCritical(env, key, NULL);
+    if (keyBytes == NULL) {
+        free(ctx);
+        throwException(env, "java/lang/OutOfMemoryError", "cipher get key");
+        return 0;
+    }
+
+    mbedtls_aes_init(&ctx->cipher);
+    int ret = mbedtls_aes_setkey_enc(&ctx->cipher, (byte*) keyBytes, keyLen * 8);
+    if (ret != 0) {
+        (*env)->ReleasePrimitiveArrayCritical(env, key, keyBytes, 0);
+        mbedtls_aes_free(&ctx->cipher);
+        free(ctx);
+
+        throwException(env, "java/security/GeneralSecurityException", "mbedtls set aes key");
+        return 0;
+    }
+
+    ctx->key = malloc(keyLen);
+    if (ctx->key == NULL) {
+        (*env)->ReleasePrimitiveArrayCritical(env, key, keyBytes, 0);
+        mbedtls_aes_free(&ctx->cipher);
+        free(ctx);
+
+        throwException(env, "java/lang/OutOfMemoryError", "cipher copy key");
+        return 0;
+    }
+    memcpy(ctx->key, keyBytes, keyLen);
+    (*env)->ReleasePrimitiveArrayCritical(env, key, keyBytes, 0);
+    return (jlong) ctx;
+}
+
+JNIEXPORT void JNICALL
+Java_com_velocitypowered_natives_encryption_MbedtlsAesImpl_free(JNIEnv *env,
+    jobject obj,
+    jlong ptr)
+{
+    velocity_cipher_context *ctx = (velocity_cipher_context*) ptr;
+    mbedtls_aes_free(&ctx->cipher);
+    free(ctx->key);
+    free(ctx);
+}
+
+JNIEXPORT void JNICALL
+Java_com_velocitypowered_natives_encryption_MbedtlsAesImpl_process(JNIEnv *env,
+    jobject obj,
+    jlong ptr,
+    jlong source,
+    jint len,
+    jlong dest,
+    jboolean encrypt)
+{
+    velocity_cipher_context *ctx = (velocity_cipher_context*) ptr;
+    mbedtls_aes_crypt_cfb8(&ctx->cipher, encrypt ? MBEDTLS_AES_ENCRYPT : MBEDTLS_AES_DECRYPT, len, ctx->key,
+        (byte*) source, (byte*) dest);
+}

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

@@ -0,0 +1,9 @@
+package com.velocitypowered.natives.encryption;
+
+public class MbedtlsAesImpl {
+    native long init(byte[] key);
+
+    native void process(long ctx, long sourceAddress, int sourceLength, long destinationAddress, boolean encrypt);
+
+    native void free(long ptr);
+}

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

@@ -0,0 +1,55 @@
+package com.velocitypowered.natives.encryption;
+
+import io.netty.buffer.ByteBuf;
+
+import javax.crypto.SecretKey;
+import javax.crypto.ShortBufferException;
+import java.security.GeneralSecurityException;
+
+public class NativeVelocityCipher implements VelocityCipher {
+    public static final VelocityCipherFactory FACTORY = new VelocityCipherFactory() {
+        @Override
+        public VelocityCipher forEncryption(SecretKey key) throws GeneralSecurityException {
+            return new NativeVelocityCipher(true, key);
+        }
+
+        @Override
+        public VelocityCipher forDecryption(SecretKey key) throws GeneralSecurityException {
+            return new NativeVelocityCipher(false, key);
+        }
+    };
+    private static final MbedtlsAesImpl impl = new MbedtlsAesImpl();
+
+    private final long ctx;
+    private final boolean encrypt;
+    private boolean disposed = false;
+
+    private NativeVelocityCipher(boolean encrypt, SecretKey key) {
+        this.encrypt = encrypt;
+        this.ctx = impl.init(key.getEncoded());
+    }
+
+    @Override
+    public void process(ByteBuf source, ByteBuf destination) throws ShortBufferException {
+        source.memoryAddress();
+        destination.memoryAddress();
+
+        // The exact amount we read in is also the amount we write out.
+        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 void dispose() {
+        if (!disposed) {
+            impl.free(ctx);
+        }
+        disposed = true;
+    }
+}

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

@@ -6,6 +6,7 @@ import com.velocitypowered.natives.compression.NativeVelocityCompressor;
 import com.velocitypowered.natives.compression.VelocityCompressor;
 import com.velocitypowered.natives.compression.VelocityCompressorFactory;
 import com.velocitypowered.natives.encryption.JavaVelocityCipher;
+import com.velocitypowered.natives.encryption.NativeVelocityCipher;
 import com.velocitypowered.natives.encryption.VelocityCipherFactory;
 
 import java.io.IOException;
@@ -51,6 +52,12 @@ public class Natives {
 
     public static final NativeCodeLoader<VelocityCipherFactory> cipher = new NativeCodeLoader<>(
             ImmutableList.of(
+                    /*new NativeCodeLoader.Variant<>(NativeCodeLoader.MACOS,
+                            copyAndLoadNative("/macosx/velocity-cipher.dylib"), "mbed TLS cipher (macOS)",
+                            NativeVelocityCipher.FACTORY),
+                    new NativeCodeLoader.Variant<>(NativeCodeLoader.LINUX,
+                            copyAndLoadNative("/linux_x64/velocity-cipher.so"), "mbed TLS cipher (Linux amd64)",
+                            NativeVelocityCipher.FACTORY),*/
                     new NativeCodeLoader.Variant<>(NativeCodeLoader.ALWAYS, () -> {}, "Java cipher", JavaVelocityCipher.FACTORY)
             )
     );

native/src/test/java/com/velocitypowered/natives/compression/VelocityCompressorTest.java

@@ -29,6 +29,7 @@ class VelocityCompressorTest {
     void nativeIntegrityCheck() throws DataFormatException {
         VelocityCompressor compressor = Natives.compressor.get().create(Deflater.DEFAULT_COMPRESSION);
         if (compressor instanceof JavaVelocityCompressor) {
+            compressor.dispose();
             fail("Loaded regular compressor");
         }
         check(compressor);
@@ -58,6 +59,7 @@ class VelocityCompressorTest {
         } finally {
             source.release();
             dest.release();
+            decompressed.release();
             compressor.dispose();
         }
     }

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

@@ -0,0 +1,73 @@
+package com.velocitypowered.natives.encryption;
+
+import com.velocitypowered.natives.util.Natives;
+import io.netty.buffer.ByteBuf;
+import io.netty.buffer.ByteBufUtil;
+import io.netty.buffer.Unpooled;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.condition.EnabledOnOs;
+
+import javax.crypto.spec.SecretKeySpec;
+import java.security.GeneralSecurityException;
+import java.util.Random;
+
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.junit.jupiter.api.Assertions.fail;
+import static org.junit.jupiter.api.condition.OS.LINUX;
+import static org.junit.jupiter.api.condition.OS.MAC;
+
+class VelocityCipherTest {
+    private static final int ENCRYPT_DATA_SIZE = 1 << 16;
+
+    @BeforeAll
+    static void checkNatives() {
+        Natives.cipher.getLoadedVariant();
+    }
+
+    @Test
+    @EnabledOnOs({ MAC, LINUX })
+    void nativeIntegrityCheck() throws GeneralSecurityException {
+        VelocityCipherFactory factory = Natives.cipher.get();
+        if (factory == JavaVelocityCipher.FACTORY) {
+            fail("Loaded regular compressor");
+        }
+        check(factory);
+    }
+
+    @Test
+    void javaIntegrityCheck() throws GeneralSecurityException {
+        check(JavaVelocityCipher.FACTORY);
+    }
+
+    private void check(VelocityCipherFactory factory) throws GeneralSecurityException {
+        // Generate a random 16-byte key.
+        Random random = new Random(1);
+        byte[] key = new byte[16];
+        random.nextBytes(key);
+
+        VelocityCipher decrypt = factory.forDecryption(new SecretKeySpec(key, "AES"));
+        VelocityCipher encrypt = factory.forEncryption(new SecretKeySpec(key, "AES"));
+
+        ByteBuf source = Unpooled.directBuffer(ENCRYPT_DATA_SIZE);
+        ByteBuf dest = Unpooled.directBuffer(ENCRYPT_DATA_SIZE);
+        ByteBuf decryptionBuf = Unpooled.directBuffer(ENCRYPT_DATA_SIZE);
+
+        byte[] randomBytes = new byte[ENCRYPT_DATA_SIZE];
+        random.nextBytes(randomBytes);
+        source.writeBytes(randomBytes);
+
+        try {
+            encrypt.process(source, dest);
+            decrypt.process(dest, decryptionBuf);
+            source.readerIndex(0);
+            assertTrue(ByteBufUtil.equals(source, decryptionBuf));
+        } finally {
+            source.release();
+            dest.release();
+            decryptionBuf.release();
+            decrypt.dispose();
+            encrypt.dispose();
+        }
+    }
+}

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

@@ -17,7 +17,7 @@ public class MinecraftCipherDecoder extends ByteToMessageDecoder {
 
     @Override
     protected void decode(ChannelHandlerContext ctx, ByteBuf in, List<Object> out) throws Exception {
-        ByteBuf decrypted = ctx.alloc().buffer();
+        ByteBuf decrypted = ctx.alloc().buffer(in.readableBytes());
         try {
             cipher.process(in, decrypted);
             out.add(decrypted);

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

@@ -18,6 +18,11 @@ public class MinecraftCipherEncoder extends MessageToByteEncoder<ByteBuf> {
         cipher.process(msg, out);
     }
 
+    @Override
+    protected ByteBuf allocateBuffer(ChannelHandlerContext ctx, ByteBuf msg, boolean preferDirect) throws Exception {
+        return ctx.alloc().directBuffer(msg.readableBytes());
+    }
+
     @Override
     public void handlerRemoved(ChannelHandlerContext ctx) throws Exception {
         cipher.dispose();