From 9f2cd71e4218a965d4bc9502a54052be0270f540 Mon Sep 17 00:00:00 2001 From: Lorenzo Dematte Date: Wed, 12 Feb 2025 15:58:33 +0100 Subject: [PATCH] [Entitlements] Fix "dynamic" instrumentation target class (#122197) --- .../impl/InstrumentationServiceImpl.java | 54 ++++++++----------- .../impl/InstrumentationServiceImplTests.java | 47 +++++++++++++++- 2 files changed, 68 insertions(+), 33 deletions(-) diff --git a/libs/entitlement/asm-provider/src/main/java/org/elasticsearch/entitlement/instrumentation/impl/InstrumentationServiceImpl.java b/libs/entitlement/asm-provider/src/main/java/org/elasticsearch/entitlement/instrumentation/impl/InstrumentationServiceImpl.java index 46b2139c79db0..05a5af374e5d9 100644 --- a/libs/entitlement/asm-provider/src/main/java/org/elasticsearch/entitlement/instrumentation/impl/InstrumentationServiceImpl.java +++ b/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() + ) + ); } } diff --git a/libs/entitlement/asm-provider/src/test/java/org/elasticsearch/entitlement/instrumentation/impl/InstrumentationServiceImplTests.java b/libs/entitlement/asm-provider/src/test/java/org/elasticsearch/entitlement/instrumentation/impl/InstrumentationServiceImplTests.java index 3c1eacdc90635..2b9b70d46c0ea 100644 --- a/libs/entitlement/asm-provider/src/test/java/org/elasticsearch/entitlement/instrumentation/impl/InstrumentationServiceImplTests.java +++ b/libs/entitlement/asm-provider/src/test/java/org/elasticsearch/entitlement/instrumentation/impl/InstrumentationServiceImplTests.java @@ -40,9 +40,16 @@ 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 +371,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",