Extract Jersey json body response schemas

dd-java-agent/agent-tooling/src/main/resources/datadog/trace/agent/tooling/bytebuddy/matcher/ignored_class_name.trie

@@ -190,6 +190,8 @@
 0 com.fasterxml.jackson.databind.util.TokenBuffer$Parser
 0 com.fasterxml.jackson.databind.ObjectMapper
 0 com.fasterxml.jackson.module.afterburner.util.MyClassLoader
+# Included for API Security response schema collection
+0 com.fasterxml.jackson.jaxrs.*
 2 com.github.mustachejava.*
 2 com.google.api.*
 0 com.google.api.client.http.HttpRequest

dd-java-agent/instrumentation/jakarta-rs-annotations-3/src/main/java/datadog/trace/instrumentation/jakarta3/MessageBodyWriterInstrumentation.java

@@ -0,0 +1,87 @@
+package datadog.trace.instrumentation.jakarta3;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface;
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static datadog.trace.api.gateway.Events.EVENTS;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+import com.google.auto.service.AutoService;
+import datadog.appsec.api.blocking.BlockingException;
+import datadog.trace.advice.ActiveRequestContext;
+import datadog.trace.advice.RequiresRequestContext;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.agent.tooling.InstrumenterModule;
+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 jakarta.ws.rs.core.MediaType;
+import java.util.function.BiFunction;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+@AutoService(InstrumenterModule.class)
+public class MessageBodyWriterInstrumentation extends InstrumenterModule.AppSec
+    implements Instrumenter.ForTypeHierarchy, Instrumenter.HasMethodAdvice {
+
+  public MessageBodyWriterInstrumentation() {
+    super("jakarta-rs");
+  }
+
+  @Override
+  public String hierarchyMarkerType() {
+    return "jakarta.ws.rs.ext.MessageBodyWriter";
+  }
+
+  @Override
+  public ElementMatcher<TypeDescription> hierarchyMatcher() {
+    return implementsInterface(named(hierarchyMarkerType()));
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
+        named("writeTo").and(takesArguments(7)), getClass().getName() + "$MessageBodyWriterAdvice");
+  }
+
+  @RequiresRequestContext(RequestContextSlot.APPSEC)
+  public static class MessageBodyWriterAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    static void before(
+        @Advice.Argument(0) Object entity,
+        @Advice.Argument(4) MediaType mediaType,
+        @ActiveRequestContext RequestContext reqCtx) {
+
+      if (!MediaType.APPLICATION_JSON_TYPE.isCompatible(mediaType)) {
+        return;
+      }
+
+      CallbackProvider cbp = AgentTracer.get().getCallbackProvider(RequestContextSlot.APPSEC);
+      BiFunction<RequestContext, Object, Flow<Void>> callback =
+          cbp.getCallback(EVENTS.responseBody());
+      if (callback == null) {
+        return;
+      }
+
+      Flow<Void> flow = callback.apply(reqCtx, entity);
+      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 MessageBodyWriter)");
+      }
+    }
+  }
+}

dd-java-agent/instrumentation/jax-rs-annotations-2/build.gradle

@@ -10,6 +10,12 @@ muzzle {
     module = "javax.ws.rs-api"
     versions = "[,]"
   }
+  pass {
+    group = "javax.ws.rs"
+    module = "javax.ws.rs-api"
+    name = 'javax-message-body-writer'
+    versions = "[,]"
+  }
 }
 
 apply from: "$rootDir/gradle/java.gradle"

dd-java-agent/instrumentation/jax-rs-annotations-2/src/main/java/datadog/trace/instrumentation/jaxrs2/MessageBodyWriterInstrumentation.java

