Add session tracking support to vert.x 3/4

DataDog · Nov 14, 2024 · 6fce896 · 6fce896
1 parent e7571fd
commit 6fce896
Show file tree

Hide file tree

Showing 8 changed files with 185 additions and 0 deletions.
diff --git a/...ava/datadog/trace/instrumentation/vertx_3_4/server/RoutingContextImplInstrumentation.java b/...ava/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");
   }
 }
diff --git a/...main/java/datadog/trace/instrumentation/vertx_3_4/server/RoutingContextSessionAdvice.java b/...main/java/datadog/trace/instrumentation/vertx_3_4/server/RoutingContextSessionAdvice.java
@@ -0,0 +1,54 @@
+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(
+      @Advice.Argument(0) final Session session,
+      @ActiveRequestContext RequestContext reqCtx,
+      @Advice.Thrown(readOnly = false) Throwable throwable) {
+    if (session == null || session.id() == null) {
+      return;
+    }
+
+    final 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 sessionId)");
+      }
+    }
+  }
+}
diff --git a/...ent/instrumentation/vertx-web-3.4/src/test/groovy/server/VertxHttpServerForkedTest.groovy b/...ent/instrumentation/vertx-web-3.4/src/test/groovy/server/VertxHttpServerForkedTest.groovy
@@ -111,6 +111,11 @@ class VertxHttpServerForkedTest extends HttpServerTest<Vertx> {
     return false
   }
 
+  @Override
+  boolean testSessionId() {
+    true
+  }
+
   @Override
   int spanCount(ServerEndpoint endpoint) {
     if (endpoint == NOT_FOUND) {

diff --git a/dd-java-agent/instrumentation/vertx-web-3.4/src/test/java/server/VertxTestServer.java b/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;
@@ -30,7 +31,11 @@
 import io.vertx.core.json.JsonObject;
 import io.vertx.ext.web.Router;
 import io.vertx.ext.web.RoutingContext;
+import io.vertx.ext.web.Session;
 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 +200,32 @@ 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());
+    final LocalSessionStore sessionStorage = LocalSessionStore.create(vertx);
+    router
+        .route(SESSION_ID.getPath())
+        .handler(
+            SessionHandler.create(sessionStorage)
+                .setCookieSecureFlag(true)
+                .setCookieHttpOnlyFlag(true));
+    router
+        .route(SESSION_ID.getPath())
+        .handler(
+            ctx ->
+                controller(
+                    ctx,
+                    SESSION_ID,
+                    () -> {
+                      final Session session = ctx.session();
+                      if (session == null) {
+                        ctx.response()
+                            .setStatusCode(500)
+                            .end("Cookie/Session handlers not present");
+                      } else {
+                        ctx.response().setStatusCode(SESSION_ID.getStatus()).end(session.id());
+                      }
+                    }));
+
     router = customizeAfterRoutes(router);
 
     vertx

diff --git a/...ava/datadog/trace/instrumentation/vertx_4_0/server/RoutingContextImplInstrumentation.java b/...ava/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");
   }
 }
diff --git a/...main/java/datadog/trace/instrumentation/vertx_4_0/server/RoutingContextSessionAdvice.java b/...main/java/datadog/trace/instrumentation/vertx_4_0/server/RoutingContextSessionAdvice.java
@@ -0,0 +1,54 @@
+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(
+      @Advice.Argument(0) final Session session,
+      @ActiveRequestContext RequestContext reqCtx,
+      @Advice.Thrown(readOnly = false) Throwable throwable) {
+    if (session == null || session.id() == null) {
+      return;
+    }
+
+    final 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 sessionId)");
+      }
+    }
+  }
+}
diff --git a/...ent/instrumentation/vertx-web-4.0/src/test/groovy/server/VertxHttpServerForkedTest.groovy b/...ent/instrumentation/vertx-web-4.0/src/test/groovy/server/VertxHttpServerForkedTest.groovy
@@ -111,6 +111,11 @@ class VertxHttpServerForkedTest extends HttpServerTest<Vertx> {
     return false
   }
 
+  @Override
+  boolean testSessionId() {
+    true
+  }
+
   @Override
   int spanCount(ServerEndpoint endpoint) {
     if (endpoint == NOT_FOUND) {

diff --git a/dd-java-agent/instrumentation/vertx-web-4.0/src/test/java/server/VertxTestServer.java b/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;
@@ -30,7 +31,10 @@
 import io.vertx.core.json.JsonObject;
 import io.vertx.ext.web.Router;
 import io.vertx.ext.web.RoutingContext;
+import io.vertx.ext.web.Session;
 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 +202,31 @@ public void start(final Promise<Void> startPromise) {
         .route(EXCEPTION.getPath())
         .handler(ctx -> controller(ctx, EXCEPTION, VertxTestServer::exception));
 
+    final LocalSessionStore sessionStorage = LocalSessionStore.create(vertx);
+    router
+        .route(SESSION_ID.getPath())
+        .handler(
+            SessionHandler.create(sessionStorage)
+                .setCookieSecureFlag(true)
+                .setCookieHttpOnlyFlag(true));
+    router
+        .route(SESSION_ID.getPath())
+        .handler(
+            ctx ->
+                controller(
+                    ctx,
+                    SESSION_ID,
+                    () -> {
+                      final Session session = ctx.session();
+                      if (session == null) {
+                        ctx.response()
+                            .setStatusCode(500)
+                            .end("Cookie/Session handlers not present");
+                      } else {
+                        ctx.response().setStatusCode(SESSION_ID.getStatus()).end(session.id());
+                      }
+                    }));
+
     router = customizeAfterRoutes(router);
 
     vertx