Skip to content

Commit 22b45ef

Browse files
manuel-alvarez-alvarezmanuel-alvarez-alvarez
committed
Extract Play json body response schemas
1 parent fe0c272 commit 22b45ef

File tree

8 files changed

+228
-7
lines changed

8 files changed

+228
-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: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,67 @@
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+
if (cbp == null) {
42+
return;
43+
}
44+
BiFunction<RequestContext, Object, Flow<Void>> callback =
45+
cbp.getCallback(EVENTS.responseBody());
46+
if (callback == null) {
47+
return;
48+
}
49+
50+
Flow<Void> flow = callback.apply(reqCtx, json);
51+
Flow.Action action = flow.getAction();
52+
if (action instanceof Flow.Action.RequestBlockingAction) {
53+
BlockResponseFunction blockResponseFunction = reqCtx.getBlockResponseFunction();
54+
if (blockResponseFunction == null) {
55+
return;
56+
}
57+
Flow.Action.RequestBlockingAction rba = (Flow.Action.RequestBlockingAction) action;
58+
blockResponseFunction.tryCommitBlockingResponse(
59+
reqCtx.getTraceSegment(),
60+
rba.getStatusCode(),
61+
rba.getBlockingContentType(),
62+
rba.getExtraHeaders());
63+
64+
throw new BlockingException("Blocked request (for StatusHeader/sendJson)");
65+
}
66+
}
67+
}

‎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: 105 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,105 @@
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+
if (cbp == null) {
79+
return;
80+
}
81+
BiFunction<RequestContext, Object, Flow<Void>> callback =
82+
cbp.getCallback(EVENTS.responseBody());
83+
if (callback == null) {
84+
return;
85+
}
86+
87+
Flow<Void> flow = callback.apply(reqCtx, json);
88+
Flow.Action action = flow.getAction();
89+
if (action instanceof Flow.Action.RequestBlockingAction) {
90+
BlockResponseFunction blockResponseFunction = reqCtx.getBlockResponseFunction();
91+
if (blockResponseFunction == null) {
92+
return;
93+
}
94+
Flow.Action.RequestBlockingAction rba = (Flow.Action.RequestBlockingAction) action;
95+
blockResponseFunction.tryCommitBlockingResponse(
96+
reqCtx.getTraceSegment(),
97+
rba.getStatusCode(),
98+
rba.getBlockingContentType(),
99+
rba.getExtraHeaders());
100+
101+
throw new BlockingException("Blocked request (for StatusHeader/sendJson)");
102+
}
103+
}
104+
}
105+
}

0 commit comments

Comments
 (0)