diff --git a/dd-java-agent/instrumentation/vertx-rx-3.5/src/test/groovy/server/VertxRxCircuitBreakerHttpServerForkedTest.groovy b/dd-java-agent/instrumentation/vertx-rx-3.5/src/test/groovy/server/VertxRxCircuitBreakerHttpServerForkedTest.groovy index 694ae1bf7c4..765ba3819d8 100644 --- a/dd-java-agent/instrumentation/vertx-rx-3.5/src/test/groovy/server/VertxRxCircuitBreakerHttpServerForkedTest.groovy +++ b/dd-java-agent/instrumentation/vertx-rx-3.5/src/test/groovy/server/VertxRxCircuitBreakerHttpServerForkedTest.groovy @@ -11,6 +11,9 @@ import io.vertx.reactivex.core.MultiMap import io.vertx.reactivex.ext.web.Router import io.vertx.reactivex.ext.web.RoutingContext import io.vertx.reactivex.ext.web.handler.BodyHandler +import io.vertx.reactivex.ext.web.handler.CookieHandler +import io.vertx.reactivex.ext.web.handler.SessionHandler +import io.vertx.reactivex.ext.web.sstore.LocalSessionStore import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.BODY_MULTIPART import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.BODY_URLENCODED @@ -22,6 +25,7 @@ import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.QUERY_ 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.USER_BLOCK import static server.VertxTestServer.CONFIG_HTTP_SERVER_PORT @@ -233,6 +237,26 @@ class VertxRxCircuitBreakerHttpServerForkedTest extends VertxHttpServerForkedTes }) } + router.route(SESSION_ID.getPath()).handler(CookieHandler.create()) + final LocalSessionStore sessionStorage = LocalSessionStore.create(super.@vertx) + router + .route(SESSION_ID.getPath()) + .handler(SessionHandler.create(sessionStorage).setCookieSecureFlag(true).setCookieHttpOnlyFlag(true)) + router + .route(SESSION_ID.getPath()) + .handler { ctx -> + controller(SESSION_ID) { + final session = ctx.session() + if (!session) { + ctx.response() + .setStatusCode(500) + .end('Cookie/Session handlers not present') + } else { + ctx.response().setStatusCode(SESSION_ID.getStatus()).end(session.id()) + } + } + } + super.@vertx.createHttpServer(new HttpServerOptions().setHandle100ContinueAutomatically(true)) .requestHandler { router.accept(it) } .listen(port) { startFuture.complete() } diff --git a/dd-java-agent/instrumentation/vertx-rx-3.5/src/test/groovy/server/VertxRxHttpServerForkedTest.groovy b/dd-java-agent/instrumentation/vertx-rx-3.5/src/test/groovy/server/VertxRxHttpServerForkedTest.groovy index 34444c8dce2..516a2ebdb22 100644 --- a/dd-java-agent/instrumentation/vertx-rx-3.5/src/test/groovy/server/VertxRxHttpServerForkedTest.groovy +++ b/dd-java-agent/instrumentation/vertx-rx-3.5/src/test/groovy/server/VertxRxHttpServerForkedTest.groovy @@ -4,6 +4,9 @@ import io.vertx.core.AbstractVerticle import io.vertx.core.Future import io.vertx.core.http.HttpServerOptions import io.vertx.ext.web.Router +import io.vertx.ext.web.handler.CookieHandler +import io.vertx.ext.web.handler.SessionHandler +import io.vertx.ext.web.sstore.LocalSessionStore import spock.lang.Ignore import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.ERROR @@ -12,6 +15,7 @@ import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.FORWAR import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.PATH_PARAM 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 server.VertxTestServer.CONFIG_HTTP_SERVER_PORT @@ -72,6 +76,24 @@ class VertxRxHttpServerForkedTest extends VertxHttpServerForkedTest { } } + 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(SESSION_ID) { + final 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()) + } + } + } + super.@vertx.createHttpServer(new HttpServerOptions().setHandle100ContinueAutomatically(true)) .requestHandler { router.accept(it) } .listen(port) { startFuture.complete() } 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