@@ -0,0 +1,92 @@
+package datadog.trace.instrumentation.jaxrs2;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface;
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static datadog.trace.api.gateway.Events.EVENTS;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+import com.google.auto.service.AutoService;
+import datadog.appsec.api.blocking.BlockingException;
+import datadog.trace.advice.ActiveRequestContext;
+import datadog.trace.advice.RequiresRequestContext;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.agent.tooling.InstrumenterModule;
+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 javax.ws.rs.core.MediaType;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+@AutoService(InstrumenterModule.class)
+public class MessageBodyWriterInstrumentation extends InstrumenterModule.AppSec
+    implements Instrumenter.ForTypeHierarchy, Instrumenter.HasMethodAdvice {
+
+  public MessageBodyWriterInstrumentation() {
+    super("jax-rs");
+  }
+
+  @Override
+  public String muzzleDirective() {
+    return "javax-message-body-writer";
+  }
+
+  @Override
+  public String hierarchyMarkerType() {
+    return "javax.ws.rs.ext.MessageBodyWriter";
+  }
+
+  @Override
+  public ElementMatcher<TypeDescription> hierarchyMatcher() {
+    return implementsInterface(named(hierarchyMarkerType()));
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
+        named("writeTo").and(takesArguments(7)), getClass().getName() + "$MessageBodyWriterAdvice");
+  }
+
+  @RequiresRequestContext(RequestContextSlot.APPSEC)
+  public static class MessageBodyWriterAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    static void before(
+        @Advice.Argument(0) Object entity,
+        @Advice.Argument(4) MediaType mediaType,
+        @ActiveRequestContext RequestContext reqCtx) {
+
+      if (!MediaType.APPLICATION_JSON_TYPE.isCompatible(mediaType)) {
+        return;
+      }
+
+      CallbackProvider cbp = AgentTracer.get().getCallbackProvider(RequestContextSlot.APPSEC);
+      BiFunction<RequestContext, Object, Flow<Void>> callback =
+          cbp.getCallback(EVENTS.responseBody());
+      if (callback == null) {
+        return;
+      }
+
+      Flow<Void> flow = callback.apply(reqCtx, entity);
+      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 MessageBodyWriter)");
+      }
+    }
+  }
+}

dd-java-agent/instrumentation/jersey/build.gradle

