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();
+ }
 }