Extract Ratpack json body response schemas

dd-java-agent/instrumentation/ratpack-1.5/src/main/java/datadog/trace/instrumentation/ratpack/JsonRendererAdvice.java

@@ -0,0 +1,53 @@
+package datadog.trace.instrumentation.ratpack;
+
+import static datadog.trace.api.gateway.Events.EVENTS;
+
+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 ratpack.jackson.JsonRender;
+
+@RequiresRequestContext(RequestContextSlot.APPSEC)
+public class JsonRendererAdvice {
+
+  // for now ignore that the parser can be configured to mix in the query string
+  @Advice.OnMethodEnter(suppress = Throwable.class)
+  static void after(
+      @Advice.Argument(1) final JsonRender render,
+      @ActiveRequestContext final RequestContext reqCtx) {
+    Object obj = render == null ? null : render.getObject();
+    if (obj == null) {
+      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, obj);
+    Flow.Action action = flow.getAction();
+    if (action instanceof Flow.Action.RequestBlockingAction) {
+      BlockResponseFunction brf = reqCtx.getBlockResponseFunction();
+      if (brf != null) {
+        Flow.Action.RequestBlockingAction rba = (Flow.Action.RequestBlockingAction) action;
+        brf.tryCommitBlockingResponse(
+            reqCtx.getTraceSegment(),
+            rba.getStatusCode(),
+            rba.getBlockingContentType(),
+            rba.getExtraHeaders());
+
+        throw new BlockingException("Blocked request (for JsonRenderer/render)");
+      }
+    }
+  }
+}

dd-java-agent/instrumentation/ratpack-1.5/src/main/java/datadog/trace/instrumentation/ratpack/JsonRendererInstrumentation.java

@@ -0,0 +1,44 @@
+package datadog.trace.instrumentation.ratpack;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+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 JsonRendererInstrumentation extends InstrumenterModule.AppSec
+    implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
+
+  // so it doesn't apply to ratpack < 1.5
+  private static final Reference FILE_IO = new Reference.Builder("ratpack.file.FileIo").build();
+
+  public JsonRendererInstrumentation() {
+    super("ratpack");
+  }
+
+  @Override
+  public String instrumentedType() {
+    return "ratpack.jackson.internal.JsonRenderer";
+  }
+
+  @Override
+  public Reference[] additionalMuzzleReferences() {
+    return new Reference[] {FILE_IO};
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
+        isMethod()
+            .and(named("render"))
+            .and(takesArguments(2))
+            .and(takesArgument(0, named("ratpack.handling.Context")))
+            .and(takesArgument(1, named("ratpack.jackson.JsonRender"))),
+        packageName + ".JsonRendererAdvice");
+  }
+}

dd-java-agent/instrumentation/ratpack-1.5/src/test/groovy/server/RatpackAsyncHttpServerTest.groovy

@@ -10,6 +10,7 @@ import ratpack.exec.Promise
 import ratpack.form.Form
 import ratpack.groovy.test.embed.GroovyEmbeddedApp
 import ratpack.handling.HandlerDecorator
+import static ratpack.jackson.Jackson.json
 
 import java.nio.charset.StandardCharsets
 
@@ -142,7 +143,8 @@ class RatpackAsyncHttpServerTest extends RatpackHttpServerTest {
             } then {endpoint ->
               controller(endpoint) {
                 context.parse(Map).then { map ->
-                  response.status(BODY_JSON.status).send('text/plain', "{\"a\":\"${map['a']}\"}")
+                  response.status(BODY_JSON.status)
+                  context.render(json(map))
                 }
               }
             }

dd-java-agent/instrumentation/ratpack-1.5/src/test/groovy/server/RatpackForkedHttpServerTest.groovy

@@ -7,6 +7,7 @@ import ratpack.exec.Promise
 import ratpack.form.Form
 import ratpack.groovy.test.embed.GroovyEmbeddedApp
 import ratpack.handling.HandlerDecorator
+import static ratpack.jackson.Jackson.json
 
 import java.nio.charset.StandardCharsets
 
