Skip to content

Commit b2cab51

Browse files
manuel-alvarez-alvarezmanuel-alvarez-alvarez
committed
Add session tracking support to vert.x 3/4
1 parent 4e74c0e commit b2cab51

File tree

8 files changed

+185
-0
lines changed

8 files changed

+185
-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: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
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+
@Advice.Argument(0) final Session session,
23+
@ActiveRequestContext RequestContext reqCtx,
24+
@Advice.Thrown(readOnly = false) Throwable throwable) {
25+
if (session == null || session.id() == null) {
26+
return;
27+
}
28+
29+
final CallbackProvider cbp = AgentTracer.get().getCallbackProvider(RequestContextSlot.APPSEC);
30+
BiFunction<RequestContext, String, Flow<Void>> callback =
31+
cbp.getCallback(EVENTS.requestSession());
32+
if (callback == null) {
33+
return;
34+
}
35+
36+
Flow<Void> flow = callback.apply(reqCtx, session.id());
37+
Flow.Action action = flow.getAction();
38+
if (action instanceof Flow.Action.RequestBlockingAction) {
39+
BlockResponseFunction blockResponseFunction = reqCtx.getBlockResponseFunction();
40+
if (blockResponseFunction == null) {
41+
return;
42+
}
43+
Flow.Action.RequestBlockingAction rba = (Flow.Action.RequestBlockingAction) action;
44+
blockResponseFunction.tryCommitBlockingResponse(
45+
reqCtx.getTraceSegment(),
46+
rba.getStatusCode(),
47+
rba.getBlockingContentType(),
48+
rba.getExtraHeaders());
49+
if (throwable == null) {
50+
throwable = new BlockingException("Blocked request (for sessionId)");
51+
}
52+
}
53+
}
54+
}

‎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
@@ -111,6 +111,11 @@ class VertxHttpServerForkedTest extends HttpServerTest<Vertx> {
111111
return false
112112
}
113113

114+
@Override
115+
boolean testSessionId() {
116+
true
117+
}
118+
114119
@Override
115120
int spanCount(ServerEndpoint endpoint) {
116121
if (endpoint == NOT_FOUND) {

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

Lines changed: 31 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;
@@ -30,7 +31,11 @@
3031
import io.vertx.core.json.JsonObject;
3132
import io.vertx.ext.web.Router;
3233
import io.vertx.ext.web.RoutingContext;
34+
import io.vertx.ext.web.Session;
3335
import io.vertx.ext.web.handler.BodyHandler;
36+
import io.vertx.ext.web.handler.CookieHandler;
37+
import io.vertx.ext.web.handler.SessionHandler;
38+
import io.vertx.ext.web.sstore.LocalSessionStore;
3439

3540
public class VertxTestServer extends AbstractVerticle {
3641
public static final String CONFIG_HTTP_SERVER_PORT = "http.server.port";
@@ -195,6 +200,32 @@ public void start(final Future<Void> startFuture) {
195200
.route(EXCEPTION.getPath())
196201
.handler(ctx -> controller(ctx, EXCEPTION, VertxTestServer::exception));
197202

203+
router.route(SESSION_ID.getPath()).handler(CookieHandler.create());
204+
final LocalSessionStore sessionStorage = LocalSessionStore.create(vertx);
205+
router
206+
.route(SESSION_ID.getPath())
207+
.handler(
208+
SessionHandler.create(sessionStorage)
209+
.setCookieSecureFlag(true)
210+
.setCookieHttpOnlyFlag(true));
211+
router
212+
.route(SESSION_ID.getPath())
213+
.handler(
214+
ctx ->
215+
controller(
216+
ctx,
217+
SESSION_ID,
218+
() -> {
219+
final Session session = ctx.session();
220+
if (session == null) {
221+
ctx.response()
222+
.setStatusCode(500)
223+
.end("Cookie/Session handlers not present");
224+
} else {
225+
ctx.response().setStatusCode(SESSION_ID.getStatus()).end(session.id());
226+
}
227+
}));
228+
198229
router = customizeAfterRoutes(router);
199230

200231
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: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
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+
@Advice.Argument(0) final Session session,
23+
@ActiveRequestContext RequestContext reqCtx,
24+
@Advice.Thrown(readOnly = false) Throwable throwable) {
25+
if (session == null || session.id() == null) {
26+
return;
27+
}
28+
29+
final CallbackProvider cbp = AgentTracer.get().getCallbackProvider(RequestContextSlot.APPSEC);
30+
BiFunction<RequestContext, String, Flow<Void>> callback =
31+
cbp.getCallback(EVENTS.requestSession());
32+
if (callback == null) {
33+
return;
34+
}
35+
36+
Flow<Void> flow = callback.apply(reqCtx, session.id());
37+
Flow.Action action = flow.getAction();
38+
if (action instanceof Flow.Action.RequestBlockingAction) {
39+
BlockResponseFunction blockResponseFunction = reqCtx.getBlockResponseFunction();
40+
if (blockResponseFunction == null) {
41+
return;
42+
}
43+
Flow.Action.RequestBlockingAction rba = (Flow.Action.RequestBlockingAction) action;
44+
blockResponseFunction.tryCommitBlockingResponse(
45+
reqCtx.getTraceSegment(),
46+
rba.getStatusCode(),
47+
rba.getBlockingContentType(),
48+
rba.getExtraHeaders());
49+
if (throwable == null) {
50+
throwable = new BlockingException("Blocked request (for sessionId)");
51+
}
52+
}
53+
}
54+
}

‎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
@@ -111,6 +111,11 @@ class VertxHttpServerForkedTest extends HttpServerTest<Vertx> {
111111
return false
112112
}
113113

114+
@Override
115+
boolean testSessionId() {
116+
true
117+
}
118+
114119
@Override
115120
int spanCount(ServerEndpoint endpoint) {
116121
if (endpoint == NOT_FOUND) {

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

Lines changed: 29 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;
@@ -30,7 +31,10 @@
3031
import io.vertx.core.json.JsonObject;
3132
import io.vertx.ext.web.Router;
3233
import io.vertx.ext.web.RoutingContext;
34+
import io.vertx.ext.web.Session;
3335
import io.vertx.ext.web.handler.BodyHandler;
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";
@@ -198,6 +202,31 @@ public void start(final Promise<Void> startPromise) {
198202
.route(EXCEPTION.getPath())
199203
.handler(ctx -> controller(ctx, EXCEPTION, VertxTestServer::exception));
200204

205+
final LocalSessionStore sessionStorage = LocalSessionStore.create(vertx);
206+
router
207+
.route(SESSION_ID.getPath())
208+
.handler(
209+
SessionHandler.create(sessionStorage)
210+
.setCookieSecureFlag(true)
211+
.setCookieHttpOnlyFlag(true));
212+
router
213+
.route(SESSION_ID.getPath())
214+
.handler(
215+
ctx ->
216+
controller(
217+
ctx,
218+
SESSION_ID,
219+
() -> {
220+
final Session session = ctx.session();
221+
if (session == null) {
222+
ctx.response()
223+
.setStatusCode(500)
224+
.end("Cookie/Session handlers not present");
225+
} else {
226+
ctx.response().setStatusCode(SESSION_ID.getStatus()).end(session.id());
227+
}
228+
}));
229+
201230
router = customizeAfterRoutes(router);
202231

203232
vertx

0 commit comments

Comments
 (0)