Handle servlet response lifecycle boundaries during RUM injection

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

@@ -244,13 +244,24 @@ public void commit() throws IOException {
     if (filter || wasDraining) {
       drain();
     }
+    matchingPos = 0;
   }
 
   @Override
   public void flush() throws IOException {
+    commit();
     downstream.flush();
   }
 
+  /** Discards buffered content and resets matching state without writing to the downstream. */
+  public void discard() {
+    pos = 0;
+    count = 0;
+    matchingPos = 0;
+    wasDraining = false;
+    bytesWritten = 0;
+  }
+
   @Override
   public void close() throws IOException {
     try {

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

@@ -245,13 +245,24 @@ public void commit() throws IOException {
     if (filter || wasDraining) {
       drain();
     }
+    matchingPos = 0;
   }
 
   @Override
   public void flush() throws IOException {
+    commit();
     downstream.flush();
   }
 
+  /** Discards buffered content and resets matching state without writing to the downstream. */
+  public void discard() {
+    pos = 0;
+    count = 0;
+    matchingPos = 0;
+    wasDraining = false;
+    bytesWritten = 0;
+  }
+
   @Override
   public void close() throws IOException {
     try {

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

@@ -196,6 +196,57 @@ class InjectingPipeOutputStreamTest extends DDSpecification {
     downstream.toByteArray() == "abc<script></script></head>0123456789</head>".getBytes("UTF-8")
   }
 
+  def 'should drain and reset a partial match on commit'() {
+    setup:
+    def downstream = new ByteArrayOutputStream()
+    def piped = new InjectingPipeOutputStream(downstream, MARKER_BYTES, CONTEXT_BYTES)
+
+    when:
+    piped.write("</he".getBytes("UTF-8"))
+    piped.commit()
+    piped.write("ad>".getBytes("UTF-8"))
+    piped.close()
+
+    then:
+    downstream.toByteArray() == "</head>".getBytes("UTF-8")
+  }
+
+  def 'should drain a partial match before flushing'() {
+    setup:
+    def downstream = new ByteArrayOutputStream()
+    def piped = new InjectingPipeOutputStream(downstream, MARKER_BYTES, CONTEXT_BYTES)
+
+    when:
+    piped.write("</he".getBytes("UTF-8"))
+    piped.flush()
+
+    then:
+    downstream.toByteArray() == "</he".getBytes("UTF-8")
+
+    when:
+    piped.write("ad>".getBytes("UTF-8"))
+    piped.close()
+
+    then:
+    downstream.toByteArray() == "</head>".getBytes("UTF-8")
+  }
+
+  def 'should discard buffered content and matching state'() {
+    setup:
+    def downstream = new ByteArrayOutputStream()
+    def piped = new InjectingPipeOutputStream(downstream, MARKER_BYTES, CONTEXT_BYTES)
+
+    when:
+    piped.write("discarded</he".getBytes("UTF-8"))
+    downstream.reset()
+    piped.discard()
+    piped.write("kept</head>".getBytes("UTF-8"))
+    piped.close()
+
+    then:
+    downstream.toByteArray() == "kept<script></script></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

@@ -211,6 +211,57 @@ class InjectingPipeWriterTest extends DDSpecification {
     downstream.toString() == "abc<script></script></head>0123456789</head>"
   }
 
+  def 'should drain and reset a partial match on commit'() {
+    setup:
+    def downstream = new StringWriter()
+    def piped = new InjectingPipeWriter(downstream, MARKER_CHARS, CONTEXT_CHARS)
+
+    when:
+    piped.write("</he".toCharArray())
+    piped.commit()
+    piped.write("ad>".toCharArray())
+    piped.close()
+
+    then:
+    downstream.toString() == "</head>"
+  }
+
+  def 'should drain a partial match before flushing'() {
+    setup:
+    def downstream = new StringWriter()
+    def piped = new InjectingPipeWriter(downstream, MARKER_CHARS, CONTEXT_CHARS)
+
+    when:
+    piped.write("</he".toCharArray())
+    piped.flush()
+
+    then:
+    downstream.toString() == "</he"
+
+    when:
+    piped.write("ad>".toCharArray())
+    piped.close()
+
+    then:
+    downstream.toString() == "</head>"
+  }
+
+  def 'should discard buffered content and matching state'() {
+    setup:
+    def downstream = new StringWriter()
+    def piped = new InjectingPipeWriter(downstream, MARKER_CHARS, CONTEXT_CHARS)
+
+    when:
+    piped.write("discarded</he".toCharArray())
+    downstream.buffer.setLength(0)
+    piped.discard()
+    piped.write("kept</head>".toCharArray())
+    piped.close()
+
+    then:
+    downstream.toString() == "kept<script></script></head>"
+  }
+
   def 'should be resilient to exceptions when onBytesWritten callback is null'() {
     setup:
     def downstream = new StringWriter()

dd-java-agent/instrumentation/servlet/jakarta-servlet-5.0/build.gradle

@@ -30,6 +30,8 @@ configurations {
   javaxClassesToRelocate
 }
 
+addTestSuiteForDir('servlet61Test', 'servlet61Test')
+
 def relocatedJavaxJar = tasks.register('relocatedJavaxJar', ShadowJar) {
   relocate 'javax.servlet', 'jakarta.servlet'
   relocate 'datadog.trace.instrumentation.servlet3', 'datadog.trace.instrumentation.servlet5'
@@ -72,6 +74,9 @@ dependencies {
   testImplementation group: 'jakarta.servlet.jsp', name: 'jakarta.servlet.jsp-api', version: '3.0.0'
   testRuntimeOnly project(':dd-java-agent:instrumentation:datadog:asm:iast-instrumenter')
 
+  servlet61TestImplementation group: 'jakarta.servlet', name: 'jakarta.servlet-api', version: '6.1.0'
+  servlet61TestImplementation libs.bundles.mockito
+
   javaxClassesToRelocate project(':dd-java-agent:instrumentation:servlet:javax-servlet:javax-servlet-iast'), {
     transitive = false
   }
@@ -90,6 +95,12 @@ dependencies {
   // tested on tomcat-5.5:latestDepTest
 }
 
+tasks.named('servlet61Test', Test) {
+  testJvmConstraints {
+    minJavaVersion = JavaVersion.VERSION_17
+  }
+}
+
 tasks.named("jar", Jar) {
   from zipTree(relocatedJavaxJarFile)
 }