Avoid multiple injections on dispatch on jboss

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/rum/RumControllableResponse.java

@@ -0,0 +1,9 @@
+package datadog.trace.bootstrap.instrumentation.rum;
+
+public interface RumControllableResponse {
+  /** Drain the held buffer. */
+  void commit();
+
+  /** Stops filtering the response. */
+  void stopFiltering();
+}

dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/RumAsyncContextInstrumentation.java

@@ -10,6 +10,7 @@
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
 import datadog.trace.api.InstrumenterConfig;
+import datadog.trace.bootstrap.instrumentation.rum.RumControllableResponse;
 import javax.servlet.AsyncContext;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
@@ -56,8 +57,8 @@ public static class CommitAdvice {
     public static void commitRumBuffer(@Advice.This final AsyncContext asyncContext) {
       final Object maybeRumWrappedResponse =
           asyncContext.getRequest().getAttribute(DD_RUM_INJECTED);
-      if (maybeRumWrappedResponse instanceof RumHttpServletResponseWrapper) {
-        ((RumHttpServletResponseWrapper) maybeRumWrappedResponse).commit();
+      if (maybeRumWrappedResponse instanceof RumControllableResponse) {
+        ((RumControllableResponse) maybeRumWrappedResponse).commit();
       }
     }
   }

dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/RumHttpServletRequestWrapper.java

@@ -2,6 +2,7 @@
 
 import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_RUM_INJECTED;
 
+import datadog.trace.bootstrap.instrumentation.rum.RumControllableResponse;
 import javax.servlet.AsyncContext;
 import javax.servlet.ServletRequest;
 import javax.servlet.ServletResponse;
@@ -11,7 +12,7 @@
 
 public class RumHttpServletRequestWrapper extends HttpServletRequestWrapper {
 
-  private final HttpServletResponse response;
+  private HttpServletResponse response;
 
   public RumHttpServletRequestWrapper(
       final HttpServletRequest request, final HttpServletResponse response) {
@@ -30,17 +31,16 @@ public AsyncContext startAsync(ServletRequest servletRequest, ServletResponse se
       throws IllegalStateException {
     // deactivate the previous wrapper
     final Object maybeRumWrappedResponse = (servletRequest.getAttribute(DD_RUM_INJECTED));
-    if (maybeRumWrappedResponse instanceof RumHttpServletResponseWrapper) {
-      ((RumHttpServletResponseWrapper) maybeRumWrappedResponse).commit();
-      ((RumHttpServletResponseWrapper) maybeRumWrappedResponse).stopFiltering();
+    if (maybeRumWrappedResponse instanceof RumControllableResponse) {
+      ((RumControllableResponse) maybeRumWrappedResponse).commit();
+      ((RumControllableResponse) maybeRumWrappedResponse).stopFiltering();
     }
-    ServletResponse actualResponse = servletResponse;
     // rewrap it
     if (servletResponse instanceof HttpServletResponse) {
-      actualResponse =
+      this.response =
           new RumHttpServletResponseWrapper(this, (HttpServletResponse) servletResponse);
-      servletRequest.setAttribute(DD_RUM_INJECTED, actualResponse);
+      servletRequest.setAttribute(DD_RUM_INJECTED, this.response);
     }
-    return super.startAsync(servletRequest, actualResponse);
+    return super.startAsync(servletRequest, this.response);
   }
 }

dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/RumHttpServletResponseWrapper.java

@@ -2,6 +2,7 @@
 
 import datadog.trace.api.rum.RumInjector;
 import datadog.trace.bootstrap.instrumentation.buffer.InjectingPipeWriter;
+import datadog.trace.bootstrap.instrumentation.rum.RumControllableResponse;
 import datadog.trace.util.MethodHandles;
 import java.io.IOException;
 import java.io.PrintWriter;
@@ -14,7 +15,8 @@
 import javax.servlet.http.HttpServletResponse;
 import javax.servlet.http.HttpServletResponseWrapper;
 
-public class RumHttpServletResponseWrapper extends HttpServletResponseWrapper {
+public class RumHttpServletResponseWrapper extends HttpServletResponseWrapper
+    implements RumControllableResponse {
   private final RumInjector rumInjector;
   private final String servletVersion;
   private WrappedServletOutputStream outputStream;
@@ -203,6 +205,7 @@ public void setContentType(String type) {
     super.setContentType(type);
   }
 
+  @Override
   public void commit() {
     if (wrappedPipeWriter != null) {
       try {
@@ -218,6 +221,7 @@ public void commit() {
     }
   }
 
+  @Override
   public void stopFiltering() {
     shouldInject = false;
     if (wrappedPipeWriter != null) {

dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet3Advice.java

@@ -17,6 +17,7 @@
 import datadog.trace.api.gateway.Flow;
 import datadog.trace.api.rum.RumInjector;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.rum.RumControllableResponse;
 import datadog.trace.instrumentation.servlet.ServletBlockingHelper;
 import java.security.Principal;
 import java.util.concurrent.atomic.AtomicBoolean;
@@ -35,7 +36,7 @@ public static boolean onEnter(
       @Advice.Local("isDispatch") boolean isDispatch,
       @Advice.Local("finishSpan") boolean finishSpan,
       @Advice.Local("contextScope") ContextScope scope,
-      @Advice.Local("rumServletWrapper") RumHttpServletResponseWrapper rumServletWrapper) {
+      @Advice.Local("rumServletWrapper") RumControllableResponse rumServletWrapper) {
     final boolean invalidRequest =
         !(request instanceof HttpServletRequest) || !(response instanceof HttpServletResponse);
     if (invalidRequest) {
@@ -47,14 +48,16 @@ public static boolean onEnter(
 
     if (RumInjector.get().isEnabled()) {
       final Object maybeRumWrapper = httpServletRequest.getAttribute(DD_RUM_INJECTED);
-      if (maybeRumWrapper instanceof RumHttpServletResponseWrapper) {
-        rumServletWrapper = (RumHttpServletResponseWrapper) maybeRumWrapper;
+      if (maybeRumWrapper instanceof RumControllableResponse) {
+        rumServletWrapper = (RumControllableResponse) maybeRumWrapper;
       } else {
         rumServletWrapper =
             new RumHttpServletResponseWrapper(httpServletRequest, (HttpServletResponse) response);
         httpServletRequest.setAttribute(DD_RUM_INJECTED, rumServletWrapper);
-        response = rumServletWrapper;
-        request = new RumHttpServletRequestWrapper(httpServletRequest, rumServletWrapper);
+        response = (ServletResponse) rumServletWrapper;
+        request =
+            new RumHttpServletRequestWrapper(
+                httpServletRequest, (HttpServletResponse) rumServletWrapper);
       }
     }
 
@@ -116,7 +119,7 @@ public static void stopSpan(
       @Advice.Local("contextScope") final ContextScope scope,
       @Advice.Local("isDispatch") boolean isDispatch,
       @Advice.Local("finishSpan") boolean finishSpan,
-      @Advice.Local("rumServletWrapper") RumHttpServletResponseWrapper rumServletWrapper,
+      @Advice.Local("rumServletWrapper") RumControllableResponse rumServletWrapper,
       @Advice.Thrown final Throwable throwable) {
     if (rumServletWrapper != null) {
       rumServletWrapper.commit();

dd-java-agent/instrumentation/servlet/request-5/src/main/java/datadog/trace/instrumentation/servlet5/JakartaServletInstrumentation.java

@@ -19,6 +19,7 @@
 import datadog.trace.api.rum.RumInjector;
 import datadog.trace.bootstrap.CallDepthThreadLocalMap;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.rum.RumControllableResponse;
 import jakarta.servlet.ServletRequest;
 import jakarta.servlet.ServletResponse;
 import jakarta.servlet.http.HttpServletRequest;
@@ -71,7 +72,7 @@ public static class JakartaServletAdvice {
     public static AgentSpan before(
         @Advice.Argument(value = 0, readOnly = false) ServletRequest request,
         @Advice.Argument(value = 1, readOnly = false) ServletResponse response,
-        @Advice.Local("rumServletWrapper") RumHttpServletResponseWrapper rumServletWrapper) {
+        @Advice.Local("rumServletWrapper") RumControllableResponse rumServletWrapper) {
       if (!(request instanceof HttpServletRequest)) {
         return null;
       }
@@ -81,15 +82,17 @@ public static AgentSpan before(
 
         if (RumInjector.get().isEnabled()) {
           final Object maybeRumWrapper = httpServletRequest.getAttribute(DD_RUM_INJECTED);
-          if (maybeRumWrapper instanceof RumHttpServletResponseWrapper) {
-            rumServletWrapper = (RumHttpServletResponseWrapper) maybeRumWrapper;
+          if (maybeRumWrapper instanceof RumControllableResponse) {
+            rumServletWrapper = (RumControllableResponse) maybeRumWrapper;
           } else {
             rumServletWrapper =
                 new RumHttpServletResponseWrapper(
                     httpServletRequest, (HttpServletResponse) response);
             httpServletRequest.setAttribute(DD_RUM_INJECTED, rumServletWrapper);
-            response = rumServletWrapper;
-            request = new RumHttpServletRequestWrapper(httpServletRequest, rumServletWrapper);
+            response = (ServletResponse) rumServletWrapper;
+            request =
+                new RumHttpServletRequestWrapper(
+                    httpServletRequest, (HttpServletResponse) rumServletWrapper);
           }
         }
       }
@@ -108,7 +111,7 @@ public static AgentSpan before(
     public static void after(
         @Advice.Enter final AgentSpan span,
         @Advice.Argument(0) final ServletRequest request,
-        @Advice.Local("rumServletWrapper") RumHttpServletResponseWrapper rumServletWrapper) {
+        @Advice.Local("rumServletWrapper") RumControllableResponse rumServletWrapper) {
       if (span == null) {
         return;
       }

dd-java-agent/instrumentation/servlet/request-5/src/main/java/datadog/trace/instrumentation/servlet5/RumAsyncContextInstrumentation.java

@@ -10,6 +10,7 @@
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
 import datadog.trace.api.InstrumenterConfig;
+import datadog.trace.bootstrap.instrumentation.rum.RumControllableResponse;
 import jakarta.servlet.AsyncContext;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
@@ -56,8 +57,8 @@ public static class CommitAdvice {
     public static void commitRumBuffer(@Advice.This final AsyncContext asyncContext) {
       final Object maybeRumWrappedResponse =
           asyncContext.getRequest().getAttribute(DD_RUM_INJECTED);
-      if (maybeRumWrappedResponse instanceof RumHttpServletResponseWrapper) {
-        ((RumHttpServletResponseWrapper) maybeRumWrappedResponse).commit();
+      if (maybeRumWrappedResponse instanceof RumControllableResponse) {
+        ((RumControllableResponse) maybeRumWrappedResponse).commit();
       }
     }
   }

dd-java-agent/instrumentation/servlet/request-5/src/main/java/datadog/trace/instrumentation/servlet5/RumHttpServletRequestWrapper.java

@@ -2,6 +2,7 @@
 
 import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_RUM_INJECTED;
 
+import datadog.trace.bootstrap.instrumentation.rum.RumControllableResponse;
 import jakarta.servlet.AsyncContext;
 import jakarta.servlet.ServletRequest;
 import jakarta.servlet.ServletResponse;
@@ -30,9 +31,9 @@ public AsyncContext startAsync(ServletRequest servletRequest, ServletResponse se
       throws IllegalStateException {
     // deactivate the previous wrapper
     final Object maybeRumWrappedResponse = (servletRequest.getAttribute(DD_RUM_INJECTED));
-    if (maybeRumWrappedResponse instanceof RumHttpServletResponseWrapper) {
-      ((RumHttpServletResponseWrapper) maybeRumWrappedResponse).commit();
-      ((RumHttpServletResponseWrapper) maybeRumWrappedResponse).stopFiltering();
+    if (maybeRumWrappedResponse instanceof RumControllableResponse) {
+      ((RumControllableResponse) maybeRumWrappedResponse).commit();
+      ((RumControllableResponse) maybeRumWrappedResponse).stopFiltering();
     }
     ServletResponse actualResponse = servletResponse;
     // rewrap it

dd-java-agent/instrumentation/servlet/request-5/src/main/java/datadog/trace/instrumentation/servlet5/RumHttpServletResponseWrapper.java

@@ -2,6 +2,7 @@
 
 import datadog.trace.api.rum.RumInjector;
 import datadog.trace.bootstrap.instrumentation.buffer.InjectingPipeWriter;
+import datadog.trace.bootstrap.instrumentation.rum.RumControllableResponse;
 import jakarta.servlet.ServletContext;
 import jakarta.servlet.ServletOutputStream;
 import jakarta.servlet.http.HttpServletRequest;
@@ -11,7 +12,8 @@
 import java.io.PrintWriter;
 import java.nio.charset.Charset;
 
-public class RumHttpServletResponseWrapper extends HttpServletResponseWrapper {
+public class RumHttpServletResponseWrapper extends HttpServletResponseWrapper
+    implements RumControllableResponse {
   private final RumInjector rumInjector;
   private final String servletVersion;
   private WrappedServletOutputStream outputStream;
@@ -180,6 +182,7 @@ public void setContentType(String type) {
     super.setContentType(type);
   }
 
+  @Override
   public void commit() {
     if (wrappedPipeWriter != null) {
       try {
@@ -195,6 +198,7 @@ public void commit() {
     }
   }
 
+  @Override
   public void stopFiltering() {
     shouldInject = false;
     if (wrappedPipeWriter != null) {

dd-smoke-tests/rum/src/main/groovy/datadog/smoketest/rum/AbstractRumServerSmokeTest.groovy

@@ -14,9 +14,13 @@ class AbstractRumServerSmokeTest extends AbstractServerSmokeTest {
     "-Ddd.rum.remote.configuration.id=12345",
   ]
 
+  String mountPoint() {
+    ""
+  }
+
   void 'test RUM SDK injection on html for path #servletPath'() {
     given:
-    def url = "http://localhost:${httpPort}/${servletPath}"
+    def url = "http://localhost:${httpPort}${mountPoint()}/${servletPath}"
     def request = new Request.Builder()
       .url(url)
       .get()
@@ -34,7 +38,7 @@ class AbstractRumServerSmokeTest extends AbstractServerSmokeTest {
 
   void 'test RUM SDK injection skip on unsupported mime type'() {
     given:
-    def url = "http://localhost:${httpPort}/xml"
+    def url = "http://localhost:${httpPort}${mountPoint()}/xml"
     def request = new Request.Builder()
       .url(url)
       .get()
@@ -52,13 +56,14 @@ class AbstractRumServerSmokeTest extends AbstractServerSmokeTest {
     assert response.header('x-datadog-rum-injected') == '1': 'RUM injected header missing'
     def content = response.body().string()
     assert content.contains('https://www.datadoghq-browser-agent.com'): 'RUM script not injected'
-    assert content.endsWith('</html>'): 'Response not fully flushed'
+    assert content.trim().endsWith('</html>'): 'Response not fully flushed'
+    assert content.indexOf("DD_RUM.init(") == content.lastIndexOf("DD_RUM.init("): 'script injected more than once'
   }
 
   static void assertRumNotInjected(Response response) {
     assert response.header('x-datadog-rum-injected') == null: 'RUM header unexpectedly injected'
     def content = response.body().string()
     assert !content.contains('https://www.datadoghq-browser-agent.com'): 'RUM script unexpectedly injected'
-    assert content.endsWith('</response>'): 'Response not fully flushed'
+    assert content.trim().endsWith('</response>'): 'Response not fully flushed'
   }
 }