Add support for session tracking in Vertx

dd-java-agent/instrumentation/vertx-web-3.4/src/main/java/datadog/trace/instrumentation/vertx_3_4/server/RoutingContextImplInstrumentation.java

@@ -3,6 +3,7 @@
 import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
 import static datadog.trace.instrumentation.vertx_3_4.server.VertxVersionMatcher.PARSABLE_HEADER_VALUE;
 import static datadog.trace.instrumentation.vertx_3_4.server.VertxVersionMatcher.VIRTUAL_HOST_HANDLER;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
 import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
 
 import com.google.auto.service.AutoService;
@@ -33,5 +34,8 @@ public void methodAdvice(MethodTransformer transformer) {
     transformer.applyAdvice(
         named("getBodyAsJson").or(named("getBodyAsJsonArray")).and(takesArguments(0)),
         packageName + ".RoutingContextJsonAdvice");
+    transformer.applyAdvice(
+        named("setSession").and(takesArgument(0, named("io.vertx.ext.web.Session"))),
+        packageName + ".RoutingContextSessionAdvice");
   }
 }

dd-java-agent/instrumentation/vertx-web-3.4/src/main/java/datadog/trace/instrumentation/vertx_3_4/server/RoutingContextSessionAdvice.java

@@ -0,0 +1,55 @@
+package datadog.trace.instrumentation.vertx_3_4.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 io.vertx.ext.web.Session;
+import java.util.function.BiFunction;
+import net.bytebuddy.asm.Advice;
+
+@RequiresRequestContext(RequestContextSlot.APPSEC)
+class RoutingContextSessionAdvice {
+  @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
+  static void after(
+      @ActiveRequestContext final RequestContext reqCtx,
+      @Advice.Argument(0) final Session session,
+      @Advice.Thrown(readOnly = false) Throwable throwable) {
+
+    if (session == null) {
+      return;
+    }
+
+    CallbackProvider cbp = AgentTracer.get().getCallbackProvider(RequestContextSlot.APPSEC);
+    BiFunction<RequestContext, String, Flow<Void>> callback =
+        cbp.getCallback(EVENTS.requestSession());
+    if (callback == null) {
+      return;
+    }
+
+    Flow<Void> flow = callback.apply(reqCtx, session.id());
+    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());
+      if (throwable == null) {
+        throwable = new BlockingException("Blocked request (for session)");
+      }
+    }
+  }
+}

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

@@ -124,6 +124,11 @@ class VertxHttpServerForkedTest extends HttpServerTest<Vertx> {
     true
   }
 
+  @Override
+  boolean testSessionId() {
+    true
+  }
+
   @Override
   Serializable expectedServerSpanRoute(ServerEndpoint endpoint) {
     switch (endpoint) {

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

@@ -13,6 +13,7 @@
 import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.QUERY_ENCODED_QUERY;
 import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.QUERY_PARAM;
 import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.REDIRECT;
+import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.SESSION_ID;
 import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.SUCCESS;
 import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.UNKNOWN;
 import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.USER_BLOCK;
@@ -31,6 +32,9 @@
 import io.vertx.ext.web.Router;
 import io.vertx.ext.web.RoutingContext;
 import io.vertx.ext.web.handler.BodyHandler;
+import io.vertx.ext.web.handler.CookieHandler;
+import io.vertx.ext.web.handler.SessionHandler;
+import io.vertx.ext.web.sstore.LocalSessionStore;
 
 public class VertxTestServer extends AbstractVerticle {
   public static final String CONFIG_HTTP_SERVER_PORT = "http.server.port";
@@ -195,6 +199,15 @@ public void start(final Future<Void> startFuture) {
         .route(EXCEPTION.getPath())
         .handler(ctx -> controller(ctx, EXCEPTION, VertxTestServer::exception));
 
+    router.route(SESSION_ID.getPath()).handler(CookieHandler.create());
+    router
+        .route(SESSION_ID.getPath())
+        .handler(SessionHandler.create(LocalSessionStore.create(vertx)));
+    router
+        .route(SESSION_ID.getPath())
+        .handler(
+            ctx -> ctx.response().setStatusCode(SESSION_ID.getStatus()).end(ctx.session().id()));
+
     router = customizeAfterRoutes(router);
 
     vertx

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

@@ -40,5 +40,8 @@ public void methodAdvice(MethodTransformer transformer) {
             .and(takesArguments(1))
             .and(takesArgument(0, int.class)),
         packageName + ".RoutingContextJsonAdvice");
+    transformer.applyAdvice(
+        named("setSession").and(takesArgument(0, named("io.vertx.ext.web.Session"))),
+        packageName + ".RoutingContextSessionAdvice");
   }
 }

dd-java-agent/instrumentation/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/server/RoutingContextSessionAdvice.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 io.vertx.ext.web.Session;
+import java.util.function.BiFunction;
+import net.bytebuddy.asm.Advice;
+
+@RequiresRequestContext(RequestContextSlot.APPSEC)
+class RoutingContextSessionAdvice {
+  @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
+  static void after(
+      @ActiveRequestContext final RequestContext reqCtx,
+      @Advice.Argument(0) final Session session,
+      @Advice.Thrown(readOnly = false) Throwable throwable) {
+
+    if (session == null) {
+      return;
+    }
+
+    CallbackProvider cbp = AgentTracer.get().getCallbackProvider(RequestContextSlot.APPSEC);
+    BiFunction<RequestContext, String, Flow<Void>> callback =
+        cbp.getCallback(EVENTS.requestSession());
+    if (callback == null) {
+      return;
+    }
+
+    Flow<Void> flow = callback.apply(reqCtx, session.id());
+    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());
+      if (throwable == null) {
+        throwable = new BlockingException("Blocked request (for session)");
+      }
+    }
+  }
+}

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

@@ -124,6 +124,11 @@ class VertxHttpServerForkedTest extends HttpServerTest<Vertx> {
     true
   }
 
+  @Override
+  boolean testSessionId() {
+    true
+  }
+
   @Override
   Serializable expectedServerSpanRoute(ServerEndpoint endpoint) {
     switch (endpoint) {

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

@@ -13,6 +13,7 @@
 import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.QUERY_ENCODED_QUERY;
 import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.QUERY_PARAM;
 import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.REDIRECT;
+import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.SESSION_ID;
 import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.SUCCESS;
 import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.UNKNOWN;
 import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.USER_BLOCK;
@@ -31,6 +32,8 @@
 import io.vertx.ext.web.Router;
 import io.vertx.ext.web.RoutingContext;
 import io.vertx.ext.web.handler.BodyHandler;
+import io.vertx.ext.web.handler.SessionHandler;
+import io.vertx.ext.web.sstore.LocalSessionStore;
 
 public class VertxTestServer extends AbstractVerticle {
   public static final String CONFIG_HTTP_SERVER_PORT = "http.server.port";
@@ -198,6 +201,14 @@ public void start(final Promise<Void> startPromise) {
         .route(EXCEPTION.getPath())
         .handler(ctx -> controller(ctx, EXCEPTION, VertxTestServer::exception));
 
+    router
+        .route(SESSION_ID.getPath())
+        .handler(SessionHandler.create(LocalSessionStore.create(vertx)));
+    router
+        .route(SESSION_ID.getPath())
+        .handler(
+            ctx -> ctx.response().setStatusCode(SESSION_ID.getStatus()).end(ctx.session().id()));
+
     router = customizeAfterRoutes(router);
 
     vertx