Fix hdfs-related IT tests for java24

build-tools-internal/src/main/java/org/elasticsearch/gradle/internal/dependencies/patches/hdfs/HdfsClassPatcher.java

@@ -0,0 +1,147 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.gradle.internal.dependencies.patches.hdfs;
+
+import org.gradle.api.artifacts.transform.CacheableTransform;
+import org.gradle.api.artifacts.transform.InputArtifact;
+import org.gradle.api.artifacts.transform.TransformAction;
+import org.gradle.api.artifacts.transform.TransformOutputs;
+import org.gradle.api.artifacts.transform.TransformParameters;
+import org.gradle.api.file.FileSystemLocation;
+import org.gradle.api.provider.Provider;
+import org.gradle.api.tasks.Classpath;
+import org.gradle.api.tasks.Input;
+import org.gradle.api.tasks.Optional;
+import org.jetbrains.annotations.NotNull;
+import org.objectweb.asm.ClassReader;
+import org.objectweb.asm.ClassVisitor;
+import org.objectweb.asm.ClassWriter;
+import java.io.File;
+import java.io.FileOutputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.util.Enumeration;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Locale;
+import java.util.Map;
+import java.util.function.Function;
+import java.util.jar.JarEntry;
+import java.util.jar.JarFile;
+import java.util.jar.JarOutputStream;
+import java.util.stream.Stream;
+
+import static java.util.Map.entry;
+
+
+@CacheableTransform
+public abstract class HdfsClassPatcher implements TransformAction<HdfsClassPatcher.Parameters> {
+
+    record JarPatchers(String artifactName, Map<String, Function<ClassWriter, ClassVisitor>> jarPatchers) {}
+
+    static final List<JarPatchers> allPatchers = List.of(
+        new JarPatchers(
+            "hadoop-common",
+            Map.ofEntries(
+                entry("org/apache/hadoop/util/ShutdownHookManager.class", ShutdownHookManagerPatcher::new),
+                entry("org/apache/hadoop/util/Shell.class", ShellPatcher::new),
+                entry("org/apache/hadoop/security/UserGroupInformation.class", SubjectGetSubjectPatcher::new)
+            )
+        ),
+        new JarPatchers(
+            "hadoop-auth",
+            Map.of(
+                "org/apache/hadoop/security/authentication/client/KerberosAuthenticator.class",
+                SubjectGetSubjectPatcher::new
+            )
+        )
+    );
+
+    interface Parameters extends TransformParameters {
+        @Input
+        @Optional
+        List<String> getMatchingArtifacts();
+
+        void setMatchingArtifacts(List<String> matchingArtifacts);
+    }
+
+    @Classpath
+    @InputArtifact
+    public abstract Provider<FileSystemLocation> getInputArtifact();
+
+    @Override
+    public void transform(@NotNull TransformOutputs outputs) {
+        File inputFile = getInputArtifact().get().getAsFile();
+
+        List<String> matchingArtifacts = getParameters().getMatchingArtifacts();
+        if (matchingArtifacts.isEmpty() == false &&
+            matchingArtifacts.stream().noneMatch(supported -> inputFile.getName().contains(supported))
+        ) {
+            outputs.file(getInputArtifact());
+        } else {
+            Stream<JarPatchers> patchersToApply = allPatchers.stream().filter(jp -> matchingArtifacts.contains(jp.artifactName()));
+            patchersToApply.forEach(patchers -> {
+                Map<String, Function<ClassWriter, ClassVisitor>> jarPatchers = new HashMap<>(patchers.jarPatchers());
+                File outputFile = outputs.file(inputFile.getName().replace(".jar", "-patched.jar"));
+
+                patchJar(inputFile, outputFile, jarPatchers);
+
+                if (jarPatchers.isEmpty() == false) {
+                    throw new IllegalArgumentException(
+                        String.format(
+                            Locale.ROOT,
+                            "error patching [%s] with [%s]: the jar does not contain [%s]",
+                            inputFile.getName(),
+                            patchers.artifactName(),
+                            String.join(", ", jarPatchers.keySet())
+                        )
+                    );
+                }
+            });
+        }
+    }
+
+    private static void patchJar(File inputFile, File outputFile, Map<String, Function<ClassWriter, ClassVisitor>> jarPatchers) {
+        try (JarFile jarFile = new JarFile(inputFile); JarOutputStream jos = new JarOutputStream(new FileOutputStream(outputFile))) {
+            Enumeration<JarEntry> entries = jarFile.entries();
+            while (entries.hasMoreElements()) {
+                JarEntry entry = entries.nextElement();
+                String entryName = entry.getName();
+                // Add the entry to the new JAR file
+                jos.putNextEntry(new JarEntry(entryName));
+                System.out.println("EntryName = " + entryName);
+
+                Function<ClassWriter, ClassVisitor> classPatcher = jarPatchers.remove(entryName);
+                if (classPatcher != null) {
+                    System.out.println("Patching " + entryName);
+                    byte[] classToPatch = jarFile.getInputStream(entry).readAllBytes();
+
+                    ClassReader classReader = new ClassReader(classToPatch);
+                    ClassWriter classWriter = new ClassWriter(classReader, 0);
+                    classReader.accept(classPatcher.apply(classWriter), 0);
+
+                    jos.write(classWriter.toByteArray());
+                } else {
+                    // Read the entry's data and write it to the new JAR
+                    try (InputStream is = jarFile.getInputStream(entry)) {
+                        byte[] buffer = new byte[1024];
+                        int bytesRead;
+                        while ((bytesRead = is.read(buffer)) != -1) {
+                            jos.write(buffer, 0, bytesRead);
+                        }
+                    }
+                }
+                jos.closeEntry();
+            }
+        } catch (IOException ex) {
+            throw new RuntimeException(ex);
+        }
+    }
+}

