From 2614d946c8cf56ecd03808d283e672cc4b4a48ac Mon Sep 17 00:00:00 2001
From: John Jiang <johnsjiang@tencent.com>
Date: Fri, 27 Dec 2024 10:59:25 +0800
Subject: [PATCH] TKSS-1016: NativeSM2KeyAgreement should be created in
 constructor

---
 .../kona/crypto/perf/SM2KeyAgreementPerfTest.java    | 12 ++++++------
 .../crypto/provider/nativeImpl/SM2KeyAgreement.java  | 11 ++++++-----
 2 files changed, 12 insertions(+), 11 deletions(-)

diff --git a/kona-crypto/src/jmh/java/com/tencent/kona/crypto/perf/SM2KeyAgreementPerfTest.java b/kona-crypto/src/jmh/java/com/tencent/kona/crypto/perf/SM2KeyAgreementPerfTest.java
index 33f2efd5..4760fd19 100644
--- a/kona-crypto/src/jmh/java/com/tencent/kona/crypto/perf/SM2KeyAgreementPerfTest.java
+++ b/kona-crypto/src/jmh/java/com/tencent/kona/crypto/perf/SM2KeyAgreementPerfTest.java
@@ -26,7 +26,6 @@
 import org.openjdk.jmh.annotations.*;
 
 import javax.crypto.KeyAgreement;
-import java.security.InvalidKeyException;
 import java.util.concurrent.TimeUnit;
 
 import static com.tencent.kona.crypto.CryptoUtils.toBytes;
@@ -72,11 +71,12 @@ public static class KeyAgreementHolder {
         @Param({"KonaCrypto", "KonaCrypto-Native"})
         String provider;
 
+        SM2KeyAgreementParamSpec paramSpec;
         KeyAgreement keyAgreement;
 
-        @Setup(Level.Invocation)
+        @Setup(Level.Trial)
         public void setup() throws Exception {
-            SM2KeyAgreementParamSpec paramSpec = new SM2KeyAgreementParamSpec(
+            paramSpec = new SM2KeyAgreementParamSpec(
                     toBytes(ID),
                     new SM2PrivateKey(toBytes(PRI_KEY)),
                     new SM2PublicKey(toBytes(PUB_KEY)),
@@ -85,13 +85,13 @@ public void setup() throws Exception {
                     true,
                     16);
             keyAgreement = KeyAgreement.getInstance("SM2", provider);
-            keyAgreement.init(
-                    new SM2PrivateKey(toBytes(TMP_PRI_KEY)), paramSpec);
         }
     }
 
     @Benchmark
-    public byte[] generateSecret(KeyAgreementHolder holder) throws InvalidKeyException {
+    public byte[] generateSecret(KeyAgreementHolder holder) throws Exception {
+        holder.keyAgreement.init(
+                new SM2PrivateKey(toBytes(TMP_PRI_KEY)), holder.paramSpec);
         holder.keyAgreement.doPhase(new SM2PublicKey(toBytes(PEER_TMP_PUB_KEY)), true);
         return holder.keyAgreement.generateSecret();
     }
diff --git a/kona-crypto/src/main/java/com/tencent/kona/crypto/provider/nativeImpl/SM2KeyAgreement.java b/kona-crypto/src/main/java/com/tencent/kona/crypto/provider/nativeImpl/SM2KeyAgreement.java
index 49e31dfc..b2c78767 100644
--- a/kona-crypto/src/main/java/com/tencent/kona/crypto/provider/nativeImpl/SM2KeyAgreement.java
+++ b/kona-crypto/src/main/java/com/tencent/kona/crypto/provider/nativeImpl/SM2KeyAgreement.java
@@ -54,7 +54,12 @@ public final class SM2KeyAgreement extends KeyAgreementSpi {
     private SM2KeyAgreementParamSpec paramSpec;
     private SM2PublicKey peerEphemeralPublicKey;
 
-    private NativeSM2KeyAgreement sm2;
+    private final NativeSM2KeyAgreement sm2;
+
+    public SM2KeyAgreement() {
+        sm2 = new NativeSM2KeyAgreement();
+        SWEEPER.register(this, new SweepNativeRef(sm2));
+    }
 
     @Override
     protected void engineInit(Key key, SecureRandom random) {
@@ -85,10 +90,6 @@ protected void engineInit(Key key, AlgorithmParameterSpec params,
         ephemeralPrivateKey = new SM2PrivateKey((ECPrivateKey) key);
         paramSpec = (SM2KeyAgreementParamSpec) params;
         peerEphemeralPublicKey = null;
-
-        sm2 = new NativeSM2KeyAgreement();
-
-        SWEEPER.register(this, new SweepNativeRef(sm2));
     }
 
     @Override