Skip to content

Commit e71c496

Browse files
manuel-alvarez-alvarezmanuel-alvarez-alvarez
committed
Add support for session tracking in jetty
1 parent 82ca377 commit e71c496

File tree

23 files changed

+308
-17
lines changed

23 files changed

+308
-17
lines changed

‎dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/AppSecRequestContext.java‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -125,6 +125,8 @@ public class AppSecRequestContext implements DataBundle, Closeable {
125125

126126
// keep a reference to the last published usr.id
127127
private volatile String userId;
128+
// keep a reference to the last published usr.session_id
129+
private volatile String sessionId;
128130

129131
private static final AtomicIntegerFieldUpdater<AppSecRequestContext> TIMEOUTS_UPDATER =
130132
AtomicIntegerFieldUpdater.newUpdater(AppSecRequestContext.class, "timeouts");
@@ -433,6 +435,14 @@ public void setUserId(String userId) {
433435
this.userId = userId;
434436
}
435437

438+
public void setSessionId(String sessionId) {
439+
this.sessionId = sessionId;
440+
}
441+
442+
public String getSessionId() {
443+
return sessionId;
444+
}
445+
436446
@Override
437447
public void close() {
438448
final AgentSpan span = AgentTracer.activeSpan();

‎dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -154,7 +154,7 @@ private TriFunction<RequestContext, UserIdCollectionMode, String, Flow<Void>> on
154154
final Address<String> address) {
155155
return (ctx_, mode, userId) -> {
156156
final AppSecRequestContext ctx = ctx_.getData(RequestContextSlot.APPSEC);
157-
if (ctx == null) {
157+
if (userId == null || ctx == null) {
158158
return NoopFlow.INSTANCE;
159159
}
160160
final TraceSegment segment = ctx_.getTraceSegment();
@@ -203,9 +203,14 @@ private TriFunction<RequestContext, UserIdCollectionMode, String, Flow<Void>> on
203203

204204
private Flow<Void> onRequestSession(final RequestContext ctx_, final String sessionId) {
205205
final AppSecRequestContext ctx = ctx_.getData(RequestContextSlot.APPSEC);
206-
if (ctx == null) {
206+
if (sessionId == null || ctx == null) {
207+
return NoopFlow.INSTANCE;
208+
}
209+
if (sessionId.equals(ctx.getSessionId())) {
207210
return NoopFlow.INSTANCE;
208211
}
212+
// unlikely that multiple threads will update the value at the same time
213+
ctx.setSessionId(sessionId);
209214
final TraceSegment segment = ctx_.getTraceSegment();
210215
segment.setTagTop("usr.session_id", sessionId);
211216
while (true) {

‎dd-java-agent/instrumentation/jetty-11/src/main/java/datadog/trace/instrumentation/jetty11/RequestInstrumentation.java‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,5 +27,8 @@ public void methodAdvice(MethodTransformer transformer) {
2727
.and(takesArgument(0, named("org.eclipse.jetty.server.handler.ContextHandler$Context")))
2828
.and(takesArgument(1, String.class)),
2929
packageName + ".SetContextPathAdvice");
30+
transformer.applyAdvice(
31+
named("setRequestedSessionId").and(takesArgument(0, String.class)),
32+
packageName + ".SetRequestedSessionIdAdvice");
3033
}
3134
}

‎dd-java-agent/instrumentation/jetty-11/src/main/java11/datadog/trace/instrumentation/jetty11/SetRequestedSessionIdAdvice.java‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
package datadog.trace.instrumentation.jetty11;
2+
3+
import datadog.trace.advice.ActiveRequestContext;
4+
import datadog.trace.advice.RequiresRequestContext;
5+
import datadog.trace.api.gateway.RequestContext;
6+
import datadog.trace.api.gateway.RequestContextSlot;
7+
import datadog.trace.instrumentation.jetty.JettyBlockingHelper;
8+
import net.bytebuddy.asm.Advice;
9+
10+
/**
11+
* Because we are processing the initial request before the requestedSessionId is set, we must
12+
* update it when it is actually set.
13+
*/
14+
@RequiresRequestContext(RequestContextSlot.APPSEC)
15+
public class SetRequestedSessionIdAdvice {
16+
@Advice.OnMethodEnter(suppress = Throwable.class)
17+
public static void setRequestedSessionId(
18+
@ActiveRequestContext RequestContext reqCtx,
19+
@Advice.Argument(0) final String requestedSessionId) {
20+
JettyBlockingHelper.maybeBlockOnSession(reqCtx, requestedSessionId);
21+
}
22+
}

‎dd-java-agent/instrumentation/jetty-11/src/test/groovy/Jetty11Test.groovy‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -84,6 +84,11 @@ abstract class Jetty11Test extends HttpServerTest<Server> {
8484
boolean hasExtraErrorInformation() {
8585
true
8686
}
87+
88+
@Override
89+
boolean testSessionId() {
90+
true
91+
}
8792
}
8893

8994
class Jetty11V0ForkedTest extends Jetty11Test implements TestingGenericHttpNamingConventions.ServerV0 {

‎dd-java-agent/instrumentation/jetty-11/src/testFixtures/groovy/JettyServer.groovy‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import org.eclipse.jetty.server.Handler
1111
import org.eclipse.jetty.server.Server
1212
import org.eclipse.jetty.server.handler.AbstractHandler
1313
import org.eclipse.jetty.server.handler.ErrorHandler
14+
import org.eclipse.jetty.server.session.SessionHandler
1415
import org.eclipse.jetty.servlet.FilterMapping
1516
import org.eclipse.jetty.servlet.ServletContextHandler
1617

@@ -19,7 +20,8 @@ class JettyServer implements HttpServer {
1920
final server = new Server(0) // select random open port
2021

2122
JettyServer(Handler handler) {
22-
server.handler = handler
23+
sessionHandler.handler = handler
24+
server.handler = sessionHandler
2325
server.addBean(errorHandler)
2426
}
2527

@@ -47,6 +49,7 @@ class JettyServer implements HttpServer {
4749

4850
static AbstractHandler servletHandler(Class<? extends Servlet> servlet) {
4951
ServletContextHandler handler = new ServletContextHandler(null, "/context-path")
52+
handler.sessionHandler = sessionHandler
5053
handler.errorHandler = errorHandler
5154
handler.servletHandler.addFilterWithMapping(EnableMultipartFilter, '/*', FilterMapping.ALL)
5255
handler.servletHandler.addServletWithMapping(servlet, '/*')
@@ -72,4 +75,6 @@ class JettyServer implements HttpServer {
7275
}
7376
}
7477
}
78+
79+
static sessionHandler = new SessionHandler()
7580
}

‎dd-java-agent/instrumentation/jetty-12/build.gradle‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,9 @@ dependencies {
3939
main_java17CompileOnly ("org.eclipse.jetty:jetty-server:12.0.0") {
4040
exclude group: 'org.slf4j', module: 'slf4j-api'
4141
}
42+
main_java17CompileOnly ("org.eclipse.jetty:jetty-session:12.0.0") {
43+
exclude group: 'org.slf4j', module: 'slf4j-api'
44+
}
4245
main_java17Implementation project(':dd-java-agent:instrumentation:jetty-common')
4346
implementation project(':dd-java-agent:instrumentation:jetty-common')
4447

@@ -62,6 +65,9 @@ dependencies {
6265
latestDepTestImplementation ("org.eclipse.jetty:jetty-server:12.+") {
6366
exclude group: 'org.slf4j', module: 'slf4j-api'
6467
}
68+
latestDepTestImplementation ("org.eclipse.jetty:jetty-session:12.+") {
69+
exclude group: 'org.slf4j', module: 'slf4j-api'
70+
}
6571

6672
ee8TestImplementation group: 'org.eclipse.jetty.ee8', name: 'jetty-ee8-servlet', version: '12.0.0' , {
6773
exclude group: 'org.slf4j', module: 'slf4j-api'

‎dd-java-agent/instrumentation/jetty-12/src/main/java/datadog/trace/instrumentation/jetty12/AbstractSessionManagerInstrumentation.java‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
package datadog.trace.instrumentation.jetty12;
2+
3+
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
4+
import static net.bytebuddy.matcher.ElementMatchers.isProtected;
5+
import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
6+
import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
7+
8+
import com.google.auto.service.AutoService;
9+
import datadog.trace.agent.tooling.Instrumenter;
10+
import datadog.trace.agent.tooling.InstrumenterModule;
11+
12+
@AutoService(InstrumenterModule.class)
13+
public final class AbstractSessionManagerInstrumentation extends InstrumenterModule.Tracing
14+
implements Instrumenter.ForSingleType {
15+
16+
public AbstractSessionManagerInstrumentation() {
17+
super("jetty");
18+
}
19+
20+
@Override
21+
public String instrumentedType() {
22+
return "org.eclipse.jetty.session.AbstractSessionManager";
23+
}
24+
25+
@Override
26+
public void methodAdvice(MethodTransformer transformer) {
27+
transformer.applyAdvice(
28+
named("resolveRequestedSessionId")
29+
.and(isProtected())
30+
.and(takesArguments(1))
31+
.and(takesArgument(0, named("org.eclipse.jetty.server.Request"))),
32+
packageName + ".ResolveRequestedSessionIdAdvice");
33+
}
34+
}

‎dd-java-agent/instrumentation/jetty-12/src/main/java17/datadog/trace/instrumentation/jetty12/ResolveRequestedSessionIdAdvice.java‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
package datadog.trace.instrumentation.jetty12;
2+
3+
import static datadog.trace.api.gateway.Events.EVENTS;
4+
import static org.eclipse.jetty.session.AbstractSessionManager.RequestedSession;
5+
6+
import datadog.appsec.api.blocking.BlockingException;
7+
import datadog.trace.advice.ActiveRequestContext;
8+
import datadog.trace.advice.RequiresRequestContext;
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 java.util.function.BiFunction;
15+
import net.bytebuddy.asm.Advice;
16+
17+
/**
18+
* Because we are processing the initial request before the requestedSession is set, we must update
19+
* it when it is actually set.
20+
*/
21+
@RequiresRequestContext(RequestContextSlot.APPSEC)
22+
public class ResolveRequestedSessionIdAdvice {
23+
@Advice.OnMethodExit(suppress = Throwable.class)
24+
public static void resolveRequestedSessionId(
25+
@ActiveRequestContext RequestContext reqCtx,
26+
@Advice.Return final RequestedSession requestedSession) {
27+
final String requestedSessionId =
28+
requestedSession == null ? null : requestedSession.sessionId();
29+
if (requestedSessionId != null && reqCtx != null) {
30+
final CallbackProvider cbp = AgentTracer.get().getCallbackProvider(RequestContextSlot.APPSEC);
31+
if (cbp == null) {
32+
return;
33+
}
34+
final BiFunction<RequestContext, String, Flow<Void>> addrCallback =
35+
cbp.getCallback(EVENTS.requestSession());
36+
if (addrCallback == null) {
37+
return;
38+
}
39+
final Flow<Void> flow = addrCallback.apply(reqCtx, requestedSessionId);
40+
Flow.Action action = flow.getAction();
41+
if (action instanceof Flow.Action.RequestBlockingAction) {
42+
throw new BlockingException("Blocked request (for sessionId)");
43+
}
44+
}
45+
}
46+
}

‎dd-java-agent/instrumentation/jetty-12/src/test/ee8/groovy/Jetty12Test.groovy‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,11 @@ abstract class Jetty12Test extends HttpServerTest<Server> implements TestingGene
4848
boolean hasExtraErrorInformation() {
4949
true
5050
}
51+
52+
@Override
53+
boolean testSessionId() {
54+
true
55+
}
5156
}
5257

5358
class Jetty12SyncTest extends Jetty12Test {

0 commit comments

Comments
 (0)