diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/WeakMap.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/WeakMap.java index 2fedd001284..85a2ac5b879 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/WeakMap.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/WeakMap.java @@ -2,21 +2,28 @@ import edu.umd.cs.findbugs.annotations.SuppressFBWarnings; import java.util.function.Function; - +import javax.annotation.Nonnull; + +/** + * Map with weakly referenced keys. + * + *

Keys must not be {@code null}: the backing {@code WeakConcurrentMap} rejects null keys and + * throws. + */ public interface WeakMap { int size(); - boolean containsKey(K target); + boolean containsKey(@Nonnull K target); - V get(K key); + V get(@Nonnull K key); - void put(K key, V value); + void put(@Nonnull K key, V value); - void putIfAbsent(K key, V value); + void putIfAbsent(@Nonnull K key, V value); - V computeIfAbsent(K key, Function supplier); + V computeIfAbsent(@Nonnull K key, Function supplier); - V remove(K key); + V remove(@Nonnull K key); abstract class Supplier { private static volatile Supplier SUPPLIER; diff --git a/dd-java-agent/instrumentation/servlet/jakarta-servlet-5.0/src/main/java/datadog/trace/instrumentation/servlet5/IastOptOutJakartaHttpServletRequestInstrumentation.java b/dd-java-agent/instrumentation/servlet/jakarta-servlet-5.0/src/main/java/datadog/trace/instrumentation/servlet5/IastOptOutJakartaHttpServletRequestInstrumentation.java index 2dbda417503..031ef810000 100644 --- a/dd-java-agent/instrumentation/servlet/jakarta-servlet-5.0/src/main/java/datadog/trace/instrumentation/servlet5/IastOptOutJakartaHttpServletRequestInstrumentation.java +++ b/dd-java-agent/instrumentation/servlet/jakarta-servlet-5.0/src/main/java/datadog/trace/instrumentation/servlet5/IastOptOutJakartaHttpServletRequestInstrumentation.java @@ -84,6 +84,11 @@ public static void onExit( return; } final ServletContext context = request.getServletContext(); + if (context == null) { + // some request copies (e.g. Wicket or Atmosphere websocket requests) have no servlet + // context + return; + } if (InstrumentationContext.get(ServletContext.class, SessionTrackingMode.class).get(context) != null) { diff --git a/dd-java-agent/instrumentation/servlet/jakarta-servlet-5.0/src/test/groovy/JakartaHttpServletRequestInstrumentationTest.groovy b/dd-java-agent/instrumentation/servlet/jakarta-servlet-5.0/src/test/groovy/JakartaHttpServletRequestInstrumentationTest.groovy index f8d0d1811bf..c3db1613c4f 100644 --- a/dd-java-agent/instrumentation/servlet/jakarta-servlet-5.0/src/test/groovy/JakartaHttpServletRequestInstrumentationTest.groovy +++ b/dd-java-agent/instrumentation/servlet/jakarta-servlet-5.0/src/test/groovy/JakartaHttpServletRequestInstrumentationTest.groovy @@ -460,6 +460,24 @@ class JakartaHttpServletRequestInstrumentationTest extends InstrumentationSpecif suite << testSuite() } + void 'getSession tolerates a null servlet context'() { + setup: + final module = Mock(ApplicationModule) + InstrumentationBridge.registerIastModule(module) + final session = Mock(HttpSession) + final delegate = Mock(HttpServletRequest) + final request = new CustomRequest(request: delegate) + + when: + final result = request.getSession() + + then: + result.is(session) + 1 * delegate.getSession() >> session + 1 * delegate.getServletContext() >> null + 0 * module._ + } + protected E runUnderIastTrace(Closure cl) { final ddctx = new TagContext().withRequestContextDataIast(iastCtx) final span = TEST_TRACER.startSpan("test", "test-iast-span", ddctx) diff --git a/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-3.0/src/main/java/datadog/trace/instrumentation/servlet3/IastOptOutHttpServletRequest3Instrumentation.java b/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-3.0/src/main/java/datadog/trace/instrumentation/servlet3/IastOptOutHttpServletRequest3Instrumentation.java index 57289e26ab7..8f7797643cc 100644 --- a/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-3.0/src/main/java/datadog/trace/instrumentation/servlet3/IastOptOutHttpServletRequest3Instrumentation.java +++ b/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-3.0/src/main/java/datadog/trace/instrumentation/servlet3/IastOptOutHttpServletRequest3Instrumentation.java @@ -92,6 +92,11 @@ public static void onExit( return; } final ServletContext context = request.getServletContext(); + if (context == null) { + // some request copies (e.g. Wicket or Atmosphere websocket requests) have no servlet + // context + return; + } if (InstrumentationContext.get(ServletContext.class, SessionTrackingMode.class).get(context) != null) { return; diff --git a/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-3.0/src/test/groovy/JettyServlet3Test.groovy b/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-3.0/src/test/groovy/JettyServlet3Test.groovy index 3271479c39b..d3373a40a3f 100644 --- a/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-3.0/src/test/groovy/JettyServlet3Test.groovy +++ b/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-3.0/src/test/groovy/JettyServlet3Test.groovy @@ -28,6 +28,7 @@ import javax.servlet.annotation.WebServlet import javax.servlet.http.HttpServlet import javax.servlet.http.HttpServletRequest import javax.servlet.http.HttpServletResponse +import javax.servlet.http.HttpSession import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.CUSTOM_EXCEPTION import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.ERROR @@ -570,6 +571,32 @@ class JettyServlet3ServeFromAsyncTimeout extends JettyServlet3Test { class IastJettyServlet3ForkedTest extends JettyServlet3TestSync { + void 'getSession tolerates a null servlet context'() { + setup: + final module = Mock(ApplicationModule) + InstrumentationBridge.registerIastModule(module) + final session = Mock(HttpSession) + final delegate = Mock(HttpServletRequest) + final request = new CustomRequest(request: delegate) + + when: + final result = request.getSession() + + then: + result.is(session) + 1 * delegate.getSession() >> session + 1 * delegate.getServletContext() >> null + 0 * module._ + + cleanup: + InstrumentationBridge.clearIastModules() + } + + private static class CustomRequest implements HttpServletRequest { + @Delegate + private HttpServletRequest request + } + @Override Class servlet() { return TestServlet3.GetSession diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/ContextStore.java b/internal-api/src/main/java/datadog/trace/bootstrap/ContextStore.java index 18dd8344a5c..0fe1e5149b9 100644 --- a/internal-api/src/main/java/datadog/trace/bootstrap/ContextStore.java +++ b/internal-api/src/main/java/datadog/trace/bootstrap/ContextStore.java @@ -1,6 +1,7 @@ package datadog.trace.bootstrap; import java.util.function.Function; +import javax.annotation.Nonnull; import javax.annotation.Nullable; /** @@ -9,6 +10,9 @@ *

Context instances are weakly referenced and will be garbage collected when their corresponding * key instance is collected. * + *

Keys must not be {@code null}: the backing weak maps reject null keys and throw. Callers that + * take a key from application code (for example a servlet request's context) must check it first. + * * @param key type to do context lookups * @param context type */ @@ -38,7 +42,7 @@ default C apply(Object key) { * @return context instance; {@code null} if the key had no context */ @Nullable - C get(K key); + C get(@Nonnull K key); /** * Unconditionally put new context instance for the given key. @@ -46,7 +50,7 @@ default C apply(Object key) { * @param key the context key * @param context context instance to save */ - void put(K key, C context); + void put(@Nonnull K key, C context); /** * Gets the context instance for the given key. If no context exists then associate it with the @@ -56,7 +60,7 @@ default C apply(Object key) { * @param context new context instance * @return existing context instance if present; otherwise new instance */ - C getOrPut(K key, C context); + C getOrPut(@Nonnull K key, C context); /** * Gets the context instance for the given key. If no context exists then create one using the @@ -66,7 +70,7 @@ default C apply(Object key) { * @param contextFactory factory instance to produce new context instances * @return existing context instance if present; otherwise new instance */ - default C getOrCreate(K key, Factory contextFactory) { + default C getOrCreate(@Nonnull K key, Factory contextFactory) { return getOrCompute(key, contextFactory); } @@ -78,7 +82,7 @@ default C getOrCreate(K key, Factory contextFactory) { * @param contextFactory factory instance to produce new context instances * @return existing context instance if present; otherwise new instance */ - C getOrCompute(K key, Function contextFactory); + C getOrCompute(@Nonnull K key, Function contextFactory); /** * Removes the context instance for the given key. @@ -87,5 +91,5 @@ default C getOrCreate(K key, Factory contextFactory) { * @return removed context instance; {@code null} if the key had no context */ @Nullable - C remove(K key); + C remove(@Nonnull K key); }