From c42d654c45c24868787a4831012e1ee78c157123 Mon Sep 17 00:00:00 2001
From: Amir Vakili <amirhosv@amazon.com>
Date: Tue, 6 Feb 2024 16:45:39 +0000
Subject: [PATCH] Benchmarking ACCP when its bundled with JDK

---
 benchmarks/README.md                          |  9 +++++
 benchmarks/lib/build.gradle.kts               | 19 ++++-------
 .../provider/benchmarks/BenchmarkUtils.java   | 34 +++++++++++++------
 3 files changed, 40 insertions(+), 22 deletions(-)

diff --git a/benchmarks/README.md b/benchmarks/README.md
index 786e0d0d..5d1d2834 100644
--- a/benchmarks/README.md
+++ b/benchmarks/README.md
@@ -25,3 +25,12 @@ Use `-PaccpLocalJar="PATH_TO_LOCAL_JAR"`:
 ```bash
  ./gradlew -PaccpLocalJar="../../build/cmake/AmazonCorrettoCryptoProvider.jar" lib:jmh
 ```
+
+### Benchmarking ACCP that is bundled in JDK
+
+Some customers bundle ACCP directly with their JDKs. To run the benchmarks with such a setup,
+one can use the following command:
+
+```bash
+ ./gradlew -PuseBundledAccp -Dorg.gradle.java.home=<PATH_TO_YOUR_CUSTOM_JDK> lib:jmh
+```
diff --git a/benchmarks/lib/build.gradle.kts b/benchmarks/lib/build.gradle.kts
index 0f2b2970..765328bb 100644
--- a/benchmarks/lib/build.gradle.kts
+++ b/benchmarks/lib/build.gradle.kts
@@ -32,21 +32,16 @@ dependencies {
         "2.+"
 
 
-    if (project.hasProperty("accpLocalJar")) {
-        jmh(files(accpLocalJar))
-    } else {
-        jmh("software.amazon.cryptools:${accpArtifactId}:$publishedAccpVersion:${osdetector.classifier}")
+    if (!project.hasProperty("useBundledAccp")) {
+        if (project.hasProperty("accpLocalJar")) {
+            jmh(files(accpLocalJar))
+        } else {
+            jmh("software.amazon.cryptools:${accpArtifactId}:$publishedAccpVersion:${osdetector.classifier}")
+        }
     }
 
 }
 
-// Apply a specific Java toolchain to ease working on different environments.
-java {
-    toolchain {
-        languageVersion.set(JavaLanguageVersion.of(11))
-    }
-}
-
 jmh {
     if (project.hasProperty("includeBenchmark")) {
         includes.add(includeBenchmark)
@@ -75,4 +70,4 @@ jmhReport {
 
 tasks.jmh {
     finalizedBy(tasks.jmhReport)
-}
\ No newline at end of file
+}
diff --git a/benchmarks/lib/src/jmh/java/com/amazon/corretto/crypto/provider/benchmarks/BenchmarkUtils.java b/benchmarks/lib/src/jmh/java/com/amazon/corretto/crypto/provider/benchmarks/BenchmarkUtils.java
index 04c1896b..1430475a 100644
--- a/benchmarks/lib/src/jmh/java/com/amazon/corretto/crypto/provider/benchmarks/BenchmarkUtils.java
+++ b/benchmarks/lib/src/jmh/java/com/amazon/corretto/crypto/provider/benchmarks/BenchmarkUtils.java
@@ -5,8 +5,10 @@
 import java.security.Provider;
 import java.security.SecureRandom;
 import java.security.Security;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.HashSet;
+import java.util.List;
 import java.util.Set;
 
 import com.amazon.corretto.crypto.provider.AmazonCorrettoCryptoProvider;
@@ -16,17 +18,29 @@ class BenchmarkUtils {
   private BenchmarkUtils() {}
 
   private static final SecureRandom sr = new SecureRandom();
-  private static final Set<String> NON_DEFAULT_PROVIDERS =
-      new HashSet(Arrays.asList("BC", "BCFIPS", "AmazonCorrettoCryptoProvider"));
-  private static final Provider[] DEFAULT_PROVIDERS;
+  private static final List<Provider> DEFAULT_PROVIDERS = new ArrayList<>();
+  private static AmazonCorrettoCryptoProvider accp = null;
+  private static BouncyCastleProvider bc = null;
 
   static {
-    DEFAULT_PROVIDERS = Security.getProviders();
-    for (Provider provider : DEFAULT_PROVIDERS) {
-      if (NON_DEFAULT_PROVIDERS.contains(provider.getName())) {
-        throw new RuntimeException("Provider prematurely (statically) registered: " + provider);
+    // For BC and ACCP, if they are installed statically, we just remove them.
+    for (Provider provider : Security.getProviders()) {
+      if ("AmazonCorrettoCryptoProvider".equals(provider.getName())) {
+        accp = (AmazonCorrettoCryptoProvider) provider;
+      } else if ("BC".equals(provider.getName())) {
+        bc = (BouncyCastleProvider) provider;
+      } else {
+        DEFAULT_PROVIDERS.add(provider);
       }
     }
+    if (accp == null) {
+      accp = AmazonCorrettoCryptoProvider.INSTANCE;
+    }
+    if (bc == null) {
+      bc = new BouncyCastleProvider();
+    }
+    removeAllProviders();
+    installDefaultProviders();
   }
 
   static byte[] getRandBytes(int n) {
@@ -45,11 +59,11 @@ static void setupProvider(String providerName) {
     switch (providerName) {
       case "AmazonCorrettoCryptoProvider":
         installDefaultProviders();
-        AmazonCorrettoCryptoProvider.install();
-        AmazonCorrettoCryptoProvider.INSTANCE.assertHealthy();
+        Security.insertProviderAt(accp, 1);
+        accp.assertHealthy();
         break;
       case "BC":
-        Security.insertProviderAt(new BouncyCastleProvider(), 1);
+        Security.insertProviderAt(bc, 1);
         break;
       case "SUN":
       case "SunEC":