Extract Vert.x json body response schemas

dd-java-agent/appsec/src/main/java/com/datadog/appsec/event/data/ObjectIntrospection.java

@@ -11,6 +11,7 @@
 import java.lang.reflect.Method;
 import java.lang.reflect.Modifier;
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.HashMap;
 import java.util.Iterator;
 import java.util.List;
@@ -211,20 +212,42 @@ private static Object doConversion(Object obj, int depth, State state) {
 
     // iterables
     if (obj instanceof Iterable) {
-      List<Object> newList;
-      if (obj instanceof List) {
-        newList = new ArrayList<>(((List<?>) obj).size());
+      final Iterator<?> it = ((Iterable<?>) obj).iterator();
+      final boolean isMap = it.hasNext() && it.next() instanceof Map.Entry;
+      // some json libraries implement objects as Iterable<Map.Entry>
+      if (isMap) {
+        Map<Object, Object> newMap;
+        if (obj instanceof Collection) {
+          newMap = new HashMap<>(((Collection<?>) obj).size());
+        } else {
+          newMap = new HashMap<>();
+        }
+        for (Map.Entry<?, ?> e : ((Iterable<Map.Entry<?, ?>>) obj)) {
+          Object key = e.getKey();
+          Object newKey = keyConversion(e.getKey(), state);
+          if (newKey == null && key != null) {
+            // probably we're out of elements anyway
+            continue;
+          }
+          newMap.put(newKey, guardedConversion(e.getValue(), depth + 1, state));
+        }
+        return newMap;
       } else {
-        newList = new ArrayList<>();
-      }
-      for (Object o : ((Iterable<?>) obj)) {
-        if (state.elemsLeft <= 0) {
-          state.listMapTooLarge = true;
-          break;
+        List<Object> newList;
+        if (obj instanceof Collection) {
+          newList = new ArrayList<>(((Collection<?>) obj).size());
+        } else {
+          newList = new ArrayList<>();
         }
-        newList.add(guardedConversion(o, depth + 1, state));
+        for (Object o : ((Iterable<?>) obj)) {
+          if (state.elemsLeft <= 0) {
+            state.listMapTooLarge = true;
+            break;
+          }
+          newList.add(guardedConversion(o, depth + 1, state));
+        }
+        return newList;
       }
-      return newList;
     }
 
     // arrays

dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/AppSecRequestContext.java

@@ -105,6 +105,7 @@ public class AppSecRequestContext implements DataBundle, Closeable {
   private boolean reqDataPublished;
   private boolean rawReqBodyPublished;
   private boolean convertedReqBodyPublished;
+  private boolean responseBodyPublished;
   private boolean respDataPublished;
   private boolean pathParamsPublished;
   private volatile Map<String, String> derivatives;
@@ -502,6 +503,14 @@ public void setConvertedReqBodyPublished(boolean convertedReqBodyPublished) {
     this.convertedReqBodyPublished = convertedReqBodyPublished;
   }
 
+  public boolean isResponseBodyPublished() {
+    return responseBodyPublished;
+  }
+
+  public void setResponseBodyPublished(final boolean responseBodyPublished) {
+    this.responseBodyPublished = responseBodyPublished;
+  }
+
   public boolean isRespDataPublished() {
     return respDataPublished;
   }

dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java

@@ -98,6 +98,7 @@ public class GatewayBridge {
   private volatile DataSubscriberInfo initialReqDataSubInfo;
   private volatile DataSubscriberInfo rawRequestBodySubInfo;
   private volatile DataSubscriberInfo requestBodySubInfo;
+  private volatile DataSubscriberInfo responseBodySubInfo;
   private volatile DataSubscriberInfo pathParamsSubInfo;
   private volatile DataSubscriberInfo respDataSubInfo;
   private volatile DataSubscriberInfo grpcServerMethodSubInfo;
@@ -137,6 +138,7 @@ public void init() {
     subscriptionService.registerCallback(EVENTS.requestMethodUriRaw(), this::onRequestMethodUriRaw);
     subscriptionService.registerCallback(EVENTS.requestBodyStart(), this::onRequestBodyStart);
     subscriptionService.registerCallback(EVENTS.requestBodyDone(), this::onRequestBodyDone);
+    subscriptionService.registerCallback(EVENTS.responseBody(), this::onResponseBody);
     subscriptionService.registerCallback(
         EVENTS.requestClientSocketAddress(), this::onRequestClientSocketAddress);
     subscriptionService.registerCallback(
@@ -177,6 +179,7 @@ public void reset() {
     initialReqDataSubInfo = null;
     rawRequestBodySubInfo = null;
     requestBodySubInfo = null;
+    responseBodySubInfo = null;
     pathParamsSubInfo = null;
     respDataSubInfo = null;
     grpcServerMethodSubInfo = null;
@@ -638,6 +641,39 @@ private Flow<Void> onRequestBodyDone(RequestContext ctx_, StoredBodySupplier sup
     }
   }
 
+  private Flow<Void> onResponseBody(RequestContext ctx_, Object obj) {
+    AppSecRequestContext ctx = ctx_.getData(RequestContextSlot.APPSEC);
+    if (ctx == null) {
+      return NoopFlow.INSTANCE;
+    }
+
+    if (ctx.isResponseBodyPublished()) {
+      log.debug(
+          "Response body already published; will ignore new value of type {}", obj.getClass());
+      return NoopFlow.INSTANCE;
+    }
+    ctx.setResponseBodyPublished(true);
+
+    while (true) {
+      DataSubscriberInfo subInfo = responseBodySubInfo;
+      if (subInfo == null) {
+        subInfo = producerService.getDataSubscribers(KnownAddresses.RESPONSE_BODY_OBJECT);
+        responseBodySubInfo = subInfo;
+      }
+      if (subInfo == null || subInfo.isEmpty()) {
+        return NoopFlow.INSTANCE;
+      }
+      Object converted = ObjectIntrospection.convert(obj, ctx);
+      DataBundle bundle = new SingletonDataBundle<>(KnownAddresses.RESPONSE_BODY_OBJECT, converted);
+      try {
+        GatewayContext gwCtx = new GatewayContext(false);
+        return producerService.publishDataEvent(subInfo, ctx, bundle, gwCtx);
+      } catch (ExpiredSubscriberInfoException e) {
+        responseBodySubInfo = null;
+      }
+    }
+  }
+
   private Flow<Void> onRequestPathParams(RequestContext ctx_, Map<String, ?> data) {
     AppSecRequestContext ctx = ctx_.getData(RequestContextSlot.APPSEC);
     if (ctx == null || ctx.isPathParamsPublished()) {

dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/event/data/ObjectIntrospectionSpecification.groovy

@@ -2,13 +2,10 @@ package com.datadog.appsec.event.data
 
 import com.datadog.appsec.gateway.AppSecRequestContext
 import com.fasterxml.jackson.databind.ObjectMapper
-import com.fasterxml.jackson.databind.node.ArrayNode
-import com.fasterxml.jackson.databind.node.ObjectNode
 import datadog.trace.api.telemetry.WafMetricCollector
 import datadog.trace.test.util.DDSpecification
 import groovy.json.JsonBuilder
 import groovy.json.JsonOutput
-import groovy.json.JsonSlurper
 import spock.lang.Shared
 
 import java.nio.CharBuffer
@@ -465,6 +462,17 @@ class ObjectIntrospectionSpecification extends DDSpecification {
     MAPPER.readTree('"unicode: \\u0041"')          || 'unicode: A'
   }
 
+  void 'iterable json objects'() {
+    setup:
+    final map = [name: 'This is just a test', list: [1, 2, 3, 4, 5]]
+
+    when:
+    final result = convert(new IterableJsonObject(map), ctx)
+
+    then:
+    result == map
+  }
+
   private static int countNesting(final Map<String, Object>object, final int levels) {
     if (object.isEmpty()) {
       return levels
@@ -475,4 +483,18 @@ class ObjectIntrospectionSpecification extends DDSpecification {
     }
     return countNesting(object.values().first() as Map, levels + 1)
   }
+
+  private static class IterableJsonObject implements Iterable<Map.Entry<String, Object>> {
+
+    private final Map<String, Object> map
+
+    IterableJsonObject(Map<String, Object> map) {
+      this.map = map
+    }
+
+    @Override
+    Iterator<Map.Entry<String, Object>> iterator() {
+      return map.entrySet().iterator()
+    }
+  }
 }

dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy

@@ -99,6 +99,7 @@ class GatewayBridgeSpecification extends DDSpecification {
   BiFunction<RequestContext, StoredBodySupplier, Void> requestBodyStartCB
   BiFunction<RequestContext, StoredBodySupplier, Flow<Void>> requestBodyDoneCB
   BiFunction<RequestContext, Object, Flow<Void>> requestBodyProcessedCB
+  BiFunction<RequestContext, Object, Flow<Void>> responseBodyCB
   BiFunction<RequestContext, Integer, Flow<Void>> responseStartedCB
   TriConsumer<RequestContext, String, String> respHeaderCB
   Function<RequestContext, Flow<Void>> respHeadersDoneCB
@@ -463,6 +464,7 @@ class GatewayBridgeSpecification extends DDSpecification {
     1 * ig.registerCallback(EVENTS.requestBodyStart(), _) >> { requestBodyStartCB = it[1]; null }
     1 * ig.registerCallback(EVENTS.requestBodyDone(), _) >> { requestBodyDoneCB = it[1]; null }
     1 * ig.registerCallback(EVENTS.requestBodyProcessed(), _) >> { requestBodyProcessedCB = it[1]; null }
+    1 * ig.registerCallback(EVENTS.responseBody(), _) >> { responseBodyCB = it[1]; null }
     1 * ig.registerCallback(EVENTS.responseStarted(), _) >> { responseStartedCB = it[1]; null }
     1 * ig.registerCallback(EVENTS.responseHeader(), _) >> { respHeaderCB = it[1]; null }
     1 * ig.registerCallback(EVENTS.responseHeaderDone(), _) >> { respHeadersDoneCB = it[1]; null }
@@ -1340,4 +1342,17 @@ class GatewayBridgeSpecification extends DDSpecification {
     arCtx.getRoute() == route
   }
 
+  void 'test on response body callback'() {
+    when:
+    responseBodyCB.apply(ctx, [test: 'this is a test'])
+
+    then:
+    1 * eventDispatcher.getDataSubscribers(KnownAddresses.RESPONSE_BODY_OBJECT) >> nonEmptyDsInfo
+    1 * eventDispatcher.publishDataEvent(_, _, _, _) >> {
+      final bundle = it[2] as DataBundle
+      final body = bundle.get(KnownAddresses.RESPONSE_BODY_OBJECT)
+      assert body['test'] == 'this is a test'
+    }
+  }
+
 }

dd-java-agent/instrumentation/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/server/RoutingContextInstrumentation.java

@@ -0,0 +1,41 @@
+package datadog.trace.instrumentation.vertx_4_0.server;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+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;
+import io.vertx.ext.web.impl.RoutingContextImpl;
+
+/**
+ * @see RoutingContextImpl#getBodyAsJson(int)
+ * @see RoutingContextImpl#getBodyAsJsonArray(int)
+ */
+@AutoService(InstrumenterModule.class)
+public class RoutingContextInstrumentation extends InstrumenterModule.AppSec
+    implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
+
+  public RoutingContextInstrumentation() {
+    super("vertx", "vertx-4.0");
+  }
+
+  @Override
+  public Reference[] additionalMuzzleReferences() {
+    return new Reference[] {VertxVersionMatcher.HTTP_1X_SERVER_RESPONSE};
+  }
+
+  @Override
+  public String instrumentedType() {
+    return "io.vertx.ext.web.RoutingContext";
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
+        named("json").and(takesArguments(1)).and(takesArgument(0, Object.class)),
+        packageName + ".RoutingContextJsonResponseAdvice");
+  }
+}

dd-java-agent/instrumentation/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/server/RoutingContextJsonResponseAdvice.java

@@ -0,0 +1,55 @@
+package datadog.trace.instrumentation.vertx_4_0.server;
+
+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;
+
+@RequiresRequestContext(RequestContextSlot.APPSEC)
+class RoutingContextJsonResponseAdvice {
+
+  @Advice.OnMethodEnter(suppress = Throwable.class)
+  static void before(
+      @Advice.Argument(0) final Object object, @ActiveRequestContext final RequestContext reqCtx) {
+
+    if (object == null) {
+      return;
+    }
+
+    CallbackProvider cbp = AgentTracer.get().getCallbackProvider(RequestContextSlot.APPSEC);
+    if (cbp == null) {
+      return;
+    }
+    BiFunction<RequestContext, Object, Flow<Void>> callback =
+        cbp.getCallback(EVENTS.responseBody());
+    if (callback == null) {
+      return;
+    }
+
+    Flow<Void> flow = callback.apply(reqCtx, object);
+    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 RoutingContext/json)");
+    }
+  }
+}

dd-java-agent/instrumentation/vertx-web-4.0/src/test/groovy/server/VertxHttpServerForkedTest.groovy

@@ -83,6 +83,11 @@ class VertxHttpServerForkedTest extends HttpServerTest<Vertx> {
     true
   }
 
+  @Override
+  boolean testResponseBodyJson() {
+    true
+  }
+
   @Override
   boolean testBlocking() {
     true

dd-java-agent/instrumentation/vertx-web-4.0/src/test/java/server/VertxTestServer.java

@@ -127,7 +127,8 @@ public void start(final Promise<Void> startPromise) {
                     BODY_JSON,
                     () -> {
                       JsonObject json = ctx.getBodyAsJson();
-                      ctx.response().setStatusCode(BODY_JSON.getStatus()).end(json.toString());
+                      ctx.response().setStatusCode(BODY_JSON.getStatus());
+                      ctx.json(json);
                     }));
     router
         .route(QUERY_ENCODED_BOTH.getRawPath())

dd-java-agent/instrumentation/vertx-web-5.0/src/test/groovy/server/VertxHttpServerForkedTest.groovy

@@ -67,6 +67,11 @@ class VertxHttpServerForkedTest extends HttpServerTest<Vertx> {
     true
   }
 
+  @Override
+  boolean testResponseBodyJson() {
+    true
+  }
+
   @Override
   boolean testBodyUrlencoded() {
     true

dd-java-agent/instrumentation/vertx-web-5.0/src/test/java/server/VertxTestServer.java

@@ -118,7 +118,8 @@ public void start(final Promise<Void> startPromise) {
                     BODY_JSON,
                     () -> {
                       JsonObject json = ctx.body().asJsonObject();
-                      ctx.response().setStatusCode(BODY_JSON.getStatus()).end(json.toString());
+                      ctx.response().setStatusCode(BODY_JSON.getStatus());
+                      ctx.json(json);
                     }));
     router
         .route(QUERY_ENCODED_BOTH.getRawPath())