Skip to content

Commit af54160

Browse files
manuel-alvarez-alvarezmanuel-alvarez-alvarez
committed
Remove update to blocking helper
1 parent 7880efa commit af54160

File tree

6 files changed

+115
-36
lines changed

6 files changed

+115
-36
lines changed

‎dd-java-agent/instrumentation/jetty-11/src/main/java11/datadog/trace/instrumentation/jetty11/SetRequestedSessionIdAdvice.java‎

Lines changed: 23 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,16 @@
11
package datadog.trace.instrumentation.jetty11;
22

3+
import static datadog.trace.api.gateway.Events.EVENTS;
4+
5+
import datadog.appsec.api.blocking.BlockingException;
36
import datadog.trace.advice.ActiveRequestContext;
47
import datadog.trace.advice.RequiresRequestContext;
8+
import datadog.trace.api.gateway.CallbackProvider;
9+
import datadog.trace.api.gateway.Flow;
510
import datadog.trace.api.gateway.RequestContext;
611
import datadog.trace.api.gateway.RequestContextSlot;
7-
import datadog.trace.instrumentation.jetty.JettyBlockingHelper;
12+
import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
13+
import java.util.function.BiFunction;
814
import net.bytebuddy.asm.Advice;
915

1016
/**
@@ -17,6 +23,21 @@ public class SetRequestedSessionIdAdvice {
1723
public static void setRequestedSessionId(
1824
@ActiveRequestContext RequestContext reqCtx,
1925
@Advice.Argument(0) final String requestedSessionId) {
20-
JettyBlockingHelper.maybeBlockOnSession(reqCtx, requestedSessionId);
26+
if (requestedSessionId != null && reqCtx != null) {
27+
final CallbackProvider cbp = AgentTracer.get().getCallbackProvider(RequestContextSlot.APPSEC);
28+
if (cbp == null) {
29+
return;
30+
}
31+
final BiFunction<RequestContext, String, Flow<Void>> addrCallback =
32+
cbp.getCallback(EVENTS.requestSession());
33+
if (addrCallback == null) {
34+
return;
35+
}
36+
final Flow<Void> flow = addrCallback.apply(reqCtx, requestedSessionId);
37+
Flow.Action action = flow.getAction();
38+
if (action instanceof Flow.Action.RequestBlockingAction) {
39+
throw new BlockingException("Blocked request (for sessionId)");
40+
}
41+
}
2142
}
2243
}

‎dd-java-agent/instrumentation/jetty-7.0/src/main/java/datadog/trace/instrumentation/jetty70/RequestInstrumentation.java‎

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

33
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
4+
import static datadog.trace.api.gateway.Events.EVENTS;
45
import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.SERVLET_CONTEXT;
56
import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.SERVLET_PATH;
67
import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_DISPATCH_SPAN_ATTRIBUTE;
@@ -10,14 +11,18 @@
1011
import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
1112

1213
import com.google.auto.service.AutoService;
14+
import datadog.appsec.api.blocking.BlockingException;
1315
import datadog.trace.advice.ActiveRequestContext;
1416
import datadog.trace.advice.RequiresRequestContext;
1517
import datadog.trace.agent.tooling.Instrumenter;
1618
import datadog.trace.agent.tooling.InstrumenterModule;
19+
import datadog.trace.api.gateway.CallbackProvider;
20+
import datadog.trace.api.gateway.Flow;
1721
import datadog.trace.api.gateway.RequestContext;
1822
import datadog.trace.api.gateway.RequestContextSlot;
1923
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
20-
import datadog.trace.instrumentation.jetty.JettyBlockingHelper;
24+
import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
25+
import java.util.function.BiFunction;
2126
import net.bytebuddy.asm.Advice;
2227
import org.eclipse.jetty.server.HttpConnection;
2328
import org.eclipse.jetty.server.Request;
@@ -100,7 +105,23 @@ public static class SetRequestedSessionId {
100105
public static void updateContextPath(
101106
@ActiveRequestContext RequestContext reqCtx,
102107
@Advice.Argument(0) final String requestedSessionId) {
103-
JettyBlockingHelper.maybeBlockOnSession(reqCtx, requestedSessionId);
108+
if (requestedSessionId != null && reqCtx != null) {
109+
final CallbackProvider cbp =
110+
AgentTracer.get().getCallbackProvider(RequestContextSlot.APPSEC);
111+
if (cbp == null) {
112+
return;
113+
}
114+
final BiFunction<RequestContext, String, Flow<Void>> addrCallback =
115+
cbp.getCallback(EVENTS.requestSession());
116+
if (addrCallback == null) {
117+
return;
118+
}
119+
final Flow<Void> flow = addrCallback.apply(reqCtx, requestedSessionId);
120+
Flow.Action action = flow.getAction();
121+
if (action instanceof Flow.Action.RequestBlockingAction) {
122+
throw new BlockingException("Blocked request (for sessionId)");
123+
}
124+
}
104125
}
105126
}
106127
}

‎dd-java-agent/instrumentation/jetty-7.6/src/main/java/datadog/trace/instrumentation/jetty76/RequestInstrumentation.java‎

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

33
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
4+
import static datadog.trace.api.gateway.Events.EVENTS;
45
import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.SERVLET_CONTEXT;
56
import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.SERVLET_PATH;
67
import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_DISPATCH_SPAN_ATTRIBUTE;
@@ -10,14 +11,18 @@
1011
import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
1112

1213
import com.google.auto.service.AutoService;
14+
import datadog.appsec.api.blocking.BlockingException;
1315
import datadog.trace.advice.ActiveRequestContext;
1416
import datadog.trace.advice.RequiresRequestContext;
1517
import datadog.trace.agent.tooling.Instrumenter;
1618
import datadog.trace.agent.tooling.InstrumenterModule;
19+
import datadog.trace.api.gateway.CallbackProvider;
20+
import datadog.trace.api.gateway.Flow;
1721
import datadog.trace.api.gateway.RequestContext;
1822
import datadog.trace.api.gateway.RequestContextSlot;
1923
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
20-
import datadog.trace.instrumentation.jetty.JettyBlockingHelper;
24+
import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
25+
import java.util.function.BiFunction;
2126
import net.bytebuddy.asm.Advice;
2227
import org.eclipse.jetty.server.AbstractHttpConnection;
2328
import org.eclipse.jetty.server.Request;
@@ -100,7 +105,23 @@ public static class SetRequestedSessionId {
100105
public static void updateContextPath(
101106
@ActiveRequestContext RequestContext reqCtx,
102107
@Advice.Argument(0) final String requestedSessionId) {
103-
JettyBlockingHelper.maybeBlockOnSession(reqCtx, requestedSessionId);
108+
if (requestedSessionId != null && reqCtx != null) {
109+
final CallbackProvider cbp =
110+
AgentTracer.get().getCallbackProvider(RequestContextSlot.APPSEC);
111+
if (cbp == null) {
112+
return;
113+
}
114+
final BiFunction<RequestContext, String, Flow<Void>> addrCallback =
115+
cbp.getCallback(EVENTS.requestSession());
116+
if (addrCallback == null) {
117+
return;
118+
}
119+
final Flow<Void> flow = addrCallback.apply(reqCtx, requestedSessionId);
120+
Flow.Action action = flow.getAction();
121+
if (action instanceof Flow.Action.RequestBlockingAction) {
122+
throw new BlockingException("Blocked request (for sessionId)");
123+
}
124+
}
104125
}
105126
}
106127
}

‎dd-java-agent/instrumentation/jetty-9/src/main/java/datadog/trace/instrumentation/jetty9/RequestInstrumentation.java‎

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

33
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
4+
import static datadog.trace.api.gateway.Events.EVENTS;
45
import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.SERVLET_CONTEXT;
56
import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.SERVLET_PATH;
67
import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_DISPATCH_SPAN_ATTRIBUTE;
@@ -10,14 +11,18 @@
1011
import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
1112

1213
import com.google.auto.service.AutoService;
14+
import datadog.appsec.api.blocking.BlockingException;
1315
import datadog.trace.advice.ActiveRequestContext;
1416
import datadog.trace.advice.RequiresRequestContext;
1517
import datadog.trace.agent.tooling.Instrumenter;
1618
import datadog.trace.agent.tooling.InstrumenterModule;
19+
import datadog.trace.api.gateway.CallbackProvider;
20+
import datadog.trace.api.gateway.Flow;
1721
import datadog.trace.api.gateway.RequestContext;
1822
import datadog.trace.api.gateway.RequestContextSlot;
1923
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
20-
import datadog.trace.instrumentation.jetty.JettyBlockingHelper;
24+
import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
25+
import java.util.function.BiFunction;
2126
import javax.servlet.http.HttpServletRequest;
2227
import net.bytebuddy.asm.Advice;
2328
import org.eclipse.jetty.http.HttpFields;
@@ -110,7 +115,23 @@ public static class SetRequestedSessionIdAdvice {
110115
public static void updateContextPath(
111116
@ActiveRequestContext RequestContext reqCtx,
112117
@Advice.Argument(0) final String requestedSessionId) {
113-
JettyBlockingHelper.maybeBlockOnSession(reqCtx, requestedSessionId);
118+
if (requestedSessionId != null && reqCtx != null) {
119+
final CallbackProvider cbp =
120+
AgentTracer.get().getCallbackProvider(RequestContextSlot.APPSEC);
121+
if (cbp == null) {
122+
return;
123+
}
124+
final BiFunction<RequestContext, String, Flow<Void>> addrCallback =
125+
cbp.getCallback(EVENTS.requestSession());
126+
if (addrCallback == null) {
127+
return;
128+
}
129+
final Flow<Void> flow = addrCallback.apply(reqCtx, requestedSessionId);
130+
Flow.Action action = flow.getAction();
131+
if (action instanceof Flow.Action.RequestBlockingAction) {
132+
throw new BlockingException("Blocked request (for sessionId)");
133+
}
134+
}
114135
}
115136
}
116137
}

‎dd-java-agent/instrumentation/jetty-9/src/main/java_jetty10/datadog/trace/instrumentation/jetty10/SetRequestedSessionIdAdvice.java‎

Lines changed: 23 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,16 @@
11
package datadog.trace.instrumentation.jetty10;
22

3+
import static datadog.trace.api.gateway.Events.EVENTS;
4+
5+
import datadog.appsec.api.blocking.BlockingException;
36
import datadog.trace.advice.ActiveRequestContext;
47
import datadog.trace.advice.RequiresRequestContext;
8+
import datadog.trace.api.gateway.CallbackProvider;
9+
import datadog.trace.api.gateway.Flow;
510
import datadog.trace.api.gateway.RequestContext;
611
import datadog.trace.api.gateway.RequestContextSlot;
7-
import datadog.trace.instrumentation.jetty.JettyBlockingHelper;
12+
import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
13+
import java.util.function.BiFunction;
814
import net.bytebuddy.asm.Advice;
915

1016
/**
@@ -17,6 +23,21 @@ public class SetRequestedSessionIdAdvice {
1723
public static void setRequestedSessionId(
1824
@ActiveRequestContext RequestContext reqCtx,
1925
@Advice.Argument(0) final String requestedSessionId) {
20-
JettyBlockingHelper.maybeBlockOnSession(reqCtx, requestedSessionId);
26+
if (requestedSessionId != null && reqCtx != null) {
27+
final CallbackProvider cbp = AgentTracer.get().getCallbackProvider(RequestContextSlot.APPSEC);
28+
if (cbp == null) {
29+
return;
30+
}
31+
final BiFunction<RequestContext, String, Flow<Void>> addrCallback =
32+
cbp.getCallback(EVENTS.requestSession());
33+
if (addrCallback == null) {
34+
return;
35+
}
36+
final Flow<Void> flow = addrCallback.apply(reqCtx, requestedSessionId);
37+
Flow.Action action = flow.getAction();
38+
if (action instanceof Flow.Action.RequestBlockingAction) {
39+
throw new BlockingException("Blocked request (for sessionId)");
40+
}
41+
}
2142
}
2243
}

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

Lines changed: 0 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,15 @@
11
package datadog.trace.instrumentation.jetty;
22

3-
import static datadog.trace.api.gateway.Events.EVENTS;
43
import static java.lang.invoke.MethodHandles.collectArguments;
54
import static java.lang.invoke.MethodHandles.lookup;
65
import static java.lang.invoke.MethodType.methodType;
76

87
import datadog.appsec.api.blocking.BlockingContentType;
98
import datadog.appsec.api.blocking.BlockingException;
10-
import datadog.trace.api.gateway.CallbackProvider;
119
import datadog.trace.api.gateway.Flow;
12-
import datadog.trace.api.gateway.RequestContext;
13-
import datadog.trace.api.gateway.RequestContextSlot;
1410
import datadog.trace.api.internal.TraceSegment;
1511
import datadog.trace.bootstrap.blocking.BlockingActionHelper;
1612
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
17-
import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
1813
import datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator;
1914
import java.io.OutputStream;
2015
import java.io.PrintWriter;
@@ -24,7 +19,6 @@
2419
import java.lang.reflect.Method;
2520
import java.nio.charset.StandardCharsets;
2621
import java.util.Map;
27-
import java.util.function.BiFunction;
2822
import org.eclipse.jetty.io.EndPoint;
2923
import org.eclipse.jetty.server.Request;
3024
import org.eclipse.jetty.server.Response;
@@ -248,24 +242,4 @@ public static void blockAndThrowOnFailure(
248242
throw new BlockingException("Throwing after being unable to commit blocking response");
249243
}
250244
}
251-
252-
public static void maybeBlockOnSession(
253-
final RequestContext reqCtx, final String requestedSessionId) {
254-
if (requestedSessionId != null && reqCtx != null) {
255-
final CallbackProvider cbp = AgentTracer.get().getCallbackProvider(RequestContextSlot.APPSEC);
256-
if (cbp == null) {
257-
return;
258-
}
259-
final BiFunction<RequestContext, String, Flow<Void>> addrCallback =
260-
cbp.getCallback(EVENTS.requestSession());
261-
if (addrCallback == null) {
262-
return;
263-
}
264-
final Flow<Void> flow = addrCallback.apply(reqCtx, requestedSessionId);
265-
Flow.Action action = flow.getAction();
266-
if (action instanceof Flow.Action.RequestBlockingAction) {
267-
throw new BlockingException("Blocked request (for sessionId)");
268-
}
269-
}
270-
}
271245
}

0 commit comments

Comments
 (0)