Skip to content

Commit f3fd311

Browse files
manuel-alvarez-alvarezmanuel-alvarez-alvarez
authored
Add support for session tracking in Vertx (#8167)
1 parent 866fc61 commit f3fd311

File tree

8 files changed

+151
-0
lines changed

8 files changed

+151
-0
lines changed

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
44
import static datadog.trace.instrumentation.vertx_3_4.server.VertxVersionMatcher.PARSABLE_HEADER_VALUE;
55
import static datadog.trace.instrumentation.vertx_3_4.server.VertxVersionMatcher.VIRTUAL_HOST_HANDLER;
6+
import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
67
import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
78

89
import com.google.auto.service.AutoService;
@@ -33,5 +34,8 @@ public void methodAdvice(MethodTransformer transformer) {
3334
transformer.applyAdvice(
3435
named("getBodyAsJson").or(named("getBodyAsJsonArray")).and(takesArguments(0)),
3536
packageName + ".RoutingContextJsonAdvice");
37+
transformer.applyAdvice(
38+
named("setSession").and(takesArgument(0, named("io.vertx.ext.web.Session"))),
39+
packageName + ".RoutingContextSessionAdvice");
3640
}
3741
}

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

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
package datadog.trace.instrumentation.vertx_3_4.server;
2+
3+
import static datadog.trace.api.gateway.Events.EVENTS;
4+
5+
import datadog.appsec.api.blocking.BlockingException;
6+
import datadog.trace.advice.ActiveRequestContext;
7+
import datadog.trace.advice.RequiresRequestContext;
8+
import datadog.trace.api.gateway.BlockResponseFunction;
9+
import datadog.trace.api.gateway.CallbackProvider;
10+
import datadog.trace.api.gateway.Flow;
11+
import datadog.trace.api.gateway.RequestContext;
12+
import datadog.trace.api.gateway.RequestContextSlot;
13+
import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
14+
import io.vertx.ext.web.Session;
15+
import java.util.function.BiFunction;
16+
import net.bytebuddy.asm.Advice;
17+
18+
@RequiresRequestContext(RequestContextSlot.APPSEC)
19+
class RoutingContextSessionAdvice {
20+
@Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
21+
static void after(
22+
@ActiveRequestContext final RequestContext reqCtx,
23+
@Advice.Argument(0) final Session session,
24+
@Advice.Thrown(readOnly = false) Throwable throwable) {
25+
26+
if (session == null) {
27+
return;
28+
}
29+
30+
CallbackProvider cbp = AgentTracer.get().getCallbackProvider(RequestContextSlot.APPSEC);
31+
BiFunction<RequestContext, String, Flow<Void>> callback =
32+
cbp.getCallback(EVENTS.requestSession());
33+
if (callback == null) {
34+
return;
35+
}
36+
37+
Flow<Void> flow = callback.apply(reqCtx, session.id());
38+
Flow.Action action = flow.getAction();
39+
if (action instanceof Flow.Action.RequestBlockingAction) {
40+
BlockResponseFunction blockResponseFunction = reqCtx.getBlockResponseFunction();
41+
if (blockResponseFunction == null) {
42+
return;
43+
}
44+
Flow.Action.RequestBlockingAction rba = (Flow.Action.RequestBlockingAction) action;
45+
blockResponseFunction.tryCommitBlockingResponse(
46+
reqCtx.getTraceSegment(),
47+
rba.getStatusCode(),
48+
rba.getBlockingContentType(),
49+
rba.getExtraHeaders());
50+
if (throwable == null) {
51+
throwable = new BlockingException("Blocked request (for session)");
52+
}
53+
}
54+
}
55+
}

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -124,6 +124,11 @@ class VertxHttpServerForkedTest extends HttpServerTest<Vertx> {
124124
true
125125
}
126126

