HTTP response schema collection and data classification for spring

Signed-off-by: sezen.leblay <[email protected]>

Author: sezen-datadog

dd-java-agent/instrumentation/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/HttpMessageConverterInstrumentation.java

@@ -71,6 +71,27 @@ public void methodAdvice(MethodTransformer transformer) {
             .and(takesArgument(1, Class.class))
             .and(takesArgument(2, named("org.springframework.http.HttpInputMessage"))),
         HttpMessageConverterInstrumentation.class.getName() + "$HttpMessageConverterReadAdvice");
+
+    transformer.applyAdvice(
+        isMethod()
+            .and(isPublic())
+            .and(named("write"))
+            .and(takesArguments(3))
+            .and(takesArgument(0, Object.class))
+            .and(takesArgument(1, named("org.springframework.http.MediaType")))
+            .and(takesArgument(2, named("org.springframework.http.HttpOutputMessage"))),
+        HttpMessageConverterInstrumentation.class.getName() + "$HttpMessageConverterWriteAdvice");
+
+    transformer.applyAdvice(
+        isMethod()
+            .and(isPublic())
+            .and(named("write"))
+            .and(takesArguments(4))
+            .and(takesArgument(0, Object.class))
+            .and(takesArgument(1, Type.class))
+            .and(takesArgument(2, named("org.springframework.http.MediaType")))
+            .and(takesArgument(3, named("org.springframework.http.HttpOutputMessage"))),
+        HttpMessageConverterInstrumentation.class.getName() + "$HttpMessageConverterWriteAdvice");
   }
 
   @RequiresRequestContext(RequestContextSlot.APPSEC)
@@ -106,4 +127,37 @@ public static void after(
       }
     }
   }
+
+  @RequiresRequestContext(RequestContextSlot.APPSEC)
+  public static class HttpMessageConverterWriteAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void before(
+        @Advice.Argument(0) final Object obj, @ActiveRequestContext RequestContext reqCtx) {
+      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) {
+        Flow.Action.RequestBlockingAction rba = (Flow.Action.RequestBlockingAction) action;
+        BlockResponseFunction brf = reqCtx.getBlockResponseFunction();
+        if (brf != null) {
+          brf.tryCommitBlockingResponse(
+              reqCtx.getTraceSegment(),
+              rba.getStatusCode(),
+              rba.getBlockingContentType(),
+              rba.getExtraHeaders());
+        }
+        throw new BlockingException("Blocked response (for HttpMessageConverter/write)");
+      }
+    }
+  }
 }

dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/test/boot/SpringBootBasedTest.groovy

@@ -77,6 +77,11 @@ class SpringBootBasedTest extends HttpServerTest<ConfigurableApplicationContext>
     return "boot-context"
   }
 
+  @Override
+  boolean testResponseBodyJson() {
+    return true
+  }
+
   @Override
   String expectedServiceName() {
     servletContext
@@ -163,8 +168,7 @@ class SpringBootBasedTest extends HttpServerTest<ConfigurableApplicationContext>
 
   @Override
   Map<String, Serializable> expectedExtraServerTags(ServerEndpoint endpoint) {
-    ["servlet.path": endpoint.path, "servlet.context": "/$servletContext"] +
-    extraServerTags
+    ["servlet.path": endpoint.path, "servlet.context": "/$servletContext"] + extraServerTags
   }
 
   @Override

dd-smoke-tests/appsec/springboot/src/main/java/datadog/smoketest/appsec/springboot/controller/WebController.java

@@ -239,6 +239,13 @@ public ResponseEntity<String> exceedResponseHeaders() {
     return new ResponseEntity<>("Custom headers added", headers, HttpStatus.OK);
   }
 
+  @PostMapping("/api_security/response")
+  public ResponseEntity<String> apiSecurityResponse(@RequestBody String body) {
+    // This endpoint is used to test API security response handling
+    // It simply returns the body received in the request
+    return ResponseEntity.ok(body);
+  }
+
   private void withProcess(final Operation<Process> op) {
     Process process = null;
     try {

dd-smoke-tests/appsec/springboot/src/test/groovy/datadog/smoketest/appsec/AppSecHttpMessageConverterSmokeTest.groovy

@@ -0,0 +1,69 @@
+package datadog.smoketest.appsec
+
+import groovy.json.JsonOutput
+import okhttp3.MediaType
+import okhttp3.Request
+import okhttp3.RequestBody
+
+class AppSecHttpMessageConverterSmokeTest extends AbstractAppSecServerSmokeTest {
+
+  @Override
+  def logLevel() {
+    'DEBUG'
+  }
+
+  @Override
+  ProcessBuilder createProcessBuilder() {
+    String springBootShadowJar = System.getProperty("datadog.smoketest.appsec.springboot.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",
+      "-jar",
+      springBootShadowJar,
+      "--server.port=${httpPort}"
+    ])
+    ProcessBuilder processBuilder = new ProcessBuilder(command)
+    processBuilder.directory(new File(buildDirectory))
+  }
+
+  @Override
+  File createTemporaryFile() {
+    return new File("${buildDirectory}/tmp/trace-structure-http-converter.out")
+  }
+
+  void 'test response schema extraction'() {
+    given:
+    def url = "http://localhost:${httpPort}/api_security/response"
+    def body = [
+      source: 'AppSecSpringSmokeTest',
+      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')
+  }
+}