Add content-length detection for InjectingPipeWriter and improve tests

sarahchen6 · sarahchen6 · commit b3fcde4cee7e · 2025-08-11T10:28:51.000-04:00
diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/buffer/InjectingPipeWriter.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/buffer/InjectingPipeWriter.java
@@ -2,6 +2,7 @@
 
 import java.io.IOException;
 import java.io.Writer;
+import java.util.function.LongConsumer;
 import javax.annotation.concurrent.NotThreadSafe;
 
 /**
@@ -23,18 +24,22 @@ public class InjectingPipeWriter extends Writer {
   private final Runnable onContentInjected;
   private final int bulkWriteThreshold;
   private final Writer downstream;
+  private final LongConsumer onBytesWritten;
+  private long bytesWritten = 0;
 
   /**
    * @param downstream the delegate writer
    * @param marker the marker to find in the stream. Must at least be one char.
    * @param contentToInject the content to inject once before the marker if found.
    * @param onContentInjected callback called when and if the content is injected.
+   * @param onBytesWritten callback called when writer is closed to report total bytes written.
    */
   public InjectingPipeWriter(
       final Writer downstream,
       final char[] marker,
       final char[] contentToInject,
-      final Runnable onContentInjected) {
+      final Runnable onContentInjected,
+      final LongConsumer onBytesWritten) {
     this.downstream = downstream;
     this.marker = marker;
     this.lookbehind = new char[marker.length];
@@ -46,11 +51,13 @@ public InjectingPipeWriter(
     this.filter = true;
     this.contentToInject = contentToInject;
     this.onContentInjected = onContentInjected;
+    this.onBytesWritten = onBytesWritten;
     this.bulkWriteThreshold = marker.length * 2 - 2;
   }
 
   @Override
   public void write(int c) throws IOException {
+    bytesWritten++;
     if (!filter) {
       if (wasDraining) {
         // continue draining
@@ -85,6 +92,7 @@ public void write(int c) throws IOException {
 
   @Override
   public void write(char[] array, int off, int len) throws IOException {
+    bytesWritten += len;
     if (!filter) {
       if (wasDraining) {
         // needs drain
@@ -113,6 +121,7 @@ public void write(char[] array, int off, int len) throws IOException {
         // we don't have a full match. write everything in a bulk except the lookbehind buffer
         // sequentially
         for (int i = off; i < off + marker.length - 1; i++) {
+          bytesWritten--; // avoid double counting
           write(array[i]);
         }
         drain();
@@ -124,12 +133,14 @@ public void write(char[] array, int off, int len) throws IOException {
         filter = wasFiltering;
 
         for (int i = len - marker.length + 1; i < len; i++) {
+          bytesWritten--; // avoid double counting
           write(array[i]);
         }
       }
     } else {
       // use slow path because the length to write is small and within the lookbehind buffer size
       for (int i = off; i < off + len; i++) {
+        bytesWritten--; // avoid double counting
         write(array[i]);
       }
     }
@@ -188,6 +199,11 @@ public void close() throws IOException {
       commit();
     } finally {
       downstream.close();
+      // report the size of the original HTTP response before injecting via callback
+      if (onBytesWritten != null) {
+        onBytesWritten.accept(bytesWritten);
+      }
+      bytesWritten = 0;
     }
   }
 
diff --git a/dd-java-agent/agent-bootstrap/src/test/groovy/datadog/trace/bootstrap/instrumentation/buffer/InjectingPipeOutputStreamTest.groovy b/dd-java-agent/agent-bootstrap/src/test/groovy/datadog/trace/bootstrap/instrumentation/buffer/InjectingPipeOutputStreamTest.groovy
@@ -87,4 +87,94 @@ class InjectingPipeOutputStreamTest extends DDSpecification {
     // expected broken since the real write happens at close (drain) being the content smaller than the buffer. And retry on close is not a common practice. Hence, we suppose loosing content
     ["<foo/>"]                                                      | "<longerThanFoo>" | "<nothing>"             | 3         | "<f"
   }
+
+  def 'should count bytes correctly when writing byte arrays'() {
+    setup:
+    def downstream = new ByteArrayOutputStream()
+    def bytesWritten = []
+    def onBytesWritten = { long bytes -> bytesWritten.add(bytes) }
+    def piped = new InjectingPipeOutputStream(downstream, "</head>".getBytes("UTF-8"), "<script></script>".getBytes("UTF-8"), null, onBytesWritten)
+
+    when:
+    piped.write("test content".getBytes("UTF-8"))
+    piped.close()
+
+    then:
+    bytesWritten.size() == 1
+    bytesWritten[0] == 12
+    downstream.toByteArray() == "test content".getBytes("UTF-8")
+  }
+
+  def 'should count bytes correctly when writing bytes individually'() {
+    setup:
+    def downstream = new ByteArrayOutputStream()
+    def bytesWritten = []
+    def onBytesWritten = { long bytes -> bytesWritten.add(bytes) }
+    def piped = new InjectingPipeOutputStream(downstream, "</head>".getBytes("UTF-8"), "<script></script>".getBytes("UTF-8"), null, onBytesWritten)
+
+    when:
+    def bytes = "test".getBytes("UTF-8")
+    for (int i = 0; i < bytes.length; i++) {
+      piped.write((int) bytes[i])
+    }
+    piped.close()
+
+    then:
+    bytesWritten.size() == 1
+    bytesWritten[0] == 4
+    downstream.toByteArray() == "test".getBytes("UTF-8")
+  }
+
+  def 'should count bytes correctly with multiple writes'() {
+    setup:
+    def downstream = new ByteArrayOutputStream()
+    def bytesWritten = []
+    def onBytesWritten = { long bytes -> bytesWritten.add(bytes) }
+    def piped = new InjectingPipeOutputStream(downstream, "</head>".getBytes("UTF-8"), "<script></script>".getBytes("UTF-8"), null, onBytesWritten)
+
+    when:
+    piped.write("test".getBytes("UTF-8"))
+    piped.write(" ".getBytes("UTF-8"))
+    piped.write("content".getBytes("UTF-8"))
+    piped.close()
+
+    then:
+    bytesWritten.size() == 1
+    bytesWritten[0] == 12
+    downstream.toByteArray() == "test content".getBytes("UTF-8")
+  }
+
+  def 'should be resilient to exceptions when onBytesWritten callback is null'() {
+    setup:
+    def downstream = new ByteArrayOutputStream()
+    def piped = new InjectingPipeOutputStream(downstream, "</head>".getBytes("UTF-8"), "<script></script>".getBytes("UTF-8"), null, null)
+
+    when:
+    piped.write("test content".getBytes("UTF-8"))
+    piped.close()
+
+    then:
+    noExceptionThrown()
+    downstream.toByteArray() == "test content".getBytes("UTF-8")
+  }
+
+  def 'should reset byte count after close'() {
+    setup:
+    def downstream = new ByteArrayOutputStream()
+    def bytesWritten = []
+    def onBytesWritten = { long bytes -> bytesWritten.add(bytes) }
+    def piped = new InjectingPipeOutputStream(downstream, "</head>".getBytes("UTF-8"), "<script></script>".getBytes("UTF-8"), null, onBytesWritten)
+
+    when:
+    piped.write("test".getBytes("UTF-8"))
+    piped.close()
+
+    piped.write("content".getBytes("UTF-8"))
+    piped.close()
+
+    then:
+    bytesWritten.size() == 2
+    bytesWritten[0] == 4
+    bytesWritten[1] == 7
+  }
 }
diff --git a/dd-java-agent/agent-bootstrap/src/test/groovy/datadog/trace/bootstrap/instrumentation/buffer/InjectingPipeWriterTest.groovy b/dd-java-agent/agent-bootstrap/src/test/groovy/datadog/trace/bootstrap/instrumentation/buffer/InjectingPipeWriterTest.groovy
@@ -36,7 +36,7 @@ class InjectingPipeWriterTest extends DDSpecification {
   def 'should filter a buffer and inject if found #found using write'() {
     setup:
     def downstream = new StringWriter()
-    def piped = new PrintWriter(new InjectingPipeWriter(downstream, marker.toCharArray(), contentToInject.toCharArray(), null))
+    def piped = new PrintWriter(new InjectingPipeWriter(downstream, marker.toCharArray(), contentToInject.toCharArray(), null, null))
     when:
     try (def closeme = piped) {
       piped.write(body)
@@ -53,7 +53,7 @@ class InjectingPipeWriterTest extends DDSpecification {
   def 'should filter a buffer and inject if found #found using append'() {
     setup:
     def downstream = new StringWriter()
-    def piped = new PrintWriter(new InjectingPipeWriter(downstream, marker.toCharArray(), contentToInject.toCharArray(), null))
+    def piped = new PrintWriter(new InjectingPipeWriter(downstream, marker.toCharArray(), contentToInject.toCharArray(), null, null))
     when:
     try (def closeme = piped) {
       piped.append(body)
@@ -71,7 +71,7 @@ class InjectingPipeWriterTest extends DDSpecification {
     setup:
     def writer = new StringWriter()
     def downstream = new GlitchedWriter(writer, glichesAt)
-    def piped = new InjectingPipeWriter(downstream, marker.toCharArray(), contentToInject.toCharArray(), null)
+    def piped = new InjectingPipeWriter(downstream, marker.toCharArray(), contentToInject.toCharArray(), null, null)
     when:
     try {
       for (String line : body) {
@@ -103,4 +103,94 @@ class InjectingPipeWriterTest extends DDSpecification {
     // expected broken since the real write happens at close (drain) being the content smaller than the buffer. And retry on close is not a common practice. Hence, we suppose loosing content
     ["<foo/>"]                                                      | "<longerThanFoo>" | "<nothing>"             | 3         | "<f"
   }
+
+  def 'should count bytes correctly when writing characters'() {
+    setup:
+    def downstream = new StringWriter()
+    def bytesWritten = []
+    def onBytesWritten = { long bytes -> bytesWritten.add(bytes) }
+    def piped = new InjectingPipeWriter(downstream, "</head>".toCharArray(), "<script></script>".toCharArray(), null, onBytesWritten)
+
+    when:
+    piped.write("test content".toCharArray())
+    piped.close()
+
+    then:
+    bytesWritten.size() == 1
+    bytesWritten[0] == 12
+    downstream.toString() == "test content"
+  }
+
+  def 'should count bytes correctly when writing characters individually'() {
+    setup:
+    def downstream = new StringWriter()
+    def bytesWritten = []
+    def onBytesWritten = { long bytes -> bytesWritten.add(bytes) }
+    def piped = new InjectingPipeWriter(downstream, "</head>".toCharArray(), "<script></script>".toCharArray(), null, onBytesWritten)
+
+    when:
+    def content = "test"
+    for (int i = 0; i < content.length(); i++) {
+      piped.write((int) content.charAt(i))
+    }
+    piped.close()
+
+    then:
+    bytesWritten.size() == 1
+    bytesWritten[0] == 4
+    downstream.toString() == "test"
+  }
+
+  def 'should count bytes correctly with multiple writes'() {
+    setup:
+    def downstream = new StringWriter()
+    def bytesWritten = []
+    def onBytesWritten = { long bytes -> bytesWritten.add(bytes) }
+    def piped = new InjectingPipeWriter(downstream, "</head>".toCharArray(), "<script></script>".toCharArray(), null, onBytesWritten)
+
+    when:
+    piped.write("test".toCharArray())
+    piped.write(" ".toCharArray())
+    piped.write("content".toCharArray())
+    piped.close()
+
+    then:
+    bytesWritten.size() == 1
+    bytesWritten[0] == 12
+    downstream.toString() == "test content"
+  }
+
+  def 'should be resilient to exceptions when onBytesWritten callback is null'() {
+    setup:
+    def downstream = new StringWriter()
+    def piped = new InjectingPipeWriter(downstream, "</head>".toCharArray(), "<script></script>".toCharArray(), null, null)
+
+    when:
+    piped.write("test content".toCharArray())
+    piped.close()
+
+    then:
+    noExceptionThrown()
+    downstream.toString() == "test content"
+  }
+
+  def 'should reset byte count after close'() {
+    setup:
+    def downstream = new StringWriter()
+    def bytesWritten = []
+    def onBytesWritten = { long bytes -> bytesWritten.add(bytes) }
+    def piped = new InjectingPipeWriter(downstream, "</head>".toCharArray(), "<script></script>".toCharArray(), null, onBytesWritten)
+
+    when:
+    piped.write("test".toCharArray())
+    piped.close()
+
+    piped.write("content".toCharArray())
+    piped.close()
+
+    then:
+    bytesWritten.size() == 2
+    bytesWritten[0] == 4
+    bytesWritten[1] == 7
+  }
 }
diff --git a/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/RumHttpServletResponseWrapper.java b/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/RumHttpServletResponseWrapper.java
@@ -94,7 +94,8 @@ public PrintWriter getWriter() throws IOException {
               super.getWriter(),
               rumInjector.getMarkerChars(),
               rumInjector.getSnippetChars(),
-              this::onInjected);
+              this::onInjected,
+              bytes -> RumInjector.getTelemetryCollector().onInjectionResponseSize("3", bytes));
       printWriter = new PrintWriter(wrappedPipeWriter);
     } catch (Exception e) {
       injectionStartTime = -1;
diff --git a/dd-java-agent/instrumentation/servlet/request-3/src/test/groovy/RumHttpServletResponseWrapperTest.groovy b/dd-java-agent/instrumentation/servlet/request-3/src/test/groovy/RumHttpServletResponseWrapperTest.groovy
@@ -159,6 +159,24 @@ class RumHttpServletResponseWrapperTest extends AgentTestRunner {
     1 * onBytesWritten.accept(11)
   }
 
+  void 'response sizes are reported by the InjectingPipeWriter callback'() {
+    setup:
+    def downstream = Mock(java.io.Writer)
+    def marker = "</head>".toCharArray()
+    def contentToInject = "<script></script>".toCharArray()
+    def onBytesWritten = Mock(java.util.function.LongConsumer)
+    def writer = new datadog.trace.bootstrap.instrumentation.buffer.InjectingPipeWriter(
+      downstream, marker, contentToInject, null, onBytesWritten)
+
+    when:
+    writer.write("test".toCharArray())
+    writer.write("content".toCharArray())
+    writer.close()
+
+    then:
+    1 * onBytesWritten.accept(11)
+  }
+
   void 'injection timing is reported when injection is successful'() {
     setup:
     // set the injection start time to simulate timing
diff --git a/dd-java-agent/instrumentation/servlet/request-5/src/main/java/datadog/trace/instrumentation/servlet5/RumHttpServletResponseWrapper.java b/dd-java-agent/instrumentation/servlet/request-5/src/main/java/datadog/trace/instrumentation/servlet5/RumHttpServletResponseWrapper.java
@@ -75,7 +75,8 @@ public PrintWriter getWriter() throws IOException {
               super.getWriter(),
               rumInjector.getMarkerChars(),
               rumInjector.getSnippetChars(),
-              this::onInjected);
+              this::onInjected,
+              bytes -> RumInjector.getTelemetryCollector().onInjectionResponseSize("5", bytes));
       printWriter = new PrintWriter(wrappedPipeWriter);
     } catch (Exception e) {
       injectionStartTime = -1;
diff --git a/dd-java-agent/instrumentation/servlet/request-5/src/test/groovy/RumHttpServletResponseWrapperTest.groovy b/dd-java-agent/instrumentation/servlet/request-5/src/test/groovy/RumHttpServletResponseWrapperTest.groovy
@@ -159,6 +159,24 @@ class RumHttpServletResponseWrapperTest extends AgentTestRunner {
     1 * onBytesWritten.accept(11)
   }
 
+  void 'response sizes are reported by the InjectingPipeWriter callback'() {
+    setup:
+    def downstream = Mock(java.io.Writer)
+    def marker = "</head>".toCharArray()
+    def contentToInject = "<script></script>".toCharArray()
+    def onBytesWritten = Mock(java.util.function.LongConsumer)
+    def writer = new datadog.trace.bootstrap.instrumentation.buffer.InjectingPipeWriter(
+      downstream, marker, contentToInject, null, onBytesWritten)
+
+    when:
+    writer.write("test".toCharArray())
+    writer.write("content".toCharArray())
+    writer.close()
+
+    then:
+    1 * onBytesWritten.accept(11)
+  }
+
   void 'injection timing is reported when injection is successful'() {
     setup:
     // set the injection start time to simulate timing
