From 328e589cdba5d549c19ccc5b7096794eb8fa0229 Mon Sep 17 00:00:00 2001
From: WilburZjh <wilburzjh@gmail.com>
Date: Sun, 3 Sep 2023 21:45:12 -0400
Subject: [PATCH] Optimize ChaCha20-Poly1305 initialization

---
 .../crypto/provider/NativeChaCha20Cipher.java | 18 ++++++++++--
 .../jdk/crypto/jniprovider/NativeCrypto.java  |  3 +-
 .../share/native/libjncrypto/NativeCrypto.c   | 28 ++++++++++++-------
 3 files changed, 36 insertions(+), 13 deletions(-)

diff --git a/closed/src/java.base/share/classes/com/sun/crypto/provider/NativeChaCha20Cipher.java b/closed/src/java.base/share/classes/com/sun/crypto/provider/NativeChaCha20Cipher.java
index 97bcee128ad..0a2ed87908d 100644
--- a/closed/src/java.base/share/classes/com/sun/crypto/provider/NativeChaCha20Cipher.java
+++ b/closed/src/java.base/share/classes/com/sun/crypto/provider/NativeChaCha20Cipher.java
@@ -24,7 +24,7 @@
  */
 /*
  * ===========================================================================
- * (c) Copyright IBM Corp. 2018, 2021 All Rights Reserved
+ * (c) Copyright IBM Corp. 2018, 2023 All Rights Reserved
  * ===========================================================================
  */
 package com.sun.crypto.provider;
@@ -93,6 +93,9 @@ abstract class NativeChaCha20Cipher extends CipherSpi {
     private static final Cleaner contextCleaner;
     private final long context;
 
+    // The previous mode, initialized to an illegal value.
+    private int prevMode = Cipher.WRAP_MODE;
+
     private final ByteArrayOutputStream aadBuf;
 
     static {
@@ -600,7 +603,18 @@ private void init(int opmode, Key key, byte[] newNonce)
         aadDone = false;
         initialized = true;
 
-        int ret = nativeCrypto.ChaCha20Init(context, ossl_mode, openssl_iv, openssl_iv.length, keyBytes, keyBytes.length);
+        // Optimize the initialization of chacha20-poly1305 when they have the same ossl_mode
+        // otherwise, we have to reinitialize based upon the ossl_mode.
+        int ret;
+        if (prevMode == ossl_mode) {
+            ret = nativeCrypto.ChaCha20Init(context, ossl_mode, openssl_iv, openssl_iv.length, keyBytes, keyBytes.length, true);
+        } else {
+            ret = nativeCrypto.ChaCha20Init(context, ossl_mode, openssl_iv, openssl_iv.length, keyBytes, keyBytes.length, false);
+            prevMode = ossl_mode;
+        }
+        if (ret == -1) {
+            throw new ProviderException("Error in Native ChaCha20Cipher");
+        }
     }
 
     /**
diff --git a/closed/src/java.base/share/classes/jdk/crypto/jniprovider/NativeCrypto.java b/closed/src/java.base/share/classes/jdk/crypto/jniprovider/NativeCrypto.java
index 6db21b437bc..4ac58f14a35 100644
--- a/closed/src/java.base/share/classes/jdk/crypto/jniprovider/NativeCrypto.java
+++ b/closed/src/java.base/share/classes/jdk/crypto/jniprovider/NativeCrypto.java
@@ -321,7 +321,8 @@ public final native int ChaCha20Init(long context,
                                     byte[] iv,
                                     int ivlen,
                                     byte[] key,
-                                    int keylen);
+                                    int keylen,
+                                    boolean doReset);
 
     public final native int ChaCha20Update(long context,
                                        byte[] input,
diff --git a/closed/src/java.base/share/native/libjncrypto/NativeCrypto.c b/closed/src/java.base/share/native/libjncrypto/NativeCrypto.c
index 747cd3b5efb..5feb00c6a3f 100644
--- a/closed/src/java.base/share/native/libjncrypto/NativeCrypto.c
+++ b/closed/src/java.base/share/native/libjncrypto/NativeCrypto.c
@@ -2153,11 +2153,11 @@ int OSSL102_RSA_set0_crt_params(RSA *r2, BIGNUM *dmp1, BIGNUM *dmq1, BIGNUM *iqm
 /*
  * Class:     jdk_crypto_jniprovider_NativeCrypto
  * Method:    ChaCha20Init
- * Signature: (JI[BI[BI)I
+ * Signature: (JI[BI[BIZ)I
  */
 JNIEXPORT jint JNICALL Java_jdk_crypto_jniprovider_NativeCrypto_ChaCha20Init
   (JNIEnv *env, jobject thisObj, jlong c, jint mode, jbyteArray iv, jint ivLen,
-  jbyteArray key, jint key_len)
+  jbyteArray key, jint key_len, jboolean doReset)
 {
     EVP_CIPHER_CTX *ctx = (EVP_CIPHER_CTX*)(intptr_t) c;
     unsigned char *ivNative = NULL;
@@ -2170,12 +2170,18 @@ JNIEXPORT jint JNICALL Java_jdk_crypto_jniprovider_NativeCrypto_ChaCha20Init
     }
 
     if ((0 == mode) || (1 == mode)) {
-        evp_cipher1 = (*OSSL_chacha20_poly1305)();
+        /* Use the existing evp_cipher? */
+        if (JNI_FALSE == doReset) {
+            evp_cipher1 = (*OSSL_chacha20_poly1305)();
+        }
         encrypt = mode;
     } else if (2 == mode) {
+        /* Use the existing evp_cipher? */
+        if (JNI_FALSE == doReset) {
+            evp_cipher1 = (*OSSL_chacha20)();
+        }
         /* encrypt or decrypt does not matter */
         encrypt = 1;
-        evp_cipher1 = (*OSSL_chacha20)();
     } else {
         return -1;
     }
@@ -2200,12 +2206,14 @@ JNIEXPORT jint JNICALL Java_jdk_crypto_jniprovider_NativeCrypto_ChaCha20Init
     }
 
     /* if using Poly1305 */
-    if (2 != mode) {
-        if (1 != (*OSSL_CIPHER_CTX_ctrl)(ctx, EVP_CTRL_AEAD_SET_IVLEN, ivLen, NULL)) {
-            printErrors();
-            (*env)->ReleaseByteArrayElements(env, iv, (jbyte*)ivNative, JNI_ABORT);
-            (*env)->ReleaseByteArrayElements(env, key, (jbyte*)keyNative, JNI_ABORT);
-            return -1;
+    if (JNI_FALSE == doReset) {
+        if (2 != mode) {
+            if (1 != (*OSSL_CIPHER_CTX_ctrl)(ctx, EVP_CTRL_AEAD_SET_IVLEN, ivLen, NULL)) {
+                printErrors();
+                (*env)->ReleaseByteArrayElements(env, iv, (jbyte*)ivNative, JNI_ABORT);
+                (*env)->ReleaseByteArrayElements(env, key, (jbyte*)keyNative, JNI_ABORT);
+                return -1;
+            }
         }
     }