Skip to content

Commit be7f9d3

Browse files
authored
Store the http.route tag value inside the iast request context in Play (#9105)
What Does This Do Store the http.route tag value inside the iast request context in Play framework instrumentation Move PlayHttpServerDecorator.onRequest to onEnter advice from onExit advice. We need to send the event on enter to have the info in the context available when vulns are detected during the requests Motivation IAST sampling algorithm requires the http.route span tag to be set on the local root span so it can be used for its sampling decision. Since Play does not use the local root span for the http.route we have to store it in the iast request context before the sampling decision is made. Additional Notes related with #8991
1 parent f6894da commit be7f9d3

File tree

32 files changed

+1093
-23
lines changed

32 files changed

+1093
-23
lines changed
Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
package com.datadog.iast;
2+
3+
import datadog.trace.api.gateway.RequestContext;
4+
import datadog.trace.api.gateway.RequestContextSlot;
5+
import java.util.function.BiConsumer;
6+
7+
public class HttpRouteHandler implements BiConsumer<RequestContext, String> {
8+
9+
@Override
10+
public void accept(final RequestContext ctx, final String route) {
11+
final IastRequestContext iastCtx = ctx.getData(RequestContextSlot.IAST);
12+
if (iastCtx != null) {
13+
iastCtx.setRoute(route);
14+
}
15+
}
16+
}

dd-java-agent/agent-iast/src/main/java/com/datadog/iast/IastRequestContext.java

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ public class IastRequestContext implements IastContext, HasMetricCollector {
3636
@Nullable private volatile String xForwardedProto;
3737
@Nullable private volatile String contentType;
3838
@Nullable private volatile String authorization;
39+
@Nullable private volatile String route;
3940

4041
/**
4142
* Use {@link IastRequestContext#IastRequestContext(TaintedObjects)} instead as we require more
@@ -121,6 +122,15 @@ public void setAuthorization(final String authorization) {
121122
this.authorization = authorization;
122123
}
123124

125+
@Nullable
126+
public String getRoute() {
127+
return route;
128+
}
129+
130+
public void setRoute(final String route) {
131+
this.route = route;
132+
}
133+
124134
public OverheadContext getOverheadContext() {
125135
return overheadContext;
126136
}

dd-java-agent/agent-iast/src/main/java/com/datadog/iast/IastSystem.java

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -122,6 +122,7 @@ public static void start(
122122
registerRequestStartedCallback(ss, addTelemetry, dependencies);
123123
registerRequestEndedCallback(ss, addTelemetry, dependencies);
124124
registerHeadersCallback(ss);
125+
registerHttpRouteCallback(ss);
125126
registerGrpcServerRequestMessageCallback(ss);
126127
maybeApplySecurityControls(instrumentation);
127128
LOGGER.debug("IAST started");
@@ -246,6 +247,10 @@ private static void registerHeadersCallback(final SubscriptionService ss) {
246247
ss.registerCallback(event, handler);
247248
}
248249

250+
private static void registerHttpRouteCallback(final SubscriptionService ss) {
251+
ss.registerCallback(Events.get().httpRoute(), new HttpRouteHandler());
252+
}
253+
249254
private static void registerGrpcServerRequestMessageCallback(final SubscriptionService ss) {
250255
ss.registerCallback(Events.get().grpcServerRequestMessage(), new GrpcRequestMessageHandler());
251256
}

dd-java-agent/agent-iast/src/main/java/com/datadog/iast/overhead/OverheadController.java

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -233,7 +233,7 @@ public boolean consumeQuota(
233233
Object methodTag = rootSpan.getTag(Tags.HTTP_METHOD);
234234
method = (methodTag == null) ? "" : methodTag.toString();
235235
Object routeTag = rootSpan.getTag(Tags.HTTP_ROUTE);
236-
path = (routeTag == null) ? "" : routeTag.toString();
236+
path = (routeTag == null) ? getHttpRouteFromRequestContext(span) : routeTag.toString();
237237
}
238238
if (!maybeSkipVulnerability(ctx, type, method, path)) {
239239
return operation.consumeQuota(ctx);
@@ -316,6 +316,19 @@ public OverheadContext getContext(@Nullable final AgentSpan span) {
316316
return globalContext;
317317
}
318318

319+
@Nullable
320+
public String getHttpRouteFromRequestContext(@Nullable final AgentSpan span) {
321+
String httpRoute = null;
322+
final RequestContext requestContext = span != null ? span.getRequestContext() : null;
323+
if (requestContext != null) {
324+
IastRequestContext iastRequestContext = requestContext.getData(RequestContextSlot.IAST);
325+
if (iastRequestContext != null) {
326+
httpRoute = iastRequestContext.getRoute();
327+
}
328+
}
329+
return httpRoute == null ? "" : httpRoute;
330+
}
331+
319332
static int computeSamplingParameter(final float pct) {
320333
if (pct >= 100) {
321334
return 100;
Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
package com.datadog.iast
2+
3+
import datadog.trace.api.gateway.RequestContext
4+
import datadog.trace.api.gateway.RequestContextSlot
5+
import datadog.trace.test.util.DDSpecification
6+
import groovy.transform.CompileDynamic
7+
8+
@CompileDynamic
9+
class HttpRouteHandlerTest extends DDSpecification {
10+
void 'route is set'() {
11+
given:
12+
final handler = new HttpRouteHandler()
13+
final iastCtx = Mock(IastRequestContext)
14+
final ctx = Mock(RequestContext)
15+
ctx.getData(RequestContextSlot.IAST) >> iastCtx
16+
17+
when:
18+
handler.accept(ctx, '/foo')
19+
20+
then:
21+
1 * ctx.getData(RequestContextSlot.IAST) >> iastCtx
22+
1 * iastCtx.setRoute('/foo')
23+
0 * _
24+
}
25+
26+
void 'does nothing when context missing'() {
27+
given:
28+
final handler = new HttpRouteHandler()
29+
final ctx = Mock(RequestContext)
30+
ctx.getData(RequestContextSlot.IAST) >> null
31+
32+
when:
33+
handler.accept(ctx, '/foo')
34+
35+
then:
36+
1 * ctx.getData(RequestContextSlot.IAST) >> null
37+
0 * _
38+
}
39+
}

dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/IastSystemTest.groovy

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,7 @@ class IastSystemTest extends DDSpecification {
6060
1 * ss.registerCallback(Events.get().requestStarted(), _)
6161
1 * ss.registerCallback(Events.get().requestEnded(), _)
6262
1 * ss.registerCallback(Events.get().requestHeader(), _)
63+
1 * ss.registerCallback(Events.get().httpRoute(), _)
6364
1 * ss.registerCallback(Events.get().grpcServerRequestMessage(), _)
6465
0 * _
6566
TestLogCollector.drainCapturedLogs().any { it.message.contains('IAST is starting') }

dd-java-agent/instrumentation/play-2.3/src/main/java/datadog/trace/instrumentation/play23/PlayHttpServerDecorator.java

Lines changed: 13 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -112,13 +112,21 @@ private void dispatchRoute(final AgentSpan span, final String route) {
112112
if (ctx == null) {
113113
return;
114114
}
115+
// Send event to AppSec provider
115116
final CallbackProvider cbp = tracer().getCallbackProvider(RequestContextSlot.APPSEC);
116-
if (cbp == null) {
117-
return;
117+
if (cbp != null) {
118+
final BiConsumer<RequestContext, String> cb = cbp.getCallback(EVENTS.httpRoute());
119+
if (cb != null) {
120+
cb.accept(ctx, route);
121+
}
118122
}
119-
final BiConsumer<RequestContext, String> cb = cbp.getCallback(EVENTS.httpRoute());
120-
if (cb != null) {
121-
cb.accept(ctx, route);
123+
// Send event to IAST provider
124+
final CallbackProvider cbpIast = tracer().getCallbackProvider(RequestContextSlot.IAST);
125+
if (cbpIast != null) {
126+
final BiConsumer<RequestContext, String> cb = cbpIast.getCallback(EVENTS.httpRoute());
127+
if (cb != null) {
128+
cb.accept(ctx, route);
129+
}
122130
}
123131
} catch (final Throwable t) {
124132
LOG.debug("Failed to dispatch route", t);

dd-java-agent/instrumentation/play-2.4/src/main/java/datadog/trace/instrumentation/play24/PlayAdvice.java

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,10 @@ public static ContextScope onEnter(@Advice.Argument(value = 0, readOnly = false)
4646

4747
req = RequestHelper.withTag(req, "_dd_HasPlayRequestSpan", "true");
4848

49+
// Moved from OnMethodExit
50+
// Call onRequest on return after tags are populated.
51+
DECORATE.onRequest(span, req, req, (AgentSpanContext.Extracted) null);
52+
4953
return scope;
5054
}
5155

@@ -63,9 +67,6 @@ public static void stopTraceOnResponse(
6367

6468
final AgentSpan playControllerSpan = spanFromContext(playControllerScope.context());
6569

66-
// Call onRequest on return after tags are populated.
67-
DECORATE.onRequest(playControllerSpan, req, req, (AgentSpanContext.Extracted) null);
68-
6970
if (throwable == null) {
7071
responseFuture.onComplete(
7172
new RequestCompleteCallback(playControllerSpan),

dd-java-agent/instrumentation/play-2.4/src/main/java/datadog/trace/instrumentation/play24/PlayHttpServerDecorator.java

Lines changed: 13 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -112,13 +112,21 @@ private void dispatchRoute(final AgentSpan span, final String route) {
112112
if (ctx == null) {
113113
return;
114114
}
115+
// Send event to AppSec provider
115116
final CallbackProvider cbp = tracer().getCallbackProvider(RequestContextSlot.APPSEC);
116-
if (cbp == null) {
117-
return;
117+
if (cbp != null) {
118+
final BiConsumer<RequestContext, String> cb = cbp.getCallback(EVENTS.httpRoute());
119+
if (cb != null) {
120+
cb.accept(ctx, route);
121+
}
118122
}
119-
final BiConsumer<RequestContext, String> cb = cbp.getCallback(EVENTS.httpRoute());
120-
if (cb != null) {
121-
cb.accept(ctx, route);
123+
// Send event to IAST provider
124+
final CallbackProvider cbpIast = tracer().getCallbackProvider(RequestContextSlot.IAST);
125+
if (cbpIast != null) {
126+
final BiConsumer<RequestContext, String> cb = cbpIast.getCallback(EVENTS.httpRoute());
127+
if (cb != null) {
128+
cb.accept(ctx, route);
129+
}
122130
}
123131
} catch (final Throwable t) {
124132
LOG.debug("Failed to dispatch route", t);

dd-java-agent/instrumentation/play-2.6/src/main/java/datadog/trace/instrumentation/play26/PlayAdvice.java

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,10 @@ public static ContextScope onEnter(
4747

4848
req = req.addAttr(HasPlayRequestSpan.KEY, HasPlayRequestSpan.INSTANCE);
4949

50+
// Moved from OnMethodExit
51+
// Call onRequest on return after tags are populated.
52+
DECORATE.onRequest(span, req, req, extractedContext);
53+
5054
return scope;
5155
}
5256

@@ -65,9 +69,6 @@ public static void stopTraceOnResponse(
6569

6670
final AgentSpan playControllerSpan = spanFromContext(playControllerScope.context());
6771

68-
// Call onRequest on return after tags are populated.
69-
DECORATE.onRequest(playControllerSpan, req, req, extractedContext);
70-
7172
if (throwable == null) {
7273
responseFuture.onComplete(
7374
new RequestCompleteCallback(playControllerSpan),

0 commit comments

Comments
 (0)