Extract Play json body response schemas

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

@@ -0,0 +1,44 @@
+package datadog.trace.instrumentation.play25.appsec;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.agent.tooling.muzzle.Reference;
+
+@AutoService(InstrumenterModule.class)
+public class ResultsStatusInstrumentation extends InstrumenterModule.AppSec
+    implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
+
+  public ResultsStatusInstrumentation() {
+    super("play");
+  }
+
+  @Override
+  public String muzzleDirective() {
+    return "play25only";
+  }
+
+  @Override
+  public Reference[] additionalMuzzleReferences() {
+    return MuzzleReferences.PLAY_25_ONLY;
+  }
+
+  @Override
+  public String instrumentedType() {
+    return "play.api.mvc.Results$Status";
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {
+      packageName + ".BodyParserHelpers",
+    };
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(named("apply"), packageName + ".ResultsStatusApplyAdvice");
+  }
+}

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

@@ -0,0 +1,40 @@
+package datadog.trace.instrumentation.play25.appsec;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.agent.tooling.muzzle.Reference;
+
+@AutoService(InstrumenterModule.class)
+public class StatusHeaderInstrumentation extends InstrumenterModule.AppSec
+    implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
+
+  public StatusHeaderInstrumentation() {
+    super("play");
+  }
+
+  @Override
+  public String muzzleDirective() {
+    return "play25only";
+  }
+
+  @Override
+  public Reference[] additionalMuzzleReferences() {
+    return MuzzleReferences.PLAY_25_ONLY;
+  }
+
+  @Override
+  public String instrumentedType() {
+    return "play.mvc.StatusHeader";
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
+        named("sendJson").and(takesArgument(0, named("com.fasterxml.jackson.databind.JsonNode"))),
+        packageName + ".StatusHeaderSendJsonAdvice");
+  }
+}

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

@@ -39,6 +39,7 @@
 import scala.collection.Iterator;
 import scala.collection.Seq;
 import scala.compat.java8.JFunction1;
