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 77114e139f2..25b3f2c9964 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 @@ -107,6 +107,7 @@ public void setHeader(String name, String value) { if (isContentLengthHeader(name)) { return; } + checkForContentType(name, value); checkForContentSecurityPolicy(name); } super.setHeader(name, value); @@ -118,6 +119,7 @@ public void addHeader(String name, String value) { if (isContentLengthHeader(name)) { return; } + checkForContentType(name, value); checkForContentSecurityPolicy(name); } super.addHeader(name, value); @@ -133,6 +135,12 @@ private void checkForContentSecurityPolicy(String name) { } } + private void checkForContentType(String name, String value) { + if ("content-type".equalsIgnoreCase(name)) { + handleContentType(value); + } + } + @Override public void setContentLength(int len) { // don't set it since we don't know if we will inject @@ -182,15 +190,20 @@ public void onInjected() { } } - @Override - public void setContentType(String type) { + private void handleContentType(String type) { + final boolean wasInjecting = shouldInject; if (shouldInject) { shouldInject = type != null && type.contains("text/html"); } - if (!shouldInject) { + if (wasInjecting && !shouldInject) { commit(); stopFiltering(); } + } + + @Override + public void setContentType(String type) { + handleContentType(type); super.setContentType(type); } 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 f2d2f049a95..595a0909f10 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 @@ -54,15 +54,36 @@ class RumHttpServletResponseWrapperTest extends InstrumentationSpecification { 1 * mockResponse.getOutputStream() } - void 'getWriter with non-HTML content reports skipped'() { - setup: + void 'getWriter with non-HTML content reports skipped (setContentType)'() { + when: wrapper.setContentType("text/plain") + wrapper.getWriter() + + then: + 1 * mockTelemetryCollector.onInjectionSkipped(SERVLET_VERSION) + 1 * mockResponse.setContentType("text/plain") + 1 * mockResponse.getWriter() + } + + void 'getWriter with non-HTML content reports skipped (setHeader)'() { + when: + wrapper.setHeader("Content-Type", "text/plain") + wrapper.getWriter() + + then: + 1 * mockTelemetryCollector.onInjectionSkipped(SERVLET_VERSION) + 1 * mockResponse.setHeader("Content-Type", "text/plain") + 1 * mockResponse.getWriter() + } + void 'getWriter with non-HTML content reports skipped (addHeader)'() { when: + wrapper.addHeader("Content-Type", "text/plain") wrapper.getWriter() then: 1 * mockTelemetryCollector.onInjectionSkipped(SERVLET_VERSION) + 1 * mockResponse.addHeader("Content-Type", "text/plain") 1 * mockResponse.getWriter() } 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 326a19007fa..cbd0b6f845b 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 @@ -126,6 +126,7 @@ public void setHeader(String name, String value) { if (isContentLengthHeader(name)) { return; } + checkForContentType(name, value); checkForContentSecurityPolicy(name); } super.setHeader(name, value); @@ -137,6 +138,7 @@ public void addHeader(String name, String value) { if (isContentLengthHeader(name)) { return; } + checkForContentType(name, value); checkForContentSecurityPolicy(name); } super.addHeader(name, value); @@ -205,15 +207,26 @@ public void onInjected() { } } - @Override - public void setContentType(String type) { + private void handleContentType(String type) { + final boolean wasInjecting = shouldInject; if (shouldInject) { shouldInject = type != null && type.contains("text/html"); } - if (!shouldInject) { + if (wasInjecting && !shouldInject) { commit(); stopFiltering(); } + } + + private void checkForContentType(String name, String value) { + if ("content-type".equalsIgnoreCase(name)) { + handleContentType(value); + } + } + + @Override + public void setContentType(String type) { + handleContentType(type); super.setContentType(type); } 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 e8ec2deaf42..341ae9cde5e 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 @@ -54,15 +54,36 @@ class RumHttpServletResponseWrapperTest extends InstrumentationSpecification { 1 * mockResponse.getOutputStream() } - void 'getWriter with non-HTML content reports skipped'() { - setup: + void 'getWriter with non-HTML content reports skipped (setContentType)'() { + when: wrapper.setContentType("text/plain") + wrapper.getWriter() + + then: + 1 * mockTelemetryCollector.onInjectionSkipped(SERVLET_VERSION) + 1 * mockResponse.setContentType("text/plain") + 1 * mockResponse.getWriter() + } + + void 'getWriter with non-HTML content reports skipped (setHeader)'() { + when: + wrapper.setHeader("Content-Type", "text/plain") + wrapper.getWriter() + + then: + 1 * mockTelemetryCollector.onInjectionSkipped(SERVLET_VERSION) + 1 * mockResponse.setHeader("Content-Type", "text/plain") + 1 * mockResponse.getWriter() + } + void 'getWriter with non-HTML content reports skipped (addHeader)'() { when: + wrapper.addHeader("Content-Type", "text/plain") wrapper.getWriter() then: 1 * mockTelemetryCollector.onInjectionSkipped(SERVLET_VERSION) + 1 * mockResponse.addHeader("Content-Type", "text/plain") 1 * mockResponse.getWriter() }