Skip to content

Commit 99bb102

Browse files
manuel-alvarez-alvarezmanuel-alvarez-alvarez
committed
Fix issue with call sites in super calls to constructor
1 parent 9bd1251 commit 99bb102

File tree

11 files changed

+218
-32
lines changed

11 files changed

+218
-32
lines changed

‎buildSrc/call-site-instrumentation-plugin/src/main/java/datadog/trace/plugin/csi/impl/AdviceGeneratorImpl.java‎

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -297,7 +297,7 @@ private static void writeStackOperations(final AdviceSpecification advice, final
297297
String mode = "COPY";
298298
if (allArgsSpec != null) {
299299
if (advice instanceof AfterSpecification) {
300-
mode = advice.isConstructor() ? "PREPEND_ARRAY_CTOR" : "PREPEND_ARRAY";
300+
mode = "PREPEND_ARRAY";
301301
} else {
302302
mode = "APPEND_ARRAY";
303303
}
@@ -344,12 +344,10 @@ private static void writeAdviceMethodCall(
344344
final MethodCallExpr invokeStatic =
345345
new MethodCallExpr()
346346
.setScope(new NameExpr("handler"))
347-
.setName("method")
348-
.addArgument(opCode("INVOKESTATIC"))
347+
.setName("advice")
349348
.addArgument(new StringLiteralExpr(method.getOwner().getInternalName()))
350349
.addArgument(new StringLiteralExpr(method.getMethodName()))
351-
.addArgument(new StringLiteralExpr(method.getMethodType().getDescriptor()))
352-
.addArgument(new BooleanLiteralExpr(false));
350+
.addArgument(new StringLiteralExpr(method.getMethodType().getDescriptor()));
353351
body.addStatement(invokeStatic);
354352
}
355353
if (requiresCast(advice)) {

‎buildSrc/call-site-instrumentation-plugin/src/test/groovy/datadog/trace/plugin/csi/impl/AdviceGeneratorTest.groovy‎

Lines changed: 12 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,7 @@ final class AdviceGeneratorTest extends BaseCsiPluginTest {
4747
pointcut('java/security/MessageDigest', 'getInstance', '(Ljava/lang/String;)Ljava/security/MessageDigest;')
4848
statements(
4949
'handler.dupParameters(descriptor, StackDupMode.COPY);',
50-
'handler.method(Opcodes.INVOKESTATIC, "datadog/trace/plugin/csi/impl/AdviceGeneratorTest$BeforeAdvice", "before", "(Ljava/lang/String;)V", false);',
50+
'handler.advice("datadog/trace/plugin/csi/impl/AdviceGeneratorTest$BeforeAdvice", "before", "(Ljava/lang/String;)V");',
5151
'handler.method(opcode, owner, name, descriptor, isInterface);'
5252
)
5353
}
@@ -78,7 +78,7 @@ final class AdviceGeneratorTest extends BaseCsiPluginTest {
7878
advices(0) {
7979
pointcut('java/lang/String', 'replaceAll', '(Ljava/lang/String;Ljava/lang/String;)Ljava/lang/String;')
8080
statements(
81-
'handler.method(Opcodes.INVOKESTATIC, "datadog/trace/plugin/csi/impl/AdviceGeneratorTest$AroundAdvice", "around", "(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)Ljava/lang/String;", false);'
81+
'handler.advice("datadog/trace/plugin/csi/impl/AdviceGeneratorTest$AroundAdvice", "around", "(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)Ljava/lang/String;");'
8282
)
8383
}
8484
}
@@ -110,7 +110,7 @@ final class AdviceGeneratorTest extends BaseCsiPluginTest {
110110
statements(
111111
'handler.dupInvoke(owner, descriptor, StackDupMode.COPY);',
112112
'handler.method(opcode, owner, name, descriptor, isInterface);',
113-
'handler.method(Opcodes.INVOKESTATIC, "datadog/trace/plugin/csi/impl/AdviceGeneratorTest$AfterAdvice", "after", "(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)Ljava/lang/String;", false);',
113+
'handler.advice("datadog/trace/plugin/csi/impl/AdviceGeneratorTest$AfterAdvice", "after", "(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)Ljava/lang/String;");',
114114
)
115115
}
116116
}
@@ -140,9 +140,9 @@ final class AdviceGeneratorTest extends BaseCsiPluginTest {
140140
advices(0) {
141141
pointcut('java/net/URL', '<init>', '(Ljava/lang/String;)V')
142142
statements(
143-
'handler.dupParameters(descriptor, StackDupMode.PREPEND_ARRAY_CTOR);',
143+
'handler.dupParameters(descriptor, StackDupMode.PREPEND_ARRAY);',
144144
'handler.method(opcode, owner, name, descriptor, isInterface);',
145-
'handler.method(Opcodes.INVOKESTATIC, "datadog/trace/plugin/csi/impl/AdviceGeneratorTest$AfterAdviceCtor", "after", "([Ljava/lang/Object;Ljava/net/URL;)Ljava/net/URL;", false);',
145+
'handler.advice("datadog/trace/plugin/csi/impl/AdviceGeneratorTest$AfterAdviceCtor", "after", "([Ljava/lang/Object;Ljava/net/URL;)Ljava/net/URL;");',
146146
)
147147
}
148148
}
@@ -208,7 +208,7 @@ final class AdviceGeneratorTest extends BaseCsiPluginTest {
208208
statements(
209209
'handler.dupParameters(descriptor, StackDupMode.PREPEND_ARRAY);',
210210
'handler.invokeDynamic(name, descriptor, bootstrapMethodHandle, bootstrapMethodArguments);',
211-
'handler.method(Opcodes.INVOKESTATIC, "datadog/trace/plugin/csi/impl/AdviceGeneratorTest$InvokeDynamicAfterAdvice", "after", "([Ljava/lang/Object;Ljava/lang/String;)Ljava/lang/String;", false);'
211+
'handler.advice("datadog/trace/plugin/csi/impl/AdviceGeneratorTest$InvokeDynamicAfterAdvice", "after", "([Ljava/lang/Object;Ljava/lang/String;)Ljava/lang/String;");'
212212
)
213213
}
214214
}
@@ -297,7 +297,7 @@ final class AdviceGeneratorTest extends BaseCsiPluginTest {
297297
'handler.dupParameters(descriptor, StackDupMode.PREPEND_ARRAY);',
298298
'handler.invokeDynamic(name, descriptor, bootstrapMethodHandle, bootstrapMethodArguments);',
299299
'handler.loadConstantArray(bootstrapMethodArguments);',
300-
'handler.method(Opcodes.INVOKESTATIC, "datadog/trace/plugin/csi/impl/AdviceGeneratorTest$InvokeDynamicWithConstantsAdvice", "after", "([Ljava/lang/Object;Ljava/lang/String;[Ljava/lang/Object;)Ljava/lang/String;", false);'
300+
'handler.advice("datadog/trace/plugin/csi/impl/AdviceGeneratorTest$InvokeDynamicWithConstantsAdvice", "after", "([Ljava/lang/Object;Ljava/lang/String;[Ljava/lang/Object;)Ljava/lang/String;");'
301301
)
302302
}
303303
}
@@ -393,7 +393,7 @@ final class AdviceGeneratorTest extends BaseCsiPluginTest {
393393
statements(
394394
'int[] parameterIndices = new int[] { 0 };',
395395
'handler.dupParameters(descriptor, parameterIndices, owner);',
396-
'handler.method(Opcodes.INVOKESTATIC, "datadog/trace/plugin/csi/impl/AdviceGeneratorTest$PartialArgumentsBeforeAdvice", "before", "(Ljava/lang/String;)V", false);',
396+
'handler.advice("datadog/trace/plugin/csi/impl/AdviceGeneratorTest$PartialArgumentsBeforeAdvice", "before", "(Ljava/lang/String;)V");',
397397
'handler.method(opcode, owner, name, descriptor, isInterface);',
398398
)
399399
}
@@ -402,7 +402,7 @@ final class AdviceGeneratorTest extends BaseCsiPluginTest {
402402
statements(
403403
'int[] parameterIndices = new int[] { 1 };',
404404
'handler.dupParameters(descriptor, parameterIndices, null);',
405-
'handler.method(Opcodes.INVOKESTATIC, "datadog/trace/plugin/csi/impl/AdviceGeneratorTest$PartialArgumentsBeforeAdvice", "before", "([Ljava/lang/Object;)V", false);',
405+
'handler.advice("datadog/trace/plugin/csi/impl/AdviceGeneratorTest$PartialArgumentsBeforeAdvice", "before", "([Ljava/lang/Object;)V");',
406406
'handler.method(opcode, owner, name, descriptor, isInterface);',
407407
)
408408
}
@@ -411,7 +411,7 @@ final class AdviceGeneratorTest extends BaseCsiPluginTest {
411411
statements(
412412
'int[] parameterIndices = new int[] { 0 };',
413413
'handler.dupInvoke(owner, descriptor, parameterIndices);',
414-
'handler.method(Opcodes.INVOKESTATIC, "datadog/trace/plugin/csi/impl/AdviceGeneratorTest$PartialArgumentsBeforeAdvice", "before", "(Ljava/lang/String;I)V", false);',
414+
'handler.advice("datadog/trace/plugin/csi/impl/AdviceGeneratorTest$PartialArgumentsBeforeAdvice", "before", "(Ljava/lang/String;I)V");',
415415
'handler.method(opcode, owner, name, descriptor, isInterface);',
416416
)
417417
}
@@ -441,9 +441,9 @@ final class AdviceGeneratorTest extends BaseCsiPluginTest {
441441
advices(0) {
442442
pointcut('java/lang/StringBuilder', '<init>', '(Ljava/lang/String;)V')
443443
statements(
444-
'handler.dupParameters(descriptor, StackDupMode.PREPEND_ARRAY_CTOR);',
444+
'handler.dupParameters(descriptor, StackDupMode.PREPEND_ARRAY);',
445445
'handler.method(opcode, owner, name, descriptor, isInterface);',
446-
'handler.method(Opcodes.INVOKESTATIC, "datadog/trace/plugin/csi/impl/AdviceGeneratorTest$SuperTypeReturnAdvice", "after", "([Ljava/lang/Object;Ljava/lang/Object;)Ljava/lang/Object;", false);',
446+
'handler.advice("datadog/trace/plugin/csi/impl/AdviceGeneratorTest$SuperTypeReturnAdvice", "after", "([Ljava/lang/Object;Ljava/lang/Object;)Ljava/lang/Object;");',
447447
'handler.instruction(Opcodes.CHECKCAST, "java/lang/StringBuilder");'
448448
)
449449
}

‎dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/csi/CallSiteTransformer.java‎

Lines changed: 44 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,8 @@
99
import datadog.trace.agent.tooling.csi.InvokeAdvice;
1010
import datadog.trace.agent.tooling.csi.InvokeDynamicAdvice;
1111
import java.security.ProtectionDomain;
12+
import java.util.Deque;
13+
import java.util.LinkedList;
1214
import javax.annotation.Nonnull;
1315
import net.bytebuddy.asm.AsmVisitorWrapper;
1416
import net.bytebuddy.description.field.FieldDescription;
@@ -118,25 +120,48 @@ public MethodVisitor visitMethod(
118120
private static class CallSiteMethodVisitor extends MethodVisitor
119121
implements CallSiteAdvice.MethodHandler {
120122
private final Advices advices;
123+
private final Deque<String> newInvocations = new LinkedList<>();
124+
private StackDupMode ctorDupMode = null;
121125

122126
private CallSiteMethodVisitor(
123127
@Nonnull final Advices advices, @Nonnull final MethodVisitor delegated) {
124128
super(ASM_API, delegated);
125129
this.advices = advices;
126130
}
127131

132+
@Override
133+
public void visitTypeInsn(final int opcode, final String type) {
134+
if (opcode == Opcodes.NEW) {
135+
newInvocations.addLast(type);
136+
}
137+
super.visitTypeInsn(opcode, type);
138+
}
139+
128140
@Override
129141
public void visitMethodInsn(
130142
final int opcode,
131143
final String owner,
132144
final String name,
133145
final String descriptor,
134146
final boolean isInterface) {
135-
CallSiteAdvice advice = advices.findAdvice(owner, name, descriptor);
136-
if (advice instanceof InvokeAdvice) {
137-
((InvokeAdvice) advice).apply(this, opcode, owner, name, descriptor, isInterface);
138-
} else {
139-
mv.visitMethodInsn(opcode, owner, name, descriptor, isInterface);
147+
try {
148+
if (opcode == Opcodes.INVOKESPECIAL && "<init>".equals(name)) {
149+
if (owner.equals(newInvocations.peekLast())) {
150+
newInvocations.removeLast();
151+
ctorDupMode = StackDupMode.PREPEND_ARRAY_ON_NEW_CTOR;
152+
} else {
153+
ctorDupMode = StackDupMode.PREPEND_ARRAY_ON_SUPER_CTOR;
154+
}
155+
}
156+
CallSiteAdvice advice = advices.findAdvice(owner, name, descriptor);
157+
// we cannot instrument calls to super in ctor
158+
if (advice instanceof InvokeAdvice) {
159+
((InvokeAdvice) advice).apply(this, opcode, owner, name, descriptor, isInterface);
160+
} else {
161+
mv.visitMethodInsn(opcode, owner, name, descriptor, isInterface);
162+
}
163+
} finally {
164+
ctorDupMode = null;
140165
}
141166
}
142167

@@ -197,6 +222,15 @@ public void method(
197222
mv.visitMethodInsn(opcode, owner, name, descriptor, isInterface);
198223
}
199224

225+
@Override
226+
public void advice(String owner, String name, String descriptor) {
227+
if (ctorDupMode == StackDupMode.PREPEND_ARRAY_ON_SUPER_CTOR) {
228+
// append this to the stack after super call
229+
mv.visitIntInsn(Opcodes.ALOAD, 0);
230+
}
231+
mv.visitMethodInsn(Opcodes.INVOKESTATIC, owner, name, descriptor, false);
232+
}
233+
200234
@Override
201235
public void invokeDynamic(
202236
final String name,
@@ -206,6 +240,11 @@ public void invokeDynamic(
206240
mv.visitInvokeDynamicInsn(name, descriptor, bootstrapMethodHandle, bootstrapMethodArguments);
207241
}
208242

243+
@Override
244+
public void dupCtorParameters(final String methodDescriptor) {
245+
dupParameters(methodDescriptor, ctorDupMode);
246+
}
247+
209248
@Override
210249
public void dupParameters(final String methodDescriptor, final StackDupMode mode) {
211250
final Type method = Type.getMethodType(methodDescriptor);

‎dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/csi/CallSiteUtils.java‎

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -150,7 +150,8 @@ public static void dup(final MethodVisitor mv, final Type[] parameters, final St
150150
dup(mv, parameters);
151151
break;
152152
case PREPEND_ARRAY:
153-
case PREPEND_ARRAY_CTOR:
153+
case PREPEND_ARRAY_ON_NEW_CTOR:
154+
case PREPEND_ARRAY_ON_SUPER_CTOR:
154155
case APPEND_ARRAY:
155156
dupN(mv, parameters, mode);
156157
break;
@@ -279,12 +280,22 @@ private static void dupN(
279280
loadArray(mv, arraySize, parameters);
280281
mv.visitInsn(POP);
281282
break;
282-
case PREPEND_ARRAY_CTOR:
283-
// move the array before the NEW and DUP opcodes
283+
case PREPEND_ARRAY_ON_NEW_CTOR:
284+
// move the array before the uninitialized entry created by NEW and DUP
285+
// stack start = [uninitialized, uninitialized, arg_0, ..., arg_n]
286+
// stack end = [array, uninitialized, uninitialized, arg_0, ..., arg_n]
284287
mv.visitInsn(DUP_X2);
285288
loadArray(mv, arraySize, parameters);
286289
mv.visitInsn(POP);
287290
break;
291+
case PREPEND_ARRAY_ON_SUPER_CTOR:
292+
// move the array before the uninitialized entry
293+
// stack start = [uninitialized, arg_0, ..., arg_n]
294+
// stack end = [array, uninitialized, arg_0, ..., arg_n]
295+
mv.visitInsn(DUP_X1);
296+
loadArray(mv, arraySize, parameters);
297+
mv.visitInsn(POP);
298+
break;
288299
case APPEND_ARRAY:
289300
loadArray(mv, arraySize, parameters);
290301
break;

‎dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/csi/CallSiteAdvice.java‎

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -27,13 +27,19 @@ interface MethodHandler {
2727
/** Performs a method invocation (static, special, virtual, interface...) */
2828
void method(int opcode, String owner, String name, String descriptor, boolean isInterface);
2929

30+
/** Performs an advice invocation (always static) */
31+
void advice(String owner, String name, String descriptor);
32+
3033
/** Performs a dynamic method invocation */
3134
void invokeDynamic(
3235
String name,
3336
String descriptor,
3437
Handle bootstrapMethodHandle,
3538
Object... bootstrapMethodArguments);
3639

40+
/** Duplicates all the ctor parameters in the stack just before the method is invoked. */
41+
void dupCtorParameters(String methodDescriptor);
42+
3743
/** Duplicates all the method parameters in the stack just before the method is invoked. */
3844
void dupParameters(String methodDescriptor, StackDupMode mode);
3945

@@ -62,11 +68,10 @@ enum StackDupMode {
6268
COPY,
6369
/** Copies the parameters in an array and prepends it */
6470
PREPEND_ARRAY,
65-
/**
66-
* Copies the parameters in an array, prepends it and swaps the array with the uninitialized
67-
* instance in a ctor
68-
*/
69-
PREPEND_ARRAY_CTOR,
71+
/** Copies the parameters in an array and adds it between NEW and DUP opcodes */
72+
PREPEND_ARRAY_ON_NEW_CTOR,
73+
/** Copies the parameters in an array and adds it before the uninitialized instance in a ctor */
74+
PREPEND_ARRAY_ON_SUPER_CTOR,
7075
/** Copies the parameters in an array and appends it */
7176
APPEND_ARRAY
7277
}

‎dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/csi/BaseCallSiteTest.groovy‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,8 @@ import net.bytebuddy.jar.asm.Type
1515
import net.bytebuddy.matcher.ElementMatcher
1616
import net.bytebuddy.utility.JavaModule
1717
import net.bytebuddy.utility.nullability.MaybeNull
18+
19+
import java.lang.reflect.Constructor
1820
import java.security.MessageDigest
1921

2022

@@ -81,6 +83,10 @@ class BaseCallSiteTest extends DDSpecification {
8183
return buildPointcut(String.getDeclaredMethod('concat', String))
8284
}
8385

86+
protected static Pointcut stringReaderPointcut() {
87+
return buildPointcut(StringReader.getDeclaredConstructor(String))
88+
}
89+
8490
protected static Pointcut messageDigestGetInstancePointcut() {
8591
return buildPointcut(MessageDigest.getDeclaredMethod('getInstance', String))
8692
}
@@ -100,6 +106,10 @@ class BaseCallSiteTest extends DDSpecification {
100106
return buildPointcut(Type.getType(executable.getDeclaringClass()).internalName, executable.name, Type.getType(executable).descriptor)
101107
}
102108

109+
protected static Pointcut buildPointcut(final Constructor<?> executable) {
110+
return buildPointcut(Type.getType(executable.getDeclaringClass()).internalName, "<init>", Type.getType(executable).descriptor)
111+
}
112+
103113
protected static Pointcut buildPointcut(final String type, final String method, final String descriptor) {
104114
return new Pointcut(type: type, method: method, descriptor: descriptor)
105115
}
@@ -157,6 +167,13 @@ class BaseCallSiteTest extends DDSpecification {
157167
return clazz.getConstructor().newInstance()
158168
}
159169

170+
protected static Class<?> loadClass(final Type type,
171+
final byte[] data,
172+
final ClassLoader loader = Thread.currentThread().contextClassLoader) {
173+
final classLoader = new ByteArrayClassLoader(loader, [(type.className): data])
174+
return classLoader.loadClass(type.className)
175+
}
176+
160177
protected static byte[] transformType(final Type source,
161178
final Type target,
162179
final CallSiteTransformer transformer,

‎dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/csi/CallSiteInstrumentationTest.groovy‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,12 @@
11
package datadog.trace.agent.tooling.csi
22

3+
import datadog.trace.agent.tooling.bytebuddy.csi.CallSiteTransformer
34
import net.bytebuddy.asm.AsmVisitorWrapper
45
import net.bytebuddy.description.type.TypeDescription
56
import net.bytebuddy.dynamic.DynamicType
7+
import net.bytebuddy.jar.asm.Type
8+
9+
import java.util.concurrent.atomic.AtomicInteger
610

711
class CallSiteInstrumentationTest extends BaseCallSiteTest {
812

@@ -60,6 +64,36 @@ class CallSiteInstrumentationTest extends BaseCallSiteTest {
6064
0 * builder.visit(_ as AsmVisitorWrapper) >> builder
6165
}
6266

67+
void 'test call site transformer with super call in ctor'() {
68+
setup:
69+
SuperInCtorExampleAdvice.CALLS.set(0)
70+
final source = Type.getType(SuperInCtorExample)
71+
final target = renameType(source, 'Test')
72+
final pointcut = stringReaderPointcut()
73+
final InvokeAdvice advice = new InvokeAdvice() {
74+
@Override
75+
void apply(CallSiteAdvice.MethodHandler handler, int opcode, String owner, String name, String descriptor, boolean isInterface) {
76+
handler.dupCtorParameters(descriptor)
77+
handler.method(opcode, owner, name, descriptor, isInterface)
78+
handler.advice(
79+
Type.getType(SuperInCtorExampleAdvice).internalName,
80+
'onInvoke',
81+
Type.getMethodType(Type.getType(StringReader), Type.getType(Object[]), Type.getType(StringReader)).getDescriptor(),
82+
)
83+
}
84+
}
85+
final callSiteTransformer = new CallSiteTransformer(mockAdvices([mockCallSites(advice, pointcut)]))
86+
87+
when:
88+
final transformedClass = transformType(source, target, callSiteTransformer)
89+
final transformed = loadClass(target, transformedClass)
90+
final reader = transformed.newInstance("test")
91+
92+
then:
93+
reader != null
94+
SuperInCtorExampleAdvice.CALLS.get() == 1
95+
}
96+
6397
static class StringCallSites implements CallSites, TestCallSites {
6498

6599
@Override
@@ -82,4 +116,14 @@ class CallSiteInstrumentationTest extends BaseCallSiteTest {
82116
handler.method(opcode, owner, name, descriptor, isInterface)
83117
}
84118
}
119+
120+
static class SuperInCtorExampleAdvice {
121+
122+
private static final AtomicInteger CALLS = new AtomicInteger(0)
123+
124+
static StringReader onInvoke(Object[] args, StringReader result) {
125+
CALLS.incrementAndGet()
126+
return result
127+
}
128+
}
85129
}

0 commit comments

Comments
 (0)