plugins/repository-hdfs/hadoop-client-api/src/patcher/java/org/elasticsearch/hdfs/patch/MethodReplacement.java → build-tools-internal/src/main/java/org/elasticsearch/gradle/internal/dependencies/patches/hdfs/MethodReplacement.java

@@ -7,7 +7,7 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
-package org.elasticsearch.hdfs.patch;
+package org.elasticsearch.gradle.internal.dependencies.patches.hdfs;
 
 import org.objectweb.asm.MethodVisitor;
 import org.objectweb.asm.Opcodes;

plugins/repository-hdfs/hadoop-client-api/src/patcher/java/org/elasticsearch/hdfs/patch/ShellPatcher.java → build-tools-internal/src/main/java/org/elasticsearch/gradle/internal/dependencies/patches/hdfs/ShellPatcher.java

@@ -7,7 +7,7 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
-package org.elasticsearch.hdfs.patch;
+package org.elasticsearch.gradle.internal.dependencies.patches.hdfs;
 
 import org.objectweb.asm.ClassVisitor;
 import org.objectweb.asm.ClassWriter;

plugins/repository-hdfs/hadoop-client-api/src/patcher/java/org/elasticsearch/hdfs/patch/ShutdownHookManagerPatcher.java → build-tools-internal/src/main/java/org/elasticsearch/gradle/internal/dependencies/patches/hdfs/ShutdownHookManagerPatcher.java

@@ -7,7 +7,7 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
-package org.elasticsearch.hdfs.patch;
+package org.elasticsearch.gradle.internal.dependencies.patches.hdfs;
 
 import org.objectweb.asm.ClassVisitor;
 import org.objectweb.asm.ClassWriter;

plugins/repository-hdfs/hadoop-client-api/src/patcher/java/org/elasticsearch/hdfs/patch/SubjectGetSubjectPatcher.java → build-tools-internal/src/main/java/org/elasticsearch/gradle/internal/dependencies/patches/hdfs/SubjectGetSubjectPatcher.java

@@ -7,7 +7,7 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
-package org.elasticsearch.hdfs.patch;
+package org.elasticsearch.gradle.internal.dependencies.patches.hdfs;
 
 import org.objectweb.asm.ClassVisitor;
 import org.objectweb.asm.ClassWriter;

plugins/repository-hdfs/build.gradle

@@ -22,12 +22,20 @@ versions << [
   'hadoop': '3.4.1'
 ]
 
