Skip to content

Commit e745dd3

Browse files
lauritlaurit
committed
rewrite loading injected classes instrumentation in asm
1 parent e35564e commit e745dd3

File tree

5 files changed

+272
-216
lines changed

5 files changed

+272
-216
lines changed

‎instrumentation/internal/internal-class-loader/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/internal/classloader/ClassLoaderInstrumentationModule.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,8 +35,9 @@ public List<String> injectedClassNames() {
3535
@Override
3636
public List<TypeInstrumentation> typeInstrumentations() {
3737
return asList(
38-
new BootDelegationInstrumentation(),
38+
// added first so it would get applied after other instrumentations
3939
new LoadInjectedClassInstrumentation(),
40+
new BootDelegationInstrumentation(),
4041
new ResourceInjectionInstrumentation(),
4142
new DefineClassInstrumentation());
4243
}

‎instrumentation/internal/internal-class-loader/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/internal/classloader/LoadInjectedClassInstrumentation.java‎

Lines changed: 207 additions & 67 deletions
Original file line numberDiff line numberDiff line change
@@ -6,26 +6,27 @@
66
package io.opentelemetry.javaagent.instrumentation.internal.classloader;
77

88
import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.extendsClass;
9-
import static io.opentelemetry.javaagent.instrumentation.internal.classloader.AdviceUtil.applyInlineAdvice;
10-
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
11-
import static net.bytebuddy.matcher.ElementMatchers.isProtected;
12-
import static net.bytebuddy.matcher.ElementMatchers.isPublic;
13-
import static net.bytebuddy.matcher.ElementMatchers.isStatic;
149
import static net.bytebuddy.matcher.ElementMatchers.named;
15-
import static net.bytebuddy.matcher.ElementMatchers.not;
16-
import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
17-
import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
1810

1911
import io.opentelemetry.javaagent.bootstrap.InjectedClassHelper;
20-
import io.opentelemetry.javaagent.bootstrap.InjectedClassHelper.HelperClassLoader;
12+
import io.opentelemetry.javaagent.bootstrap.InjectedClassHelper.HelperClassInfo;
2113
import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
2214
import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
23-
import java.lang.invoke.MethodHandles;
24-
import net.bytebuddy.asm.Advice;
25-
import net.bytebuddy.asm.Advice.AssignReturned;
26-
import net.bytebuddy.description.method.MethodDescription;
15+
import io.opentelemetry.javaagent.extension.instrumentation.internal.AsmApi;
16+
import net.bytebuddy.asm.AsmVisitorWrapper;
17+
import net.bytebuddy.description.field.FieldDescription;
18+
import net.bytebuddy.description.field.FieldList;
19+
import net.bytebuddy.description.method.MethodList;
2720
import net.bytebuddy.description.type.TypeDescription;
21+
import net.bytebuddy.implementation.Implementation;
2822
import net.bytebuddy.matcher.ElementMatcher;
23+
import net.bytebuddy.pool.TypePool;
24+
import org.objectweb.asm.ClassVisitor;
25+
import org.objectweb.asm.ClassWriter;
26+
import org.objectweb.asm.Label;
27+
import org.objectweb.asm.MethodVisitor;
28+
import org.objectweb.asm.Opcodes;
29+
import org.objectweb.asm.Type;
2930

3031
/**
3132
* This instrumentation inserts loading of our injected helper classes at the start of {@code
@@ -40,67 +41,206 @@ public ElementMatcher<TypeDescription> typeMatcher() {
4041

4142
@Override
4243
public void transform(TypeTransformer transformer) {
43-
ElementMatcher.Junction<MethodDescription> methodMatcher =
44-
isMethod()
45-
.and(named("loadClass"))
46-
.and(
47-
takesArguments(1)
48-
.and(takesArgument(0, String.class))
49-
.or(
50-
takesArguments(2)
51-
.and(takesArgument(0, String.class))
52-
.and(takesArgument(1, boolean.class))))
53-
.and(isPublic().or(isProtected()))
54-
.and(not(isStatic()));
55-
boolean useLookup = Double.parseDouble(System.getProperty("java.specification.version")) >= 11;
56-
// Inline instrumentation to prevent problems with invokedynamic-recursion
57-
applyInlineAdvice(
58-
transformer,
59-
methodMatcher,
60-
this.getClass().getName()
61-
+ (useLookup ? "$LoadClassAdvice" : "$LoadClassWithoutLookupAdvice"));
44+
transformer.applyTransformer(
45+
(builder, typeDescription, classLoader, module, protectionDomain) ->
46+
builder.visit(
47+
new AsmVisitorWrapper() {
48+
@Override
49+
public int mergeWriter(int flags) {
50+
return flags | ClassWriter.COMPUTE_MAXS;
51+
}
52+
53+
@Override
54+
public int mergeReader(int flags) {
55+
return flags;
56+
}
57+
58+
@Override
59+
public ClassVisitor wrap(
60+
TypeDescription instrumentedType,
61+
ClassVisitor classVisitor,
62+
Implementation.Context implementationContext,
63+
TypePool typePool,
64+
FieldList<FieldDescription.InDefinedShape> fields,
65+
MethodList<?> methods,
66+
int writerFlags,
67+
int readerFlags) {
68+
return new ClassLoaderClassVisitor(classVisitor);
69+
}
70+
}));
6271
}
6372

64-
@SuppressWarnings("unused")
65-
public static class LoadClassAdvice {
66-
67-
@Advice.OnMethodEnter(skipOn = Advice.OnNonDefaultValue.class)
68-
public static Class<?> onEnter(
69-
@Advice.This ClassLoader classLoader, @Advice.Argument(0) String name) throws Throwable {
70-
HelperClassLoader helperClassLoader =
71-
InjectedClassHelper.getHelperClassLoader(classLoader, name);
72-
return helperClassLoader != null
73-
? helperClassLoader.loadHelperClass(MethodHandles.lookup())
74-
: null;
75-
}
73+
private static class ClassLoaderClassVisitor extends ClassVisitor implements Opcodes {
74+
private String internalClassName;
7675

77-
@AssignReturned.ToReturned
78-
@Advice.OnMethodExit(onThrowable = Throwable.class)
79-
public static Class<?> onExit(
80-
@Advice.Return Class<?> originalResult, @Advice.Enter Class<?> loadedClass) {
81-
return loadedClass != null ? loadedClass : originalResult;
76+
ClassLoaderClassVisitor(ClassVisitor classVisitor) {
77+
super(AsmApi.VERSION, classVisitor);
8278
}
83-
}
8479

85-
@SuppressWarnings("unused")
86-
public static class LoadClassWithoutLookupAdvice {
87-
88-
@Advice.OnMethodEnter(skipOn = Advice.OnNonDefaultValue.class)
89-
public static Class<?> onEnter(
90-
@Advice.This ClassLoader classLoader, @Advice.Argument(0) String name) throws Throwable {
91-
HelperClassLoader helperClassLoader =
92-
InjectedClassHelper.getHelperClassLoader(classLoader, name);
93-
// on jdk8 we can't use MethodHandles.lookup() because it fails when called from
94-
// java.lang.ClassLoader with java.lang.IllegalArgumentException: illegal lookupClass: class
95-
// java.lang.ClassLoader
96-
return helperClassLoader != null ? helperClassLoader.loadHelperClass(null) : null;
80+
@Override
81+
public void visit(
82+
int version,
83+
int access,
84+
String name,
85+
String signature,
86+
String superName,
87+
String[] interfaces) {
88+
super.visit(version, access, name, signature, superName, interfaces);
89+
internalClassName = name;
9790
}
9891

99-
@AssignReturned.ToReturned
100-
@Advice.OnMethodExit(onThrowable = Throwable.class)
101-
public static Class<?> onExit(
102-
@Advice.Return Class<?> originalResult, @Advice.Enter Class<?> loadedClass) {
103-
return loadedClass != null ? loadedClass : originalResult;
92+
@Override
93+
public MethodVisitor visitMethod(
94+
int access, String name, String descriptor, String signature, String[] exceptions) {
95+
MethodVisitor mv = super.visitMethod(access, name, descriptor, signature, exceptions);
96+
if ("loadClass".equals(name)
97+
&& ("(Ljava/lang/String;)Ljava/lang/Class;".equals(descriptor)
98+
|| "(Ljava/lang/String;Z)Ljava/lang/Class;".equals(descriptor))) {
99+
100+
int argumentCount = Type.getArgumentTypes(descriptor).length;
101+
return new MethodVisitor(api, mv) {
102+
@Override
103+
public void visitCode() {
104+
super.visitCode();
105+
106+
// inserts the following at the start of the loadClass method:
107+
/*
108+
InjectedClassHelper.HelperClassInfo helperClassInfo = InjectedClassHelper.getHelperClassInfo(this, name);
109+
if (helperClassInfo != null) {
110+
Class<?> clazz = findLoadedClass(name);
111+
if (clazz != null) {
112+
return clazz;
113+
}
114+
try {
115+
byte[] bytes = helperClassInfo.getClassBytes();
116+
return defineClass(name, bytes, 0, bytes.length, helperClassInfo.getProtectionDomain());
117+
} catch (LinkageError error) {
118+
clazz = findLoadedClass(name);
119+
if (clazz != null) {
120+
return clazz;
121+
}
122+
throw error;
123+
}
124+
}
125+
*/
126+
127+
Label startTry = new Label();
128+
Label endTry = new Label();
129+
Label handler = new Label();
130+
mv.visitTryCatchBlock(startTry, endTry, handler, "java/lang/LinkageError");
131+
// InjectedClassHelper.HelperClassInfo helperClassInfo =
132+
// InjectedClassHelper.getHelperClassInfo(this, name);
133+
mv.visitVarInsn(ALOAD, 0);
134+
mv.visitVarInsn(ALOAD, 1);
135+
mv.visitMethodInsn(
136+
INVOKESTATIC,
137+
Type.getInternalName(InjectedClassHelper.class),
138+
"getHelperClassInfo",
139+
"(Ljava/lang/ClassLoader;Ljava/lang/String;)"
140+
+ Type.getDescriptor(HelperClassInfo.class),
141+
false);
142+
mv.visitVarInsn(ASTORE, argumentCount + 1); // store helperClassInfo
143+
mv.visitVarInsn(ALOAD, argumentCount + 1);
144+
Label notHelperClass = new Label();
145+
mv.visitJumpInsn(IFNULL, notHelperClass);
146+
147+
// getHelperClassInfo returned non-null
148+
// Class<?> clazz = findLoadedClass(name);
149+
mv.visitVarInsn(ALOAD, 0);
150+
mv.visitVarInsn(ALOAD, 1);
151+
mv.visitMethodInsn(
152+
INVOKEVIRTUAL,
153+
internalClassName,
154+
"findLoadedClass",
155+
"(Ljava/lang/String;)Ljava/lang/Class;",
156+
false);
157+
mv.visitVarInsn(ASTORE, argumentCount + 2); // store clazz
158+
mv.visitVarInsn(ALOAD, argumentCount + 2);
159+
mv.visitJumpInsn(IFNULL, startTry);
160+
161+
// findLoadedClass returned non-null
162+
// return clazz
163+
mv.visitVarInsn(ALOAD, argumentCount + 2);
164+
mv.visitInsn(ARETURN);
165+
166+
mv.visitLabel(startTry);
167+
mv.visitFrame(
168+
Opcodes.F_APPEND,
169+
2,
170+
new Object[] {Type.getInternalName(HelperClassInfo.class), "java/lang/Class"},
171+
0,
172+
null);
173+
// byte[] bytes = helperClassInfo.getClassBytes();
174+
mv.visitVarInsn(ALOAD, argumentCount + 1);
175+
mv.visitMethodInsn(
176+
INVOKEINTERFACE,
177+
Type.getInternalName(HelperClassInfo.class),
178+
"getClassBytes",
179+
"()[B",
180+
true);
181+
mv.visitVarInsn(ASTORE, argumentCount + 3); // store bytes
182+
183+
// return defineClass(name, bytes, 0, bytes.length,
184+
// helperClassInfo.getProtectionDomain());
185+
mv.visitVarInsn(ALOAD, 0);
186+
mv.visitVarInsn(ALOAD, 1);
187+
mv.visitVarInsn(ALOAD, argumentCount + 3);
188+
mv.visitInsn(ICONST_0);
189+
mv.visitVarInsn(ALOAD, argumentCount + 3);
190+
mv.visitInsn(ARRAYLENGTH);
191+
mv.visitVarInsn(ALOAD, argumentCount + 1);
192+
mv.visitMethodInsn(
193+
INVOKEINTERFACE,
194+
Type.getInternalName(HelperClassInfo.class),
195+
"getProtectionDomain",
196+
"()Ljava/security/ProtectionDomain;",
197+
true);
198+
mv.visitMethodInsn(
199+
INVOKEVIRTUAL,
200+
internalClassName,
201+
"defineClass",
202+
"(Ljava/lang/String;[BIILjava/security/ProtectionDomain;)Ljava/lang/Class;",
203+
false);
204+
mv.visitLabel(endTry);
205+
mv.visitInsn(ARETURN);
206+
207+
mv.visitLabel(handler);
208+
mv.visitFrame(Opcodes.F_SAME1, 0, null, 1, new Object[] {"java/lang/LinkageError"});
209+
mv.visitVarInsn(ASTORE, argumentCount + 3); // store LinkageError
210+
// clazz = findLoadedClass(name);
211+
mv.visitVarInsn(ALOAD, 0);
212+
mv.visitVarInsn(ALOAD, 1);
213+
mv.visitMethodInsn(
214+
INVOKEVIRTUAL,
215+
internalClassName,
216+
"findLoadedClass",
217+
"(Ljava/lang/String;)Ljava/lang/Class;",
218+
false);
219+
mv.visitVarInsn(ASTORE, argumentCount + 2); // score clazz
220+
mv.visitVarInsn(ALOAD, argumentCount + 2);
221+
Label throwError = new Label();
222+
mv.visitJumpInsn(IFNULL, throwError);
223+
// return clazz
224+
mv.visitVarInsn(ALOAD, argumentCount + 2);
225+
mv.visitInsn(ARETURN);
226+
mv.visitLabel(throwError);
227+
mv.visitFrame(Opcodes.F_APPEND, 1, new Object[] {"java/lang/LinkageError"}, 0, null);
228+
// throw error
229+
mv.visitVarInsn(ALOAD, argumentCount + 3);
230+
mv.visitInsn(ATHROW);
231+
232+
mv.visitLabel(notHelperClass);
233+
mv.visitFrame(Opcodes.F_CHOP, 3, null, 0, null);
234+
}
235+
236+
@Override
237+
public void visitMaxs(int maxStack, int maxLocals) {
238+
// minimally we have argumentCount parameters + this + 3 locals added by us
239+
super.visitMaxs(maxStack, Math.max(maxLocals, argumentCount + 1 + 3));
240+
}
241+
};
242+
}
243+
return mv;
104244
}
105245
}
106246
}

