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