Skip to content

Commit cb58b47

Browse files
manuel-alvarez-alvarezmanuel-alvarez-alvarez
committed
Extract Play json body response schemas
1 parent d33bf35 commit cb58b47

File tree

9 files changed

+226
-7
lines changed

9 files changed

+226
-7
lines changed

‎dd-java-agent/instrumentation/play-2.4/src/main/java/datadog/trace/instrumentation/play25/appsec/StatusHeaderInstrumentation.java

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
package datadog.trace.instrumentation.play25.appsec;
2+
3+
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
4+
import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
5+
6+
import com.google.auto.service.AutoService;
7+
import datadog.trace.agent.tooling.Instrumenter;
8+
import datadog.trace.agent.tooling.InstrumenterModule;
9+
import datadog.trace.agent.tooling.muzzle.Reference;
10+
11+
@AutoService(InstrumenterModule.class)
12+
public class StatusHeaderInstrumentation extends InstrumenterModule.AppSec
13+
implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
14+
15+
public StatusHeaderInstrumentation() {
16+
super("play");
17+
}
18+
19+
@Override
20+
public String muzzleDirective() {
21+
return "play25only";
22+
}
23+
24+
@Override
25+
public Reference[] additionalMuzzleReferences() {
26+
return MuzzleReferences.PLAY_25_ONLY;
27+
}
28+
29+
@Override
30+
public String instrumentedType() {
31+
return "play.mvc.StatusHeader";
32+
}
33+
34+
@Override
35+
public void methodAdvice(MethodTransformer transformer) {
36+
transformer.applyAdvice(
37+
named("sendJson").and(takesArgument(0, named("com.fasterxml.jackson.databind.JsonNode"))),
38+
packageName + ".StatusHeaderSendJsonAdvice");
39+
}
40+
}

‎dd-java-agent/instrumentation/play-2.4/src/main/java_play25/datadog/trace/instrumentation/play25/appsec/StatusHeaderSendJsonAdvice.java

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
package datadog.trace.instrumentation.play25.appsec;
2+
3+
import static datadog.trace.api.gateway.Events.EVENTS;
4+
5+
import com.fasterxml.jackson.databind.JsonNode;
6+
import datadog.appsec.api.blocking.BlockingException;
7+
import datadog.trace.advice.ActiveRequestContext;
8+
import datadog.trace.advice.RequiresRequestContext;
9+
import datadog.trace.api.gateway.BlockResponseFunction;
10+
import datadog.trace.api.gateway.CallbackProvider;
11+
import datadog.trace.api.gateway.Flow;
12+
import datadog.trace.api.gateway.RequestContext;
13+
import datadog.trace.api.gateway.RequestContextSlot;
14+
import datadog.trace.bootstrap.CallDepthThreadLocalMap;
15+
import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
16+
import java.util.function.BiFunction;
17+
import net.bytebuddy.asm.Advice;
18+
import play.mvc.StatusHeader;
19+
20+
@RequiresRequestContext(RequestContextSlot.APPSEC)
21+
public class StatusHeaderSendJsonAdvice {
22+
23+
@Advice.OnMethodEnter(suppress = Throwable.class)
24+
static void before() {
25+
CallDepthThreadLocalMap.incrementCallDepth(StatusHeader.class);
26+
}
27+
28+
@Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
29+
static void after(
30+
@Advice.Argument(0) final JsonNode json, @ActiveRequestContext RequestContext reqCtx) {
31+
final int depth = CallDepthThreadLocalMap.decrementCallDepth(StatusHeader.class);
32+
if (depth > 0) {
33+
return;
34+
}
35+
36+
if (json == null) {
37+
return;
38+
}
39+
40+
CallbackProvider cbp = AgentTracer.get().getCallbackProvider(RequestContextSlot.APPSEC);
41+
BiFunction<RequestContext, Object, Flow<Void>> callback =
42+
cbp.getCallback(EVENTS.responseBody());
43+
if (callback == null) {
44+
return;
45+
}
46+
47+
Flow<Void> flow = callback.apply(reqCtx, json);
48+
Flow.Action action = flow.getAction();
49+
if (action instanceof Flow.Action.RequestBlockingAction) {
50+
BlockResponseFunction blockResponseFunction = reqCtx.getBlockResponseFunction();
51+
if (blockResponseFunction == null) {
52+
return;
53+
}
54+
Flow.Action.RequestBlockingAction rba = (Flow.Action.RequestBlockingAction) action;
55+
blockResponseFunction.tryCommitBlockingResponse(
56+
reqCtx.getTraceSegment(),
57+
rba.getStatusCode(),
58+
rba.getBlockingContentType(),
59+
rba.getExtraHeaders());
60+
61+
throw new BlockingException("Blocked request (for StatusHeader/sendJson)");
62+
}
63+
}
64+
}