+import scala.math.BigDecimal;
 
 public class BodyParserHelpers {
 
@@ -231,15 +232,20 @@ private static void handleException(Exception e, String logMessage) {
     log.warn(logMessage, e);
   }
 
-  private static Object jsValueToJavaObject(JsValue value, int maxRecursion) {
+  public static Object jsValueToJavaObject(JsValue value) {
+    return jsValueToJavaObject(value, MAX_RECURSION);
+  }
+
+  public static Object jsValueToJavaObject(JsValue value, int maxRecursion) {
     if (value == null || maxRecursion <= 0) {
       return null;
     }
 
     if (value instanceof JsString) {
       return ((JsString) value).value();
     } else if (value instanceof JsNumber) {
-      return ((JsNumber) value).value();
+      final BigDecimal number = ((JsNumber) value).value();
+      return number == null ? null : number.bigDecimal();
     } else if (value instanceof JsBoolean) {
       return ((JsBoolean) value).value();
     } else if (value instanceof JsObject) {

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

@@ -0,0 +1,57 @@
+package datadog.trace.instrumentation.play25.appsec;
+
+import static datadog.trace.api.gateway.Events.EVENTS;
+import static datadog.trace.instrumentation.play25.appsec.BodyParserHelpers.jsValueToJavaObject;
+
+import datadog.appsec.api.blocking.BlockingException;
+import datadog.trace.advice.ActiveRequestContext;
+import datadog.trace.advice.RequiresRequestContext;
+import datadog.trace.api.gateway.BlockResponseFunction;
+import datadog.trace.api.gateway.CallbackProvider;
+import datadog.trace.api.gateway.Flow;
+import datadog.trace.api.gateway.RequestContext;
+import datadog.trace.api.gateway.RequestContextSlot;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
+import java.util.function.BiFunction;
+import net.bytebuddy.asm.Advice;
+import play.api.libs.json.JsValue;
+
+@RequiresRequestContext(RequestContextSlot.APPSEC)
+public class ResultsStatusApplyAdvice {
+
+  @Advice.OnMethodEnter(suppress = Throwable.class)
+  static void after(
+      @Advice.Argument(0) final Object content, @ActiveRequestContext RequestContext reqCtx) {
+
+    if (!(content instanceof JsValue)) {
+      return;
+    }
+
+    CallbackProvider cbp = AgentTracer.get().getCallbackProvider(RequestContextSlot.APPSEC);
+    if (cbp == null) {
+      return;
+    }
+    BiFunction<RequestContext, Object, Flow<Void>> callback =
+        cbp.getCallback(EVENTS.responseBody());
+    if (callback == null) {
+      return;
+    }
+
+    Flow<Void> flow = callback.apply(reqCtx, jsValueToJavaObject((JsValue) content));
+    Flow.Action action = flow.getAction();
+    if (action instanceof Flow.Action.RequestBlockingAction) {
+      BlockResponseFunction blockResponseFunction = reqCtx.getBlockResponseFunction();
+      if (blockResponseFunction == null) {
+        return;
+      }
+      Flow.Action.RequestBlockingAction rba = (Flow.Action.RequestBlockingAction) action;
+      blockResponseFunction.tryCommitBlockingResponse(
+          reqCtx.getTraceSegment(),
+          rba.getStatusCode(),
+          rba.getBlockingContentType(),
+          rba.getExtraHeaders());
+
+      throw new BlockingException("Blocked request (for Results$Status/apply)");
+    }
+  }
+}

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

@@ -0,0 +1,67 @@
+package datadog.trace.instrumentation.play25.appsec;
+
+import static datadog.trace.api.gateway.Events.EVENTS;
+
+import com.fasterxml.jackson.databind.JsonNode;
+import datadog.appsec.api.blocking.BlockingException;
+import datadog.trace.advice.ActiveRequestContext;
+import datadog.trace.advice.RequiresRequestContext;
+import datadog.trace.api.gateway.BlockResponseFunction;
+import datadog.trace.api.gateway.CallbackProvider;
+import datadog.trace.api.gateway.Flow;
+import datadog.trace.api.gateway.RequestContext;
+import datadog.trace.api.gateway.RequestContextSlot;
+import datadog.trace.bootstrap.CallDepthThreadLocalMap;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
+import java.util.function.BiFunction;
+import net.bytebuddy.asm.Advice;
+import play.mvc.StatusHeader;
+
+@RequiresRequestContext(RequestContextSlot.APPSEC)
+public class StatusHeaderSendJsonAdvice {
+
+  @Advice.OnMethodEnter(suppress = Throwable.class)
+  static void before() {
+    CallDepthThreadLocalMap.incrementCallDepth(StatusHeader.class);
+  }
+
+  @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
+  static void after(
+      @Advice.Argument(0) final JsonNode json, @ActiveRequestContext RequestContext reqCtx) {
+    final int depth = CallDepthThreadLocalMap.decrementCallDepth(StatusHeader.class);
+    if (depth > 0) {
+      return;
+    }
+
+    if (json == null) {
+      return;
+    }
+
+    CallbackProvider cbp = AgentTracer.get().getCallbackProvider(RequestContextSlot.APPSEC);
+    if (cbp == null) {
+      return;
+    }
+    BiFunction<RequestContext, Object, Flow<Void>> callback =
+        cbp.getCallback(EVENTS.responseBody());
+    if (callback == null) {
+      return;
+    }
+
+    Flow<Void> flow = callback.apply(reqCtx, json);
+    Flow.Action action = flow.getAction();
+    if (action instanceof Flow.Action.RequestBlockingAction) {
+      BlockResponseFunction blockResponseFunction = reqCtx.getBlockResponseFunction();
+      if (blockResponseFunction == null) {
+        return;
+      }
+      Flow.Action.RequestBlockingAction rba = (Flow.Action.RequestBlockingAction) action;
+      blockResponseFunction.tryCommitBlockingResponse(
+          reqCtx.getTraceSegment(),
+          rba.getStatusCode(),
+          rba.getBlockingContentType(),
+          rba.getExtraHeaders());
+
+      throw new BlockingException("Blocked request (for StatusHeader/sendJson)");
+    }
+  }
+}

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

@@ -1,7 +1,6 @@
 package datadog.trace.instrumentation.play25.server
 
 import com.fasterxml.jackson.databind.JsonNode
-import com.fasterxml.jackson.databind.ObjectMapper
 import datadog.appsec.api.blocking.Blocking
 import datadog.trace.agent.test.base.HttpServerTest
 import groovy.transform.CompileStatic
@@ -130,7 +129,7 @@ class PlayRouters {
         ->
         JsonNode json = body().asJson()
         controller(BODY_JSON) {
-          Results.status(BODY_JSON.status, new ObjectMapper().writeValueAsString(json))
+          Results.status(BODY_JSON.status, json)
         }
       } as Supplier)
       .build()
@@ -253,7 +252,7 @@ class PlayRouters {
         CompletableFuture.supplyAsync({
           ->
           controller(BODY_JSON) {
-            Results.status(BODY_JSON.status, new ObjectMapper().writeValueAsString(json))
+            Results.status(BODY_JSON.status, json)
           }
         }, execContext)
       } as Supplier)

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

@@ -93,6 +93,11 @@ class PlayServerTest extends HttpServerTest<Server> {
     true
   }
 
+  @Override
+  boolean testResponseBodyJson() {
+    true
+  }
+
   @Override
   String testPathParam() {
     '/path/?/param'

dd-java-agent/instrumentation/play-2.4/src/test/scala/datadog/trace/instrumentation/play25/PlayController.scala

@@ -77,7 +77,7 @@ class PlayController(implicit ec: ExecutionContext) extends Controller {
 
   def bodyJson = controller(ServerEndpoint.BODY_JSON) { request =>
     val body: JsValue = request.body.asJson.getOrElse(JsNull)
-    Results.Ok(Json.stringify(body))
+    Results.Ok(body)
   }
 
   private def controller(endpoint: ServerEndpoint)(block: Request[AnyContent] => Result) : Action[AnyContent] = {

dd-java-agent/instrumentation/play-2.4/src/test/scala/datadog/trace/instrumentation/play25/PlayRoutersScala.scala

@@ -143,7 +143,7 @@ object PlayRoutersScala {
       case POST(p"/body-json") => Action.async(parser) { request =>
         controller(BODY_JSON) {
           val body: JsValue = request.body.asJson.getOrElse(JsNull)
-          Results.Ok(Json.stringify(body))
+          Results.Ok(body)
         }
       }
     }

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

@@ -1,11 +1,11 @@
 package datadog.trace.instrumentation.play26.server
 
+import com.fasterxml.jackson.databind.ObjectMapper
 import datadog.appsec.api.blocking.Blocking
 import datadog.trace.agent.test.base.HttpServerTest
 import groovy.transform.CompileStatic
 import play.BuiltInComponents
 import play.api.libs.json.JsValue
-import play.api.libs.json.Json$
 import play.api.mvc.AnyContent
 import play.libs.concurrent.HttpExecution
 import play.mvc.Http
@@ -150,7 +150,7 @@ class PlayRouters {
         ->
         controller(BODY_JSON) {
           JsValue json = body().asJson().get()
-          Results.status(BODY_JSON.status, Json$.MODULE$.stringify(json))
+          Results.status(BODY_JSON.status, new ObjectMapper().readTree(json.toString()))
         }
       } as Supplier)
       .POST(BODY_XML.path).routeTo({
@@ -286,7 +286,7 @@ class PlayRouters {
         CompletableFuture.supplyAsync({
           ->
           controller(BODY_JSON) {
-            Results.status(BODY_JSON.status, Json$.MODULE$.stringify(json))
+            Results.status(BODY_JSON.status, new ObjectMapper().readTree(json.toString()))
           }
         }, execContext)
       } as Supplier)

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

@@ -96,6 +96,11 @@ class PlayServerTest extends HttpServerTest<Server> {
     true
   }
 
+  @Override
+  boolean testResponseBodyJson() {
+    true
+  }
+
   @Override
   String testPathParam() {
     '/path/?/param'

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

@@ -1,7 +1,6 @@
 package datadog.trace.instrumentation.play26.server.latestdep
 
 import com.fasterxml.jackson.databind.JsonNode
-import com.fasterxml.jackson.databind.ObjectMapper
 import datadog.appsec.api.blocking.Blocking
 import datadog.trace.agent.test.base.HttpServerTest
 import datadog.trace.instrumentation.play26.server.TestHttpErrorHandler
@@ -121,7 +120,7 @@ class PlayRouters {
       .POST(BODY_JSON.path).routingTo({ Http.Request req ->
         controller(BODY_JSON) {
           JsonNode json = req.body().asJson()
-          Results.status(BODY_JSON.status, new ObjectMapper().writeValueAsString(json))
+          Results.status(BODY_JSON.status, json)
         }
       } as RequestFunctions.Params0<Result>)
       .POST(BODY_XML.path).routingTo({ Http.Request req ->
@@ -254,7 +253,7 @@ class PlayRouters {
         CompletableFuture.supplyAsync({
           ->
           controller(BODY_JSON) {
-            Results.status(BODY_JSON.status, new ObjectMapper().writeValueAsString(json))
+            Results.status(BODY_JSON.status, json)
           }
         }, execContext)
       } as RequestFunctions.Params0<? extends CompletionStage<Result>>)

dd-java-agent/instrumentation/play-2.6/src/latestDepTest/scala/datadog/trace/instrumentation/play26/server/latestdep/PlayController.scala

@@ -79,7 +79,7 @@ class PlayController(cc: ControllerComponents)(implicit ec: ExecutionContext) ex
 
   def bodyJson = controller(ServerEndpoint.BODY_JSON) { request =>
     val body: JsValue = request.body.asJson.getOrElse(JsNull)
-    Results.Ok(Json.stringify(body))
+    Results.Ok(body)
   }
 
   def bodyXml = controller(ServerEndpoint.BODY_XML) { request =>

dd-java-agent/instrumentation/play-2.6/src/latestDepTest/scala/datadog/trace/instrumentation/play26/server/latestdep/PlayRoutersScala.scala

@@ -147,7 +147,7 @@ object PlayRoutersScala {
       case POST(p"/body-json") => defaultActionBuilder.async(parser) { request =>
         controller(BODY_JSON) {
           val body: JsValue = request.body.asJson.getOrElse(JsNull)
-          Results.Ok(Json.stringify(body))
+          Results.Ok(body)
         }
       }