Skip to content

Commit fdbd079

Browse files
ldematteldematte
authored
[Entitlements] Fix Entitlement initialization to work across multiple versions (#121192)
1 parent 8886594 commit fdbd079

File tree

3 files changed

+162
-30
lines changed

3 files changed

+162
-30
lines changed

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

Lines changed: 60 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -23,49 +23,86 @@
2323
import java.io.IOException;
2424
import java.lang.reflect.Method;
2525
import java.lang.reflect.Modifier;
26+
import java.util.ArrayDeque;
2627
import java.util.Arrays;
28+
import java.util.Collections;
2729
import java.util.HashMap;
30+
import java.util.HashSet;
2831
import java.util.List;
2932
import java.util.Locale;
3033
import java.util.Map;
34+
import java.util.Set;
3135
import java.util.stream.Collectors;
3236
import java.util.stream.Stream;
3337

3438
public class InstrumentationServiceImpl implements InstrumentationService {
3539

40+
private static final String OBJECT_INTERNAL_NAME = Type.getInternalName(Object.class);
41+
3642
@Override
3743
public Instrumenter newInstrumenter(Class<?> clazz, Map<MethodKey, CheckMethod> methods) {
3844
return InstrumenterImpl.create(clazz, methods);
3945
}
4046

4147
@Override
4248
public Map<MethodKey, CheckMethod> lookupMethods(Class<?> checkerClass) throws IOException {
43-
var methodsToInstrument = new HashMap<MethodKey, CheckMethod>();
44-
var classFileInfo = InstrumenterImpl.getClassFileInfo(checkerClass);
45-
ClassReader reader = new ClassReader(classFileInfo.bytecodes());
46-
ClassVisitor visitor = new ClassVisitor(Opcodes.ASM9) {
47-
@Override
48-
public MethodVisitor visitMethod(
49-
int access,
50-
String checkerMethodName,
51-
String checkerMethodDescriptor,
52-
String signature,
53-
String[] exceptions
54-
) {
55-
var mv = super.visitMethod(access, checkerMethodName, checkerMethodDescriptor, signature, exceptions);
56-
if (checkerMethodName.startsWith(InstrumentationService.CHECK_METHOD_PREFIX)) {
57-
var checkerMethodArgumentTypes = Type.getArgumentTypes(checkerMethodDescriptor);
58-
var methodToInstrument = parseCheckerMethodSignature(checkerMethodName, checkerMethodArgumentTypes);
49+
Map<MethodKey, CheckMethod> methodsToInstrument = new HashMap<>();
5950

60-
var checkerParameterDescriptors = Arrays.stream(checkerMethodArgumentTypes).map(Type::getDescriptor).toList();
61-
var checkMethod = new CheckMethod(Type.getInternalName(checkerClass), checkerMethodName, checkerParameterDescriptors);
51+
Set<Class<?>> visitedClasses = new HashSet<>();
52+
ArrayDeque<Class<?>> classesToVisit = new ArrayDeque<>(Collections.singleton(checkerClass));
53+
while (classesToVisit.isEmpty() == false) {
54+
var currentClass = classesToVisit.remove();
55+
if (visitedClasses.contains(currentClass)) {
56+
continue;
57+
}
58+
visitedClasses.add(currentClass);
6259

63-
methodsToInstrument.put(methodToInstrument, checkMethod);
60+
var classFileInfo = InstrumenterImpl.getClassFileInfo(currentClass);
61+
ClassReader reader = new ClassReader(classFileInfo.bytecodes());
62+
ClassVisitor visitor = new ClassVisitor(Opcodes.ASM9) {
63+
64+
@Override
65+
public void visit(int version, int access, String name, String signature, String superName, String[] interfaces) {
66+
super.visit(version, access, name, signature, superName, interfaces);
67+
try {
68+
if (OBJECT_INTERNAL_NAME.equals(superName) == false) {
69+
classesToVisit.add(Class.forName(Type.getObjectType(superName).getClassName()));
70+
}
71+
for (var interfaceName : interfaces) {
72+
classesToVisit.add(Class.forName(Type.getObjectType(interfaceName).getClassName()));
73+
}
74+
} catch (ClassNotFoundException e) {
75+
throw new IllegalArgumentException("Cannot inspect checker class " + checkerClass.getName(), e);
76+
}
6477
}
65-
return mv;
66-
}
67-
};
68-
reader.accept(visitor, 0);
78+
79+
@Override
80+
public MethodVisitor visitMethod(
81+
int access,
82+
String checkerMethodName,
83+
String checkerMethodDescriptor,
84+
String signature,
85+
String[] exceptions
86+
) {
87+
var mv = super.visitMethod(access, checkerMethodName, checkerMethodDescriptor, signature, exceptions);
88+
if (checkerMethodName.startsWith(InstrumentationService.CHECK_METHOD_PREFIX)) {
89+
var checkerMethodArgumentTypes = Type.getArgumentTypes(checkerMethodDescriptor);
90+
var methodToInstrument = parseCheckerMethodSignature(checkerMethodName, checkerMethodArgumentTypes);
91+
92+
var checkerParameterDescriptors = Arrays.stream(checkerMethodArgumentTypes).map(Type::getDescriptor).toList();
93+
var checkMethod = new CheckMethod(
94+
Type.getInternalName(currentClass),
95+
checkerMethodName,
96+
checkerParameterDescriptors
97+
);
98+
99+
methodsToInstrument.putIfAbsent(methodToInstrument, checkMethod);
100+
}
101+
return mv;
102+
}
103+
};
104+
reader.accept(visitor, 0);
105+
}
69106
return methodsToInstrument;
70107
}
71108

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

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,14 @@ interface TestChecker {
5555
void check$org_example_TestTargetClass$instanceMethodWithArgs(Class<?> clazz, TestTargetClass that, int x, int y);
5656
}
5757

58+
interface TestCheckerDerived extends TestChecker {
59+
void check$org_example_TestTargetClass$instanceMethodNoArgs(Class<?> clazz, TestTargetClass that);
60+
61+
void check$org_example_TestTargetClass$differentInstanceMethod(Class<?> clazz, TestTargetClass that);
62+
}
63+
64+
interface TestCheckerDerived2 extends TestCheckerDerived, TestChecker {}
65+
5866
interface TestCheckerOverloads {
5967
void check$org_example_TestTargetClass$$staticMethodWithOverload(Class<?> clazz, int x, int y);
6068

@@ -160,6 +168,75 @@ public void testInstrumentationTargetLookupWithOverloads() throws IOException {
160168
);
161169
}
162170

171+
public void testInstrumentationTargetLookupWithDerivedClass() throws IOException {
172+
Map<MethodKey, CheckMethod> checkMethods = instrumentationService.lookupMethods(TestCheckerDerived2.class);
173+
174+
assertThat(checkMethods, aMapWithSize(4));
175+
assertThat(
176+
checkMethods,
177+
hasEntry(
178+
equalTo(new MethodKey("org/example/TestTargetClass", "staticMethod", List.of("I", "java/lang/String", "java/lang/Object"))),
179+
equalTo(
180+
new CheckMethod(
181+
"org/elasticsearch/entitlement/instrumentation/impl/InstrumentationServiceImplTests$TestChecker",
182+
"check$org_example_TestTargetClass$$staticMethod",
183+
List.of("Ljava/lang/Class;", "I", "Ljava/lang/String;", "Ljava/lang/Object;")
184+
)
185+
)
186+
)
187+
);
188+
assertThat(
189+
checkMethods,
190+
hasEntry(
191+
equalTo(new MethodKey("org/example/TestTargetClass", "instanceMethodNoArgs", List.of())),
192+
equalTo(
193+
new CheckMethod(
194+
"org/elasticsearch/entitlement/instrumentation/impl/InstrumentationServiceImplTests$TestCheckerDerived",
195+
"check$org_example_TestTargetClass$instanceMethodNoArgs",
196+
List.of(
197+
"Ljava/lang/Class;",
198+
"Lorg/elasticsearch/entitlement/instrumentation/impl/InstrumentationServiceImplTests$TestTargetClass;"
199+
)
200+
)
201+
)
202+
)
203+
);
204+
assertThat(
205+
checkMethods,
206+
hasEntry(
207+
equalTo(new MethodKey("org/example/TestTargetClass", "instanceMethodWithArgs", List.of("I", "I"))),
208+
equalTo(
209+
new CheckMethod(
210+
"org/elasticsearch/entitlement/instrumentation/impl/InstrumentationServiceImplTests$TestChecker",
211+
"check$org_example_TestTargetClass$instanceMethodWithArgs",
212+
List.of(
213+
"Ljava/lang/Class;",
214+
"Lorg/elasticsearch/entitlement/instrumentation/impl/InstrumentationServiceImplTests$TestTargetClass;",
215+
"I",
216+
"I"
217+
)
218+
)
219+
)
220+
)
221+
);
222+
assertThat(
223+
checkMethods,
224+
hasEntry(
225+
equalTo(new MethodKey("org/example/TestTargetClass", "differentInstanceMethod", List.of())),
226+
equalTo(
227+
new CheckMethod(
228+
"org/elasticsearch/entitlement/instrumentation/impl/InstrumentationServiceImplTests$TestCheckerDerived",
229+
"check$org_example_TestTargetClass$differentInstanceMethod",
230+
List.of(
231+
"Ljava/lang/Class;",
232+
"Lorg/elasticsearch/entitlement/instrumentation/impl/InstrumentationServiceImplTests$TestTargetClass;"
233+
)
234+
)
235+
)
236+
)
237+
);
238+
}
239+
163240
public void testInstrumentationTargetLookupWithCtors() throws IOException {
164241
Map<MethodKey, CheckMethod> checkMethods = instrumentationService.lookupMethods(TestCheckerCtors.class);
165242

‎libs/entitlement/src/main/java/org/elasticsearch/entitlement/initialization/EntitlementInitialization.java‎

Lines changed: 25 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -66,8 +66,9 @@ public static EntitlementChecker checker() {
6666
public static void initialize(Instrumentation inst) throws Exception {
6767
manager = initChecker();
6868

69-
Map<MethodKey, CheckMethod> checkMethods = new HashMap<>(INSTRUMENTATION_SERVICE.lookupMethods(EntitlementChecker.class));
69+
var latestCheckerInterface = getVersionSpecificCheckerClass(EntitlementChecker.class);
7070

71+
Map<MethodKey, CheckMethod> checkMethods = new HashMap<>(INSTRUMENTATION_SERVICE.lookupMethods(latestCheckerInterface));
7172
var fileSystemProviderClass = FileSystems.getDefault().provider().getClass();
7273
Stream.of(
7374
INSTRUMENTATION_SERVICE.lookupImplementationMethod(
@@ -83,7 +84,7 @@ public static void initialize(Instrumentation inst) throws Exception {
8384

8485
var classesToTransform = checkMethods.keySet().stream().map(MethodKey::className).collect(Collectors.toSet());
8586

86-
Instrumenter instrumenter = INSTRUMENTATION_SERVICE.newInstrumenter(EntitlementChecker.class, checkMethods);
87+
Instrumenter instrumenter = INSTRUMENTATION_SERVICE.newInstrumenter(latestCheckerInterface, checkMethods);
8788
inst.addTransformer(new Transformer(instrumenter, classesToTransform), true);
8889
inst.retransformClasses(findClassesToRetransform(inst.getAllLoadedClasses(), classesToTransform));
8990
}
@@ -130,23 +131,40 @@ private static PolicyManager createPolicyManager() {
130131
return new PolicyManager(serverPolicy, agentEntitlements, pluginPolicies, resolver, AGENTS_PACKAGE_NAME, ENTITLEMENTS_MODULE);
131132
}
132133

133-
private static ElasticsearchEntitlementChecker initChecker() {
134-
final PolicyManager policyManager = createPolicyManager();
135-
134+
/**
135+
* Returns the "most recent" checker class compatible with the current runtime Java version.
136+
* For checkers, we have (optionally) version specific classes, each with a prefix (e.g. Java23).
137+
* The mapping cannot be automatic, as it depends on the actual presence of these classes in the final Jar (see
138+
* the various mainXX source sets).
139+
*/
140+
private static Class<?> getVersionSpecificCheckerClass(Class<?> baseClass) {
141+
String packageName = baseClass.getPackageName();
142+
String baseClassName = baseClass.getSimpleName();
136143
int javaVersion = Runtime.version().feature();
144+
137145
final String classNamePrefix;
138146
if (javaVersion >= 23) {
147+
// All Java version from 23 onwards will be able to use che checks in the Java23EntitlementChecker interface and implementation
139148
classNamePrefix = "Java23";
140149
} else {
150+
// For any other Java version, the basic EntitlementChecker interface and implementation contains all the supported checks
141151
classNamePrefix = "";
142152
}
143-
final String className = "org.elasticsearch.entitlement.runtime.api." + classNamePrefix + "ElasticsearchEntitlementChecker";
153+
final String className = packageName + "." + classNamePrefix + baseClassName;
144154
Class<?> clazz;
145155
try {
146156
clazz = Class.forName(className);
147157
} catch (ClassNotFoundException e) {
148-
throw new AssertionError("entitlement lib cannot find entitlement impl", e);
158+
throw new AssertionError("entitlement lib cannot find entitlement class " + className, e);
149159
}
160+
return clazz;
161+
}
162+
163+
private static ElasticsearchEntitlementChecker initChecker() {
164+
final PolicyManager policyManager = createPolicyManager();
165+
166+
final Class<?> clazz = getVersionSpecificCheckerClass(ElasticsearchEntitlementChecker.class);
167+
150168
Constructor<?> constructor;
151169
try {
152170
constructor = clazz.getConstructor(PolicyManager.class);

0 commit comments

Comments
 (0)