‎dd-java-agent/instrumentation/play-2.4/src/test/groovy/datadog/trace/instrumentation/play25/server/PlayRouters.groovy

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -130,7 +130,7 @@ class PlayRouters {
130130
->
131131
JsonNode json = body().asJson()
132132
controller(BODY_JSON) {
133-
Results.status(BODY_JSON.status, new ObjectMapper().writeValueAsString(json))
133+
Results.status(BODY_JSON.status, json)
134134
}
135135
} as Supplier)
136136
.build()

‎dd-java-agent/instrumentation/play-2.4/src/test/groovy/datadog/trace/instrumentation/play25/server/PlayServerTest.groovy

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -93,6 +93,11 @@ class PlayServerTest extends HttpServerTest<Server> {
9393
true
9494
}
9595

96+
@Override
97+
boolean testResponseBodyJson() {
98+
true
99+
}
100+
96101
@Override
97102
String testPathParam() {
98103
'/path/?/param'

‎dd-java-agent/instrumentation/play-2.6/src/baseTest/groovy/datadog/trace/instrumentation/play26/server/PlayRouters.groovy

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,11 @@
11
package datadog.trace.instrumentation.play26.server
22

3+
import com.fasterxml.jackson.databind.ObjectMapper
34
import datadog.appsec.api.blocking.Blocking
45
import datadog.trace.agent.test.base.HttpServerTest
56
import groovy.transform.CompileStatic
67
import play.BuiltInComponents
78
import play.api.libs.json.JsValue
8-
import play.api.libs.json.Json$
99
import play.api.mvc.AnyContent
1010
import play.libs.concurrent.HttpExecution
1111
import play.mvc.Http
@@ -150,7 +150,7 @@ class PlayRouters {
150150
->
151151
controller(BODY_JSON) {
152152
JsValue json = body().asJson().get()
153-
Results.status(BODY_JSON.status, Json$.MODULE$.stringify(json))
153+
Results.status(BODY_JSON.status, new ObjectMapper().readTree(json.toString()))
154154
}
155155
} as Supplier)
156156
.POST(BODY_XML.path).routeTo({
@@ -286,7 +286,7 @@ class PlayRouters {
286286
CompletableFuture.supplyAsync({
287287
->
288288
controller(BODY_JSON) {
289-
Results.status(BODY_JSON.status, Json$.MODULE$.stringify(json))
289+
Results.status(BODY_JSON.status, new ObjectMapper().readTree(json.toString()))
290290
}
291291
}, execContext)
292292
} as Supplier)

‎dd-java-agent/instrumentation/play-2.6/src/baseTest/groovy/datadog/trace/instrumentation/play26/server/PlayServerTest.groovy

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -96,6 +96,11 @@ class PlayServerTest extends HttpServerTest<Server> {
9696
true
9797
}
9898

99+
@Override
100+
boolean testResponseBodyJson() {
101+
true
102+
}
103+
99104
@Override
100105
String testPathParam() {
101106
'/path/?/param'

‎dd-java-agent/instrumentation/play-2.6/src/latestDepTest/groovy/datadog/trace/instrumentation/play26/server/latestdep/PlayRouters.groovy

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
package datadog.trace.instrumentation.play26.server.latestdep
22

33
import com.fasterxml.jackson.databind.JsonNode
4-
import com.fasterxml.jackson.databind.ObjectMapper
54
import datadog.appsec.api.blocking.Blocking
65
import datadog.trace.agent.test.base.HttpServerTest
76
import datadog.trace.instrumentation.play26.server.TestHttpErrorHandler
@@ -121,7 +120,7 @@ class PlayRouters {
121120
.POST(BODY_JSON.path).routingTo({ Http.Request req ->
122121
controller(BODY_JSON) {
123122
JsonNode json = req.body().asJson()
124-
Results.status(BODY_JSON.status, new ObjectMapper().writeValueAsString(json))
123+
Results.status(BODY_JSON.status, json)
125124
}
126125
} as RequestFunctions.Params0<Result>)
127126
.POST(BODY_XML.path).routingTo({ Http.Request req ->
@@ -254,7 +253,7 @@ class PlayRouters {
254253
CompletableFuture.supplyAsync({
255254
->
256255
controller(BODY_JSON) {
257-
Results.status(BODY_JSON.status, new ObjectMapper().writeValueAsString(json))
256+
Results.status(BODY_JSON.status, json)
258257
}
259258
}, execContext)
260259
} as RequestFunctions.Params0<? extends CompletionStage<Result>>)

‎dd-java-agent/instrumentation/play-2.6/src/main/java/datadog/trace/instrumentation/play26/appsec/StatusHeaderInstrumentation.java

Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,102 @@
1+
package datadog.trace.instrumentation.play26.appsec;
2+
3+
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
4+
import static datadog.trace.api.gateway.Events.EVENTS;
5+
import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
6+
7+
import com.fasterxml.jackson.databind.JsonNode;
8+
import com.google.auto.service.AutoService;
9+
import datadog.appsec.api.blocking.BlockingException;
10+
import datadog.trace.advice.ActiveRequestContext;
11+
import datadog.trace.advice.RequiresRequestContext;
12+
import datadog.trace.agent.tooling.Instrumenter;
13+
import datadog.trace.agent.tooling.InstrumenterModule;
14+
import datadog.trace.agent.tooling.muzzle.Reference;
15+
import datadog.trace.api.gateway.BlockResponseFunction;
16+
import datadog.trace.api.gateway.CallbackProvider;
17+
import datadog.trace.api.gateway.Flow;
18+
import datadog.trace.api.gateway.RequestContext;
19+
import datadog.trace.api.gateway.RequestContextSlot;
20+
import datadog.trace.bootstrap.CallDepthThreadLocalMap;
21+
import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
22+
import datadog.trace.instrumentation.play26.MuzzleReferences;
23+
import java.util.function.BiFunction;
24+
import net.bytebuddy.asm.Advice;
25+
import play.mvc.StatusHeader;
26+
27+
@AutoService(InstrumenterModule.class)
28+
public class StatusHeaderInstrumentation extends InstrumenterModule.AppSec
29+
implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
30+
31+
public StatusHeaderInstrumentation() {
32+
super("play");
33+
}
34+
35+
@Override
36+
public String muzzleDirective() {
37+
return "play26Plus";
38+
}
39+
40+
@Override
41+
public Reference[] additionalMuzzleReferences() {
42+
return MuzzleReferences.PLAY_26_PLUS; // force failure in <2.6
43+
}
44+
45+
@Override
46+
public String instrumentedType() {
47+
return "play.mvc.StatusHeader";
48+
}
49+
50+
@Override
51+
public void methodAdvice(MethodTransformer transformer) {
52+
transformer.applyAdvice(
53+
named("sendJson").and(takesArgument(0, named("com.fasterxml.jackson.databind.JsonNode"))),
54+
StatusHeaderInstrumentation.class.getName() + "$StatusHeaderSendJsonAdvice");
55+
}
56+
57+
@RequiresRequestContext(RequestContextSlot.APPSEC)
58+
public static class StatusHeaderSendJsonAdvice {
59+
60+
@Advice.OnMethodEnter(suppress = Throwable.class)
61+
static void before() {
62+
CallDepthThreadLocalMap.incrementCallDepth(StatusHeader.class);
63+
}
64+
65+
@Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
66+
static void after(
67+
@Advice.Argument(0) final JsonNode json, @ActiveRequestContext RequestContext reqCtx) {
68+
final int depth = CallDepthThreadLocalMap.decrementCallDepth(StatusHeader.class);
69+
if (depth > 0) {
70+
return;
71+
}
72+
73+
if (json == null) {
74+
return;
75+
}
76+
77+
CallbackProvider cbp = AgentTracer.get().getCallbackProvider(RequestContextSlot.APPSEC);
78+
BiFunction<RequestContext, Object, Flow<Void>> callback =
79+
cbp.getCallback(EVENTS.responseBody());
80+
if (callback == null) {
81+
return;
82+
}
83+
84+
Flow<Void> flow = callback.apply(reqCtx, json);
85+
Flow.Action action = flow.getAction();
86+
if (action instanceof Flow.Action.RequestBlockingAction) {
87+
BlockResponseFunction blockResponseFunction = reqCtx.getBlockResponseFunction();
88+
if (blockResponseFunction == null) {
89+
return;
90+
}
91+
Flow.Action.RequestBlockingAction rba = (Flow.Action.RequestBlockingAction) action;
92+
blockResponseFunction.tryCommitBlockingResponse(
93+
reqCtx.getTraceSegment(),
94+
rba.getStatusCode(),
95+
rba.getBlockingContentType(),
96+
rba.getExtraHeaders());
97+
98+
throw new BlockingException("Blocked request (for StatusHeader/sendJson)");
99+
}
100+
}
101+
}
102+
}

‎dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/base/HttpServerTest.groovy

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@ import datadog.trace.bootstrap.instrumentation.api.URIUtils
3939
import datadog.trace.core.DDSpan
4040
import datadog.trace.core.datastreams.StatsGroup
4141
import datadog.trace.test.util.Flaky
42+
import groovy.json.JsonSlurper
4243
import groovy.transform.Canonical
4344
import groovy.transform.CompileStatic
4445
import net.bytebuddy.utility.RandomString
@@ -2497,6 +2498,9 @@ abstract class HttpServerTest<SERVER> extends WithHttpServer<SERVER> {
24972498

24982499
final BiFunction<RequestContext, Object, Flow<Void>> responseBodyObjectCb =
24992500
({ RequestContext rqCtxt, Object obj ->
2501+
if (obj.class.package.name.startsWith('com.fasterxml.jackson.databind')) {
2502+
obj = new JsonSlurper().parseText(obj.toString())
2503+
}
25002504
if (obj instanceof Map) {
25012505
obj = obj.collectEntries {
25022506
[

0 commit comments

Comments
 (0)