Add getSubject -> current method replacement patcher (elastic#119779)

idegtiarenko · Jan 9, 2025 · 169cb44 · 169cb44
1 parent 16d8aaf
commit 169cb44
Show file tree

Hide file tree

Showing 2 changed files with 65 additions and 1 deletion.
diff --git a/...dfs/hadoop-client-api/src/patcher/java/org/elasticsearch/hdfs/patch/HdfsClassPatcher.java b/...dfs/hadoop-client-api/src/patcher/java/org/elasticsearch/hdfs/patch/HdfsClassPatcher.java
@@ -27,7 +27,11 @@ public class HdfsClassPatcher {
         "org/apache/hadoop/util/ShutdownHookManager.class",
         ShutdownHookManagerPatcher::new,
         "org/apache/hadoop/util/Shell.class",
-        ShellPatcher::new
+        ShellPatcher::new,
+        "org/apache/hadoop/security/UserGroupInformation.class",
+        SubjectGetSubjectPatcher::new,
+        "org/apache/hadoop/security/authentication/client/KerberosAuthenticator.class",
+        SubjectGetSubjectPatcher::new
     );
 
     public static void main(String[] args) throws Exception {

diff --git a/...op-client-api/src/patcher/java/org/elasticsearch/hdfs/patch/SubjectGetSubjectPatcher.java b/...op-client-api/src/patcher/java/org/elasticsearch/hdfs/patch/SubjectGetSubjectPatcher.java
@@ -0,0 +1,60 @@
+/*
+ * 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.hdfs.patch;
+
+import org.objectweb.asm.ClassVisitor;
+import org.objectweb.asm.ClassWriter;
+import org.objectweb.asm.MethodVisitor;
+import org.objectweb.asm.Type;
+
+import static org.objectweb.asm.Opcodes.ASM9;
+import static org.objectweb.asm.Opcodes.INVOKESTATIC;
+import static org.objectweb.asm.Opcodes.POP;
+
+class SubjectGetSubjectPatcher extends ClassVisitor {
+    SubjectGetSubjectPatcher(ClassWriter classWriter) {
+        super(ASM9, classWriter);
+    }
+
+    @Override
+    public MethodVisitor visitMethod(int access, String name, String descriptor, String signature, String[] exceptions) {
+        return new ReplaceCallMethodVisitor(super.visitMethod(access, name, descriptor, signature, exceptions), name, access, descriptor);
+    }
+
+    /**
+     * Replaces calls to Subject.getSubject(context); with calls to Subject.current();
+     */
+    private static class ReplaceCallMethodVisitor extends MethodVisitor {
+        private static final String SUBJECT_CLASS_INTERNAL_NAME = "javax/security/auth/Subject";
+        private static final String METHOD_NAME = "getSubject";
+
+        ReplaceCallMethodVisitor(MethodVisitor methodVisitor, String name, int access, String descriptor) {
+            super(ASM9, methodVisitor);
+        }
+
+        @Override
+        public void visitMethodInsn(int opcode, String owner, String name, String descriptor, boolean isInterface) {
+            if (opcode == INVOKESTATIC && SUBJECT_CLASS_INTERNAL_NAME.equals(owner) && METHOD_NAME.equals(name)) {
+                // Get rid of the extra arg on the stack
+                mv.visitInsn(POP);
+                // Call Subject.current()
+                mv.visitMethodInsn(
+                    INVOKESTATIC,
+                    SUBJECT_CLASS_INTERNAL_NAME,
+                    "current",
+                    Type.getMethodDescriptor(Type.getObjectType(SUBJECT_CLASS_INTERNAL_NAME)),
+                    false
+                );
+            } else {
+                super.visitMethodInsn(opcode, owner, name, descriptor, isInterface);
+            }
+        }
+    }
+}