package com.github.abrasha.crypto;

import org.openjdk.jmh.annotations.*;
import org.openjdk.jmh.infra.Blackhole;

import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.spec.SecretKeySpec;
import java.security.Key;
import java.util.concurrent.TimeUnit;

import static java.nio.charset.StandardCharsets.UTF_8;

// Run complete.
//
// Benchmark                       Mode  Cnt     Score     Error  Units
// CryptoInit.withAvailableCipher  avgt    9   134.556 ±   6.433  ns/op
// CryptoInit.withNewCipher        avgt    9  6483.068 ± 234.472  ns/op
@BenchmarkMode(Mode.AverageTime)
@OutputTimeUnit(TimeUnit.NANOSECONDS)
public class CryptoInit {
    
    private static byte[] MESSAGE_TO_ENCRYPT = "catch me if you can".getBytes(UTF_8);
    
    private static class KeyContainer {
        static final Key key;
        
        static {
            try {
                KeyGenerator generator = KeyGenerator.getInstance("AES");
                generator.init(256);
                
                key = new SecretKeySpec(generator.generateKey().getEncoded(), "AES");
            } catch (Exception e) {
                // ignored
                throw new RuntimeException(e);
            }
        }
        
    }
    
    @State(Scope.Benchmark)
    public static class CipherContainer {
        
        private ThreadLocal<Cipher> cipherThreadLocal = new ThreadLocal<>();
        
        private static Cipher initCipher(){
            try {
    
                Cipher cipher = Cipher.getInstance("AES/CBC/PKCS5Padding");
                cipher.init(Cipher.ENCRYPT_MODE, KeyContainer.key);
                return cipher;
            } catch (Exception e){
                throw new RuntimeException(e);
            }
        }
        
        public Cipher getCipher(){
            Cipher cached = cipherThreadLocal.get();
            if (cached == null){
                cached = initCipher();
                cipherThreadLocal.set(cached);
            }
            return cached;
        }
    }
    
    @Benchmark
    public void withAvailableCipher(CipherContainer container, Blackhole bh) throws Exception {
        bh.consume(container.getCipher().doFinal(MESSAGE_TO_ENCRYPT));
    }
    
    @Benchmark
    public void withNewCipher(CipherContainer container, Blackhole bh) throws Exception {
        Cipher cipher = Cipher.getInstance("AES/CBC/PKCS5Padding");
        cipher.init(Cipher.ENCRYPT_MODE, KeyContainer.key);
        bh.consume(cipher.doFinal(MESSAGE_TO_ENCRYPT));
    }
    
}

@@ -170,16 +170,58 @@
 erator;%0A
+import javax.crypto.spec.IvParameterSpec;%0A
 import j
@@ -278,16 +278,51 @@
 ty.Key;%0A
+import java.security.SecureRandom;%0A
 import j
@@ -907,16 +907,86 @@
 ey key;%0A
+        static final SecureRandom SECURE_RANDOM = new SecureRandom();%0A
         
@@ -1374,24 +1374,221 @@
  %7D%0A        %0A
+        static IvParameterSpec generateIV() %7B%0A            byte%5B%5D bytes = new byte%5B16%5D;%0A            SECURE_RANDOM.nextBytes(bytes);%0A            return new IvParameterSpec(bytes);%0A        %7D%0A        %0A
     %7D%0A    %0A 
@@ -1720,20 +1720,16 @@
 dLocal =
- new
  ThreadL
@@ -1736,65 +1736,27 @@
 ocal
-%3C%3E();%0A        %0A        private static Cipher initCipher()
+.withInitial(() -%3E 
 %7B%0A  
@@ -1771,25 +1771,24 @@
   try %7B%0A    
-%0A
             
@@ -1787,35 +1787,22 @@
         
-    Cipher cipher =
+return
  Cipher.
@@ -1842,107 +1842,8 @@
 %22);%0A
-                cipher.init(Cipher.ENCRYPT_MODE, KeyContainer.key);%0A                return cipher;%0A
     
@@ -1871,16 +1871,17 @@
 ption e)
+ 
 %7B%0A      
@@ -1936,32 +1936,34 @@
      %7D%0A        %7D
+);
 %0A        %0A      
@@ -1964,23 +1964,16 @@
         
-public 
 Cipher g
@@ -1982,16 +1982,34 @@
 Cipher()
+ throws Exception 
 %7B%0A      
@@ -2022,21 +2022,21 @@
 Cipher c
-ached
+ipher
  = ciphe
@@ -2071,129 +2071,87 @@
     
-if (cached == null)%7B%0A                cached = 
+cipher.
 init
+(
 Cipher
-();%0A                cipherThreadLocal.set(cached
+.ENCRYPT_MODE, KeyContainer.key, KeyContainer.generateIV()
 );%0A
-            %7D%0A
     
@@ -2170,13 +2170,13 @@
 rn c
-ached
+ipher
 ;%0A  
@@ -2436,35 +2436,8 @@
 her(
-CipherContainer container, 
 Blac
@@ -2590,16 +2590,43 @@
 iner.key
+, KeyContainer.generateIV()
 );%0A     

Fix CryptoInit benchmark: should be initialized with different IV every time