snippet inject support like <head lang="en"> (open-telemetry#8736)

Co-authored-by: Mateusz Rzeszutek <[email protected]> Co-authored-by: Lauri Tulmin <[email protected]>
breedx-splk · Aug 15, 2023 · faf59b6 · faf59b6
1 parent 6c21a09
commit faf59b6
Show file tree

Hide file tree

Showing 11 changed files with 121 additions and 6 deletions.
diff --git a/.../opentelemetry/javaagent/instrumentation/servlet/v3_0/snippet/SnippetPrintWriterTest.java b/.../opentelemetry/javaagent/instrumentation/servlet/v3_0/snippet/SnippetPrintWriterTest.java
@@ -124,6 +124,22 @@ void testWriteWithOffset() throws IOException {
     assertThat(response.getStringContent()).isEqualTo(expectedHtml);
   }
 
+  @Test
+  void testInjectToTextHtmlWithOtherHeadStyle() throws IOException {
+    String snippet = "\n  <script type=\"text/javascript\"> Test </script>";
+    String html = readFileAsString("beforeSnippetInjectionWithOtherHeadStyle.html");
+
+    InMemoryHttpServletResponse response = createInMemoryHttpServletResponse("text/html");
+    Servlet3SnippetInjectingResponseWrapper responseWrapper =
+        new Servlet3SnippetInjectingResponseWrapper(response, snippet);
+
+    responseWrapper.getWriter().write(html);
+    responseWrapper.getWriter().flush();
+
+    String expectedHtml = readFileAsString("afterSnippetInjectionWithOtherHeadStyle.html");
+    assertThat(response.getStringContent()).isEqualTo(expectedHtml);
+  }
+
   private static InMemoryHttpServletResponse createInMemoryHttpServletResponse(String contentType) {
     HttpServletResponse response = mock(HttpServletResponse.class);
     when(response.getContentType()).thenReturn(contentType);

diff --git a/...emetry/javaagent/instrumentation/servlet/v3_0/snippet/SnippetServletOutputStreamTest.java b/...emetry/javaagent/instrumentation/servlet/v3_0/snippet/SnippetServletOutputStreamTest.java
@@ -125,6 +125,25 @@ void testHeadTagSplitAcrossTwoWrites() throws IOException {
     assertThat(out.getBytes()).isEqualTo(expectedSecondPart.getBytes(UTF_8));
   }
 
+  @Test
+  void testInjectionWithOtherHeadStyle() throws IOException {
+    String snippet = "\n  <script type=\"text/javascript\"> Test </script>";
+    byte[] html = readFileAsBytes("beforeSnippetInjectionWithOtherHeadStyle.html");
+
+    InjectionState obj = createInjectionStateForTesting(snippet, UTF_8);
+    InMemoryServletOutputStream out = new InMemoryServletOutputStream();
+
+    Supplier<String> stringSupplier = snippet::toString;
+    OutputStreamSnippetInjectionHelper helper =
+        new OutputStreamSnippetInjectionHelper(stringSupplier);
+    boolean injected = helper.handleWrite(obj, out, html, 0, html.length);
+    assertThat(obj.getHeadTagBytesSeen()).isEqualTo(-1);
+    assertThat(injected).isEqualTo(true);
+
+    byte[] expectedHtml = readFileAsBytes("afterSnippetInjectionWithOtherHeadStyle.html");
+    assertThat(out.getBytes()).isEqualTo(expectedHtml);
+  }
+
   private static InjectionState createInjectionStateForTesting(String snippet, Charset charset) {
     HttpServletResponse response = mock(HttpServletResponse.class);
     when(response.isCommitted()).thenReturn(false);

diff --git a/...-3.0/javaagent-unit-tests/src/test/resources/afterSnippetInjectionWithOtherHeadStyle.html b/...-3.0/javaagent-unit-tests/src/test/resources/afterSnippetInjectionWithOtherHeadStyle.html
@@ -0,0 +1,11 @@
+<!DOCTYPE html>
+<html lang="en">
+<head lang="en">
+  <script type="text/javascript"> Test </script>
+  <meta charset="UTF-8">
+  <title>Title</title>
+</head>
+<body>
+
+</body>
+</html>
diff --git a/...3.0/javaagent-unit-tests/src/test/resources/beforeSnippetInjectionWithOtherHeadStyle.html b/...3.0/javaagent-unit-tests/src/test/resources/beforeSnippetInjectionWithOtherHeadStyle.html
@@ -0,0 +1,10 @@
+<!DOCTYPE html>
+<html lang="en">
+<head lang="en">
+  <meta charset="UTF-8">
+  <title>Title</title>
+</head>
+<body>
+
+</body>
+</html>
diff --git a/...vaagent/instrumentation/servlet/v3_0/snippet/Servlet3SnippetInjectingResponseWrapper.java b/...vaagent/instrumentation/servlet/v3_0/snippet/Servlet3SnippetInjectingResponseWrapper.java
@@ -150,7 +150,8 @@ public boolean isContentTypeTextHtml() {
     if (contentType == null) {
       contentType = super.getHeader("content-type");
     }
-    return contentType != null && contentType.startsWith("text/html");
+    return contentType != null
+        && (contentType.startsWith("text/html") || contentType.startsWith("application/xhtml+xml"));
   }
 
   @Override

diff --git a/.../opentelemetry/javaagent/instrumentation/servlet/v5_0/snippet/SnippetPrintWriterTest.java b/.../opentelemetry/javaagent/instrumentation/servlet/v5_0/snippet/SnippetPrintWriterTest.java
@@ -123,6 +123,22 @@ void testWriteWithOffset() throws IOException {
     assertThat(response.getStringContent()).isEqualTo(expectedHtml);
   }
 
+  @Test
+  void testInjectToTextHtmlWithOtherHeadStyle() throws IOException {
+    String snippet = "\n  <script type=\"text/javascript\"> Test </script>";
+    String html = TestUtil.readFileAsString("beforeSnippetInjectionWithOtherHeadStyle.html");
+
+    InMemoryHttpServletResponse response = createInMemoryHttpServletResponse("text/html");
+    Servlet5SnippetInjectingResponseWrapper responseWrapper =
+        new Servlet5SnippetInjectingResponseWrapper(response, snippet);
+
+    responseWrapper.getWriter().write(html);
+    responseWrapper.getWriter().flush();
+
+    String expectedHtml = TestUtil.readFileAsString("afterSnippetInjectionWithOtherHeadStyle.html");
+    assertThat(response.getStringContent()).isEqualTo(expectedHtml);
+  }
+
   private static InMemoryHttpServletResponse createInMemoryHttpServletResponse(String contentType) {
     HttpServletResponse response = mock(HttpServletResponse.class);
     when(response.getContentType()).thenReturn(contentType);

diff --git a/...emetry/javaagent/instrumentation/servlet/v5_0/snippet/SnippetServletOutputStreamTest.java b/...emetry/javaagent/instrumentation/servlet/v5_0/snippet/SnippetServletOutputStreamTest.java
@@ -127,6 +127,25 @@ void testHeadTagSplitAcrossTwoWrites() throws IOException {
     assertThat(out.getBytes()).isEqualTo(expectedSecondPart.getBytes(UTF_8));
   }
 
+  @Test
+  void testInjectionWithOtherHeadStyle() throws IOException {
+    String snippet = "\n  <script type=\"text/javascript\"> Test </script>";
+    byte[] html = readFileAsBytes("beforeSnippetInjectionWithOtherHeadStyle.html");
+
+    InjectionState obj = createInjectionStateForTesting(snippet, UTF_8);
+    InMemoryServletOutputStream out = new InMemoryServletOutputStream();
+
+    Supplier<String> stringSupplier = snippet::toString;
+    OutputStreamSnippetInjectionHelper helper =
+        new OutputStreamSnippetInjectionHelper(stringSupplier);
+    boolean injected = helper.handleWrite(obj, out, html, 0, html.length);
+    assertThat(obj.getHeadTagBytesSeen()).isEqualTo(-1);
+    assertThat(injected).isEqualTo(true);
+
+    byte[] expectedHtml = readFileAsBytes("afterSnippetInjectionWithOtherHeadStyle.html");
+    assertThat(out.getBytes()).isEqualTo(expectedHtml);
+  }
+
   private static InjectionState createInjectionStateForTesting(String snippet, Charset charset) {
     HttpServletResponse response = mock(HttpServletResponse.class);
     when(response.isCommitted()).thenReturn(false);

diff --git a/...-5.0/javaagent-unit-tests/src/test/resources/afterSnippetInjectionWithOtherHeadStyle.html b/...-5.0/javaagent-unit-tests/src/test/resources/afterSnippetInjectionWithOtherHeadStyle.html
@@ -0,0 +1,11 @@
+<!DOCTYPE html>
+<html lang="en">
+<head lang="en">
+  <script type="text/javascript"> Test </script>
+  <meta charset="UTF-8">
+  <title>Title</title>
+</head>
+<body>
+
+</body>
+</html>
diff --git a/...5.0/javaagent-unit-tests/src/test/resources/beforeSnippetInjectionWithOtherHeadStyle.html b/...5.0/javaagent-unit-tests/src/test/resources/beforeSnippetInjectionWithOtherHeadStyle.html
@@ -0,0 +1,10 @@
+<!DOCTYPE html>
+<html lang="en">
+<head lang="en">
+  <meta charset="UTF-8">
+  <title>Title</title>
+</head>
+<body>
+
+</body>
+</html>
diff --git a/...vaagent/instrumentation/servlet/v5_0/snippet/Servlet5SnippetInjectingResponseWrapper.java b/...vaagent/instrumentation/servlet/v5_0/snippet/Servlet5SnippetInjectingResponseWrapper.java
@@ -142,7 +142,8 @@ public boolean isContentTypeTextHtml() {
     if (contentType == null) {
       contentType = super.getHeader("content-type");
     }
-    return contentType != null && contentType.startsWith("text/html");
+    return contentType != null
+        && (contentType.startsWith("text/html") || contentType.startsWith("application/xhtml+xml"));
   }
 
   @Override

diff --git a/.../bootstrap/src/main/java/io/opentelemetry/javaagent/bootstrap/servlet/InjectionState.java b/.../bootstrap/src/main/java/io/opentelemetry/javaagent/bootstrap/servlet/InjectionState.java
@@ -8,7 +8,7 @@
 // this is shared by both ServletOutputStream and PrintWriter injection
 public class InjectionState {
   private static final int HEAD_TAG_WRITTEN_FAKE_VALUE = -1;
-  private static final int HEAD_TAG_LENGTH = "<head>".length();
+  private static final int HEAD_TAG_PREFIX_LENGTH = "<head".length();
   private final SnippetInjectingResponseWrapper wrapper;
   private int headTagBytesSeen = 0;
 
@@ -45,7 +45,7 @@ public boolean processByte(int b) {
     } else {
       headTagBytesSeen = 0;
     }
-    if (headTagBytesSeen == HEAD_TAG_LENGTH) {
+    if (headTagBytesSeen > HEAD_TAG_PREFIX_LENGTH && b == '>') {
       setHeadTagWritten();
       return true;
     } else {
@@ -64,10 +64,11 @@ private boolean inHeadTag(int b) {
       return true;
     } else if (headTagBytesSeen == 4 && b == 'd') {
       return true;
-    } else if (headTagBytesSeen == 5 && b == '>') {
+    } else if (headTagBytesSeen == 5 && (b == '>' || Character.isWhitespace(b))) {
       return true;
+    } else {
+      return headTagBytesSeen > 5;
     }
-    return false;
   }
 
   public SnippetInjectingResponseWrapper getWrapper() {