package net.md_5.bungee;

import net.md_5.bungee.jni.cipher.NativeCipher;
import net.md_5.bungee.jni.cipher.JavaCipher;
import net.md_5.bungee.jni.cipher.BungeeCipher;
import io.netty.buffer.ByteBuf;
import io.netty.buffer.Unpooled;
import org.junit.Assert;
import org.junit.Test;
import org.junit.FixMethodOrder;
import org.junit.runners.MethodSorters;
import javax.crypto.SecretKey;
import javax.crypto.spec.SecretKeySpec;
import net.md_5.bungee.jni.NativeCode;

@FixMethodOrder(MethodSorters.NAME_ASCENDING)
public class NativeCipherTest
{

    private final byte[] plainBytes = "This is a test".getBytes();
    private final byte[] cipheredBytes = new byte[]
    {
        50, -7, 89, 1, -11, -32, -118, -48, -2, -72, 105, 97, -70, -81
    };
    private final SecretKey secret = new SecretKeySpec( new byte[ 16 ], "AES" );
    private static final int BENCHMARK_COUNT = 50000;
    //
    private static final NativeCode<BungeeCipher> factory = new NativeCode( "native-cipher", JavaCipher.class, NativeCipher.class );

    @Test
    public void testOpenSSL() throws Exception
    {
        if ( NativeCode.isSupported() )
        {
            boolean loaded = factory.load();
            Assert.assertTrue( "Native cipher failed to load!", loaded );

            NativeCipher cipher = new NativeCipher();
            System.out.println( "Testing OpenSSL cipher..." );
            testACipher( cipher );
        }
    }

    @Test
    public void testOpenSSLBenchmark() throws Exception
    {
        if ( NativeCode.isSupported() )
        {
            boolean loaded = factory.load();
            Assert.assertTrue( "Native cipher failed to load!", loaded );

            NativeCipher cipher = new NativeCipher();

            System.out.println( "Benchmarking OpenSSL cipher..." );
            testBenchmark( cipher );
        }
    }

    @Test
    public void testJDK() throws Exception
    {
        // Create JDK cipher
        BungeeCipher cipher = new JavaCipher();

        System.out.println( "Testing Java cipher..." );
        testACipher( cipher );
    }

    @Test
    public void testJDKBenchmark() throws Exception
    {
        // Create JDK cipher
        BungeeCipher cipher = new JavaCipher();

        System.out.println( "Benchmarking Java cipher..." );
        testBenchmark( cipher );
    }

    /**
     * Hackish test which can test both native and fallback ciphers using direct
     * buffers.
     */
    public void testACipher(BungeeCipher cipher) throws Exception
    {
        // Create input buf
        ByteBuf nativePlain = Unpooled.directBuffer( plainBytes.length );
        nativePlain.writeBytes( plainBytes );
        // Create expected buf
        ByteBuf nativeCiphered = Unpooled.directBuffer( cipheredBytes.length );
        nativeCiphered.writeBytes( cipheredBytes );
        // Create output buf
        ByteBuf out = Unpooled.directBuffer( plainBytes.length );

        // Encrypt
        cipher.init( true, secret );
        cipher.cipher( nativePlain, out );
        Assert.assertEquals( nativeCiphered, out );

        out.clear();

        // Decrypt
        cipher.init( false, secret );
        cipher.cipher( nativeCiphered, out );
        nativePlain.resetReaderIndex();
        Assert.assertEquals( nativePlain, out );

        System.out.println( "This cipher works correctly!" );
    }

    public void testBenchmark(BungeeCipher cipher) throws Exception
    {
        // Create input buf
        ByteBuf nativePlain = Unpooled.directBuffer( plainBytes.length );
        nativePlain.writeBytes( plainBytes );
        // Create expected buf
        ByteBuf nativeCiphered = Unpooled.directBuffer( cipheredBytes.length );
        nativeCiphered.writeBytes( cipheredBytes );
        // Create output buf
        ByteBuf out = Unpooled.directBuffer( plainBytes.length );

        // Encrypt
        cipher.init( true, secret );
        long start = System.currentTimeMillis();
        for ( int i = 0; i < BENCHMARK_COUNT; i++ )
        {
            cipher.cipher( nativePlain, out );
            nativePlain.readerIndex( 0 );
            out.clear();
        }
        System.out.println( String.format( "Encryption Iteration: %d, Elapsed: %d ms", BENCHMARK_COUNT, System.currentTimeMillis() - start ) );

        // Decrypt
        cipher.init( false, secret );
        start = System.currentTimeMillis();
        for ( int i = 0; i < BENCHMARK_COUNT; i++ )
        {
            cipher.cipher( nativeCiphered, out );
            nativeCiphered.readerIndex( 0 );
            out.clear();
        }
        System.out.println( String.format( "Decryption Iteration: %d, Elapsed: %d ms", BENCHMARK_COUNT, System.currentTimeMillis() - start ) );
    }
}

@@ -226,16 +226,41 @@
 pooled;%0A
+import java.util.Random;%0A
 import o
@@ -897,13 +897,12 @@
 T = 
-50000
+4096
 ;%0A  
@@ -3481,150 +3481,86 @@
     
-ByteBuf nativePlain = Unpooled.directBuffer( plainBytes.length );%0A        nativePlain.writeBytes( plainBytes );%0A        // Create expected buf
+byte%5B%5D random = new byte%5B 1 %3C%3C 12 %5D;%0A        new Random().nextBytes( random );
 %0A   
@@ -3574,32 +3574,29 @@
 teBuf native
-Ciphered
+Plain
  = Unpooled.
@@ -3608,38 +3608,16 @@
 tBuffer(
- cipheredBytes.length 
 );%0A     
@@ -3621,32 +3621,29 @@
       native
-Ciphered
+Plain
 .writeBytes(
@@ -3643,33 +3643,27 @@
 eBytes( 
-cipheredBytes
+random
  );%0A
+%0A
         
@@ -3691,35 +3691,46 @@
         ByteBuf 
-out
+nativeCiphered
  = Unpooled.dire
@@ -3920,32 +3920,68 @@
 i++ )%0A        %7B%0A
+            nativeCiphered.clear();%0A
             ciph
@@ -3996,35 +3996,46 @@
 r( nativePlain, 
-out
+nativeCiphered
  );%0A            
@@ -4068,33 +4068,8 @@
  );%0A
-            out.clear();%0A
     
@@ -4215,24 +4215,120 @@
 start ) );%0A%0A
+        // Create output buf%0A        ByteBuf out = Unpooled.directBuffer( plainBytes.length );%0A%0A
         // D

Use more realistic cipher test sizes and counts.