127+
@Override
128+
boolean testSessionId() {
129+
true
130+
}
131+
127132
@Override
128133
Serializable expectedServerSpanRoute(ServerEndpoint endpoint) {
129134
switch (endpoint) {

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

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.QUERY_ENCODED_QUERY;
1414
import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.QUERY_PARAM;
1515
import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.REDIRECT;
16+
import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.SESSION_ID;
1617
import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.SUCCESS;
1718
import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.UNKNOWN;
1819
import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.USER_BLOCK;
@@ -31,6 +32,9 @@
3132
import io.vertx.ext.web.Router;
3233
import io.vertx.ext.web.RoutingContext;
3334
import io.vertx.ext.web.handler.BodyHandler;
35+
import io.vertx.ext.web.handler.CookieHandler;
36+
import io.vertx.ext.web.handler.SessionHandler;
37+
import io.vertx.ext.web.sstore.LocalSessionStore;
3438

3539
public class VertxTestServer extends AbstractVerticle {
3640
public static final String CONFIG_HTTP_SERVER_PORT = "http.server.port";
@@ -195,6 +199,15 @@ public void start(final Future<Void> startFuture) {
195199
.route(EXCEPTION.getPath())
196200
.handler(ctx -> controller(ctx, EXCEPTION, VertxTestServer::exception));
197201

202+
router.route(SESSION_ID.getPath()).handler(CookieHandler.create());
203+
router
204+
.route(SESSION_ID.getPath())
205+
.handler(SessionHandler.create(LocalSessionStore.create(vertx)));
206+
router
207+
.route(SESSION_ID.getPath())
208+
.handler(
209+
ctx -> ctx.response().setStatusCode(SESSION_ID.getStatus()).end(ctx.session().id()));
210+
198211
router = customizeAfterRoutes(router);
199212

200213
vertx

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,5 +40,8 @@ public void methodAdvice(MethodTransformer transformer) {
4040
.and(takesArguments(1))
4141
.and(takesArgument(0, int.class)),
4242
packageName + ".RoutingContextJsonAdvice");
43+
transformer.applyAdvice(
44+
named("setSession").and(takesArgument(0, named("io.vertx.ext.web.Session"))),
45+
packageName + ".RoutingContextSessionAdvice");
4346
}
4447
}

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

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
package datadog.trace.instrumentation.vertx_4_0.server;
2+
3+
import static datadog.trace.api.gateway.Events.EVENTS;
4+
5+
import datadog.appsec.api.blocking.BlockingException;
6+
import datadog.trace.advice.ActiveRequestContext;
7+
import datadog.trace.advice.RequiresRequestContext;
8+
import datadog.trace.api.gateway.BlockResponseFunction;
9+
import datadog.trace.api.gateway.CallbackProvider;
10+
import datadog.trace.api.gateway.Flow;
11+
import datadog.trace.api.gateway.RequestContext;
12+
import datadog.trace.api.gateway.RequestContextSlot;
13+
import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
14+
import io.vertx.ext.web.Session;
15+
import java.util.function.BiFunction;
16+
import net.bytebuddy.asm.Advice;
17+
18+
@RequiresRequestContext(RequestContextSlot.APPSEC)
19+
class RoutingContextSessionAdvice {
20+
@Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
21+
static void after(
22+
@ActiveRequestContext final RequestContext reqCtx,
23+
@Advice.Argument(0) final Session session,
24+
@Advice.Thrown(readOnly = false) Throwable throwable) {
25+
26+
if (session == null) {
27+
return;
28+
}
29+
30+
CallbackProvider cbp = AgentTracer.get().getCallbackProvider(RequestContextSlot.APPSEC);
31+
BiFunction<RequestContext, String, Flow<Void>> callback =
32+
cbp.getCallback(EVENTS.requestSession());
33+
if (callback == null) {
34+
return;
35+
}
36+
37+
Flow<Void> flow = callback.apply(reqCtx, session.id());
38+
Flow.Action action = flow.getAction();
39+
if (action instanceof Flow.Action.RequestBlockingAction) {
40+
BlockResponseFunction blockResponseFunction = reqCtx.getBlockResponseFunction();
41+
if (blockResponseFunction == null) {
42+
return;
43+
}
44+
Flow.Action.RequestBlockingAction rba = (Flow.Action.RequestBlockingAction) action;
45+
blockResponseFunction.tryCommitBlockingResponse(
46+
reqCtx.getTraceSegment(),
47+
rba.getStatusCode(),
48+
rba.getBlockingContentType(),
49+
rba.getExtraHeaders());
50+
if (throwable == null) {
51+
throwable = new BlockingException("Blocked request (for session)");
52+
}
53+
}
54+
}
55+
}

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -124,6 +124,11 @@ class VertxHttpServerForkedTest extends HttpServerTest<Vertx> {
124124
true
125125
}
126126

127+
@Override
128+
boolean testSessionId() {
129+
true
130+
}
131+
127132
@Override
128133
Serializable expectedServerSpanRoute(ServerEndpoint endpoint) {
129134
switch (endpoint) {

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

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.QUERY_ENCODED_QUERY;
1414
import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.QUERY_PARAM;
1515
import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.REDIRECT;
16+
import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.SESSION_ID;
1617
import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.SUCCESS;
1718
import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.UNKNOWN;
1819
import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.USER_BLOCK;
@@ -31,6 +32,8 @@
3132
import io.vertx.ext.web.Router;
3233
import io.vertx.ext.web.RoutingContext;
3334
import io.vertx.ext.web.handler.BodyHandler;
35+
import io.vertx.ext.web.handler.SessionHandler;
36+
import io.vertx.ext.web.sstore.LocalSessionStore;
3437

3538
public class VertxTestServer extends AbstractVerticle {
3639
public static final String CONFIG_HTTP_SERVER_PORT = "http.server.port";
@@ -198,6 +201,14 @@ public void start(final Promise<Void> startPromise) {
198201
.route(EXCEPTION.getPath())
199202
.handler(ctx -> controller(ctx, EXCEPTION, VertxTestServer::exception));
200203

204+
router
205+
.route(SESSION_ID.getPath())
206+
.handler(SessionHandler.create(LocalSessionStore.create(vertx)));
207+
router
208+
.route(SESSION_ID.getPath())
209+
.handler(
210+
ctx -> ctx.response().setStatusCode(SESSION_ID.getStatus()).end(ctx.session().id()));
211+
201212
router = customizeAfterRoutes(router);
202213

203214
vertx

0 commit comments

Comments
 (0)