Simplify instrumenter and tests (elastic#118493)

This commit simplifies the entitlements instrumentation service and instrumenter a bit. It especially removes some repetition in the instrumenter tests.
rjernst · Dec 13, 2024 · bf619a5 · bf619a5
1 parent d68d355
commit bf619a5
Show file tree

Hide file tree

Showing 8 changed files with 222 additions and 327 deletions.
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
@@ -58,30 +58,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();
-        String handleClass = checkerClass + "Handle";
-        String getCheckerClassMethodDescriptor = Type.getMethodDescriptor(Type.getObjectType(checkerClass));
+    public static InstrumenterImpl create(Class<?> checkerClass, Map<MethodKey, CheckMethod> checkMethods) {
+        Type checkerClassType = Type.getType(checkerClass);
+        String handleClass = checkerClassType.getInternalName() + "Handle";
+        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;
@@ -306,5 +290,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(