From f5f0ac0f25a8af88e09e7b3fcd5fa42856c79b7e Mon Sep 17 00:00:00 2001
From: Finn Carroll <carrofin@amazon.com>
Date: Mon, 23 Dec 2024 15:01:45 -0800
Subject: [PATCH] Extract jars to sub dirs during thirdPartyAudit task.

Signed-off-by: Finn Carroll <carrofin@amazon.com>
---
 .../gradle/precommit/ThirdPartyAuditTask.java | 33 ++++++++++++++-----
 1 file changed, 24 insertions(+), 9 deletions(-)

diff --git a/buildSrc/src/main/java/org/opensearch/gradle/precommit/ThirdPartyAuditTask.java b/buildSrc/src/main/java/org/opensearch/gradle/precommit/ThirdPartyAuditTask.java
index a74781ac44720..78304a4c7704d 100644
--- a/buildSrc/src/main/java/org/opensearch/gradle/precommit/ThirdPartyAuditTask.java
+++ b/buildSrc/src/main/java/org/opensearch/gradle/precommit/ThirdPartyAuditTask.java
@@ -229,8 +229,7 @@ public Set<File> getJarsToScan() {
     @TaskAction
     public void runThirdPartyAudit() throws IOException {
         Set<File> jars = getJarsToScan();
-
-        extractJars(jars);
+        Set<File> extractedJars = extractJars(jars);
 
         final String forbiddenApisOutput = runForbiddenAPIsCli();
 
@@ -248,7 +247,7 @@ public void runThirdPartyAudit() throws IOException {
 
         Set<String> jdkJarHellClasses = null;
         if (this.jarHellEnabled) {
-            jdkJarHellClasses = runJdkJarHellCheck();
+            jdkJarHellClasses = runJdkJarHellCheck(extractedJars);
         }
 
         if (missingClassExcludes != null) {
@@ -301,16 +300,26 @@ private void logForbiddenAPIsOutput(String forbiddenApisOutput) {
         getLogger().error("Forbidden APIs output:\n{}==end of forbidden APIs==", forbiddenApisOutput);
     }
 
-    private void extractJars(Set<File> jars) {
+    /**
+     * Extract project jars to build directory as specified by getJarExpandDir.
+     * Handle multi release jars by keeping versions closest to `targetCompatibility` version.
+     * @param jars to extract to build dir
+     * @return File set of extracted jars
+     */
+    private Set<File> extractJars(Set<File> jars) {
+        Set<File> extractedJars = new TreeSet<>();
         File jarExpandDir = getJarExpandDir();
         // We need to clean up to make sure old dependencies don't linger
         getProject().delete(jarExpandDir);
 
         jars.forEach(jar -> {
+            String jarPrefix = jar.getName().replaceAll("-\\d+\\.\\d+\\.\\d+\\.jar$", "");
+            File jarSubDir = new File(jarExpandDir, jarPrefix);
+            extractedJars.add(jarSubDir);
             FileTree jarFiles = getProject().zipTree(jar);
             getProject().copy(spec -> {
                 spec.from(jarFiles);
-                spec.into(jarExpandDir);
+                spec.into(jarSubDir);
                 // exclude classes from multi release jars
                 spec.exclude("META-INF/versions/**");
             });
@@ -329,7 +338,7 @@ private void extractJars(Set<File> jars) {
                 Integer.parseInt(targetCompatibility.get().getMajorVersion())
             ).forEach(majorVersion -> getProject().copy(spec -> {
                 spec.from(getProject().zipTree(jar));
-                spec.into(jarExpandDir);
+                spec.into(jarSubDir);
                 String metaInfPrefix = "META-INF/versions/" + majorVersion;
                 spec.include(metaInfPrefix + "/**");
                 // Drop the version specific prefix
@@ -337,6 +346,8 @@ private void extractJars(Set<File> jars) {
                 spec.setIncludeEmptyDirs(false);
             }));
         });
+
+        return extractedJars;
     }
 
     private void assertNoJarHell(Set<String> jdkJarHellClasses) {
@@ -398,7 +409,12 @@ private String runForbiddenAPIsCli() throws IOException {
         return forbiddenApisOutput;
     }
 
-    private Set<String> runJdkJarHellCheck() throws IOException {
+    /**
+     * Execute java with JDK_JAR_HELL_MAIN_CLASS against provided jars with OpenSearch core in the classpath.
+     * @param jars to scan for jarHell violations.
+     * @return standard out of jarHell process.
+     */
+    private Set<String> runJdkJarHellCheck(Set<File> jars) throws IOException {
         ByteArrayOutputStream standardOut = new ByteArrayOutputStream();
         InjectedExecOps execOps = getProject().getObjects().newInstance(InjectedExecOps.class);
         ExecResult execResult = execOps.getExecOps().javaexec(spec -> {
@@ -407,9 +423,8 @@ private Set<String> runJdkJarHellCheck() throws IOException {
                 getRuntimeConfiguration(),
                 getProject().getConfigurations().getByName(CompileOnlyResolvePlugin.RESOLVEABLE_COMPILE_ONLY_CONFIGURATION_NAME)
             );
-
             spec.getMainClass().set(JDK_JAR_HELL_MAIN_CLASS);
-            spec.args(getJarExpandDir());
+            spec.args(jars);
             spec.setIgnoreExitValue(true);
             if (javaHome != null) {
                 spec.setExecutable(javaHome + "/bin/java");