DataDog · sarahchen6 · Jul 28, 2025 · Jul 28, 2025 · Jul 28, 2025 · Jul 28, 2025
@@ -51,7 +51,7 @@ public class InjectingPipeOutputStreamBenchmark {
   public void withPipe() throws Exception {
     try (final PrintWriter out =
         new PrintWriter(
-            new InjectingPipeOutputStream(new ByteArrayOutputStream(), marker, content, null))) {
+            new InjectingPipeOutputStream(new ByteArrayOutputStream(), marker, content))) {
       htmlContent.forEach(out::println);
     }
   }

@@ -2,6 +2,7 @@
 
 import java.io.IOException;
 import java.io.OutputStream;
+import java.util.function.LongConsumer;
 import javax.annotation.concurrent.NotThreadSafe;
 
 /**
@@ -23,18 +24,37 @@ public class InjectingPipeOutputStream extends OutputStream {
   private final Runnable onContentInjected;
   private final int bulkWriteThreshold;
   private final OutputStream downstream;
+  private final LongConsumer onBytesWritten;
+  private long bytesWritten = 0;
 
   /**
+   * This constructor is typically used for testing where we care about the logic and not the
+   * telemetry.
+   *
+   * @param downstream the delegate output stream
+   * @param marker the marker to find in the stream. Must at least be one byte.
+   * @param contentToInject the content to inject once before the marker if found.
+   */
+  public InjectingPipeOutputStream(
+      final OutputStream downstream, final byte[] marker, final byte[] contentToInject) {
+    this(downstream, marker, contentToInject, null, null);
+  }
+
+  /**
+   * This constructor contains the full set of parameters.
+   *
    * @param downstream the delegate output stream
    * @param marker the marker to find in the stream. Must at least be one byte.
    * @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 stream is closed to report total bytes written.
    */
   public InjectingPipeOutputStream(
       final OutputStream downstream,
       final byte[] marker,
       final byte[] contentToInject,
-      final Runnable onContentInjected) {
+      final Runnable onContentInjected,
+      final LongConsumer onBytesWritten) {
     this.downstream = downstream;
     this.marker = marker;
     this.lookbehind = new byte[marker.length];
@@ -46,6 +66,7 @@ public InjectingPipeOutputStream(
     this.filter = true;
     this.contentToInject = contentToInject;
     this.onContentInjected = onContentInjected;
+    this.onBytesWritten = onBytesWritten;
     this.bulkWriteThreshold = marker.length * 2 - 2;
   }
 
@@ -57,11 +78,13 @@ public void write(int b) throws IOException {
         drain();
       }
       downstream.write(b);
+      bytesWritten++;
       return;
     }
 
     if (count == lookbehind.length) {
       downstream.write(lookbehind[pos]);
+      bytesWritten++;
     } else {
       count++;
     }
@@ -91,6 +114,7 @@ public void write(byte[] array, int off, int len) throws IOException {
         drain();
       }
       downstream.write(array, off, len);
+      bytesWritten += len;
       return;
     }
 
