Fix RUM injection for offset response writes

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

@@ -118,6 +118,7 @@ public void write(int b) throws IOException {
 
   @Override
   public void write(byte[] array, int off, int len) throws IOException {
+    final int end = off + len;
     if (!filter) {
       if (wasDraining) {
         // needs drain
@@ -129,15 +130,25 @@ public void write(byte[] array, int off, int len) throws IOException {
     }
 
     if (len > bulkWriteThreshold) {
+      // A match that started in the previous write precedes every match wholly in this array.
+      if (matchingPos > 0 && arrayCompletesPendingMatch(array, off, end)) {
+        int pendingMatchLength = marker.length - matchingPos;
+        for (int i = off; i < off + pendingMatchLength; i++) {
+          write(array[i]);
+        }
+        write(array, off + pendingMatchLength, len - pendingMatchLength);
+        return;
+      }
+
       // if the content is large enough, we can bulk write everything but the N trail and tail.
       // This because the buffer can already contain some byte from a previous single write.
       // Also we need to fill the buffer with the tail since we don't know about the next write.
-      int idx = arrayContains(array, off, len, marker);
+      int idx = arrayContains(array, off, end, marker);
       if (idx >= 0) {
         // we have a full match. just write everything
         filter = false;
         drain();
-        int bytesToWrite = idx;
+        int bytesToWrite = idx - off;
         downstream.write(array, off, bytesToWrite);
         bytesWritten += bytesToWrite;
         long injectionStart = System.nanoTime();
@@ -149,8 +160,8 @@ public void write(byte[] array, int off, int len) throws IOException {
         if (onContentInjected != null) {
           onContentInjected.run();
         }
-        bytesToWrite = len - idx;
-        downstream.write(array, off + idx, bytesToWrite);
+        bytesToWrite = end - idx;
+        downstream.write(array, idx, bytesToWrite);
         bytesWritten += bytesToWrite;
       } else {
         // we don't have a full match. write everything in a bulk except the lookbehind buffer
@@ -167,20 +178,33 @@ public void write(byte[] array, int off, int len) throws IOException {
         downstream.write(array, off + marker.length - 1, bytesToWrite);
         bytesWritten += bytesToWrite;
         filter = wasFiltering;
-        for (int i = len - marker.length + 1; i < len; i++) {
+        for (int i = end - marker.length + 1; i < end; i++) {
           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++) {
+      for (int i = off; i < end; i++) {
         write(array[i]);
       }
     }
   }
 
-  private int arrayContains(byte[] array, int off, int len, byte[] search) {
-    for (int i = off; i < len - search.length; i++) {
+  private boolean arrayCompletesPendingMatch(byte[] array, int off, int end) {
+    int pendingMatchLength = marker.length - matchingPos;
+    if (off + pendingMatchLength > end) {
+      return false;
+    }
+    for (int i = 0; i < pendingMatchLength; i++) {
+      if (array[off + i] != marker[matchingPos + i]) {
+        return false;
+      }
+    }
+    return true;
+  }
+
+  private int arrayContains(byte[] array, int off, int end, byte[] search) {
+    for (int i = off; i <= end - search.length; i++) {
       if (array[i] == search[0]) {
         boolean found = true;
         int k = i;

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

@@ -118,6 +118,7 @@ public void write(int c) throws IOException {
 
   @Override
   public void write(char[] array, int off, int len) throws IOException {
+    final int end = off + len;
     if (!filter) {
       if (wasDraining) {
         // needs drain
@@ -129,15 +130,25 @@ public void write(char[] array, int off, int len) throws IOException {
     }
 
     if (len > bulkWriteThreshold) {
+      // A match that started in the previous write precedes every match wholly in this array.
+      if (matchingPos > 0 && arrayCompletesPendingMatch(array, off, end)) {
+        int pendingMatchLength = marker.length - matchingPos;
+        for (int i = off; i < off + pendingMatchLength; i++) {
+          write(array[i]);
+        }
+        write(array, off + pendingMatchLength, len - pendingMatchLength);
+        return;
+      }
+
       // if the content is large enough, we can bulk write everything but the N trail and tail.
       // This because the buffer can already contain some byte from a previous single write.
       // Also we need to fill the buffer with the tail since we don't know about the next write.
-      int idx = arrayContains(array, off, len, marker);
+      int idx = arrayContains(array, off, end, marker);
       if (idx >= 0) {
         // we have a full match. just write everything
         filter = false;
         drain();
-        int bytesToWrite = idx;
+        int bytesToWrite = idx - off;
         downstream.write(array, off, bytesToWrite);
         bytesWritten += bytesToWrite;
         long injectionStart = System.nanoTime();
@@ -149,8 +160,8 @@ public void write(char[] array, int off, int len) throws IOException {
         if (onContentInjected != null) {
           onContentInjected.run();
         }
-        bytesToWrite = len - idx;
-        downstream.write(array, off + idx, bytesToWrite);
+        bytesToWrite = end - idx;
+        downstream.write(array, idx, bytesToWrite);
         bytesWritten += bytesToWrite;
       } else {
         // we don't have a full match. write everything in a bulk except the lookbehind buffer
@@ -168,20 +179,33 @@ public void write(char[] array, int off, int len) throws IOException {
         bytesWritten += bytesToWrite;
         filter = wasFiltering;
 
-        for (int i = len - marker.length + 1; i < len; i++) {
+        for (int i = end - marker.length + 1; i < end; i++) {
           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++) {
+      for (int i = off; i < end; i++) {
         write(array[i]);
       }
     }
   }
 
-  private int arrayContains(char[] array, int off, int len, char[] search) {
-    for (int i = off; i < len - search.length; i++) {
+  private boolean arrayCompletesPendingMatch(char[] array, int off, int end) {
+    int pendingMatchLength = marker.length - matchingPos;
+    if (off + pendingMatchLength > end) {
+      return false;
+    }
+    for (int i = 0; i < pendingMatchLength; i++) {
+      if (array[off + i] != marker[matchingPos + i]) {
+        return false;
+      }
+    }
+    return true;
+  }
+
+  private int arrayContains(char[] array, int off, int end, char[] search) {
+    for (int i = off; i <= end - search.length; i++) {
       if (array[i] == search[0]) {
         boolean found = true;
         int k = i;

dd-java-agent/agent-bootstrap/src/test/groovy/datadog/trace/bootstrap/instrumentation/buffer/InjectingPipeOutputStreamTest.groovy

@@ -1,6 +1,7 @@
 package datadog.trace.bootstrap.instrumentation.buffer
 
 import datadog.trace.test.util.DDSpecification
+import java.util.concurrent.atomic.AtomicInteger
 import java.util.function.LongConsumer
 
 class InjectingPipeOutputStreamTest extends DDSpecification {
@@ -152,6 +153,49 @@ class InjectingPipeOutputStreamTest extends DDSpecification {
     downstream.toByteArray() == testBytes.getBytes("UTF-8")
   }
 
+  def 'should honor non-zero offsets in bulk writes: #scenario'() {
+    setup:
+    def prefix = "ignored-prefix"
+    def payloadBytes = payload.getBytes("UTF-8")
+    def source = (prefix + payload + "ignored-suffix").getBytes("UTF-8")
+    def downstream = new ByteArrayOutputStream()
+    def counter = new Counter()
+    def injections = new AtomicInteger()
+    def piped = new InjectingPipeOutputStream(downstream, MARKER_BYTES, CONTEXT_BYTES, injections.&incrementAndGet, { long bytes -> counter.incr(bytes) }, null)
+
+    when:
+    piped.write(source, prefix.getBytes("UTF-8").length, payloadBytes.length)
+    piped.close()
+
+    then:
+    downstream.toByteArray() == expected.getBytes("UTF-8")
+    injections.get() == expectedInjections
+    counter.value == payloadBytes.length
+
+    where:
+    scenario                  | payload                                                     | expected                                                                      | expectedInjections
+    "without a marker"        | "<html><body>safe</body></html>"                           | "<html><body>safe</body></html>"                                             | 0
+    "with a marker"           | "<html><head>dynamic</head><body>safe</body></html>"        | "<html><head>dynamic<script></script></head><body>safe</body></html>"          | 1
+    "with a marker at the end" | "<html><head>dynamic-content</head>"                        | "<html><head>dynamic-content<script></script></head>"                          | 1
+  }
+
+  def 'should prioritize a marker spanning a bulk write boundary'() {
+    setup:
+    def prefix = "ignored-prefix"
+    def payload = ">0123456789</head>"
+    def source = (prefix + payload + "ignored-suffix").getBytes("UTF-8")
+    def downstream = new ByteArrayOutputStream()
+    def piped = new InjectingPipeOutputStream(downstream, MARKER_BYTES, CONTEXT_BYTES)
+
+    when:
+    piped.write("abc</head".getBytes("UTF-8"))
+    piped.write(source, prefix.getBytes("UTF-8").length, payload.getBytes("UTF-8").length)
+    piped.close()
+
+    then:
+    downstream.toByteArray() == "abc<script></script></head>0123456789</head>".getBytes("UTF-8")
+  }
+
   def 'should be resilient to exceptions when onBytesWritten callback is null'() {
     setup:
     def testBytes = "test content".getBytes("UTF-8")

dd-java-agent/agent-bootstrap/src/test/groovy/datadog/trace/bootstrap/instrumentation/buffer/InjectingPipeWriterTest.groovy

@@ -1,6 +1,7 @@
 package datadog.trace.bootstrap.instrumentation.buffer
 
 import datadog.trace.test.util.DDSpecification
+import java.util.concurrent.atomic.AtomicInteger
 import java.util.function.LongConsumer
 
 class InjectingPipeWriterTest extends DDSpecification {
@@ -168,6 +169,48 @@ class InjectingPipeWriterTest extends DDSpecification {
     downstream.toString() == testBytes
   }
 
+  def 'should honor non-zero offsets in bulk writes: #scenario'() {
+    setup:
+    def prefix = "ignored-prefix"
+    def source = (prefix + payload + "ignored-suffix").toCharArray()
+    def downstream = new StringWriter()
+    def counter = new Counter()
+    def injections = new AtomicInteger()
+    def piped = new InjectingPipeWriter(downstream, MARKER_CHARS, CONTEXT_CHARS, injections.&incrementAndGet, { long bytes -> counter.incr(bytes) }, null)
+
+    when:
+    piped.write(source, prefix.length(), payload.length())
+    piped.close()
+
+    then:
+    downstream.toString() == expected
+    injections.get() == expectedInjections
+    counter.value == payload.length()
+
+    where:
+    scenario                  | payload                                                     | expected                                                                      | expectedInjections
+    "without a marker"        | "<html><body>safe</body></html>"                           | "<html><body>safe</body></html>"                                             | 0
+    "with a marker"           | "<html><head>dynamic</head><body>safe</body></html>"        | "<html><head>dynamic<script></script></head><body>safe</body></html>"          | 1
+    "with a marker at the end" | "<html><head>dynamic-content</head>"                        | "<html><head>dynamic-content<script></script></head>"                          | 1
+  }
+
+  def 'should prioritize a marker spanning a bulk write boundary'() {
+    setup:
+    def prefix = "ignored-prefix"
+    def payload = ">0123456789</head>"
+    def source = (prefix + payload + "ignored-suffix").toCharArray()
+    def downstream = new StringWriter()
+    def piped = new InjectingPipeWriter(downstream, MARKER_CHARS, CONTEXT_CHARS)
+
+    when:
+    piped.write("abc</head".toCharArray())
+    piped.write(source, prefix.length(), payload.length())
+    piped.close()
+
+    then:
+    downstream.toString() == "abc<script></script></head>0123456789</head>"
+  }
+
   def 'should be resilient to exceptions when onBytesWritten callback is null'() {
     setup:
     def downstream = new StringWriter()

dd-smoke-tests/spring-boot-3.3-webmvc/application/build.gradle

@@ -24,6 +24,7 @@ java {
 
 dependencies {
   implementation 'org.springframework.boot:spring-boot-starter-web'
+  implementation 'org.springframework.boot:spring-boot-starter-thymeleaf'
   implementation 'org.springframework.boot:spring-boot-starter-data-jpa'
   implementation group: 'com.h2database', name: 'h2', version: '2.1.214'
   compileOnly group:"com.google.code.findbugs", name:"jsr305", version:"3.0.2"

dd-smoke-tests/spring-boot-3.3-webmvc/application/src/main/java/datadog/smoketest/springboot/controller/FormResponseController.java

@@ -0,0 +1,36 @@
+package datadog.smoketest.springboot.controller;
+
+import jakarta.servlet.http.HttpServletResponse;
+import java.io.IOException;
+import java.util.Locale;
+import org.springframework.stereotype.Controller;
+import org.springframework.web.bind.annotation.GetMapping;
+import org.thymeleaf.context.Context;
+import org.thymeleaf.spring6.SpringTemplateEngine;
+
+@Controller
+public class FormResponseController {
+  private static final int RESPONSE_OFFSET = 6;
+
+  private final SpringTemplateEngine templateEngine;
+
+  public FormResponseController(SpringTemplateEngine templateEngine) {
+    this.templateEngine = templateEngine;
+  }
+
+  @GetMapping("/form-response")
+  public void formResponse(HttpServletResponse response) throws IOException {
+    Context context = new Context(Locale.ROOT);
+    context.setVariable("returnUrl", "https://app.example.test/flows/complete?request=request-123");
+    context.setVariable("formAction", "https://provider.example.test/flows/continue");
+    context.setVariable("requestId", "request-123");
+
+    String content = templateEngine.process("form-response", context);
+    // Exercise response writers that receive a slice of a reusable buffer.
+    char[] responseBuffer = new char[RESPONSE_OFFSET + content.length()];
+    content.getChars(0, content.length(), responseBuffer, RESPONSE_OFFSET);
+
+    response.setContentType("text/html");
+    response.getWriter().write(responseBuffer, RESPONSE_OFFSET, content.length());
+  }
+}

dd-smoke-tests/spring-boot-3.3-webmvc/application/src/main/resources/templates/form-response.html

@@ -0,0 +1,16 @@
+<!DOCTYPE html>
+<html lang="en" xmlns:th="http://www.thymeleaf.org">
+<head>
+  <meta charset="UTF-8">
+  <title>Continue request</title>
+  <script th:inline="javascript">
+    const returnUrl = /*[[${returnUrl}]]*/ "https://app.example.test/flows/complete";
+    window.formResponse = { returnUrl: returnUrl };
+  </script>
+</head>
+<body onload="document.getElementById('response-form').submit()">
+  <form id="response-form" method="post" th:action="${formAction}">
+    <input type="hidden" name="requestId" th:value="${requestId}">
+  </form>
+</body>
+</html>