Simplify instrumenter and tests

This commit simplifies the entitlements instrumentation service and instrumenter a bit. It especially removes some repetition in the instrumenter tests.
rjernst · Dec 11, 2024 · 0a2c86e · 0a2c86e
1 parent 9837e78
commit 0a2c86e
Show file tree

Hide file tree

Showing 9 changed files with 231 additions and 326 deletions.
diff --git a/libs/entitlement/asm-provider/src/main/java/module-info.java b/libs/entitlement/asm-provider/src/main/java/module-info.java
@@ -13,6 +13,7 @@
 module org.elasticsearch.entitlement.instrumentation {
     requires org.objectweb.asm;
     requires org.elasticsearch.entitlement;
+    requires java.logging;
 
     provides InstrumentationService with InstrumentationServiceImpl;
 }
diff --git a/...n/java/org/elasticsearch/entitlement/instrumentation/impl/InstrumentationServiceImpl.java b/...n/java/org/elasticsearch/entitlement/instrumentation/impl/InstrumentationServiceImpl.java
@@ -29,15 +29,13 @@
 public class InstrumentationServiceImpl implements InstrumentationService {
 
     @Override
-    public Instrumenter newInstrumenter(Map<MethodKey, CheckMethod> checkMethods) {
-        return InstrumenterImpl.create(checkMethods);
+    public Instrumenter newInstrumenter(Class<?> clazz, Map<MethodKey, CheckMethod> methods) {
+        return InstrumenterImpl.create(clazz, methods);
     }
 
     @Override
-    public Map<MethodKey, CheckMethod> lookupMethodsToInstrument(String entitlementCheckerClassName) throws ClassNotFoundException,
-        IOException {
+    public Map<MethodKey, CheckMethod> lookupMethods(Class<?> checkerClass) throws IOException {
         var methodsToInstrument = new HashMap<MethodKey, CheckMethod>();
-        var checkerClass = Class.forName(entitlementCheckerClassName);
         var classFileInfo = InstrumenterImpl.getClassFileInfo(checkerClass);
         ClassReader reader = new ClassReader(classFileInfo.bytecodes());
         ClassVisitor visitor = new ClassVisitor(Opcodes.ASM9) {

diff --git a/...er/src/main/java/org/elasticsearch/entitlement/instrumentation/impl/InstrumenterImpl.java b/...er/src/main/java/org/elasticsearch/entitlement/instrumentation/impl/InstrumenterImpl.java
@@ -25,6 +25,7 @@
 import java.io.IOException;
 import java.io.InputStream;
 import java.util.Map;
+import java.util.logging.LogManager;
 import java.util.stream.Stream;
 
 import static org.objectweb.asm.ClassWriter.COMPUTE_FRAMES;
@@ -58,30 +59,14 @@ public class InstrumenterImpl implements Instrumenter {
         this.checkMethods = checkMethods;
     }
 
-    static String getCheckerClassName() {
-        int javaVersion = Runtime.version().feature();
-        final String classNamePrefix;
-        if (javaVersion >= 23) {
-            classNamePrefix = "Java23";
-        } else {
-            classNamePrefix = "";
-        }
-        return "org/elasticsearch/entitlement/bridge/" + classNamePrefix + "EntitlementChecker";
-    }
-
-    public static InstrumenterImpl create(Map<MethodKey, CheckMethod> checkMethods) {
-        String checkerClass = getCheckerClassName();
+    public static InstrumenterImpl create(Class<?> checkerClass, Map<MethodKey, CheckMethod> checkMethods) {
+        Type checkerClassType = Type.getType(checkerClass);
         String handleClass = checkerClass + "Handle";
-        String getCheckerClassMethodDescriptor = Type.getMethodDescriptor(Type.getObjectType(checkerClass));
+        String getCheckerClassMethodDescriptor = Type.getMethodDescriptor(checkerClassType);
         return new InstrumenterImpl(handleClass, getCheckerClassMethodDescriptor, "", checkMethods);
     }
 
-    public ClassFileInfo instrumentClassFile(Class<?> clazz) throws IOException {
-        ClassFileInfo initial = getClassFileInfo(clazz);
-        return new ClassFileInfo(initial.fileName(), instrumentClass(Type.getInternalName(clazz), initial.bytecodes()));
-    }
-
-    public static ClassFileInfo getClassFileInfo(Class<?> clazz) throws IOException {
+    static ClassFileInfo getClassFileInfo(Class<?> clazz) throws IOException {
         String internalName = Type.getInternalName(clazz);
         String fileName = "/" + internalName + ".class";
         byte[] originalBytecodes;
@@ -171,7 +156,7 @@ public MethodVisitor visitMethod(int access, String name, String descriptor, Str
                 var key = new MethodKey(className, name, Stream.of(Type.getArgumentTypes(descriptor)).map(Type::getInternalName).toList());
                 var instrumentationMethod = checkMethods.get(key);
                 if (instrumentationMethod != null) {
-                    // LOGGER.debug("Will instrument method {}", key);
+                    System.out.println("Will instrument method " + name);
                     return new EntitlementMethodVisitor(Opcodes.ASM9, mv, isStatic, isCtor, descriptor, instrumentationMethod);
                 } else {
                     // LOGGER.trace("Will not instrument method {}", key);
@@ -306,5 +291,5 @@ protected void pushEntitlementChecker(MethodVisitor mv) {
         mv.visitMethodInsn(INVOKESTATIC, handleClass, "instance", getCheckerClassMethodDescriptor, false);
     }
 
-    public record ClassFileInfo(String fileName, byte[] bytecodes) {}
+    record ClassFileInfo(String fileName, byte[] bytecodes) {}
 }
diff --git a/...a/org/elasticsearch/entitlement/instrumentation/impl/InstrumentationServiceImplTests.java b/...a/org/elasticsearch/entitlement/instrumentation/impl/InstrumentationServiceImplTests.java
@@ -51,8 +51,8 @@ interface TestCheckerCtors {
         void check$org_example_TestTargetClass$(Class<?> clazz, int x, String y);
     }
 
-    public void testInstrumentationTargetLookup() throws IOException, ClassNotFoundException {
-        Map<MethodKey, CheckMethod> checkMethods = instrumentationService.lookupMethodsToInstrument(TestChecker.class.getName());
+    public void testInstrumentationTargetLookup() throws IOException {
+        Map<MethodKey, CheckMethod> checkMethods = instrumentationService.lookupMethods(TestChecker.class);
 
         assertThat(checkMethods, aMapWithSize(3));
         assertThat(
@@ -116,8 +116,8 @@ public void testInstrumentationTargetLookup() throws IOException, ClassNotFoundE
         );
     }
 
-    public void testInstrumentationTargetLookupWithOverloads() throws IOException, ClassNotFoundException {
-        Map<MethodKey, CheckMethod> checkMethods = instrumentationService.lookupMethodsToInstrument(TestCheckerOverloads.class.getName());
+    public void testInstrumentationTargetLookupWithOverloads() throws IOException {
+        Map<MethodKey, CheckMethod> checkMethods = instrumentationService.lookupMethods(TestCheckerOverloads.class);
 
         assertThat(checkMethods, aMapWithSize(2));
         assertThat(
@@ -148,8 +148,8 @@ public void testInstrumentationTargetLookupWithOverloads() throws IOException, C
         );
     }
 
-    public void testInstrumentationTargetLookupWithCtors() throws IOException, ClassNotFoundException {
-        Map<MethodKey, CheckMethod> checkMethods = instrumentationService.lookupMethodsToInstrument(TestCheckerCtors.class.getName());
+    public void testInstrumentationTargetLookupWithCtors() throws IOException {
+        Map<MethodKey, CheckMethod> checkMethods = instrumentationService.lookupMethods(TestCheckerCtors.class);
 
         assertThat(checkMethods, aMapWithSize(2));
         assertThat(