‎javaagent-bootstrap/src/main/java/io/opentelemetry/javaagent/bootstrap/InjectedClassHelper.java‎

Lines changed: 13 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55

66
package io.opentelemetry.javaagent.bootstrap;
77

8-
import java.lang.invoke.MethodHandles;
8+
import java.security.ProtectionDomain;
99
import java.util.function.BiFunction;
1010
import java.util.function.BiPredicate;
1111
import java.util.function.Function;
@@ -38,25 +38,27 @@ public static boolean isHelperClass(ClassLoader classLoader, String className) {
3838
return helperClassDetector.test(classLoader, className);
3939
}
4040

41-
private static volatile BiFunction<ClassLoader, String, HelperClassLoader> helperClassLoader;
41+
private static volatile BiFunction<ClassLoader, String, HelperClassInfo> helperClassInfo;
4242

43-
public static void internalSetHelperClassLoader(
44-
BiFunction<ClassLoader, String, HelperClassLoader> helperClassLoader) {
45-
if (InjectedClassHelper.helperClassLoader != null) {
43+
public static void internalSetHelperClassInfo(
44+
BiFunction<ClassLoader, String, HelperClassInfo> helperClassInfo) {
45+
if (InjectedClassHelper.helperClassInfo != null) {
4646
// Only possible by misuse of this API, just ignore.
4747
return;
4848
}
49-
InjectedClassHelper.helperClassLoader = helperClassLoader;
49+
InjectedClassHelper.helperClassInfo = helperClassInfo;
5050
}
5151

52-
public static HelperClassLoader getHelperClassLoader(ClassLoader classLoader, String className) {
53-
if (helperClassLoader == null) {
52+
public static HelperClassInfo getHelperClassInfo(ClassLoader classLoader, String className) {
53+
if (helperClassInfo == null) {
5454
return null;
5555
}
56-
return helperClassLoader.apply(classLoader, className);
56+
return helperClassInfo.apply(classLoader, className);
5757
}
5858

59-
public interface HelperClassLoader {
60-
Class<?> loadHelperClass(MethodHandles.Lookup lookup) throws Throwable;
59+
public interface HelperClassInfo {
60+
byte[] getClassBytes();
61+
62+
ProtectionDomain getProtectionDomain();
6163
}
6264
}

0 commit comments

Comments
 (0)