+def patched = Attribute.of('patched', Boolean)
+
 configurations {
   hdfsFixture2
   hdfsFixture3
+  hdfsClientApi {
+    attributes {
+      attribute(patched, true)
+    }
+  }
 }
 
 dependencies {
+  hdfsClientApi "org.apache.hadoop:hadoop-client-api:${versions.hadoop}"
   api project(path: 'hadoop-client-api', configuration: 'default')
   if (isEclipse) {
     /*
@@ -69,6 +77,20 @@ dependencies {
 
   hdfsFixture2 project(path: ':test:fixtures:hdfs-fixture', configuration: 'shadowedHdfs2')
   hdfsFixture3 project(path: ':test:fixtures:hdfs-fixture', configuration: 'shadow')
+
+  attributesSchema {
+    attribute(patched)
+  }
+  artifactTypes.getByName("jar") {
+    attributes.attribute(patched, false)
+  }
+  registerTransform(org.elasticsearch.gradle.internal.dependencies.patches.hdfs.HdfsClassPatcher) {
+    from.attribute(patched, false)
+    to.attribute(patched, true)
+    parameters {
+      matchingArtifacts = ["hadoop-common", "hadoop-auth"]
+    }
+  }
 }
 
 restResources {

plugins/repository-hdfs/hadoop-client-api/build.gradle

@@ -2,53 +2,32 @@ import org.gradle.api.file.ArchiveOperations
 
 apply plugin: 'elasticsearch.java'
 
-sourceSets {
-  patcher
-}
+def patched = Attribute.of('patched', Boolean)
 
 configurations {
   thejar {
     canBeResolved = true
+    attributes {
+      attribute(patched, true)
+    }
   }
 }
 
 dependencies {
   thejar("org.apache.hadoop:hadoop-client-api:${project.parent.versions.hadoop}") {
     transitive = false
   }
-
-  patcherImplementation 'org.ow2.asm:asm:9.7.1'
-  patcherImplementation 'org.ow2.asm:asm-tree:9.7.1'
-}
-
-def outputDir = layout.buildDirectory.dir("patched-classes")
-
-def patchTask = tasks.register("patchClasses", JavaExec) {
-  inputs.files(configurations.thejar).withPathSensitivity(PathSensitivity.RELATIVE)
-  inputs.files(sourceSets.patcher.output).withPathSensitivity(PathSensitivity.RELATIVE)
-  outputs.dir(outputDir)
-  classpath = sourceSets.patcher.runtimeClasspath
-  mainClass = 'org.elasticsearch.hdfs.patch.HdfsClassPatcher'
-  def thejar = configurations.thejar
-  doFirst {
-    args(thejar.singleFile, outputDir.get().asFile)
+  attributesSchema {
+    attribute(patched)
   }
-}
-
-interface InjectedArchiveOps {
-  @Inject ArchiveOperations getArchiveOperations()
-}
-
-tasks.named('jar').configure {
-  dependsOn(configurations.thejar)
-  def injected = project.objects.newInstance(InjectedArchiveOps)
-  def thejar = configurations.thejar
-  from(patchTask)
-  from({ injected.getArchiveOperations().zipTree(thejar.singleFile) }) {
-    eachFile {
-      if (outputDir.get().file(it.relativePath.pathString).asFile.exists()) {
-        it.exclude()
-      }
+  artifactTypes.getByName("jar") {
+    attributes.attribute(patched, false)
+  }
+  registerTransform(org.elasticsearch.gradle.internal.dependencies.patches.hdfs.HdfsClassPatcher) {
+    from.attribute(patched, false)
+    to.attribute(patched, true)
+    parameters {
+      matchingArtifacts = ["hadoop-common", "hadoop-auth"]
     }
   }
 }

test/fixtures/hdfs-fixture/build.gradle

@@ -13,8 +13,29 @@ apply plugin: 'com.gradleup.shadow'
 
 import com.github.jengelman.gradle.plugins.shadow.tasks.ShadowJar
 
+def patched = Attribute.of('patched', Boolean)
+dependencies {
+  attributesSchema {
+    attribute(patched)
+  }
+  artifactTypes.getByName("jar") {
+    attributes.attribute(patched, false)
+  }
+  registerTransform(org.elasticsearch.gradle.internal.dependencies.patches.hdfs.HdfsClassPatcher) {
+    from.attribute(patched, false)
+    to.attribute(patched, true)
+    parameters {
+      matchingArtifacts = ["hadoop-common"]
+    }
+  }
+}
+
 configurations {
-  hdfs2
+  hdfs2 {
+    attributes {
+      attribute(patched, true)
+    }
+  }
   hdfs3
   consumable("shadowedHdfs2")
 }