From 57f6c18c89f21a9795f709d83db528900bab1aa4 Mon Sep 17 00:00:00 2001 From: Andrea Marziali Date: Tue, 15 Sep 2026 11:56:15 +0200 Subject: [PATCH 1/4] Handle servlet response lifecycle boundaries during RUM injection --- .../buffer/InjectingPipeOutputStream.java | 11 ++ .../buffer/InjectingPipeWriter.java | 11 ++ .../InjectingPipeOutputStreamTest.groovy | 51 +++++++++ .../buffer/InjectingPipeWriterTest.groovy | 51 +++++++++ .../RumHttpServletResponseWrapper.java | 65 +++++++++-- .../servlet5/WrappedServletOutputStream.java | 4 + .../RumHttpServletResponseWrapperTest.groovy | 106 ++++++++++++++++++ .../RumHttpServletResponseWrapper.java | 65 +++++++++-- .../servlet3/WrappedServletOutputStream.java | 4 + .../RumHttpServletResponseWrapperTest.groovy | 106 ++++++++++++++++++ 10 files changed, 460 insertions(+), 14 deletions(-) diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/buffer/InjectingPipeOutputStream.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/buffer/InjectingPipeOutputStream.java index 9d49166bd1b..8a0db21f1e7 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/buffer/InjectingPipeOutputStream.java +++ b/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 { 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 index a7439abaacf..87cf779402d 100644 --- 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 @@ -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 { 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 index ed5525bbb4e..9c43eadb7b9 100644 --- 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 @@ -196,6 +196,57 @@ class InjectingPipeOutputStreamTest extends DDSpecification { downstream.toByteArray() == "abc0123456789".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("".getBytes("UTF-8")) + piped.close() + + then: + downstream.toByteArray() == "".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("".getBytes("UTF-8")) + piped.close() + + then: + downstream.toByteArray() == "".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".getBytes("UTF-8")) + piped.close() + + then: + downstream.toByteArray() == "kept".getBytes("UTF-8") + } + def 'should be resilient to exceptions when onBytesWritten callback is null'() { setup: def testBytes = "test content".getBytes("UTF-8") 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 index 19307f61c2b..6f0aafd1ce3 100644 --- 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 @@ -211,6 +211,57 @@ class InjectingPipeWriterTest extends DDSpecification { downstream.toString() == "abc0123456789" } + 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("".toCharArray()) + piped.close() + + then: + downstream.toString() == "" + } + + 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("".toCharArray()) + piped.close() + + then: + downstream.toString() == "" + } + + 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".toCharArray()) + piped.close() + + then: + downstream.toString() == "kept" + } + def 'should be resilient to exceptions when onBytesWritten callback is null'() { setup: def downstream = new StringWriter() diff --git a/dd-java-agent/instrumentation/servlet/jakarta-servlet-5.0/src/main/java/datadog/trace/instrumentation/servlet5/RumHttpServletResponseWrapper.java b/dd-java-agent/instrumentation/servlet/jakarta-servlet-5.0/src/main/java/datadog/trace/instrumentation/servlet5/RumHttpServletResponseWrapper.java index 25b3f2c9964..b49649093b6 100644 --- a/dd-java-agent/instrumentation/servlet/jakarta-servlet-5.0/src/main/java/datadog/trace/instrumentation/servlet5/RumHttpServletResponseWrapper.java +++ b/dd-java-agent/instrumentation/servlet/jakarta-servlet-5.0/src/main/java/datadog/trace/instrumentation/servlet5/RumHttpServletResponseWrapper.java @@ -166,19 +166,48 @@ public void setCharacterEncoding(String charset) { @Override public void reset() { + super.reset(); + discardBufferedContent(); + setActiveFilters(false); this.outputStream = null; this.wrappedPipeWriter = null; this.printWriter = null; - this.shouldInject = false; - super.reset(); + this.shouldInject = true; + this.contentEncoding = null; } @Override public void resetBuffer() { - this.outputStream = null; - this.wrappedPipeWriter = null; - this.printWriter = null; super.resetBuffer(); + discardBufferedContent(); + setActiveFilters(shouldInject); + } + + @Override + public void flushBuffer() throws IOException { + flushBufferedContent(); + super.flushBuffer(); + } + + @Override + public void sendError(int sc) throws IOException { + super.sendError(sc); + discardBufferedContent(); + stopFiltering(); + } + + @Override + public void sendError(int sc, String msg) throws IOException { + super.sendError(sc, msg); + discardBufferedContent(); + stopFiltering(); + } + + @Override + public void sendRedirect(String location) throws IOException { + super.sendRedirect(location); + discardBufferedContent(); + stopFiltering(); } public void onInjected() { @@ -226,11 +255,33 @@ public void commit() { @Override public void stopFiltering() { shouldInject = false; + setActiveFilters(false); + } + + private void flushBufferedContent() throws IOException { + if (wrappedPipeWriter != null) { + wrappedPipeWriter.commit(); + } + if (outputStream != null) { + outputStream.commit(); + } + } + + private void discardBufferedContent() { + if (wrappedPipeWriter != null) { + wrappedPipeWriter.discard(); + } + if (outputStream != null) { + outputStream.discard(); + } + } + + private void setActiveFilters(boolean filter) { if (wrappedPipeWriter != null) { - wrappedPipeWriter.setFilter(false); + wrappedPipeWriter.setFilter(filter); } if (outputStream != null) { - outputStream.setFilter(false); + outputStream.setFilter(filter); } } } diff --git a/dd-java-agent/instrumentation/servlet/jakarta-servlet-5.0/src/main/java/datadog/trace/instrumentation/servlet5/WrappedServletOutputStream.java b/dd-java-agent/instrumentation/servlet/jakarta-servlet-5.0/src/main/java/datadog/trace/instrumentation/servlet5/WrappedServletOutputStream.java index fe2a87774ac..c42e38193f6 100644 --- a/dd-java-agent/instrumentation/servlet/jakarta-servlet-5.0/src/main/java/datadog/trace/instrumentation/servlet5/WrappedServletOutputStream.java +++ b/dd-java-agent/instrumentation/servlet/jakarta-servlet-5.0/src/main/java/datadog/trace/instrumentation/servlet5/WrappedServletOutputStream.java @@ -42,6 +42,10 @@ public void commit() throws IOException { filtered.commit(); } + public void discard() { + filtered.discard(); + } + @Override public void flush() throws IOException { filtered.flush(); diff --git a/dd-java-agent/instrumentation/servlet/jakarta-servlet-5.0/src/test/groovy/RumHttpServletResponseWrapperTest.groovy b/dd-java-agent/instrumentation/servlet/jakarta-servlet-5.0/src/test/groovy/RumHttpServletResponseWrapperTest.groovy index 595a0909f10..cb0e0679a44 100644 --- a/dd-java-agent/instrumentation/servlet/jakarta-servlet-5.0/src/test/groovy/RumHttpServletResponseWrapperTest.groovy +++ b/dd-java-agent/instrumentation/servlet/jakarta-servlet-5.0/src/test/groovy/RumHttpServletResponseWrapperTest.groovy @@ -181,6 +181,112 @@ class RumHttpServletResponseWrapperTest extends InstrumentationSpecification { 1 * mockTelemetryCollector.onInjectionFailed(SERVLET_VERSION, null) } + void 'flushBuffer drains and resets a partial marker'() { + setup: + def downstream = new StringWriter() + def writer = attachWriter(downstream) + + when: + writer.write("") + wrapper.commit() + + then: + downstream.toString() == "" + 0 * mockTelemetryCollector.onInjectionSucceed(_) + } + + void 'resetBuffer discards content and re-arms injection'() { + setup: + def downstream = new StringWriter() + def writer = attachWriter(downstream) + + when: + writer.write("") + wrapper.commit() + + then: + !downstream.toString().startsWith("" + 1 * mockResponse.resetBuffer() + 1 * mockTelemetryCollector.onInjectionSucceed(SERVLET_VERSION) + } + + void 'reset retires the old writer and restores injection for a new writer'() { + setup: + def downstream = new StringWriter() + def oldWriter = attachWriter(downstream) + + when: + oldWriter.write("") + wrapper.commit() + + then: + !downstream.toString().startsWith("" + 1 * mockResponse.reset() + 1 * mockTelemetryCollector.onInjectionSucceed(SERVLET_VERSION) + } + + void 'sendError discards buffered content and stops filtering'() { + setup: + def downstream = new StringWriter() + def writer = attachWriter(downstream) + + when: + writer.write("") + wrapper.commit() + + then: + downstream.toString() == "ad>" + 1 * mockResponse.sendError(500) + 0 * mockTelemetryCollector.onInjectionSucceed(_) + } + + void 'sendRedirect discards buffered content and stops filtering'() { + setup: + def downstream = new StringWriter() + def writer = attachWriter(downstream) + + when: + writer.write("") + wrapper.commit() + + then: + downstream.toString() == "ad>" + 1 * mockResponse.sendRedirect("/other") + 0 * mockTelemetryCollector.onInjectionSucceed(_) + } + + private PrintWriter attachWriter(StringWriter downstream) { + def pipe = new InjectingPipeWriter( + downstream, + "".toCharArray(), + "".toCharArray(), + wrapper.&onInjected, + null, + null) + def writer = new PrintWriter(pipe) + wrapper.@wrappedPipeWriter = pipe + wrapper.@printWriter = writer + return writer + } + // Callback is created in the RumHttpServletResponseWrapper and passed to InjectingPipeOutputStream via WrappedServletOutputStream. // When the stream is closed, the callback is called with the number of bytes written to the stream and the time taken to write the injection content. void 'response sizes are reported to the telemetry collector via the WrappedServletOutputStream callback'() { diff --git a/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-3.0/src/main/java/datadog/trace/instrumentation/servlet3/RumHttpServletResponseWrapper.java b/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-3.0/src/main/java/datadog/trace/instrumentation/servlet3/RumHttpServletResponseWrapper.java index cbd0b6f845b..3b25a1c874f 100644 --- a/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-3.0/src/main/java/datadog/trace/instrumentation/servlet3/RumHttpServletResponseWrapper.java +++ b/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-3.0/src/main/java/datadog/trace/instrumentation/servlet3/RumHttpServletResponseWrapper.java @@ -183,19 +183,48 @@ public void setCharacterEncoding(String charset) { @Override public void reset() { + super.reset(); + discardBufferedContent(); + setActiveFilters(false); this.outputStream = null; this.wrappedPipeWriter = null; this.printWriter = null; - this.shouldInject = false; - super.reset(); + this.shouldInject = true; + this.contentEncoding = null; } @Override public void resetBuffer() { - this.outputStream = null; - this.wrappedPipeWriter = null; - this.printWriter = null; super.resetBuffer(); + discardBufferedContent(); + setActiveFilters(shouldInject); + } + + @Override + public void flushBuffer() throws IOException { + flushBufferedContent(); + super.flushBuffer(); + } + + @Override + public void sendError(int sc) throws IOException { + super.sendError(sc); + discardBufferedContent(); + stopFiltering(); + } + + @Override + public void sendError(int sc, String msg) throws IOException { + super.sendError(sc, msg); + discardBufferedContent(); + stopFiltering(); + } + + @Override + public void sendRedirect(String location) throws IOException { + super.sendRedirect(location); + discardBufferedContent(); + stopFiltering(); } public void onInjected() { @@ -249,11 +278,33 @@ public void commit() { @Override public void stopFiltering() { shouldInject = false; + setActiveFilters(false); + } + + private void flushBufferedContent() throws IOException { + if (wrappedPipeWriter != null) { + wrappedPipeWriter.commit(); + } + if (outputStream != null) { + outputStream.commit(); + } + } + + private void discardBufferedContent() { + if (wrappedPipeWriter != null) { + wrappedPipeWriter.discard(); + } + if (outputStream != null) { + outputStream.discard(); + } + } + + private void setActiveFilters(boolean filter) { if (wrappedPipeWriter != null) { - wrappedPipeWriter.setFilter(false); + wrappedPipeWriter.setFilter(filter); } if (outputStream != null) { - outputStream.setFilter(false); + outputStream.setFilter(filter); } } } diff --git a/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-3.0/src/main/java/datadog/trace/instrumentation/servlet3/WrappedServletOutputStream.java b/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-3.0/src/main/java/datadog/trace/instrumentation/servlet3/WrappedServletOutputStream.java index c4b575836ff..f5a20a656db 100644 --- a/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-3.0/src/main/java/datadog/trace/instrumentation/servlet3/WrappedServletOutputStream.java +++ b/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-3.0/src/main/java/datadog/trace/instrumentation/servlet3/WrappedServletOutputStream.java @@ -95,6 +95,10 @@ public void commit() throws IOException { filtered.commit(); } + public void discard() { + filtered.discard(); + } + public void setFilter(boolean filter) { filtered.setFilter(filter); } diff --git a/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-3.0/src/test/groovy/RumHttpServletResponseWrapperTest.groovy b/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-3.0/src/test/groovy/RumHttpServletResponseWrapperTest.groovy index 341ae9cde5e..4964d22f762 100644 --- a/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-3.0/src/test/groovy/RumHttpServletResponseWrapperTest.groovy +++ b/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-3.0/src/test/groovy/RumHttpServletResponseWrapperTest.groovy @@ -181,6 +181,112 @@ class RumHttpServletResponseWrapperTest extends InstrumentationSpecification { 1 * mockTelemetryCollector.onInjectionFailed(SERVLET_VERSION, null) } + void 'flushBuffer drains and resets a partial marker'() { + setup: + def downstream = new StringWriter() + def writer = attachWriter(downstream) + + when: + writer.write("") + wrapper.commit() + + then: + downstream.toString() == "" + 0 * mockTelemetryCollector.onInjectionSucceed(_) + } + + void 'resetBuffer discards content and re-arms injection'() { + setup: + def downstream = new StringWriter() + def writer = attachWriter(downstream) + + when: + writer.write("") + wrapper.commit() + + then: + !downstream.toString().startsWith("" + 1 * mockResponse.resetBuffer() + 1 * mockTelemetryCollector.onInjectionSucceed(SERVLET_VERSION) + } + + void 'reset retires the old writer and restores injection for a new writer'() { + setup: + def downstream = new StringWriter() + def oldWriter = attachWriter(downstream) + + when: + oldWriter.write("") + wrapper.commit() + + then: + !downstream.toString().startsWith("" + 1 * mockResponse.reset() + 1 * mockTelemetryCollector.onInjectionSucceed(SERVLET_VERSION) + } + + void 'sendError discards buffered content and stops filtering'() { + setup: + def downstream = new StringWriter() + def writer = attachWriter(downstream) + + when: + writer.write("") + wrapper.commit() + + then: + downstream.toString() == "ad>" + 1 * mockResponse.sendError(500) + 0 * mockTelemetryCollector.onInjectionSucceed(_) + } + + void 'sendRedirect discards buffered content and stops filtering'() { + setup: + def downstream = new StringWriter() + def writer = attachWriter(downstream) + + when: + writer.write("") + wrapper.commit() + + then: + downstream.toString() == "ad>" + 1 * mockResponse.sendRedirect("/other") + 0 * mockTelemetryCollector.onInjectionSucceed(_) + } + + private PrintWriter attachWriter(StringWriter downstream) { + def pipe = new InjectingPipeWriter( + downstream, + "".toCharArray(), + "".toCharArray(), + wrapper.&onInjected, + null, + null) + def writer = new PrintWriter(pipe) + wrapper.@wrappedPipeWriter = pipe + wrapper.@printWriter = writer + return writer + } + // Callback is created in the RumHttpServletResponseWrapper and passed to InjectingPipeOutputStream via WrappedServletOutputStream. // When the stream is closed, the callback is called with the number of bytes written to the stream and the time taken to write the injection content. void 'response sizes are reported to the telemetry collector via the WrappedServletOutputStream callback'() { From 1b514ff0de17df71de73bf51cc1682fb98845839 Mon Sep 17 00:00:00 2001 From: Andrea Marziali Date: Tue, 15 Sep 2026 13:31:52 +0200 Subject: [PATCH 2/4] Prevent retired RUM wrappers from injecting --- .../servlet5/RumHttpServletResponseWrapper.java | 10 ++++++++-- .../RumHttpServletResponseWrapperTest.groovy | 14 ++++++++++++++ .../servlet3/RumHttpServletResponseWrapper.java | 10 ++++++++-- .../RumHttpServletResponseWrapperTest.groovy | 14 ++++++++++++++ 4 files changed, 44 insertions(+), 4 deletions(-) diff --git a/dd-java-agent/instrumentation/servlet/jakarta-servlet-5.0/src/main/java/datadog/trace/instrumentation/servlet5/RumHttpServletResponseWrapper.java b/dd-java-agent/instrumentation/servlet/jakarta-servlet-5.0/src/main/java/datadog/trace/instrumentation/servlet5/RumHttpServletResponseWrapper.java index b49649093b6..ee15e0554a0 100644 --- a/dd-java-agent/instrumentation/servlet/jakarta-servlet-5.0/src/main/java/datadog/trace/instrumentation/servlet5/RumHttpServletResponseWrapper.java +++ b/dd-java-agent/instrumentation/servlet/jakarta-servlet-5.0/src/main/java/datadog/trace/instrumentation/servlet5/RumHttpServletResponseWrapper.java @@ -20,6 +20,7 @@ public class RumHttpServletResponseWrapper extends HttpServletResponseWrapper private InjectingPipeWriter wrappedPipeWriter; private PrintWriter printWriter; private boolean shouldInject = true; + private boolean retired; private String contentEncoding = null; public RumHttpServletResponseWrapper(HttpServletRequest request, HttpServletResponse response) { @@ -172,7 +173,7 @@ public void reset() { this.outputStream = null; this.wrappedPipeWriter = null; this.printWriter = null; - this.shouldInject = true; + this.shouldInject = !retired; this.contentEncoding = null; } @@ -226,7 +227,7 @@ private void handleContentType(String type) { } if (wasInjecting && !shouldInject) { commit(); - stopFiltering(); + disableFiltering(); } } @@ -254,6 +255,11 @@ public void commit() { @Override public void stopFiltering() { + retired = true; + disableFiltering(); + } + + private void disableFiltering() { shouldInject = false; setActiveFilters(false); } diff --git a/dd-java-agent/instrumentation/servlet/jakarta-servlet-5.0/src/test/groovy/RumHttpServletResponseWrapperTest.groovy b/dd-java-agent/instrumentation/servlet/jakarta-servlet-5.0/src/test/groovy/RumHttpServletResponseWrapperTest.groovy index cb0e0679a44..94de6935035 100644 --- a/dd-java-agent/instrumentation/servlet/jakarta-servlet-5.0/src/test/groovy/RumHttpServletResponseWrapperTest.groovy +++ b/dd-java-agent/instrumentation/servlet/jakarta-servlet-5.0/src/test/groovy/RumHttpServletResponseWrapperTest.groovy @@ -239,6 +239,20 @@ class RumHttpServletResponseWrapperTest extends InstrumentationSpecification { 1 * mockTelemetryCollector.onInjectionSucceed(SERVLET_VERSION) } + void 'reset does not reactivate a retired nested wrapper'() { + setup: + wrapper.stopFiltering() + def outerWrapper = new RumHttpServletResponseWrapper(mockRequest, wrapper) + + when: + outerWrapper.reset() + + then: + !wrapper.@shouldInject + outerWrapper.@shouldInject + 1 * mockResponse.reset() + } + void 'sendError discards buffered content and stops filtering'() { setup: def downstream = new StringWriter() diff --git a/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-3.0/src/main/java/datadog/trace/instrumentation/servlet3/RumHttpServletResponseWrapper.java b/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-3.0/src/main/java/datadog/trace/instrumentation/servlet3/RumHttpServletResponseWrapper.java index 3b25a1c874f..d6dd8a3284e 100644 --- a/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-3.0/src/main/java/datadog/trace/instrumentation/servlet3/RumHttpServletResponseWrapper.java +++ b/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-3.0/src/main/java/datadog/trace/instrumentation/servlet3/RumHttpServletResponseWrapper.java @@ -23,6 +23,7 @@ public class RumHttpServletResponseWrapper extends HttpServletResponseWrapper private PrintWriter printWriter; private InjectingPipeWriter wrappedPipeWriter; private boolean shouldInject = true; + private boolean retired; private String contentEncoding = null; private static final MethodHandle SET_CONTENT_LENGTH_LONG = getMh("setContentLengthLong"); @@ -189,7 +190,7 @@ public void reset() { this.outputStream = null; this.wrappedPipeWriter = null; this.printWriter = null; - this.shouldInject = true; + this.shouldInject = !retired; this.contentEncoding = null; } @@ -243,7 +244,7 @@ private void handleContentType(String type) { } if (wasInjecting && !shouldInject) { commit(); - stopFiltering(); + disableFiltering(); } } @@ -277,6 +278,11 @@ public void commit() { @Override public void stopFiltering() { + retired = true; + disableFiltering(); + } + + private void disableFiltering() { shouldInject = false; setActiveFilters(false); } diff --git a/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-3.0/src/test/groovy/RumHttpServletResponseWrapperTest.groovy b/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-3.0/src/test/groovy/RumHttpServletResponseWrapperTest.groovy index 4964d22f762..a516a863264 100644 --- a/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-3.0/src/test/groovy/RumHttpServletResponseWrapperTest.groovy +++ b/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-3.0/src/test/groovy/RumHttpServletResponseWrapperTest.groovy @@ -239,6 +239,20 @@ class RumHttpServletResponseWrapperTest extends InstrumentationSpecification { 1 * mockTelemetryCollector.onInjectionSucceed(SERVLET_VERSION) } + void 'reset does not reactivate a retired nested wrapper'() { + setup: + wrapper.stopFiltering() + def outerWrapper = new RumHttpServletResponseWrapper(mockRequest, wrapper) + + when: + outerWrapper.reset() + + then: + !wrapper.@shouldInject + outerWrapper.@shouldInject + 1 * mockResponse.reset() + } + void 'sendError discards buffered content and stops filtering'() { setup: def downstream = new StringWriter() From 52946eff9ed7530bfd1129b67ae47c0874a89f87 Mon Sep 17 00:00:00 2001 From: Andrea Marziali Date: Tue, 15 Sep 2026 14:57:22 +0200 Subject: [PATCH 3/4] Handle RUM cleanup for terminal servlet responses --- .../servlet/jakarta-servlet-5.0/build.gradle | 11 +++ .../RumHttpServletResponseWrapper.java | 80 ++++++++++++++++--- ...tpServletResponseWrapperServlet61Test.java | 54 +++++++++++++ .../RumHttpServletResponseWrapperTest.groovy | 38 +++++++++ .../RumHttpServletResponseWrapper.java | 27 ++++--- .../RumHttpServletResponseWrapperTest.groovy | 38 +++++++++ 6 files changed, 230 insertions(+), 18 deletions(-) create mode 100644 dd-java-agent/instrumentation/servlet/jakarta-servlet-5.0/src/servlet61Test/java/datadog/trace/instrumentation/servlet5/RumHttpServletResponseWrapperServlet61Test.java diff --git a/dd-java-agent/instrumentation/servlet/jakarta-servlet-5.0/build.gradle b/dd-java-agent/instrumentation/servlet/jakarta-servlet-5.0/build.gradle index 4a54ab04a0a..dadadf1c874 100644 --- a/dd-java-agent/instrumentation/servlet/jakarta-servlet-5.0/build.gradle +++ b/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) } diff --git a/dd-java-agent/instrumentation/servlet/jakarta-servlet-5.0/src/main/java/datadog/trace/instrumentation/servlet5/RumHttpServletResponseWrapper.java b/dd-java-agent/instrumentation/servlet/jakarta-servlet-5.0/src/main/java/datadog/trace/instrumentation/servlet5/RumHttpServletResponseWrapper.java index ee15e0554a0..0d1d805b243 100644 --- a/dd-java-agent/instrumentation/servlet/jakarta-servlet-5.0/src/main/java/datadog/trace/instrumentation/servlet5/RumHttpServletResponseWrapper.java +++ b/dd-java-agent/instrumentation/servlet/jakarta-servlet-5.0/src/main/java/datadog/trace/instrumentation/servlet5/RumHttpServletResponseWrapper.java @@ -3,6 +3,7 @@ import datadog.trace.api.rum.RumInjector; import datadog.trace.bootstrap.instrumentation.buffer.InjectingPipeWriter; import datadog.trace.bootstrap.instrumentation.rum.RumControllableResponse; +import datadog.trace.util.MethodHandles; import jakarta.servlet.ServletContext; import jakarta.servlet.ServletOutputStream; import jakarta.servlet.http.HttpServletRequest; @@ -10,6 +11,7 @@ import jakarta.servlet.http.HttpServletResponseWrapper; import java.io.IOException; import java.io.PrintWriter; +import java.lang.invoke.MethodHandle; import java.nio.charset.Charset; public class RumHttpServletResponseWrapper extends HttpServletResponseWrapper @@ -23,6 +25,21 @@ public class RumHttpServletResponseWrapper extends HttpServletResponseWrapper private boolean retired; private String contentEncoding = null; + private static final MethodHandles METHOD_HANDLES = + new MethodHandles(HttpServletResponse.class.getClassLoader()); + private static final MethodHandle SEND_REDIRECT_WITH_BUFFER_6_1 = + METHOD_HANDLES.method(HttpServletResponse.class, "sendRedirect", String.class, boolean.class); + private static final MethodHandle SEND_REDIRECT_WITH_STATUS_6_1 = + METHOD_HANDLES.method(HttpServletResponse.class, "sendRedirect", String.class, int.class); + private static final MethodHandle SEND_REDIRECT_6_1 = + METHOD_HANDLES.method( + HttpServletResponse.class, "sendRedirect", String.class, int.class, boolean.class); + + @SuppressWarnings("unchecked") + private static void sneakyThrow(Throwable e) throws E { + throw (E) e; + } + public RumHttpServletResponseWrapper(HttpServletRequest request, HttpServletResponse response) { super(response); this.rumInjector = RumInjector.get(); @@ -192,23 +209,68 @@ public void flushBuffer() throws IOException { @Override public void sendError(int sc) throws IOException { - super.sendError(sc); - discardBufferedContent(); - stopFiltering(); + try { + super.sendError(sc); + } finally { + discardBufferedContent(); + stopFiltering(); + } } @Override public void sendError(int sc, String msg) throws IOException { - super.sendError(sc, msg); - discardBufferedContent(); - stopFiltering(); + try { + super.sendError(sc, msg); + } finally { + discardBufferedContent(); + stopFiltering(); + } } @Override public void sendRedirect(String location) throws IOException { - super.sendRedirect(location); - discardBufferedContent(); - stopFiltering(); + try { + super.sendRedirect(location); + } finally { + discardBufferedContent(); + stopFiltering(); + } + } + + public void sendRedirect(String location, boolean clearBuffer) throws IOException { + sendRedirect( + location, clearBuffer, SEND_REDIRECT_WITH_BUFFER_6_1, getResponse(), location, clearBuffer); + } + + public void sendRedirect(String location, int sc) throws IOException { + sendRedirect(location, true, SEND_REDIRECT_WITH_STATUS_6_1, getResponse(), location, sc); + } + + public void sendRedirect(String location, int sc, boolean clearBuffer) throws IOException { + sendRedirect( + location, clearBuffer, SEND_REDIRECT_6_1, getResponse(), location, sc, clearBuffer); + } + + private void sendRedirect( + String location, boolean clearBuffer, MethodHandle method, Object... arguments) + throws IOException { + try { + if (!clearBuffer) { + commit(); + } + if (method == null) { + super.sendRedirect(location); + } else { + method.invokeWithArguments(arguments); + } + } catch (Throwable t) { + sneakyThrow(t); + } finally { + if (clearBuffer) { + discardBufferedContent(); + } + stopFiltering(); + } } public void onInjected() { diff --git a/dd-java-agent/instrumentation/servlet/jakarta-servlet-5.0/src/servlet61Test/java/datadog/trace/instrumentation/servlet5/RumHttpServletResponseWrapperServlet61Test.java b/dd-java-agent/instrumentation/servlet/jakarta-servlet-5.0/src/servlet61Test/java/datadog/trace/instrumentation/servlet5/RumHttpServletResponseWrapperServlet61Test.java new file mode 100644 index 00000000000..21fc46a4071 --- /dev/null +++ b/dd-java-agent/instrumentation/servlet/jakarta-servlet-5.0/src/servlet61Test/java/datadog/trace/instrumentation/servlet5/RumHttpServletResponseWrapperServlet61Test.java @@ -0,0 +1,54 @@ +package datadog.trace.instrumentation.servlet5; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import datadog.trace.bootstrap.instrumentation.buffer.InjectingPipeWriter; +import jakarta.servlet.ServletContext; +import jakarta.servlet.http.HttpServletRequest; +import jakarta.servlet.http.HttpServletResponse; +import java.io.IOException; +import java.io.StringWriter; +import java.lang.reflect.Field; +import org.tabletest.junit.TableTest; + +class RumHttpServletResponseWrapperServlet61Test { + + @TableTest({ + "scenario | clearBuffer", + "clear buffer | true ", + "retain buffer | false " + }) + void servlet61RedirectRespectsClearBuffer(boolean clearBuffer) + throws IOException, ReflectiveOperationException { + ServletContext servletContext = mock(ServletContext.class); + when(servletContext.getEffectiveMajorVersion()).thenReturn(6); + HttpServletRequest request = mock(HttpServletRequest.class); + when(request.getServletContext()).thenReturn(servletContext); + HttpServletResponse response = mock(HttpServletResponse.class); + RumHttpServletResponseWrapper wrapper = new RumHttpServletResponseWrapper(request, response); + + StringWriter downstream = new StringWriter(); + InjectingPipeWriter pipe = + new InjectingPipeWriter( + downstream, "".toCharArray(), "".toCharArray()); + setWrappedPipeWriter(wrapper, pipe); + pipe.write("> { throw new IOException("error response failed") } + + when: + wrapper.commit() + + then: + downstream.toString().isEmpty() + } + void 'sendRedirect discards buffered content and stops filtering'() { setup: def downstream = new StringWriter() @@ -287,6 +306,25 @@ class RumHttpServletResponseWrapperTest extends InstrumentationSpecification { 0 * mockTelemetryCollector.onInjectionSucceed(_) } + void 'sendRedirect discards buffered content when the delegate throws'() { + setup: + def downstream = new StringWriter() + attachWriter(downstream).write("> { throw new IOException("redirect failed") } + + when: + wrapper.commit() + + then: + downstream.toString().isEmpty() + } + private PrintWriter attachWriter(StringWriter downstream) { def pipe = new InjectingPipeWriter( downstream, diff --git a/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-3.0/src/main/java/datadog/trace/instrumentation/servlet3/RumHttpServletResponseWrapper.java b/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-3.0/src/main/java/datadog/trace/instrumentation/servlet3/RumHttpServletResponseWrapper.java index d6dd8a3284e..2c060917070 100644 --- a/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-3.0/src/main/java/datadog/trace/instrumentation/servlet3/RumHttpServletResponseWrapper.java +++ b/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-3.0/src/main/java/datadog/trace/instrumentation/servlet3/RumHttpServletResponseWrapper.java @@ -209,23 +209,32 @@ public void flushBuffer() throws IOException { @Override public void sendError(int sc) throws IOException { - super.sendError(sc); - discardBufferedContent(); - stopFiltering(); + try { + super.sendError(sc); + } finally { + discardBufferedContent(); + stopFiltering(); + } } @Override public void sendError(int sc, String msg) throws IOException { - super.sendError(sc, msg); - discardBufferedContent(); - stopFiltering(); + try { + super.sendError(sc, msg); + } finally { + discardBufferedContent(); + stopFiltering(); + } } @Override public void sendRedirect(String location) throws IOException { - super.sendRedirect(location); - discardBufferedContent(); - stopFiltering(); + try { + super.sendRedirect(location); + } finally { + discardBufferedContent(); + stopFiltering(); + } } public void onInjected() { diff --git a/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-3.0/src/test/groovy/RumHttpServletResponseWrapperTest.groovy b/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-3.0/src/test/groovy/RumHttpServletResponseWrapperTest.groovy index a516a863264..646a92da1f7 100644 --- a/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-3.0/src/test/groovy/RumHttpServletResponseWrapperTest.groovy +++ b/dd-java-agent/instrumentation/servlet/javax-servlet/javax-servlet-3.0/src/test/groovy/RumHttpServletResponseWrapperTest.groovy @@ -270,6 +270,25 @@ class RumHttpServletResponseWrapperTest extends InstrumentationSpecification { 0 * mockTelemetryCollector.onInjectionSucceed(_) } + void 'sendError discards buffered content when the delegate throws'() { + setup: + def downstream = new StringWriter() + attachWriter(downstream).write("> { throw new IOException("error response failed") } + + when: + wrapper.commit() + + then: + downstream.toString().isEmpty() + } + void 'sendRedirect discards buffered content and stops filtering'() { setup: def downstream = new StringWriter() @@ -287,6 +306,25 @@ class RumHttpServletResponseWrapperTest extends InstrumentationSpecification { 0 * mockTelemetryCollector.onInjectionSucceed(_) } + void 'sendRedirect discards buffered content when the delegate throws'() { + setup: + def downstream = new StringWriter() + attachWriter(downstream).write("> { throw new IOException("redirect failed") } + + when: + wrapper.commit() + + then: + downstream.toString().isEmpty() + } + private PrintWriter attachWriter(StringWriter downstream) { def pipe = new InjectingPipeWriter( downstream, From c07186bfc755288b80d48c4f009e144f07fad610 Mon Sep 17 00:00:00 2001 From: Andrea Marziali Date: Tue, 15 Sep 2026 15:56:44 +0200 Subject: [PATCH 4/4] fix codex findings --- .../RumHttpServletResponseWrapper.java | 42 +++++++++++++++---- ...tpServletResponseWrapperServlet61Test.java | 23 +++++++--- .../RumHttpServletResponseWrapperTest.groovy | 19 +++++++++ .../RumHttpServletResponseWrapper.java | 30 ++++++++++--- .../RumHttpServletResponseWrapperTest.groovy | 19 +++++++++ 5 files changed, 112 insertions(+), 21 deletions(-) diff --git a/dd-java-agent/instrumentation/servlet/jakarta-servlet-5.0/src/main/java/datadog/trace/instrumentation/servlet5/RumHttpServletResponseWrapper.java b/dd-java-agent/instrumentation/servlet/jakarta-servlet-5.0/src/main/java/datadog/trace/instrumentation/servlet5/RumHttpServletResponseWrapper.java index 0d1d805b243..9d779386da0 100644 --- a/dd-java-agent/instrumentation/servlet/jakarta-servlet-5.0/src/main/java/datadog/trace/instrumentation/servlet5/RumHttpServletResponseWrapper.java +++ b/dd-java-agent/instrumentation/servlet/jakarta-servlet-5.0/src/main/java/datadog/trace/instrumentation/servlet5/RumHttpServletResponseWrapper.java @@ -209,31 +209,49 @@ public void flushBuffer() throws IOException { @Override public void sendError(int sc) throws IOException { + boolean rejected = false; try { super.sendError(sc); + } catch (IllegalStateException e) { + rejected = true; + throw e; } finally { - discardBufferedContent(); - stopFiltering(); + if (!rejected) { + discardBufferedContent(); + stopFiltering(); + } } } @Override public void sendError(int sc, String msg) throws IOException { + boolean rejected = false; try { super.sendError(sc, msg); + } catch (IllegalStateException e) { + rejected = true; + throw e; } finally { - discardBufferedContent(); - stopFiltering(); + if (!rejected) { + discardBufferedContent(); + stopFiltering(); + } } } @Override public void sendRedirect(String location) throws IOException { + boolean rejected = false; try { super.sendRedirect(location); + } catch (IllegalStateException e) { + rejected = true; + throw e; } finally { - discardBufferedContent(); - stopFiltering(); + if (!rejected) { + discardBufferedContent(); + stopFiltering(); + } } } @@ -254,6 +272,7 @@ public void sendRedirect(String location, int sc, boolean clearBuffer) throws IO private void sendRedirect( String location, boolean clearBuffer, MethodHandle method, Object... arguments) throws IOException { + boolean rejected = false; try { if (!clearBuffer) { commit(); @@ -263,13 +282,18 @@ private void sendRedirect( } else { method.invokeWithArguments(arguments); } + } catch (IllegalStateException e) { + rejected = true; + throw e; } catch (Throwable t) { sneakyThrow(t); } finally { - if (clearBuffer) { - discardBufferedContent(); + if (!rejected) { + if (clearBuffer) { + discardBufferedContent(); + } + stopFiltering(); } - stopFiltering(); } } diff --git a/dd-java-agent/instrumentation/servlet/jakarta-servlet-5.0/src/servlet61Test/java/datadog/trace/instrumentation/servlet5/RumHttpServletResponseWrapperServlet61Test.java b/dd-java-agent/instrumentation/servlet/jakarta-servlet-5.0/src/servlet61Test/java/datadog/trace/instrumentation/servlet5/RumHttpServletResponseWrapperServlet61Test.java index 21fc46a4071..6f6a8179b26 100644 --- a/dd-java-agent/instrumentation/servlet/jakarta-servlet-5.0/src/servlet61Test/java/datadog/trace/instrumentation/servlet5/RumHttpServletResponseWrapperServlet61Test.java +++ b/dd-java-agent/instrumentation/servlet/jakarta-servlet-5.0/src/servlet61Test/java/datadog/trace/instrumentation/servlet5/RumHttpServletResponseWrapperServlet61Test.java @@ -1,6 +1,8 @@ package datadog.trace.instrumentation.servlet5; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -17,11 +19,12 @@ class RumHttpServletResponseWrapperServlet61Test { @TableTest({ - "scenario | clearBuffer", - "clear buffer | true ", - "retain buffer | false " + "scenario | clearBuffer | rejected", + "clear buffer | true | false ", + "retain buffer | false | false ", + "rejected call | true | true " }) - void servlet61RedirectRespectsClearBuffer(boolean clearBuffer) + void servlet61RedirectRespectsClearBuffer(boolean clearBuffer, boolean rejected) throws IOException, ReflectiveOperationException { ServletContext servletContext = mock(ServletContext.class); when(servletContext.getEffectiveMajorVersion()).thenReturn(6); @@ -37,10 +40,18 @@ void servlet61RedirectRespectsClearBuffer(boolean clearBuffer) setWrappedPipeWriter(wrapper, pipe); pipe.write(" wrapper.sendRedirect("/other", 307, clearBuffer)); + } else { + wrapper.sendRedirect("/other", 307, clearBuffer); + } wrapper.commit(); - assertEquals(clearBuffer ? "" : "> { throw new IllegalStateException("already committed") } + + when: + wrapper.commit() + + then: + downstream.toString() == "> { throw new IllegalStateException("already committed") } + + when: + wrapper.commit() + + then: + downstream.toString() == "