@@ -103,12 +127,16 @@ public void write(byte[] array, int off, int len) throws IOException {
         // we have a full match. just write everything
         filter = false;
         drain();
-        downstream.write(array, off, idx);
+        int bytesToWrite = idx;
+        downstream.write(array, off, bytesToWrite);
+        bytesWritten += bytesToWrite;
         downstream.write(contentToInject);
         if (onContentInjected != null) {
           onContentInjected.run();
         }
-        downstream.write(array, off + idx, len - idx);
+        bytesToWrite = len - idx;
+        downstream.write(array, off + idx, bytesToWrite);
+        bytesWritten += bytesToWrite;
       } else {
         // we don't have a full match. write everything in a bulk except the lookbehind buffer
         // sequentially
@@ -120,7 +148,9 @@ public void write(byte[] array, int off, int len) throws IOException {
 
         // will be reset if no errors after the following write
         filter = false;
-        downstream.write(array, off + marker.length - 1, len - bulkWriteThreshold);
+        int bytesToWrite = len - bulkWriteThreshold;
+        downstream.write(array, off + marker.length - 1, bytesToWrite);
+        bytesWritten += bytesToWrite;
         filter = wasFiltering;
         for (int i = len - marker.length + 1; i < len; i++) {
           write(array[i]);
@@ -163,6 +193,7 @@ private void drain() throws IOException {
       int cnt = count;
       for (int i = 0; i < cnt; i++) {
         downstream.write(lookbehind[(start + i) % lookbehind.length]);
+        bytesWritten++;
         count--;
       }
       filter = wasFiltering;
@@ -185,6 +216,11 @@ public void flush() throws IOException {
   public void close() throws IOException {
     try {
       commit();
+      // report the size of the original HTTP response before injecting via callback
+      if (onBytesWritten != null) {
+        onBytesWritten.accept(bytesWritten);
+      }
+      bytesWritten = 0;
     } finally {
       downstream.close();
     }

@@ -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,37 @@ 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;
 
   /**
+   * This constructor is typically used for testing where we care about the logic and not the
+   * telemetry.
+   *
+   * @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.
+   */
+  public InjectingPipeWriter(
+      final Writer downstream, final char[] marker, final char[] contentToInject) {
+    this(downstream, marker, contentToInject, null, null);
+  }
+
+  /**
+   * This constructor contains the full set of parameters.
+   *
    * @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,6 +66,7 @@ public InjectingPipeWriter(
     this.filter = true;
     this.contentToInject = contentToInject;
     this.onContentInjected = onContentInjected;
+    this.onBytesWritten = onBytesWritten;
     this.bulkWriteThreshold = marker.length * 2 - 2;
   }
 
@@ -57,11 +78,13 @@ public void write(int c) throws IOException {
         drain();
       }
       downstream.write(c);
+      bytesWritten++;
       return;
     }
 
     if (count == lookbehind.length) {
       downstream.write(lookbehind[pos]);
+      bytesWritten++;
     } else {
       count++;
     }
@@ -91,6 +114,7 @@ public void write(char[] array, int off, int len) throws IOException {
         drain();
       }
       downstream.write(array, off, len);
+      bytesWritten += len;
       return;
     }
 
@@ -103,12 +127,16 @@ public void write(char[] array, int off, int len) throws IOException {
         // we have a full match. just write everything
         filter = false;
         drain();
-        downstream.write(array, off, idx);
+        int bytesToWrite = idx;
+        downstream.write(array, off, bytesToWrite);
+        bytesWritten += bytesToWrite;
         downstream.write(contentToInject);
         if (onContentInjected != null) {
           onContentInjected.run();
         }
-        downstream.write(array, off + idx, len - idx);
+        bytesToWrite = len - idx;
+        downstream.write(array, off + idx, bytesToWrite);
+        bytesWritten += bytesToWrite;
       } else {
         // we don't have a full match. write everything in a bulk except the lookbehind buffer
         // sequentially
@@ -120,7 +148,9 @@ public void write(char[] array, int off, int len) throws IOException {
 
         // will be reset if no errors after the following write
         filter = false;
-        downstream.write(array, off + marker.length - 1, len - bulkWriteThreshold);
+        int bytesToWrite = len - bulkWriteThreshold;
+        downstream.write(array, off + marker.length - 1, bytesToWrite);
+        bytesWritten += bytesToWrite;
         filter = wasFiltering;
 
         for (int i = len - marker.length + 1; i < len; i++) {
@@ -164,6 +194,7 @@ private void drain() throws IOException {
       int cnt = count;
       for (int i = 0; i < cnt; i++) {
         downstream.write(lookbehind[(start + i) % lookbehind.length]);
+        bytesWritten++;
         count--;
       }
       filter = wasFiltering;
@@ -188,6 +219,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;
     }
   }
 

@@ -3,6 +3,9 @@ package datadog.trace.bootstrap.instrumentation.buffer
 import datadog.trace.test.util.DDSpecification
 
 class InjectingPipeOutputStreamTest extends DDSpecification {
+  static final byte[] MARKER_BYTES = "</head>".getBytes("UTF-8")
+  static final byte[] CONTEXT_BYTES = "<script></script>".getBytes("UTF-8")
+
   static class GlitchedOutputStream extends FilterOutputStream {
     int glitchesPos
     int count
@@ -33,10 +36,18 @@ class InjectingPipeOutputStreamTest extends DDSpecification {
     }
   }
 
+  static class Counter {
+    int value = 0
+
+    def incr(long count) {
+      this.value += count
+    }
+  }
+
   def 'should filter a buffer and inject if found #found'() {
     setup:
     def downstream = new ByteArrayOutputStream()
-    def piped = new OutputStreamWriter(new InjectingPipeOutputStream(downstream, marker.getBytes("UTF-8"), contentToInject.getBytes("UTF-8"), null),
+    def piped = new OutputStreamWriter(new InjectingPipeOutputStream(downstream, marker.getBytes("UTF-8"), contentToInject.getBytes("UTF-8")),
     "UTF-8")
     when:
     try (def closeme = piped) {
@@ -55,7 +66,7 @@ class InjectingPipeOutputStreamTest extends DDSpecification {
     setup:
     def baos = new ByteArrayOutputStream()
     def downstream = new GlitchedOutputStream(baos, glichesAt)
-    def piped = new InjectingPipeOutputStream(downstream, marker.getBytes("UTF-8"), contentToInject.getBytes("UTF-8"), null)
+    def piped = new InjectingPipeOutputStream(downstream, marker.getBytes("UTF-8"), contentToInject.getBytes("UTF-8"))
     when:
     try {
       for (String line : body) {
@@ -87,4 +98,74 @@ 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 testBytes = "test content".getBytes("UTF-8")
+    def downstream = new ByteArrayOutputStream()
+    def counter = new Counter()
+    def piped = new InjectingPipeOutputStream(downstream, MARKER_BYTES, CONTEXT_BYTES, null, { long bytes -> counter.incr(bytes) })
+
+    when:
+    piped.write(testBytes)
+    piped.close()
+
+    then:
+    counter.value == 12
+    downstream.toByteArray() == testBytes
+  }
+
+  def 'should count bytes correctly when writing bytes individually'() {
+    setup:
+    def testBytes = "test".getBytes("UTF-8")
+    def downstream = new ByteArrayOutputStream()
+    def counter = new Counter()
+    def piped = new InjectingPipeOutputStream(downstream, MARKER_BYTES, CONTEXT_BYTES, null, { long bytes -> counter.incr(bytes) })
+
+    when:
+    for (int i = 0; i < testBytes.length; i++) {
+      piped.write((int) testBytes[i])
+    }
+    piped.close()
+
+    then:
+    counter.value == 4
+    downstream.toByteArray() == testBytes
+  }
+
+  def 'should count bytes correctly with multiple writes'() {
+    setup:
+    def part1 = "test".getBytes("UTF-8")
+    def part2 = " ".getBytes("UTF-8")
+    def part3 = "content".getBytes("UTF-8")
+    def testBytes = "test content".getBytes("UTF-8")
+    def downstream = new ByteArrayOutputStream()
+    def counter = new Counter()
+    def piped = new InjectingPipeOutputStream(downstream, MARKER_BYTES, CONTEXT_BYTES, null, { long bytes -> counter.incr(bytes) })
+
+    when:
+    piped.write(part1)
+    piped.write(part2)
+    piped.write(part3)
+    piped.close()
+
+    then:
+    counter.value == 12
+    downstream.toByteArray() == testBytes
+  }
+
+  def 'should be resilient to exceptions when onBytesWritten callback is null'() {
+    setup:
+    def testBytes = "test content".getBytes("UTF-8")
+    def downstream = new ByteArrayOutputStream()
+    def piped = new InjectingPipeOutputStream(downstream, MARKER_BYTES, CONTEXT_BYTES)
+
+    when:
+    piped.write(testBytes)
+    piped.close()
+
+    then:
+    noExceptionThrown()
+    downstream.toByteArray() == testBytes
+  }
 }