Improve RUM injection matching and avoid truncating responses (#9323)

* Make rum injector stream/writer more resilient to errors

* fix tests

* Improve RUM injection matching and avoid truncating responses

* Use MH for setContentLengthLong since available from 3.1

* Fix tomcat11 tests

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

Co-authored-by: Bruce Bujon <[email protected]>

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

Co-authored-by: Bruce Bujon <[email protected]>

* remove leftover

---------

Co-authored-by: Bruce Bujon <[email protected]>

Author: amarziali

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/buffer/InjectingPipeOutputStream.java

@@ -2,13 +2,15 @@
 
 import java.io.IOException;
 import java.io.OutputStream;
+import javax.annotation.concurrent.NotThreadSafe;
 
 /**
  * An OutputStream containing a circular buffer with a lookbehind buffer of n bytes. The first time
  * that the latest n bytes matches the marker, a content is injected before. In case of IOException
  * thrown by the downstream, the buffer will be lost unless the error occurred when draining it. In
  * this case the draining will be resumed.
  */
+@NotThreadSafe
 public class InjectingPipeOutputStream extends OutputStream {
   private final byte[] lookbehind;
   private int pos;
@@ -168,6 +170,13 @@ private void drain() throws IOException {
     }
   }
 
+  public void commit() throws IOException {
+    if (filter || wasDraining) {
+      filter = false;
+      drain();
+    }
+  }
+
   @Override
   public void flush() throws IOException {
     downstream.flush();
@@ -176,9 +185,7 @@ public void flush() throws IOException {
   @Override
   public void close() throws IOException {
     try {
-      if (filter || wasDraining) {
-        drain();
-      }
+      commit();
     } finally {
       downstream.close();
     }

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/buffer/InjectingPipeWriter.java

@@ -2,13 +2,15 @@
 
 import java.io.IOException;
 import java.io.Writer;
+import javax.annotation.concurrent.NotThreadSafe;
 
 /**
  * A Writer containing a circular buffer with a lookbehind buffer of n bytes. The first time that
  * the latest n bytes matches the marker, a content is injected before. In case of IOException
  * thrown by the downstream, the buffer will be lost unless the error occurred when draining it. In
  * this case the draining will be resumed.
  */
+@NotThreadSafe
 public class InjectingPipeWriter extends Writer {
   private final char[] lookbehind;
   private int pos;
@@ -169,6 +171,13 @@ private void drain() throws IOException {
     }
   }
 
+  public void commit() throws IOException {
+    if (filter || wasDraining) {
+      filter = false;
+      drain();
+    }
+  }
+
   @Override
   public void flush() throws IOException {
     downstream.flush();
@@ -177,9 +186,7 @@ public void flush() throws IOException {
   @Override
   public void close() throws IOException {
     try {
-      if (filter || wasDraining) {
-        drain();
-      }
+      commit();
     } finally {
       downstream.close();
     }

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

@@ -2,18 +2,37 @@
 
 import datadog.trace.api.rum.RumInjector;
 import datadog.trace.bootstrap.instrumentation.buffer.InjectingPipeWriter;
+import datadog.trace.util.MethodHandles;
 import java.io.IOException;
 import java.io.PrintWriter;
+import java.lang.invoke.MethodHandle;
 import java.nio.charset.Charset;
 import javax.servlet.ServletOutputStream;
+import javax.servlet.ServletResponse;
 import javax.servlet.http.HttpServletResponse;
 import javax.servlet.http.HttpServletResponseWrapper;
 
 public class RumHttpServletResponseWrapper extends HttpServletResponseWrapper {
   private final RumInjector rumInjector;
-  private ServletOutputStream outputStream;
+  private WrappedServletOutputStream outputStream;
   private PrintWriter printWriter;
-  private boolean shouldInject = false;
+  private InjectingPipeWriter wrappedPipeWriter;
+  private boolean shouldInject = true;
+
+  private static final MethodHandle SET_CONTENT_LENGTH_LONG = getMh("setContentLengthLong");
+
+  private static MethodHandle getMh(final String name) {
+    try {
+      return new MethodHandles(ServletResponse.class.getClassLoader())
+          .method(ServletResponse.class, name);
+    } catch (Throwable ignored) {
+      return null;
+    }
+  }
+
+  private static <E extends Throwable> void sneakyThrow(Throwable e) throws E {
+    throw (E) e;
+  }
 
   public RumHttpServletResponseWrapper(HttpServletResponse response) {
     super(response);
@@ -22,50 +41,68 @@ public RumHttpServletResponseWrapper(HttpServletResponse response) {
 
   @Override
   public ServletOutputStream getOutputStream() throws IOException {
+    if (outputStream != null) {
+      return outputStream;
+    }
     if (!shouldInject) {
       return super.getOutputStream();
     }
-    if (outputStream == null) {
-      String encoding = getCharacterEncoding();
-      if (encoding == null) {
-        encoding = Charset.defaultCharset().name();
-      }
-      outputStream =
-          new WrappedServletOutputStream(
-              super.getOutputStream(),
-              rumInjector.getMarkerBytes(encoding),
-              rumInjector.getSnippetBytes(encoding),
-              this::onInjected);
+    String encoding = getCharacterEncoding();
+    if (encoding == null) {
+      encoding = Charset.defaultCharset().name();
     }
+    outputStream =
+        new WrappedServletOutputStream(
+            super.getOutputStream(),
+            rumInjector.getMarkerBytes(encoding),
+            rumInjector.getSnippetBytes(encoding),
+            this::onInjected);
+
     return outputStream;
   }
 
   @Override
   public PrintWriter getWriter() throws IOException {
-    final PrintWriter delegate = super.getWriter();
-    if (!shouldInject) {
-      return delegate;
+    if (printWriter != null) {
+      return printWriter;
     }
-    if (printWriter == null) {
-      printWriter =
-          new PrintWriter(
-              new InjectingPipeWriter(
-                  delegate,
-                  rumInjector.getMarkerChars(),
-                  rumInjector.getSnippetChars(),
-                  this::onInjected));
+    if (!shouldInject) {
+      return super.getWriter();
     }
+    wrappedPipeWriter =
+        new InjectingPipeWriter(
+            super.getWriter(),
+            rumInjector.getMarkerChars(),
+            rumInjector.getSnippetChars(),
+            this::onInjected);
+    printWriter = new PrintWriter(wrappedPipeWriter);
+
     return printWriter;
   }
 
   @Override
   public void setContentLength(int len) {
     // don't set it since we don't know if we will inject
+    if (!shouldInject) {
+      super.setContentLength(len);
+    }
+  }
+
+  @Override
+  public void setContentLengthLong(long len) {
+    if (!shouldInject && SET_CONTENT_LENGTH_LONG != null) {
+      try {
+        SET_CONTENT_LENGTH_LONG.invoke(getResponse(), len);
+      } catch (Throwable t) {
+        sneakyThrow(t);
+      }
+    }
   }
 
   @Override
   public void reset() {
     this.outputStream = null;
+    this.wrappedPipeWriter = null;
     this.printWriter = null;
     this.shouldInject = false;
     super.reset();
@@ -74,8 +111,8 @@ public void reset() {
   @Override
   public void resetBuffer() {
     this.outputStream = null;
+    this.wrappedPipeWriter = null;
     this.printWriter = null;
-    this.shouldInject = false;
     super.resetBuffer();
   }
 
@@ -89,7 +126,27 @@ public void onInjected() {
 
   @Override
   public void setContentType(String type) {
-    shouldInject = type != null && type.contains("text/html");
+    if (shouldInject) {
+      shouldInject = type != null && type.contains("text/html");
+    }
+    if (!shouldInject) {
+      commit();
+    }
     super.setContentType(type);
   }
+
+  public void commit() {
+    if (wrappedPipeWriter != null) {
+      try {
+        wrappedPipeWriter.commit();
+      } catch (Throwable ignored) {
+      }
+    }
+    if (outputStream != null) {
+      try {
+        outputStream.commit();
+      } catch (Throwable ignored) {
+      }
+    }
+  }
 }

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

@@ -35,7 +35,8 @@ public static boolean onEnter(
       @Advice.Argument(value = 1, readOnly = false) ServletResponse response,
       @Advice.Local("isDispatch") boolean isDispatch,
       @Advice.Local("finishSpan") boolean finishSpan,
-      @Advice.Local("contextScope") ContextScope scope) {
+      @Advice.Local("contextScope") ContextScope scope,
+      @Advice.Local("rumServletWrapper") RumHttpServletResponseWrapper rumServletWrapper) {
     final boolean invalidRequest =
         !(request instanceof HttpServletRequest) || !(response instanceof HttpServletResponse);
     if (invalidRequest) {
@@ -47,7 +48,8 @@ public static boolean onEnter(
 
     if (RumInjector.get().isEnabled() && httpServletRequest.getAttribute(DD_RUM_INJECTED) == null) {
       httpServletRequest.setAttribute(DD_RUM_INJECTED, Boolean.TRUE);
-      httpServletResponse = new RumHttpServletResponseWrapper(httpServletResponse);
+      rumServletWrapper = new RumHttpServletResponseWrapper(httpServletResponse);
+      httpServletResponse = rumServletWrapper;
       response = httpServletResponse;
     }
 
@@ -108,7 +110,11 @@ 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.Thrown final Throwable throwable) {
+    if (rumServletWrapper != null) {
+      rumServletWrapper.commit();
+    }
     // Set user.principal regardless of who created this span.
     final Object spanAttr = request.getAttribute(DD_SPAN_ATTRIBUTE);
     if (Config.get().isServletPrincipalEnabled()

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

@@ -3,13 +3,12 @@
 import datadog.trace.bootstrap.instrumentation.buffer.InjectingPipeOutputStream;
 import datadog.trace.util.MethodHandles;
 import java.io.IOException;
-import java.io.OutputStream;
 import java.lang.invoke.MethodHandle;
 import javax.servlet.ServletOutputStream;
 import javax.servlet.WriteListener;
 
 public class WrappedServletOutputStream extends ServletOutputStream {
-  private final OutputStream filtered;
+  private final InjectingPipeOutputStream filtered;
   private final ServletOutputStream delegate;
 
   private static final MethodHandle IS_READY_MH = getMh("isReady");
@@ -83,4 +82,8 @@ public void setWriteListener(WriteListener writeListener) {
       sneakyThrow(e);
     }
   }
+
+  public void commit() throws IOException {
+    filtered.commit();
+  }
 }

dd-java-agent/instrumentation/servlet/request-3/src/testFixtures/groovy/datadog/trace/instrumentation/servlet3/RumServlet.groovy

@@ -14,8 +14,8 @@ class RumServlet extends HttpServlet {
 
   @Override
   protected void service(HttpServletRequest req, HttpServletResponse resp) throws ServletException, IOException {
-    resp.setContentType(mimeType)
     try (def writer = resp.getWriter()) {
+      resp.setContentType(mimeType)
       writer.println("\n" +
       "<!doctype html>\n" +
       "<html>\n" +

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

@@ -68,7 +68,8 @@ public static class JakartaServletAdvice {
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static AgentSpan before(
         @Advice.Argument(0) final ServletRequest request,
-        @Advice.Argument(value = 1, readOnly = false) ServletResponse response) {
+        @Advice.Argument(value = 1, readOnly = false) ServletResponse response,
+        @Advice.Local("rumServletWrapper") RumHttpServletResponseWrapper rumServletWrapper) {
       if (!(request instanceof HttpServletRequest)) {
         return null;
       }
@@ -79,7 +80,8 @@ public static AgentSpan before(
         if (RumInjector.get().isEnabled()
             && httpServletRequest.getAttribute(DD_RUM_INJECTED) == null) {
           httpServletRequest.setAttribute(DD_RUM_INJECTED, Boolean.TRUE);
-          response = new RumHttpServletResponseWrapper((HttpServletResponse) response);
+          rumServletWrapper = new RumHttpServletResponseWrapper((HttpServletResponse) response);
+          response = rumServletWrapper;
         }
       }
 
@@ -95,10 +97,15 @@ public static AgentSpan before(
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void after(
-        @Advice.Enter final AgentSpan span, @Advice.Argument(0) final ServletRequest request) {
+        @Advice.Enter final AgentSpan span,
+        @Advice.Argument(0) final ServletRequest request,
+        @Advice.Local("rumServletWrapper") RumHttpServletResponseWrapper rumServletWrapper) {
       if (span == null) {
         return;
       }
+      if (rumServletWrapper != null) {
+        rumServletWrapper.commit();
+      }
 
       CallDepthThreadLocalMap.reset(HttpServletRequest.class);
       final HttpServletRequest httpServletRequest =