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/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 25b3f2c9964..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 @@ -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 @@ -20,8 +22,24 @@ public class RumHttpServletResponseWrapper extends HttpServletResponseWrapper private InjectingPipeWriter wrappedPipeWriter; private PrintWriter printWriter; private boolean shouldInject = true; + 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(); @@ -166,19 +184,117 @@ 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 = !retired; + 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 { + boolean rejected = false; + try { + super.sendError(sc); + } catch (IllegalStateException e) { + rejected = true; + throw e; + } finally { + 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 { + 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 { + if (!rejected) { + 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 { + boolean rejected = false; + try { + if (!clearBuffer) { + commit(); + } + if (method == null) { + super.sendRedirect(location); + } else { + method.invokeWithArguments(arguments); + } + } catch (IllegalStateException e) { + rejected = true; + throw e; + } catch (Throwable t) { + sneakyThrow(t); + } finally { + if (!rejected) { + if (clearBuffer) { + discardBufferedContent(); + } + stopFiltering(); + } + } } public void onInjected() { @@ -197,7 +313,7 @@ private void handleContentType(String type) { } if (wasInjecting && !shouldInject) { commit(); - stopFiltering(); + disableFiltering(); } } @@ -225,12 +341,39 @@ public void commit() { @Override public void stopFiltering() { + retired = true; + disableFiltering(); + } + + private void disableFiltering() { 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/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..6f6a8179b26 --- /dev/null +++ b/dd-java-agent/instrumentation/servlet/jakarta-servlet-5.0/src/servlet61Test/java/datadog/trace/instrumentation/servlet5/RumHttpServletResponseWrapperServlet61Test.java @@ -0,0 +1,65 @@ +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; + +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 | rejected", + "clear buffer | true | false ", + "retain buffer | false | false ", + "rejected call | true | true " + }) + void servlet61RedirectRespectsClearBuffer(boolean clearBuffer, boolean rejected) + 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(" wrapper.sendRedirect("/other", 307, clearBuffer)); + } else { + wrapper.sendRedirect("/other", 307, clearBuffer); + } + wrapper.commit(); + + assertEquals(clearBuffer && !rejected ? "" : "") + 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 '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() + def writer = attachWriter(downstream) + + when: + writer.write("") + wrapper.commit() + + then: + downstream.toString() == "ad>" + 1 * mockResponse.sendError(500) + 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 'sendError preserves buffered content when the delegate rejects the call'() { + setup: + def downstream = new StringWriter() + attachWriter(downstream).write("> { throw new IllegalStateException("already committed") } + + when: + wrapper.commit() + + then: + downstream.toString() == "") + wrapper.commit() + + then: + downstream.toString() == "ad>" + 1 * mockResponse.sendRedirect("/other") + 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, + "".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..4339dfd5fb4 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"); @@ -183,19 +184,75 @@ 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 = !retired; + 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 { + boolean rejected = false; + try { + super.sendError(sc); + } catch (IllegalStateException e) { + rejected = true; + throw e; + } finally { + 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 { + 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 { + if (!rejected) { + discardBufferedContent(); + stopFiltering(); + } + } } public void onInjected() { @@ -214,7 +271,7 @@ private void handleContentType(String type) { } if (wasInjecting && !shouldInject) { commit(); - stopFiltering(); + disableFiltering(); } } @@ -248,12 +305,39 @@ public void commit() { @Override public void stopFiltering() { + retired = true; + disableFiltering(); + } + + private void disableFiltering() { 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..cc5186e9e95 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,183 @@ 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 '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() + def writer = attachWriter(downstream) + + when: + writer.write("") + wrapper.commit() + + then: + downstream.toString() == "ad>" + 1 * mockResponse.sendError(500) + 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 'sendError preserves buffered content when the delegate rejects the call'() { + setup: + def downstream = new StringWriter() + attachWriter(downstream).write("> { throw new IllegalStateException("already committed") } + + when: + wrapper.commit() + + then: + downstream.toString() == "") + wrapper.commit() + + then: + downstream.toString() == "ad>" + 1 * mockResponse.sendRedirect("/other") + 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, + "".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'() {