From 8dc58579828c09ad9750dd4653eb0192f8928a9a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Manuel=20=C3=81lvarez=20=C3=81lvarez?= Date: Wed, 8 Jan 2025 10:49:00 +0100 Subject: [PATCH] Add support for session tracking in Vertx --- .../RoutingContextImplInstrumentation.java | 4 ++ .../server/RoutingContextSessionAdvice.java | 55 +++++++++++++++++++ .../server/VertxHttpServerForkedTest.groovy | 5 ++ .../src/test/java/server/VertxTestServer.java | 13 +++++ .../RoutingContextImplInstrumentation.java | 3 + .../server/RoutingContextSessionAdvice.java | 55 +++++++++++++++++++ .../server/VertxHttpServerForkedTest.groovy | 5 ++ .../src/test/java/server/VertxTestServer.java | 11 ++++ 8 files changed, 151 insertions(+) create mode 100644 dd-java-agent/instrumentation/vertx-web-3.4/src/main/java/datadog/trace/instrumentation/vertx_3_4/server/RoutingContextSessionAdvice.java create mode 100644 dd-java-agent/instrumentation/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/server/RoutingContextSessionAdvice.java diff --git a/dd-java-agent/instrumentation/vertx-web-3.4/src/main/java/datadog/trace/instrumentation/vertx_3_4/server/RoutingContextImplInstrumentation.java b/dd-java-agent/instrumentation/vertx-web-3.4/src/main/java/datadog/trace/instrumentation/vertx_3_4/server/RoutingContextImplInstrumentation.java index 4c9b1d1a6fc..bc52bbfdbed 100644 --- a/dd-java-agent/instrumentation/vertx-web-3.4/src/main/java/datadog/trace/instrumentation/vertx_3_4/server/RoutingContextImplInstrumentation.java +++ b/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"); } } diff --git a/dd-java-agent/instrumentation/vertx-web-3.4/src/main/java/datadog/trace/instrumentation/vertx_3_4/server/RoutingContextSessionAdvice.java b/dd-java-agent/instrumentation/vertx-web-3.4/src/main/java/datadog/trace/instrumentation/vertx_3_4/server/RoutingContextSessionAdvice.java new file mode 100644 index 00000000000..866e9c71ffa --- /dev/null +++ b/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> callback = + cbp.getCallback(EVENTS.requestSession()); + if (callback == null) { + return; + } + + Flow 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)"); + } + } + } +} diff --git a/dd-java-agent/instrumentation/vertx-web-3.4/src/test/groovy/server/VertxHttpServerForkedTest.groovy b/dd-java-agent/instrumentation/vertx-web-3.4/src/test/groovy/server/VertxHttpServerForkedTest.groovy index 54cf85bff5d..4b585252db3 100644 --- a/dd-java-agent/instrumentation/vertx-web-3.4/src/test/groovy/server/VertxHttpServerForkedTest.groovy +++ b/dd-java-agent/instrumentation/vertx-web-3.4/src/test/groovy/server/VertxHttpServerForkedTest.groovy @@ -124,6 +124,11 @@ class VertxHttpServerForkedTest extends HttpServerTest { true } + @Override + boolean testSessionId() { + true + } + @Override Serializable expectedServerSpanRoute(ServerEndpoint endpoint) { switch (endpoint) { 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 index 8d0d64b7614..c795d7860b6 100644 --- 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; @@ -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 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 diff --git a/dd-java-agent/instrumentation/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/server/RoutingContextImplInstrumentation.java b/dd-java-agent/instrumentation/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/server/RoutingContextImplInstrumentation.java index 5306e71bf85..7a1e157e174 100644 --- a/dd-java-agent/instrumentation/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/server/RoutingContextImplInstrumentation.java +++ b/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"); } } diff --git a/dd-java-agent/instrumentation/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/server/RoutingContextSessionAdvice.java b/dd-java-agent/instrumentation/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/server/RoutingContextSessionAdvice.java new file mode 100644 index 00000000000..2824822b3be --- /dev/null +++ b/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> callback = + cbp.getCallback(EVENTS.requestSession()); + if (callback == null) { + return; + } + + Flow 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)"); + } + } + } +} diff --git a/dd-java-agent/instrumentation/vertx-web-4.0/src/test/groovy/server/VertxHttpServerForkedTest.groovy b/dd-java-agent/instrumentation/vertx-web-4.0/src/test/groovy/server/VertxHttpServerForkedTest.groovy index a79fe77403c..d1b99d55b55 100644 --- a/dd-java-agent/instrumentation/vertx-web-4.0/src/test/groovy/server/VertxHttpServerForkedTest.groovy +++ b/dd-java-agent/instrumentation/vertx-web-4.0/src/test/groovy/server/VertxHttpServerForkedTest.groovy @@ -124,6 +124,11 @@ class VertxHttpServerForkedTest extends HttpServerTest { true } + @Override + boolean testSessionId() { + true + } + @Override Serializable expectedServerSpanRoute(ServerEndpoint endpoint) { switch (endpoint) { 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 index daa4562cc63..743f7bf83df 100644 --- 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; @@ -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 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