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 c25647aed74..81c11342088 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..048e4f6c791 --- /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,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> 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 sessionId)"); + } + } + } +} 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..bba468d16a0 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 @@ -111,6 +111,11 @@ class VertxHttpServerForkedTest extends HttpServerTest { 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 index 8d0d64b7614..7f3298f5310 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; @@ -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 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/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 ffaa0bf10bb..e574404cef1 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..bfe12e5a80b --- /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,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> 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 sessionId)"); + } + } + } +} 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..41fabb89d66 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 @@ -111,6 +111,11 @@ class VertxHttpServerForkedTest extends HttpServerTest { 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 index daa4562cc63..4feee63d559 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; @@ -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 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