Skip to content

Commit 089d457

Browse files
PerfectSlayerPerfectSlayer
committed
wip(asm): Fix startSpan hook
1 parent 0106a4c commit 089d457

File tree

3 files changed

+33
-81
lines changed

3 files changed

+33
-81
lines changed

‎dd-java-agent/instrumentation/jetty-common/src/main/java/datadog/trace/instrumentation/jetty/HandleRequestVisitor.java

Lines changed: 9 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
package datadog.trace.instrumentation.jetty;
22

3-
import datadog.trace.api.gateway.Flow;
4-
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
3+
import datadog.context.Context;
54
import java.util.List;
65
import net.bytebuddy.jar.asm.Label;
76
import net.bytebuddy.jar.asm.MethodVisitor;
@@ -15,23 +14,20 @@
1514
* and replaces it with:
1615
*
1716
* <pre>
18-
* if (span != null && span.getRequestBlockingAction() &&
19-
* JettyBlockingHelper.block(
20-
* this.getRequest(), this.getResponse(),
21-
* span.getRequestBlockingAction(), span) {
17+
* if (JettyBlockingHelper.block(this.getRequest(), this.getResponse(), context) {
2218
* // nothing
2319
* } else {
2420
* server.handle(this);
2521
* }
2622
* </pre>
2723
*
28-
* <p>It needs first to get the index of the span variable that's set when a new span is created.
24+
* <p>It needs first to get the index of the context variable that's set when a new span is created.
2925
*/
3026
public class HandleRequestVisitor extends MethodVisitor {
3127
private static final Logger log = LoggerFactory.getLogger(HandleRequestVisitor.class);
3228

3329
private boolean lookForStore;
34-
private int agentSpanVar = -1;
30+
private int contextVar = -1;
3531
private final int classVersion;
3632
private final String connClassInternalName;
3733
private boolean success;
@@ -53,12 +49,12 @@ DelayLoadsMethodVisitor delayVisitorDelegate() {
5349
@Override
5450
public void visitMethodInsn(
5551
int opcode, String owner, String name, String descriptor, boolean isInterface) {
56-
if (agentSpanVar == -1) {
52+
if (contextVar == -1) {
5753
lookForStore =
5854
!lookForStore
5955
&& opcode == Opcodes.INVOKEVIRTUAL
6056
&& name.equals("startSpan")
61-
&& descriptor.endsWith("Ldatadog/trace/bootstrap/instrumentation/api/AgentSpan;");
57+
&& descriptor.endsWith("Ldatadog.context.Context;");
6258
} else if (opcode == Opcodes.INVOKEVIRTUAL
6359
&& owner.equals("org/eclipse/jetty/server/Server")
6460
&& name.equals("handle")
@@ -71,26 +67,8 @@ public void visitMethodInsn(
7167
return;
7268
}
7369

74-
Label doBlockLabel = new Label();
75-
Label beforeHandle = new Label();
7670
Label afterHandle = new Label();
7771

78-
super.visitVarInsn(Opcodes.ALOAD, agentSpanVar);
79-
super.visitJumpInsn(Opcodes.IFNULL, beforeHandle);
80-
super.visitVarInsn(Opcodes.ALOAD, agentSpanVar);
81-
super.visitMethodInsn(
82-
Opcodes.INVOKEINTERFACE,
83-
"datadog/trace/bootstrap/instrumentation/api/AgentSpan",
84-
"getRequestBlockingAction",
85-
"()" + Type.getDescriptor(Flow.Action.RequestBlockingAction.class),
86-
true);
87-
super.visitJumpInsn(Opcodes.IFNONNULL, doBlockLabel);
88-
super.visitJumpInsn(Opcodes.GOTO, beforeHandle);
89-
90-
super.visitLabel(doBlockLabel);
91-
if (needsStackFrames()) {
92-
super.visitFrame(Opcodes.F_SAME, 0, null, 0, null);
93-
}
9472
super.visitVarInsn(Opcodes.ALOAD, 0);
9573
super.visitMethodInsn(
9674
Opcodes.INVOKEVIRTUAL,
@@ -105,29 +83,17 @@ public void visitMethodInsn(
10583
"getResponse",
10684
"()Lorg/eclipse/jetty/server/Response;",
10785
false);
108-
super.visitVarInsn(Opcodes.ALOAD, agentSpanVar);
109-
super.visitMethodInsn(
110-
Opcodes.INVOKEINTERFACE,
111-
"datadog/trace/bootstrap/instrumentation/api/AgentSpan",
112-
"getRequestBlockingAction",
113-
"()" + Type.getDescriptor(Flow.Action.RequestBlockingAction.class),
114-
true);
115-
super.visitVarInsn(Opcodes.ALOAD, agentSpanVar);
86+
super.visitVarInsn(Opcodes.ALOAD, contextVar);
11687
super.visitMethodInsn(
11788
Opcodes.INVOKESTATIC,
11889
Type.getInternalName(JettyBlockingHelper.class),
11990
"block",
12091
"(Lorg/eclipse/jetty/server/Request;Lorg/eclipse/jetty/server/Response;"
121-
+ Type.getDescriptor(Flow.Action.RequestBlockingAction.class)
122-
+ Type.getDescriptor(AgentSpan.class)
92+
+ Type.getDescriptor(Context.class)
12393
+ ")Z",
12494
false);
12595
super.visitJumpInsn(Opcodes.IFNE, afterHandle);
12696

127-
super.visitLabel(beforeHandle);
128-
if (needsStackFrames()) {
129-
super.visitFrame(Opcodes.F_SAME, 0, null, 0, null);
130-
}
13197
mv.commitLoads(savedLoads);
13298
super.visitMethodInsn(opcode, owner, name, descriptor, isInterface);
13399
super.visitLabel(afterHandle);
@@ -144,7 +110,7 @@ public void visitMethodInsn(
144110
@Override
145111
public void visitVarInsn(int opcode, int varIndex) {
146112
if (lookForStore && opcode == Opcodes.ASTORE) {
147-
agentSpanVar = varIndex;
113+
contextVar = varIndex;
148114
lookForStore = false;
149115
}
150116

‎dd-java-agent/instrumentation/jetty-common/src/main/java/datadog/trace/instrumentation/jetty/JettyBlockingHelper.java

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,10 +6,12 @@
66

77
import datadog.appsec.api.blocking.BlockingContentType;
88
import datadog.appsec.api.blocking.BlockingException;
9+
import datadog.context.Context;
910
import datadog.trace.api.gateway.Flow;
1011
import datadog.trace.api.internal.TraceSegment;
1112
import datadog.trace.bootstrap.blocking.BlockingActionHelper;
1213
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
14+
import datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge;
1315
import datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator;
1416
import java.io.OutputStream;
1517
import java.io.PrintWriter;
@@ -225,6 +227,15 @@ public static boolean block(
225227
return true;
226228
}
227229

230+
public static boolean block(Request request, Response response, Context context) {
231+
AgentSpan span = Java8BytecodeBridge.spanFromContext(context);
232+
Flow.Action.RequestBlockingAction rba;
233+
if (span == null || (rba = span.getRequestBlockingAction()) == null) {
234+
return false;
235+
}
236+
return block(request, response, rba, span);
237+
}
238+
228239
public static boolean block(
229240
Request request, Response response, Flow.Action.RequestBlockingAction rba, AgentSpan span) {
230241
return block(

‎dd-java-agent/instrumentation/jetty-common/src/main/java/datadog/trace/instrumentation/jetty9/HandleVisitor.java

Lines changed: 13 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package datadog.trace.instrumentation.jetty9;
22

3+
import datadog.context.Context;
34
import datadog.trace.api.gateway.Flow;
45
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
56
import datadog.trace.instrumentation.jetty.JettyBlockingHelper;
@@ -25,8 +26,7 @@
2526
* </code> is replaced with: <code>
2627
* case REQUEST_DISPATCH:
2728
* // ...
28-
* if (span != null && span.getBlockingAction() != null &&
29-
* JettyBlockingHelper.block(this.getRequest(), this.getResponse())) {
29+
* if (JettyBlockingHelper.block(this.getRequest(), this.getResponse(), context)) {
3030
* // nothing
3131
* } else {
3232
* getServer().handle(this);
@@ -82,7 +82,7 @@ public class HandleVisitor extends MethodVisitor {
8282
private static final Logger log = LoggerFactory.getLogger(HandleVisitor.class);
8383

8484
private boolean lookForStore;
85-
private int agentSpanVar = -1;
85+
private int contextVar = -1;
8686
private boolean success;
8787
private final String methodName;
8888

@@ -98,12 +98,13 @@ DelayCertainInsMethodVisitor delayVisitorDelegate() {
9898
@Override
9999
public void visitMethodInsn(
100100
int opcode, String owner, String name, String descriptor, boolean isInterface) {
101-
if (agentSpanVar == -1) {
101+
if (contextVar == -1) {
102102
lookForStore =
103103
!lookForStore
104104
&& opcode == Opcodes.INVOKEVIRTUAL
105105
&& name.equals("startSpan")
106-
&& descriptor.endsWith("Ldatadog/trace/bootstrap/instrumentation/api/AgentSpan;");
106+
&& descriptor.endsWith("Ldatadog.context.Context;");
107+
&& descriptor.endsWith("Ldatadog/context/Context;");
107108
} else if (!success
108109
&& opcode == Opcodes.INVOKEVIRTUAL
109110
&& owner.equals("org/eclipse/jetty/server/Server")
@@ -124,24 +125,8 @@ public void visitMethodInsn(
124125
return;
125126
}
126127

127-
Label doBlockLabel = new Label();
128-
Label beforeHandle = new Label();
129128
Label afterHandle = new Label();
130129

131-
super.visitVarInsn(Opcodes.ALOAD, agentSpanVar);
132-
super.visitJumpInsn(Opcodes.IFNULL, beforeHandle);
133-
super.visitVarInsn(Opcodes.ALOAD, agentSpanVar);
134-
super.visitMethodInsn(
135-
Opcodes.INVOKEINTERFACE,
136-
"datadog/trace/bootstrap/instrumentation/api/AgentSpan",
137-
"getRequestBlockingAction",
138-
"()" + Type.getDescriptor(Flow.Action.RequestBlockingAction.class),
139-
true);
140-
super.visitJumpInsn(Opcodes.IFNONNULL, doBlockLabel);
141-
super.visitJumpInsn(Opcodes.GOTO, beforeHandle);
142-
143-
super.visitLabel(doBlockLabel);
144-
super.visitFrame(Opcodes.F_SAME, 0, null, 0, null);
145130
super.visitVarInsn(Opcodes.ALOAD, 0);
146131
super.visitMethodInsn(
147132
Opcodes.INVOKEVIRTUAL,
@@ -156,27 +141,17 @@ public void visitMethodInsn(
156141
"getResponse",
157142
"()Lorg/eclipse/jetty/server/Response;",
158143
false);
159-
super.visitVarInsn(Opcodes.ALOAD, agentSpanVar);
160-
super.visitMethodInsn(
161-
Opcodes.INVOKEINTERFACE,
162-
"datadog/trace/bootstrap/instrumentation/api/AgentSpan",
163-
"getRequestBlockingAction",
164-
"()" + Type.getDescriptor(Flow.Action.RequestBlockingAction.class),
165-
true);
166-
super.visitVarInsn(Opcodes.ALOAD, agentSpanVar);
144+
super.visitVarInsn(Opcodes.ALOAD, contextVar);
167145
super.visitMethodInsn(
168146
Opcodes.INVOKESTATIC,
169147
Type.getInternalName(JettyBlockingHelper.class),
170148
"block",
171149
"(Lorg/eclipse/jetty/server/Request;Lorg/eclipse/jetty/server/Response;"
172-
+ Type.getDescriptor(Flow.Action.RequestBlockingAction.class)
173-
+ Type.getDescriptor(AgentSpan.class)
150+
+ Type.getDescriptor(Context.class)
174151
+ ")Z",
175152
false);
176153
super.visitJumpInsn(Opcodes.IFNE, afterHandle);
177154

178-
super.visitLabel(beforeHandle);
179-
super.visitFrame(Opcodes.F_SAME, 0, null, 0, null);
180155
mv.commitVisitations(savedVisitations);
181156
super.visitMethodInsn(opcode, owner, name, descriptor, isInterface);
182157
super.visitLabel(afterHandle);
@@ -214,9 +189,9 @@ public void visitMethodInsn(
214189
Label beforeRegularDispatch = new Label();
215190
Label afterRegularDispatch = new Label();
216191

217-
super.visitVarInsn(Opcodes.ALOAD, agentSpanVar);
192+
super.visitVarInsn(Opcodes.ALOAD, contextVar);
218193
super.visitJumpInsn(Opcodes.IFNULL, beforeRegularDispatch);
219-
super.visitVarInsn(Opcodes.ALOAD, agentSpanVar);
194+
super.visitVarInsn(Opcodes.ALOAD, contextVar);
220195
super.visitMethodInsn(
221196
Opcodes.INVOKEINTERFACE,
222197
"datadog/trace/bootstrap/instrumentation/api/AgentSpan",
@@ -248,14 +223,14 @@ public void visitMethodInsn(
248223
"getResponse",
249224
"()Lorg/eclipse/jetty/server/Response;",
250225
false);
251-
super.visitVarInsn(Opcodes.ALOAD, agentSpanVar);
226+
super.visitVarInsn(Opcodes.ALOAD, contextVar);
252227
super.visitMethodInsn(
253228
Opcodes.INVOKEINTERFACE,
254229
"datadog/trace/bootstrap/instrumentation/api/AgentSpan",
255230
"getRequestBlockingAction",
256231
"()" + Type.getDescriptor(Flow.Action.RequestBlockingAction.class),
257232
true);
258-
super.visitVarInsn(Opcodes.ALOAD, agentSpanVar);
233+
super.visitVarInsn(Opcodes.ALOAD, contextVar);
259234

260235
// create the lambda
261236
super.visitInvokeDynamicInsn(
@@ -337,7 +312,7 @@ private boolean checkDispatchMethodState(final List<Function> savedVisitations)
337312
@Override
338313
public void visitVarInsn(int opcode, int varIndex) {
339314
if (lookForStore && opcode == Opcodes.ASTORE) {
340-
agentSpanVar = varIndex;
315+
contextVar = varIndex;
341316
lookForStore = false;
342317
}
343318

0 commit comments

Comments
 (0)