Fix: methodkey for target function must use the method declaring class

Author: ldematte

libs/entitlement/asm-provider/src/main/java/org/elasticsearch/entitlement/instrumentation/impl/InstrumentationServiceImpl.java

@@ -118,7 +118,8 @@ public InstrumentationInfo lookupImplementationMethod(
     ) throws NoSuchMethodException, ClassNotFoundException {
 
         var targetMethod = targetSuperclass.getDeclaredMethod(methodName, parameterTypes);
-        validateTargetMethod(implementationClass, targetMethod);
+        var implementationMethod = implementationClass.getMethod(targetMethod.getName(), targetMethod.getParameterTypes());
+        validateTargetMethod(implementationClass, targetMethod, implementationMethod);
 
         var checkerAdditionalArguments = Stream.of(Class.class, targetSuperclass);
         var checkMethodArgumentTypes = Stream.concat(checkerAdditionalArguments, Arrays.stream(parameterTypes))
@@ -169,15 +170,15 @@ public MethodVisitor visitMethod(
 
         return new InstrumentationInfo(
             new MethodKey(
-                Type.getInternalName(implementationClass),
-                targetMethod.getName(),
+                Type.getInternalName(implementationMethod.getDeclaringClass()),
+                implementationMethod.getName(),
                 Arrays.stream(parameterTypes).map(c -> Type.getType(c).getInternalName()).toList()
             ),
             checkMethod[0]
         );
     }
 
-    private static void validateTargetMethod(Class<?> implementationClass, Method targetMethod) {
+    private static void validateTargetMethod(Class<?> implementationClass, Method targetMethod, Method implementationMethod) {
         if (targetMethod.getDeclaringClass().isAssignableFrom(implementationClass) == false) {
             throw new IllegalArgumentException(
                 String.format(
@@ -209,37 +210,26 @@ private static void validateTargetMethod(Class<?> implementationClass, Method ta
                 )
             );
         }
-        try {
-            var implementationMethod = implementationClass.getMethod(targetMethod.getName(), targetMethod.getParameterTypes());
-            var methodModifiers = implementationMethod.getModifiers();
-            if (Modifier.isAbstract(methodModifiers)) {
-                throw new IllegalArgumentException(
-                    String.format(
-                        Locale.ROOT,
-                        "Not a valid instrumentation method: %s is abstract in %s",
-                        targetMethod.getName(),
-                        implementationClass.getName()
-                    )
-                );
-            }
-            if (Modifier.isPublic(methodModifiers) == false) {
-                throw new IllegalArgumentException(
-                    String.format(
-                        Locale.ROOT,
-                        "Not a valid instrumentation method: %s is not public in %s",
-                        targetMethod.getName(),
-                        implementationClass.getName()
-                    )
-                );
-            }
-        } catch (NoSuchMethodException e) {
-            assert false
-                : String.format(
+        var methodModifiers = implementationMethod.getModifiers();
+        if (Modifier.isAbstract(methodModifiers)) {
+            throw new IllegalArgumentException(
+                String.format(
                     Locale.ROOT,
-                    "Not a valid instrumentation method: %s cannot be found in %s",
+                    "Not a valid instrumentation method: %s is abstract in %s",
                     targetMethod.getName(),
                     implementationClass.getName()
-                );
+                )
+            );
+        }
+        if (Modifier.isPublic(methodModifiers) == false) {
+            throw new IllegalArgumentException(
+                String.format(
+                    Locale.ROOT,
+                    "Not a valid instrumentation method: %s is not public in %s",
+                    targetMethod.getName(),
+                    implementationClass.getName()
+                )
+            );
         }
     }
 

libs/entitlement/asm-provider/src/test/java/org/elasticsearch/entitlement/instrumentation/impl/InstrumentationServiceImplTests.java

@@ -40,9 +40,15 @@ public void instanceMethod(int x, String y) {}
 
     abstract static class TestTargetBaseClass {
         abstract void instanceMethod(int x, String y);
+        abstract void instanceMethod2(int x, String y);
     }
 
-    static class TestTargetImplementationClass extends TestTargetBaseClass {
+    abstract static class TestTargetIntermediateClass extends TestTargetBaseClass {
+        @Override
+        public void instanceMethod2(int x, String y) {}
+    }
+
+    static class TestTargetImplementationClass extends TestTargetIntermediateClass {
         @Override
         public void instanceMethod(int x, String y) {}
     }
@@ -364,6 +370,44 @@ public void testLookupImplementationMethodWithBaseClass() throws ClassNotFoundEx
         );
     }
 
+    public void testLookupImplementationMethodWithInheritance() throws ClassNotFoundException, NoSuchMethodException {
+        var info = instrumentationService.lookupImplementationMethod(
+            TestTargetBaseClass.class,
+            "instanceMethod2",
+            TestTargetImplementationClass.class,
+            TestCheckerMixed.class,
+            "checkInstanceMethodManual",
+            int.class,
+            String.class
+        );
+
+        assertThat(
+            info.targetMethod(),
+            equalTo(
+                new MethodKey(
+                    "org/elasticsearch/entitlement/instrumentation/impl/InstrumentationServiceImplTests$TestTargetIntermediateClass",
+                    "instanceMethod2",
+                    List.of("I", "java/lang/String")
+                )
+            )
+        );
+        assertThat(
+            info.checkMethod(),
+            equalTo(
+                new CheckMethod(
+                    "org/elasticsearch/entitlement/instrumentation/impl/InstrumentationServiceImplTests$TestCheckerMixed",
+                    "checkInstanceMethodManual",
+                    List.of(
+                        "Ljava/lang/Class;",
+                        "Lorg/elasticsearch/entitlement/instrumentation/impl/InstrumentationServiceImplTests$TestTargetBaseClass;",
+                        "I",
+                        "Ljava/lang/String;"
+                    )
+                )
+            )
+        );
+    }
+
     public void testParseCheckerMethodSignatureStaticMethod() {
         var methodKey = InstrumentationServiceImpl.parseCheckerMethodSignature(
             "check$org_example_TestClass$$staticMethod",