Skip to content

Commit 613f58c

Browse files
ldematteldematte
authored
[Entitlements] Fix Entitlement initialization to work across multiple versions (#121192) (#121738) (#121742)
1 parent 3222a64 commit 613f58c

File tree

3 files changed

+169
-59
lines changed

3 files changed

+169
-59
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: 32 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,6 @@
3737
import java.nio.file.spi.FileSystemProvider;
3838
import java.util.ArrayList;
3939
import java.util.HashMap;
40-
import java.util.HashSet;
4140
import java.util.List;
4241
import java.util.Map;
4342
import java.util.Set;
@@ -67,16 +66,9 @@ public static EntitlementChecker checker() {
6766
public static void initialize(Instrumentation inst) throws Exception {
6867
manager = initChecker();
6968

70-
Map<MethodKey, CheckMethod> checkMethods = new HashMap<>();
71-
int javaVersion = Runtime.version().feature();
72-
Set<Class<?>> interfaces = new HashSet<>();
73-
for (int i = 17; i <= javaVersion; ++i) {
74-
interfaces.add(getVersionSpecificCheckerClass(i, "org.elasticsearch.entitlement.bridge", "EntitlementChecker"));
75-
}
76-
for (var checkerInterface : interfaces) {
77-
checkMethods.putAll(INSTRUMENTATION_SERVICE.lookupMethods(checkerInterface));
78-
}
69+
var latestCheckerInterface = getVersionSpecificCheckerClass(EntitlementChecker.class);
7970

71+
Map<MethodKey, CheckMethod> checkMethods = new HashMap<>(INSTRUMENTATION_SERVICE.lookupMethods(latestCheckerInterface));
8072
var fileSystemProviderClass = FileSystems.getDefault().provider().getClass();
8173
Stream.of(
8274
INSTRUMENTATION_SERVICE.lookupImplementationMethod(
@@ -90,11 +82,6 @@ public static void initialize(Instrumentation inst) throws Exception {
9082
)
9183
).forEach(instrumentation -> checkMethods.put(instrumentation.targetMethod(), instrumentation.checkMethod()));
9284

93-
var latestCheckerInterface = getVersionSpecificCheckerClass(
94-
javaVersion,
95-
"org.elasticsearch.entitlement.bridge",
96-
"EntitlementChecker"
97-
);
9885
var classesToTransform = checkMethods.keySet().stream().map(MethodKey::className).collect(Collectors.toSet());
9986

10087
Instrumenter instrumenter = INSTRUMENTATION_SERVICE.newInstrumenter(latestCheckerInterface, checkMethods);
@@ -144,36 +131,27 @@ private static PolicyManager createPolicyManager() {
144131
return new PolicyManager(serverPolicy, agentEntitlements, pluginPolicies, resolver, AGENTS_PACKAGE_NAME, ENTITLEMENTS_MODULE);
145132
}
146133

147-
private static ElasticsearchEntitlementChecker initChecker() {
148-
final PolicyManager policyManager = createPolicyManager();
149-
150-
Class<?> clazz = getVersionSpecificCheckerClass(
151-
Runtime.version().feature(),
152-
"org.elasticsearch.entitlement.runtime.api",
153-
"ElasticsearchEntitlementChecker"
154-
);
155-
Constructor<?> constructor;
156-
try {
157-
constructor = clazz.getConstructor(PolicyManager.class);
158-
} catch (NoSuchMethodException e) {
159-
throw new AssertionError("entitlement impl is missing no arg constructor", e);
160-
}
161-
try {
162-
return (ElasticsearchEntitlementChecker) constructor.newInstance(policyManager);
163-
} catch (IllegalAccessException | InvocationTargetException | InstantiationException e) {
164-
throw new AssertionError(e);
165-
}
166-
}
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();
143+
int javaVersion = Runtime.version().feature();
167144

168-
private static Class<?> getVersionSpecificCheckerClass(int javaVersion, String packageName, String baseClassName) {
169145
final String classNamePrefix;
170146
if (javaVersion == 21) {
171147
classNamePrefix = "Java21";
172148
} else if (javaVersion == 22) {
173149
classNamePrefix = "Java22";
174150
} else if (javaVersion >= 23) {
151+
// All Java version from 23 onwards will be able to use che checks in the Java23EntitlementChecker interface and implementation
175152
classNamePrefix = "Java23";
176153
} else {
154+
// For any other Java version, the basic EntitlementChecker interface and implementation contains all the supported checks
177155
classNamePrefix = "";
178156
}
179157
final String className = packageName + "." + classNamePrefix + baseClassName;
@@ -186,6 +164,24 @@ private static Class<?> getVersionSpecificCheckerClass(int javaVersion, String p
186164
return clazz;
187165
}
188166

167+
private static ElasticsearchEntitlementChecker initChecker() {
168+
final PolicyManager policyManager = createPolicyManager();
169+
170+
final Class<?> clazz = getVersionSpecificCheckerClass(ElasticsearchEntitlementChecker.class);
171+
172+
Constructor<?> constructor;
173+
try {
174+
constructor = clazz.getConstructor(PolicyManager.class);
175+
} catch (NoSuchMethodException e) {
176+
throw new AssertionError("entitlement impl is missing no arg constructor", e);
177+
}
178+
try {
179+
return (ElasticsearchEntitlementChecker) constructor.newInstance(policyManager);
180+
} catch (IllegalAccessException | InvocationTargetException | InstantiationException e) {
181+
throw new AssertionError(e);
182+
}
183+
}
184+
189185
private static final InstrumentationService INSTRUMENTATION_SERVICE = new ProviderLocator<>(
190186
"entitlement",
191187
InstrumentationService.class,

0 commit comments

Comments
 (0)