@@ -105,10 +106,11 @@ class RatpackForkedHttpServerTest extends RatpackHttpServerTest {
           all {
             Promise.sync {
               BODY_JSON
-            }.fork().then {endpoint ->
+            }.fork().then { endpoint ->
               controller(endpoint) {
                 context.parse(Map).then { map ->
-                  response.status(BODY_JSON.status).send('text/plain', "{\"a\":\"${map['a']}\"}")
+                  response.status(BODY_JSON.status)
+                  context.render(json(map))
                 }
               }
             }

dd-java-agent/instrumentation/ratpack-1.5/src/test/groovy/server/RatpackHttpServerTest.groovy

@@ -77,6 +77,11 @@ class RatpackHttpServerTest extends HttpServerTest<EmbeddedApp> {
     true
   }
 
+  @Override
+  boolean testResponseBodyJson() {
+    true
+  }
+
   @Override
   String testPathParam() {
     true

dd-java-agent/instrumentation/ratpack-1.5/src/test/groovy/server/SyncRatpackApp.groovy

@@ -5,6 +5,7 @@ import ratpack.form.Form
 import ratpack.groovy.test.embed.GroovyEmbeddedApp
 import ratpack.handling.HandlerDecorator
 import ratpack.test.embed.EmbeddedApp
+import static ratpack.jackson.Jackson.json
 
 import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.BODY_JSON
 import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.BODY_MULTIPART
@@ -46,9 +47,10 @@ enum SyncRatpackApp implements EmbeddedApp {
       prefix(CREATED.relativeRawPath()) {
         all {
           controller(CREATED) {
-            request.body.then { typedData ->
+            request.body.then {
+              typedData ->
               response.status(CREATED.status)
-                .send('text/plain', "${CREATED.body}: ${typedData.text}")
+              .send('text/plain', "${CREATED.body}: ${typedData.text}")
             }
           }
         }
@@ -61,9 +63,14 @@ enum SyncRatpackApp implements EmbeddedApp {
       prefix(BODY_URLENCODED.relativeRawPath()) {
         all {
           controller(BODY_URLENCODED) {
-            context.parse(Form).then { form ->
-              def text = form.findAll { it.key != 'ignore'}
-              .collectEntries {[it.key, it.value as List]} as String
+            context.parse(Form).then {
+              form ->
+              def text = form.findAll {
+                it.key != 'ignore'
+              }
+              .collectEntries {
+                [it.key, it.value as List]
+              } as String
               response.status(BODY_URLENCODED.status).send('text/plain', text)
             }
           }
@@ -72,8 +79,11 @@ enum SyncRatpackApp implements EmbeddedApp {
       prefix(BODY_MULTIPART.relativeRawPath()) {
         all {
           controller(BODY_MULTIPART) {
-            context.parse(Form).then { form ->
-              def text = form.collectEntries {[it.key, it.value as List]} as String
+            context.parse(Form).then {
+              form ->
+              def text = form.collectEntries {
+                [it.key, it.value as List]
+              } as String
               response.status(BODY_MULTIPART.status).send('text/plain', text)
             }
           }
@@ -82,8 +92,11 @@ enum SyncRatpackApp implements EmbeddedApp {
       prefix(BODY_JSON.relativeRawPath()) {
         all {
           controller(BODY_JSON) {
-            context.parse(Map).then { map ->
-              response.status(BODY_JSON.status).send('text/plain', "{\"a\":\"${map['a']}\"}")
+            context.parse(Map).then {
+              map -> {
+                response.status(BODY_JSON.status)
+                context.render(json(map))
+              }
             }
           }
         }

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

@@ -2503,7 +2503,12 @@ abstract class HttpServerTest<SERVER> extends WithHttpServer<SERVER> {
 
     final BiFunction<RequestContext, Object, Flow<Void>> responseBodyObjectCb =
     ({ RequestContext rqCtxt, Object obj ->
-      String body = obj.toString()
+      String body
+      if (obj instanceof Map) {
+        body = JsonOutput.toJson(obj as Map)
+      } else {
+        body = obj.toString()
+      }
       Context context = rqCtxt.getData(RequestContextSlot.APPSEC)
       if (context.responseBodyTag) {
         rqCtxt.traceSegment.setTagTop('response.body', body)

dd-smoke-tests/ratpack-1.5/build.gradle

@@ -0,0 +1,23 @@
+plugins {
+  id "com.gradleup.shadow"
+}
+
+apply from: "$rootDir/gradle/java.gradle"
+
+jar {
+  manifest {
+    attributes('Main-Class': 'datadog.smoketest.ratpack.RatpackApp')
+  }
+}
+dependencies {
+  implementation 'io.ratpack:ratpack-core:1.5.0'
+
+  testImplementation project(':dd-smoke-tests')
+  testImplementation project(':dd-smoke-tests:appsec')
+}
+
+tasks.withType(Test).configureEach {
+  dependsOn "shadowJar"
+  jvmArgs "-Ddatadog.smoketest.ratpack.shadowJar.path=${tasks.shadowJar.archiveFile.get()}"
+}
+

dd-smoke-tests/ratpack-1.5/src/main/java/datadog/smoketest/ratpack/RatpackApp.java

@@ -0,0 +1,41 @@
+package datadog.smoketest.ratpack;
+
+import static ratpack.jackson.Jackson.json;
+
+import com.fasterxml.jackson.databind.JsonNode;
+import ratpack.server.RatpackServer;
+
+public class RatpackApp {
+
+  public static void main(String[] args) throws Exception {
+    int port = Integer.parseInt(System.getProperty("ratpack.http.port", "8080"));
+    RatpackServer.start(
+        server ->
+            server
+                .serverConfig(config -> config.port(port))
+                .handlers(
+                    chain ->
+                        chain
+                            .path(
+                                "api_security/sampling/:status_code",
+                                ctx -> {
+                                  ctx.getResponse()
+                                      .status(
+                                          Integer.parseInt(ctx.getPathTokens().get("status_code")))
+                                      .send("EXECUTED");
+                                })
+                            .path(
+                                "api_security/response",
+                                ctx ->
+                                    ctx.parse(JsonNode.class)
+                                        .then(
+                                            node -> {
+                                              ctx.getResponse().status(200);
+                                              ctx.render(json(node));
+                                            }))
+                            .all(
+                                ctx -> {
+                                  ctx.getResponse().status(404).send("Not Found");
+                                })));
+  }
+}

dd-smoke-tests/ratpack-1.5/src/test/groovy/AppSecRatpackSmokeTest.groovy

@@ -0,0 +1,100 @@
+import datadog.smoketest.appsec.AbstractAppSecServerSmokeTest
+import datadog.trace.agent.test.utils.OkHttpUtils
+import groovy.json.JsonOutput
+import okhttp3.MediaType
+import okhttp3.Request
+import okhttp3.RequestBody
+import okhttp3.Response
+
+class AppSecRatpackSmokeTest extends AbstractAppSecServerSmokeTest {
+
+  @Override
+  ProcessBuilder createProcessBuilder() {
+    String ratpackUberJar = System.getProperty("datadog.smoketest.ratpack.shadowJar.path")
+
+    List<String> command = new ArrayList<>()
+    command.add(javaPath())
+    command.addAll(defaultJavaProperties)
+    command.addAll(defaultAppSecProperties)
+    command.addAll((String[]) [
+      "-Ddd.writer.type=MultiWriter:TraceStructureWriter:${output.getAbsolutePath()},DDAgentWriter",
+      "-Dratpack.http.port=${httpPort}",
+      "-jar",
+      ratpackUberJar
+    ])
+    ProcessBuilder processBuilder = new ProcessBuilder(command)
+    processBuilder.directory(new File(buildDirectory))
+  }
+
+  @Override
+  File createTemporaryFile() {
+    return new File("${buildDirectory}/tmp/trace-structure-ratpack.out")
+  }
+
+  void 'API Security samples only one request per endpoint'() {
+    given:
+    def url = "http://localhost:${httpPort}/api_security/sampling/200?test=value"
+    def client = OkHttpUtils.clientBuilder().build()
+    def request = new Request.Builder()
+      .url(url)
+      .addHeader('X-My-Header', "value")
+      .get()
+      .build()
+
+    when:
+    List<Response> responses = (1..3).collect {
+      client.newCall(request).execute()
+    }
+
+    then:
+    responses.each {
+      assert it.code() == 200
+    }
+    waitForTraceCount(3)
+    def spans = rootSpans.toList().toSorted { it.span.duration }
+    spans.size() == 3
+    def sampledSpans = spans.findAll {
+      it.meta.keySet().any {
+        it.startsWith('_dd.appsec.s.req.')
+      }
+    }
+    sampledSpans.size() == 1
+    def span = sampledSpans[0]
+    span.meta.containsKey('_dd.appsec.s.req.query')
+    span.meta.containsKey('_dd.appsec.s.req.params')
+    span.meta.containsKey('_dd.appsec.s.req.headers')
+  }
+
+  void 'test response schema extraction'() {
+    given:
+    def url = "http://localhost:${httpPort}/api_security/response"
+    def client = OkHttpUtils.clientBuilder().build()
+    def body = [
+      source: 'AppSecRatpackSmokeTest',
+      tests : [
+        [
+          name  : 'API Security samples only one request per endpoint',
+          status: 'SUCCESS'
+        ],
+        [
+          name  : 'test response schema extraction',
+          status: 'FAILED'
+        ]
+      ]
+    ]
+    def request = new Request.Builder()
+      .url(url)
+      .post(RequestBody.create(MediaType.get('application/json'), JsonOutput.toJson(body)))
+      .build()
+
+    when:
+    final response = client.newCall(request).execute()
+    waitForTraceCount(1)
+
+    then:
+    response.code() == 200
+    def span = rootSpans.first()
+    span.meta.containsKey('_dd.appsec.s.res.headers')
+    span.meta.containsKey('_dd.appsec.s.res.body')
+  }
+}