Skip to content

Commit c77f09e

Browse files
ldematteldematte
authored
[Entitlements] Refactor InstrumenterImpl tests (#117688)
Following up #117332 (comment), I refactored `InstrumenterImpl` tests, splitting them into 2 suites: - `SyntheticInstrumenterImplTests`, which tests the mechanics of instrumentation using ad-hoc test cases. This should see little change now that we have our Instrumenter working as intended - `InstrumenterImplTests`, which is back to its original intent to make sure (1) the right arguments make it all the way to the check methods, and (2) if the check method throws, that exception correctly bubbles up through the instrumented method. The PR also includes a little change to `InstrumenterImpl` construction to clean it up a bit and make it more testable.
1 parent c74c06d commit c77f09e

File tree

11 files changed

+646
-381
lines changed

11 files changed

+646
-381
lines changed

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

Lines changed: 7 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99

1010
package org.elasticsearch.entitlement.instrumentation.impl;
1111

12-
import org.elasticsearch.entitlement.instrumentation.CheckerMethod;
12+
import org.elasticsearch.entitlement.instrumentation.CheckMethod;
1313
import org.elasticsearch.entitlement.instrumentation.InstrumentationService;
1414
import org.elasticsearch.entitlement.instrumentation.Instrumenter;
1515
import org.elasticsearch.entitlement.instrumentation.MethodKey;
@@ -20,37 +20,23 @@
2020
import org.objectweb.asm.Type;
2121

2222
import java.io.IOException;
23-
import java.lang.reflect.Method;
2423
import java.util.Arrays;
2524
import java.util.HashMap;
2625
import java.util.List;
2726
import java.util.Locale;
2827
import java.util.Map;
29-
import java.util.stream.Stream;
3028

3129
public class InstrumentationServiceImpl implements InstrumentationService {
3230

3331
@Override
34-
public Instrumenter newInstrumenter(String classNameSuffix, Map<MethodKey, CheckerMethod> instrumentationMethods) {
35-
return new InstrumenterImpl(classNameSuffix, instrumentationMethods);
36-
}
37-
38-
/**
39-
* @return a {@link MethodKey} suitable for looking up the given {@code targetMethod} in the entitlements trampoline
40-
*/
41-
public MethodKey methodKeyForTarget(Method targetMethod) {
42-
Type actualType = Type.getMethodType(Type.getMethodDescriptor(targetMethod));
43-
return new MethodKey(
44-
Type.getInternalName(targetMethod.getDeclaringClass()),
45-
targetMethod.getName(),
46-
Stream.of(actualType.getArgumentTypes()).map(Type::getInternalName).toList()
47-
);
32+
public Instrumenter newInstrumenter(Map<MethodKey, CheckMethod> checkMethods) {
33+
return InstrumenterImpl.create(checkMethods);
4834
}
4935

5036
@Override
51-
public Map<MethodKey, CheckerMethod> lookupMethodsToInstrument(String entitlementCheckerClassName) throws ClassNotFoundException,
37+
public Map<MethodKey, CheckMethod> lookupMethodsToInstrument(String entitlementCheckerClassName) throws ClassNotFoundException,
5238
IOException {
53-
var methodsToInstrument = new HashMap<MethodKey, CheckerMethod>();
39+
var methodsToInstrument = new HashMap<MethodKey, CheckMethod>();
5440
var checkerClass = Class.forName(entitlementCheckerClassName);
5541
var classFileInfo = InstrumenterImpl.getClassFileInfo(checkerClass);
5642
ClassReader reader = new ClassReader(classFileInfo.bytecodes());
@@ -69,9 +55,9 @@ public MethodVisitor visitMethod(
6955
var methodToInstrument = parseCheckerMethodSignature(checkerMethodName, checkerMethodArgumentTypes);
7056

7157
var checkerParameterDescriptors = Arrays.stream(checkerMethodArgumentTypes).map(Type::getDescriptor).toList();
72-
var checkerMethod = new CheckerMethod(Type.getInternalName(checkerClass), checkerMethodName, checkerParameterDescriptors);
58+
var checkMethod = new CheckMethod(Type.getInternalName(checkerClass), checkerMethodName, checkerParameterDescriptors);
7359

74-
methodsToInstrument.put(methodToInstrument, checkerMethod);
60+
methodsToInstrument.put(methodToInstrument, checkMethod);
7561

7662
return mv;
7763
}

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

Lines changed: 37 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99

1010
package org.elasticsearch.entitlement.instrumentation.impl;
1111

12-
import org.elasticsearch.entitlement.instrumentation.CheckerMethod;
12+
import org.elasticsearch.entitlement.instrumentation.CheckMethod;
1313
import org.elasticsearch.entitlement.instrumentation.Instrumenter;
1414
import org.elasticsearch.entitlement.instrumentation.MethodKey;
1515
import org.objectweb.asm.AnnotationVisitor;
@@ -37,30 +37,43 @@
3737

3838
public class InstrumenterImpl implements Instrumenter {
3939

40-
private static final String checkerClassDescriptor;
41-
private static final String handleClass;
42-
static {
40+
private final String getCheckerClassMethodDescriptor;
41+
private final String handleClass;
42+
43+
/**
44+
* To avoid class name collisions during testing without an agent to replace classes in-place.
45+
*/
46+
private final String classNameSuffix;
47+
private final Map<MethodKey, CheckMethod> checkMethods;
48+
49+
InstrumenterImpl(
50+
String handleClass,
51+
String getCheckerClassMethodDescriptor,
52+
String classNameSuffix,
53+
Map<MethodKey, CheckMethod> checkMethods
54+
) {
55+
this.handleClass = handleClass;
56+
this.getCheckerClassMethodDescriptor = getCheckerClassMethodDescriptor;
57+
this.classNameSuffix = classNameSuffix;
58+
this.checkMethods = checkMethods;
59+
}
60+
61+
static String getCheckerClassName() {
4362
int javaVersion = Runtime.version().feature();
4463
final String classNamePrefix;
4564
if (javaVersion >= 23) {
4665
classNamePrefix = "Java23";
4766
} else {
4867
classNamePrefix = "";
4968
}
50-
String checkerClass = "org/elasticsearch/entitlement/bridge/" + classNamePrefix + "EntitlementChecker";
51-
handleClass = checkerClass + "Handle";
52-
checkerClassDescriptor = Type.getObjectType(checkerClass).getDescriptor();
69+
return "org/elasticsearch/entitlement/bridge/" + classNamePrefix + "EntitlementChecker";
5370
}
5471

55-
/**
56-
* To avoid class name collisions during testing without an agent to replace classes in-place.
57-
*/
58-
private final String classNameSuffix;
59-
private final Map<MethodKey, CheckerMethod> instrumentationMethods;
60-
61-
public InstrumenterImpl(String classNameSuffix, Map<MethodKey, CheckerMethod> instrumentationMethods) {
62-
this.classNameSuffix = classNameSuffix;
63-
this.instrumentationMethods = instrumentationMethods;
72+
public static InstrumenterImpl create(Map<MethodKey, CheckMethod> checkMethods) {
73+
String checkerClass = getCheckerClassName();
74+
String handleClass = checkerClass + "Handle";
75+
String getCheckerClassMethodDescriptor = Type.getMethodDescriptor(Type.getObjectType(checkerClass));
76+
return new InstrumenterImpl(handleClass, getCheckerClassMethodDescriptor, "", checkMethods);
6477
}
6578

6679
public ClassFileInfo instrumentClassFile(Class<?> clazz) throws IOException {
@@ -156,7 +169,7 @@ public MethodVisitor visitMethod(int access, String name, String descriptor, Str
156169
boolean isStatic = (access & ACC_STATIC) != 0;
157170
boolean isCtor = "<init>".equals(name);
158171
var key = new MethodKey(className, name, Stream.of(Type.getArgumentTypes(descriptor)).map(Type::getInternalName).toList());
159-
var instrumentationMethod = instrumentationMethods.get(key);
172+
var instrumentationMethod = checkMethods.get(key);
160173
if (instrumentationMethod != null) {
161174
// LOGGER.debug("Will instrument method {}", key);
162175
return new EntitlementMethodVisitor(Opcodes.ASM9, mv, isStatic, isCtor, descriptor, instrumentationMethod);
@@ -190,7 +203,7 @@ class EntitlementMethodVisitor extends MethodVisitor {
190203
private final boolean instrumentedMethodIsStatic;
191204
private final boolean instrumentedMethodIsCtor;
192205
private final String instrumentedMethodDescriptor;
193-
private final CheckerMethod instrumentationMethod;
206+
private final CheckMethod checkMethod;
194207
private boolean hasCallerSensitiveAnnotation = false;
195208

196209
EntitlementMethodVisitor(
@@ -199,13 +212,13 @@ class EntitlementMethodVisitor extends MethodVisitor {
199212
boolean instrumentedMethodIsStatic,
200213
boolean instrumentedMethodIsCtor,
201214
String instrumentedMethodDescriptor,
202-
CheckerMethod instrumentationMethod
215+
CheckMethod checkMethod
203216
) {
204217
super(api, methodVisitor);
205218
this.instrumentedMethodIsStatic = instrumentedMethodIsStatic;
206219
this.instrumentedMethodIsCtor = instrumentedMethodIsCtor;
207220
this.instrumentedMethodDescriptor = instrumentedMethodDescriptor;
208-
this.instrumentationMethod = instrumentationMethod;
221+
this.checkMethod = checkMethod;
209222
}
210223

211224
@Override
@@ -278,19 +291,19 @@ private void forwardIncomingArguments() {
278291
private void invokeInstrumentationMethod() {
279292
mv.visitMethodInsn(
280293
INVOKEINTERFACE,
281-
instrumentationMethod.className(),
282-
instrumentationMethod.methodName(),
294+
checkMethod.className(),
295+
checkMethod.methodName(),
283296
Type.getMethodDescriptor(
284297
Type.VOID_TYPE,
285-
instrumentationMethod.parameterDescriptors().stream().map(Type::getType).toArray(Type[]::new)
298+
checkMethod.parameterDescriptors().stream().map(Type::getType).toArray(Type[]::new)
286299
),
287300
true
288301
);
289302
}
290303
}
291304

292305
protected void pushEntitlementChecker(MethodVisitor mv) {
293-
mv.visitMethodInsn(INVOKESTATIC, handleClass, "instance", "()" + checkerClassDescriptor, false);
306+
mv.visitMethodInsn(INVOKESTATIC, handleClass, "instance", getCheckerClassMethodDescriptor, false);
294307
}
295308

296309
public record ClassFileInfo(String fileName, byte[] bytecodes) {}

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

Lines changed: 21 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99

1010
package org.elasticsearch.entitlement.instrumentation.impl;
1111

12-
import org.elasticsearch.entitlement.instrumentation.CheckerMethod;
12+
import org.elasticsearch.entitlement.instrumentation.CheckMethod;
1313
import org.elasticsearch.entitlement.instrumentation.InstrumentationService;
1414
import org.elasticsearch.entitlement.instrumentation.MethodKey;
1515
import org.elasticsearch.test.ESTestCase;
@@ -52,15 +52,15 @@ interface TestCheckerCtors {
5252
}
5353

5454
public void testInstrumentationTargetLookup() throws IOException, ClassNotFoundException {
55-
Map<MethodKey, CheckerMethod> methodsMap = instrumentationService.lookupMethodsToInstrument(TestChecker.class.getName());
55+
Map<MethodKey, CheckMethod> checkMethods = instrumentationService.lookupMethodsToInstrument(TestChecker.class.getName());
5656

57-
assertThat(methodsMap, aMapWithSize(3));
57+
assertThat(checkMethods, aMapWithSize(3));
5858
assertThat(
59-
methodsMap,
59+
checkMethods,
6060
hasEntry(
6161
equalTo(new MethodKey("org/example/TestTargetClass", "staticMethod", List.of("I", "java/lang/String", "java/lang/Object"))),
6262
equalTo(
63-
new CheckerMethod(
63+
new CheckMethod(
6464
"org/elasticsearch/entitlement/instrumentation/impl/InstrumentationServiceImplTests$TestChecker",
6565
"check$org_example_TestTargetClass$staticMethod",
6666
List.of("Ljava/lang/Class;", "I", "Ljava/lang/String;", "Ljava/lang/Object;")
@@ -69,7 +69,7 @@ public void testInstrumentationTargetLookup() throws IOException, ClassNotFoundE
6969
)
7070
);
7171
assertThat(
72-
methodsMap,
72+
checkMethods,
7373
hasEntry(
7474
equalTo(
7575
new MethodKey(
@@ -79,7 +79,7 @@ public void testInstrumentationTargetLookup() throws IOException, ClassNotFoundE
7979
)
8080
),
8181
equalTo(
82-
new CheckerMethod(
82+
new CheckMethod(
8383
"org/elasticsearch/entitlement/instrumentation/impl/InstrumentationServiceImplTests$TestChecker",
8484
"check$$instanceMethodNoArgs",
8585
List.of(
@@ -91,7 +91,7 @@ public void testInstrumentationTargetLookup() throws IOException, ClassNotFoundE
9191
)
9292
);
9393
assertThat(
94-
methodsMap,
94+
checkMethods,
9595
hasEntry(
9696
equalTo(
9797
new MethodKey(
@@ -101,7 +101,7 @@ public void testInstrumentationTargetLookup() throws IOException, ClassNotFoundE
101101
)
102102
),
103103
equalTo(
104-
new CheckerMethod(
104+
new CheckMethod(
105105
"org/elasticsearch/entitlement/instrumentation/impl/InstrumentationServiceImplTests$TestChecker",
106106
"check$$instanceMethodWithArgs",
107107
List.of(
@@ -117,15 +117,15 @@ public void testInstrumentationTargetLookup() throws IOException, ClassNotFoundE
117117
}
118118

119119
public void testInstrumentationTargetLookupWithOverloads() throws IOException, ClassNotFoundException {
120-
Map<MethodKey, CheckerMethod> methodsMap = instrumentationService.lookupMethodsToInstrument(TestCheckerOverloads.class.getName());
120+
Map<MethodKey, CheckMethod> checkMethods = instrumentationService.lookupMethodsToInstrument(TestCheckerOverloads.class.getName());
121121

122-
assertThat(methodsMap, aMapWithSize(2));
122+
assertThat(checkMethods, aMapWithSize(2));
123123
assertThat(
124-
methodsMap,
124+
checkMethods,
125125
hasEntry(
126126
equalTo(new MethodKey("org/example/TestTargetClass", "staticMethodWithOverload", List.of("I", "java/lang/String"))),
127127
equalTo(
128-
new CheckerMethod(
128+
new CheckMethod(
129129
"org/elasticsearch/entitlement/instrumentation/impl/InstrumentationServiceImplTests$TestCheckerOverloads",
130130
"check$org_example_TestTargetClass$staticMethodWithOverload",
131131
List.of("Ljava/lang/Class;", "I", "Ljava/lang/String;")
@@ -134,11 +134,11 @@ public void testInstrumentationTargetLookupWithOverloads() throws IOException, C
134134
)
135135
);
136136
assertThat(
137-
methodsMap,
137+
checkMethods,
138138
hasEntry(
139139
equalTo(new MethodKey("org/example/TestTargetClass", "staticMethodWithOverload", List.of("I", "I"))),
140140
equalTo(
141-
new CheckerMethod(
141+
new CheckMethod(
142142
"org/elasticsearch/entitlement/instrumentation/impl/InstrumentationServiceImplTests$TestCheckerOverloads",
143143
"check$org_example_TestTargetClass$staticMethodWithOverload",
144144
List.of("Ljava/lang/Class;", "I", "I")
@@ -149,15 +149,15 @@ public void testInstrumentationTargetLookupWithOverloads() throws IOException, C
149149
}
150150

151151
public void testInstrumentationTargetLookupWithCtors() throws IOException, ClassNotFoundException {
152-
Map<MethodKey, CheckerMethod> methodsMap = instrumentationService.lookupMethodsToInstrument(TestCheckerCtors.class.getName());
152+
Map<MethodKey, CheckMethod> checkMethods = instrumentationService.lookupMethodsToInstrument(TestCheckerCtors.class.getName());
153153

154-
assertThat(methodsMap, aMapWithSize(2));
154+
assertThat(checkMethods, aMapWithSize(2));
155155
assertThat(
156-
methodsMap,
156+
checkMethods,
157157
hasEntry(
158158
equalTo(new MethodKey("org/example/TestTargetClass", "<init>", List.of("I", "java/lang/String"))),
159159
equalTo(
160-
new CheckerMethod(
160+
new CheckMethod(
161161
"org/elasticsearch/entitlement/instrumentation/impl/InstrumentationServiceImplTests$TestCheckerCtors",
162162
"check$org_example_TestTargetClass$",
163163
List.of("Ljava/lang/Class;", "I", "Ljava/lang/String;")
@@ -166,11 +166,11 @@ public void testInstrumentationTargetLookupWithCtors() throws IOException, Class
166166
)
167167
);
168168
assertThat(
169-
methodsMap,
169+
checkMethods,
170170
hasEntry(
171171
equalTo(new MethodKey("org/example/TestTargetClass", "<init>", List.of())),
172172
equalTo(
173-
new CheckerMethod(
173+
new CheckMethod(
174174
"org/elasticsearch/entitlement/instrumentation/impl/InstrumentationServiceImplTests$TestCheckerCtors",
175175
"check$org_example_TestTargetClass$",
176176
List.of("Ljava/lang/Class;")

0 commit comments

Comments
 (0)