@@ -56,6 +56,7 @@ dependencies {
   jersey2JettyTestRuntimeOnly group: 'javax.xml.bind', name: 'jaxb-api', version: '2.2.3'
   jersey2JettyTestRuntimeOnly project(':dd-java-agent:instrumentation:jetty-9')
   jersey2JettyTestRuntimeOnly project(':dd-java-agent:instrumentation:jersey-2-appsec')
+  jersey2JettyTestRuntimeOnly project(':dd-java-agent:instrumentation:jax-rs-annotations-2')
 
   jersey3JettyTestImplementation project(':dd-java-agent:testing'), {
     exclude group: 'org.eclipse.jetty', module: 'jetty-server'
@@ -72,6 +73,7 @@ dependencies {
   jersey3JettyTestRuntimeOnly project(':dd-java-agent:instrumentation:jetty-11')
   jersey3JettyTestRuntimeOnly project(':dd-java-agent:instrumentation:jersey-2-appsec')
   jersey3JettyTestRuntimeOnly project(':dd-java-agent:instrumentation:jersey-3-appsec')
+  jersey3JettyTestRuntimeOnly project(':dd-java-agent:instrumentation:jakarta-rs-annotations-3')
 }
 
 configurations.getByName('jersey3JettyTestRuntimeClasspath').resolutionStrategy {

dd-java-agent/instrumentation/jersey/src/jersey2JettyTest/groovy/datadog/trace/instrumentation/jersey2/ClassToConvertBodyTo.groovy

@@ -1,5 +1,12 @@
 package datadog.trace.instrumentation.jersey2
 
+import groovy.json.JsonBuilder
+
 class ClassToConvertBodyTo {
   String a
+
+  @Override
+  String toString() {
+    new JsonBuilder([a: a]).toString()
+  }
 }

dd-java-agent/instrumentation/jersey/src/jersey2JettyTest/groovy/datadog/trace/instrumentation/jersey2/Jersey2JettyTest.groovy

@@ -9,6 +9,11 @@ import javax.ws.rs.ext.ExceptionMapper
 
 class Jersey2JettyTest extends HttpServerTest<JettyServer> {
 
+  @Override
+  boolean testResponseBodyJson() {
+    return true
+  }
+
   @Override
   HttpServer server() {
     new JettyServer()

dd-java-agent/instrumentation/jersey/src/jersey2JettyTest/groovy/datadog/trace/instrumentation/jersey2/ServiceResource.groovy

@@ -10,6 +10,7 @@ import javax.ws.rs.HeaderParam
 import javax.ws.rs.POST
 import javax.ws.rs.Path
 import javax.ws.rs.PathParam
+import javax.ws.rs.Produces
 import javax.ws.rs.QueryParam
 import javax.ws.rs.core.MediaType
 import javax.ws.rs.core.Response
@@ -87,10 +88,14 @@ class ServiceResource {
 
   @POST
   @Path("body-json")
+  @Produces(MediaType.APPLICATION_JSON)
   Response bodyJson(ClassToConvertBodyTo obj) {
-    controller(BODY_JSON) {
-      Response.status(BODY_JSON.status).entity("""{"a":"${obj.a}"}""" as String).build()
-    }
+    return controller(BODY_JSON, () -> {
+      Response response = Response.status(BODY_JSON.status)
+      .entity(obj)
+      .build()
+      return response
+    })
   }
 
   @GET

dd-java-agent/instrumentation/jersey/src/jersey3JettyTest/groovy/datadog/trace/instrumentation/jersey3/ClassToConvertBodyTo.groovy

@@ -1,5 +1,12 @@
 package datadog.trace.instrumentation.jersey3
 
+import groovy.json.JsonBuilder
+
 class ClassToConvertBodyTo {
   String a
+
+  @Override
+  String toString() {
+    new JsonBuilder([a: a]).toString()
+  }
 }

dd-java-agent/instrumentation/jersey/src/jersey3JettyTest/groovy/datadog/trace/instrumentation/jersey3/Jersey3JettyTest.groovy

@@ -8,6 +8,11 @@ import jakarta.ws.rs.ext.ExceptionMapper
 
 class Jersey3JettyTest extends HttpServerTest<JettyServer> {
 
+  @Override
+  boolean testResponseBodyJson() {
+    return true
+  }
+
   @Override
   HttpServer server() {
     new JettyServer()

dd-java-agent/instrumentation/jersey/src/jersey3JettyTest/groovy/datadog/trace/instrumentation/jersey3/ServiceResource.groovy

@@ -1,6 +1,7 @@
 package datadog.trace.instrumentation.jersey3
 
 import datadog.appsec.api.blocking.Blocking
+import jakarta.ws.rs.Produces
 import org.glassfish.jersey.media.multipart.FormDataParam
 
 import jakarta.ws.rs.Consumes
@@ -87,10 +88,13 @@ class ServiceResource {
 
   @POST
   @Path("body-json")
+  @Produces(MediaType.APPLICATION_JSON)
   Response bodyJson(ClassToConvertBodyTo obj) {
-    controller(BODY_JSON) {
-      Response.status(BODY_JSON.status).entity("""{"a":"${obj.a}"}""" as String).build()
-    }
+    controller(BODY_JSON, () ->
+    Response.status(BODY_JSON.status)
+    .entity(obj)
+    .build()
+    )
   }
 
   @GET

dd-smoke-tests/jersey-2/build.gradle

@@ -20,6 +20,7 @@ dependencies {
   implementation group: 'javax.xml', name: 'jaxb-api', version:'2.1'
   testImplementation project(':dd-smoke-tests')
   testImplementation(testFixtures(project(":dd-smoke-tests:iast-util")))
+  testImplementation project(':dd-smoke-tests:appsec')
 }
 
 tasks.withType(Test).configureEach {

dd-smoke-tests/jersey-2/src/main/java/com/restserver/RequestBody.java

@@ -0,0 +1,45 @@
+package com.restserver;
+
+import java.util.List;
+
+public class RequestBody {
+  private List<KeyValue> main;
+  private Object nullable;
+
+  public List<KeyValue> getMain() {
+    return main;
+  }
+
+  public void setMain(List<KeyValue> main) {
+    this.main = main;
+  }
+
+  public Object getNullable() {
+    return nullable;
+  }
+
+  public void setNullable(Object nullable) {
+    this.nullable = nullable;
+  }
+
+  public static class KeyValue {
+    private String key;
+    private Double value;
+
+    public String getKey() {
+      return key;
+    }
+
+    public void setKey(String key) {
+      this.key = key;
+    }
+
+    public Double getValue() {
+      return value;
+    }
+
+    public void setValue(Double value) {
+      this.value = value;
+    }
+  }
+}

dd-smoke-tests/jersey-2/src/main/java/com/restserver/Resource.java

@@ -139,4 +139,18 @@ public Response responseLocation(@QueryParam("param") String param) throws URISy
   public Response getCookie() throws SQLException {
     return Response.ok().cookie(new NewCookie("user-id", "7")).build();
   }
+
+  @Path("/api_security/response")
+  @POST
+  @Produces(MediaType.APPLICATION_JSON)
+  @Consumes(MediaType.APPLICATION_JSON)
+  public Response bodyJson(RequestBody input) {
+    return Response.ok(input).build();
+  }
+
+  @GET
+  @Path("/api_security/sampling/{i}")
+  public Response apiSecuritySamplingWithStatus(@PathParam("i") int i) {
+    return Response.status(i).header("content-type", "text/plain").entity("Hello!\n").build();
+  }
 }