Rum injection: filter setContentType headers

amarziali · amarziali · commit baac90500332 · 2025-12-10T11:40:16.000+01:00
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
@@ -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
@@ -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
@@ -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
@@ -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()
   }
 
@@ -80,6 +101,20 @@ class RumHttpServletResponseWrapperTest extends InstrumentationSpecification {
     1 * mockTelemetryCollector.onInjectionFailed(SERVLET_VERSION, null)
   }
 
+  void 'getOutputStream exception reports failure'() {
+    setup:
+    wrapper.setContentType("text/html")
+    mockResponse.getOutputStream() >> { throw new IOException("stream error") }
+
+    when:
+    try {
+      wrapper.getOutputStream()
+    } catch (IOException ignored) {}
+
+    then:
+    1 * mockTelemetryCollector.onInjectionFailed(SERVLET_VERSION, null)
+  }
+
   void 'getWriter exception reports failure'() {
     setup:
     wrapper.setContentType("text/html")

Original file line number	Diff line number	Diff line change
`@@ -107,6 +107,7 @@ public void setHeader(String name, String value) {`
`107`	`107`	`if (isContentLengthHeader(name)) {`
`108`	`108`	`return;`
`109`	`109`	`}`
	`110`	`+ checkForContentType(name, value);`
`110`	`111`	`checkForContentSecurityPolicy(name);`
`111`	`112`	`}`
`112`	`113`	`super.setHeader(name, value);`
`@@ -118,6 +119,7 @@ public void addHeader(String name, String value) {`
`118`	`119`	`if (isContentLengthHeader(name)) {`
`119`	`120`	`return;`
`120`	`121`	`}`
	`122`	`+ checkForContentType(name, value);`
`121`	`123`	`checkForContentSecurityPolicy(name);`
`122`	`124`	`}`
`123`	`125`	`super.addHeader(name, value);`
`@@ -133,6 +135,12 @@ private void checkForContentSecurityPolicy(String name) {`
`133`	`135`	`}`
`134`	`136`	`}`
`135`	`137`
	`138`	`+ private void checkForContentType(String name, String value) {`
	`139`	`+ if ("content-type".equalsIgnoreCase(name)) {`
	`140`	`+ handleContentType(value);`
	`141`	`+ }`
	`142`	`+ }`
	`143`	`+`
`136`	`144`	`@Override`
`137`	`145`	`public void setContentLength(int len) {`
`138`	`146`	`// don't set it since we don't know if we will inject`
`@@ -182,15 +190,20 @@ public void onInjected() {`
`182`	`190`	`}`
`183`	`191`	`}`
`184`	`192`
`185`		`- @Override`
`186`		`- public void setContentType(String type) {`
	`193`	`+ private void handleContentType(String type) {`
	`194`	`+ final boolean wasInjecting = shouldInject;`
`187`	`195`	`if (shouldInject) {`
`188`	`196`	`shouldInject = type != null && type.contains("text/html");`
`189`	`197`	`}`
`190`		`- if (!shouldInject) {`
	`198`	`+ if (wasInjecting && !shouldInject) {`
`191`	`199`	`commit();`
`192`	`200`	`stopFiltering();`
`193`	`201`	`}`
	`202`	`+ }`
	`203`	`+`
	`204`	`+ @Override`
	`205`	`+ public void setContentType(String type) {`
	`206`	`+ handleContentType(type);`
`194`	`207`	`super.setContentType(type);`
`195`	`208`	`}`
`196`	`209`
Original file line number	Diff line number	Diff line change
`@@ -126,6 +126,7 @@ public void setHeader(String name, String value) {`
`126`	`126`	`if (isContentLengthHeader(name)) {`
`127`	`127`	`return;`
`128`	`128`	`}`
	`129`	`+ checkForContentType(name, value);`
`129`	`130`	`checkForContentSecurityPolicy(name);`
`130`	`131`	`}`
`131`	`132`	`super.setHeader(name, value);`
`@@ -137,6 +138,7 @@ public void addHeader(String name, String value) {`
`137`	`138`	`if (isContentLengthHeader(name)) {`
`138`	`139`	`return;`
`139`	`140`	`}`
	`141`	`+ checkForContentType(name, value);`
`140`	`142`	`checkForContentSecurityPolicy(name);`
`141`	`143`	`}`
`142`	`144`	`super.addHeader(name, value);`
`@@ -205,15 +207,26 @@ public void onInjected() {`
`205`	`207`	`}`
`206`	`208`	`}`
`207`	`209`
`208`		`- @Override`
`209`		`- public void setContentType(String type) {`
	`210`	`+ private void handleContentType(String type) {`
	`211`	`+ final boolean wasInjecting = shouldInject;`
`210`	`212`	`if (shouldInject) {`
`211`	`213`	`shouldInject = type != null && type.contains("text/html");`
`212`	`214`	`}`
`213`		`- if (!shouldInject) {`
	`215`	`+ if (wasInjecting && !shouldInject) {`
`214`	`216`	`commit();`
`215`	`217`	`stopFiltering();`
`216`	`218`	`}`
	`219`	`+ }`
	`220`	`+`
	`221`	`+ private void checkForContentType(String name, String value) {`
	`222`	`+ if ("content-type".equalsIgnoreCase(name)) {`
	`223`	`+ handleContentType(value);`
	`224`	`+ }`
	`225`	`+ }`
	`226`	`+`
	`227`	`+ @Override`
	`228`	`+ public void setContentType(String type) {`
	`229`	`+ handleContentType(type);`
`217`	`230`	`super.setContentType(type);`
`218`	`231`	`}`
`219`	`232`