JavaScript Snippet Injection

instrumentation/servlet/servlet-3.0/javaagent-unit-tests/build.gradle.kts

@@ -0,0 +1,9 @@
+plugins {
+  id("otel.java-conventions")
+}
+
+dependencies {
+  testImplementation("javax.servlet:javax.servlet-api:3.0.1")
+  testImplementation(project(":instrumentation:servlet:servlet-common:bootstrap"))
+  testImplementation(project(":instrumentation:servlet:servlet-3.0:javaagent"))
+}

instrumentation/servlet/servlet-3.0/javaagent-unit-tests/src/test/java/io/opentelemetry/javaagent/instrumentation/servlet/v3_0/snippet/InjectionTest.java

@@ -0,0 +1,184 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.servlet.v3_0.snippet;
+
+import static io.opentelemetry.javaagent.instrumentation.servlet.v3_0.snippet.TestUtil.readFile;
+import static org.assertj.core.api.AssertionsForClassTypes.assertThat;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import io.opentelemetry.javaagent.bootstrap.servlet.ExperimentalSnippetHolder;
+import java.io.IOException;
+import java.io.StringWriter;
+import java.nio.charset.StandardCharsets;
+import javax.servlet.ServletOutputStream;
+import org.junit.jupiter.api.Disabled;
+import org.junit.jupiter.api.Test;
+
+class InjectionTest {
+
+  @Test
+  void testInjectionForStringContainHeadTag() throws IOException {
+    String testSnippet = "\n  <script type=\"text/javascript\"> Test </script>";
+    ExperimentalSnippetHolder.setSnippet(testSnippet);
+    // read the originalFile
+    String original = readFile("staticHtmlOrigin.html");
+    // read the correct answer
+    String correct = readFile("staticHtmlAfter.html");
+    byte[] originalBytes = original.getBytes(StandardCharsets.UTF_8);
+    SnippetInjectingResponseWrapper response = mock(SnippetInjectingResponseWrapper.class);
+    when(response.isCommitted()).thenReturn(false);
+    when(response.getCharacterEncoding()).thenReturn(StandardCharsets.UTF_8.name());
+    InjectionState obj = new InjectionState(response);
+
+    StringWriter writer = new StringWriter();
+
+    ServletOutputStream sp =
+        new ServletOutputStream() {
+          @Override
+          public void write(int b) throws IOException {
+            writer.write(b);
+          }
+        };
+    boolean injected =
+        ServletOutputStreamInjectionHelper.handleWrite(
+            originalBytes, 0, originalBytes.length, obj, sp);
+    assertThat(obj.getHeadTagBytesSeen()).isEqualTo(-1);
+    assertThat(injected).isEqualTo(true);
+    writer.flush();
+
+    String result = writer.toString();
+    writer.close();
+    assertThat(result).isEqualTo(correct);
+  }
+
+  @Test
+  @Disabled
+  void testInjectionForChinese() throws IOException {
+    String testSnippet = "\n  <script type=\"text/javascript\"> Test </script>";
+    ExperimentalSnippetHolder.setSnippet(testSnippet);
+    // read the originalFile
+    String original = readFile("staticHtmlChineseOrigin.html");
+    // read the correct answer
+    String correct = readFile("staticHtmlChineseAfter.html");
+    byte[] originalBytes = original.getBytes(StandardCharsets.UTF_8);
+    SnippetInjectingResponseWrapper response = mock(SnippetInjectingResponseWrapper.class);
+    when(response.isCommitted()).thenReturn(false);
+    when(response.getCharacterEncoding()).thenReturn(StandardCharsets.UTF_8.name());
+    InjectionState obj = new InjectionState(response);
+
+    StringWriter writer = new StringWriter();
+
+    ServletOutputStream sp =
+        new ServletOutputStream() {
+          @Override
+          public void write(int b) throws IOException {
+            writer.write(b);
+          }
+        };
+    boolean injected =
+        ServletOutputStreamInjectionHelper.handleWrite(
+            originalBytes, 0, originalBytes.length, obj, sp);
+    assertThat(obj.getHeadTagBytesSeen()).isEqualTo(-1);
+    assertThat(injected).isEqualTo(true);
+    writer.flush();
+
+    String result = writer.toString();
+    writer.close();
+    assertThat(result).isEqualTo(correct);
+  }
+
+  @Test
+  void testInjectionForStringWithoutHeadTag() throws IOException {
+    String testSnippet = "\n  <script type=\"text/javascript\"> Test </script>";
+    ExperimentalSnippetHolder.setSnippet(testSnippet);
+    // read the originalFile
+    String original = readFile("htmlWithoutHeadTag.html");
+
+    byte[] originalBytes = original.getBytes(StandardCharsets.UTF_8);
+    SnippetInjectingResponseWrapper response = mock(SnippetInjectingResponseWrapper.class);
+    when(response.isCommitted()).thenReturn(false);
+    when(response.getCharacterEncoding()).thenReturn(StandardCharsets.UTF_8.name());
+    InjectionState obj = new InjectionState(response);
+    StringWriter writer = new StringWriter();
+
+    ServletOutputStream sp =
+        new ServletOutputStream() {
+          @Override
+          public void write(int b) throws IOException {
+            writer.write(b);
+          }
+        };
+    boolean injected =
+        ServletOutputStreamInjectionHelper.handleWrite(
+            originalBytes, 0, originalBytes.length, obj, sp);
+    assertThat(obj.getHeadTagBytesSeen()).isEqualTo(0);
+    assertThat(injected).isEqualTo(false);
+    writer.flush();
+    String result = writer.toString();
+    writer.close();
+    assertThat(result).isEqualTo("");
+  }
+
+  @Test
+  void testHalfHeadTag() throws IOException {
+    String testSnippet = "\n  <script type=\"text/javascript\"> Test </script>";
+    ExperimentalSnippetHolder.setSnippet(testSnippet);
+    // read the original string
+    String originalFirstPart = "<!DOCTYPE html>\n" + "<html lang=\"en\">\n" + "<he";
+    byte[] originalFirstPartBytes = originalFirstPart.getBytes(StandardCharsets.UTF_8);
+    SnippetInjectingResponseWrapper response = mock(SnippetInjectingResponseWrapper.class);
+    when(response.isCommitted()).thenReturn(false);
+    when(response.getCharacterEncoding()).thenReturn(StandardCharsets.UTF_8.name());
+    InjectionState obj = new InjectionState(response);
+    StringWriter writer = new StringWriter();
+
+    ServletOutputStream sp =
+        new ServletOutputStream() {
+          @Override
+          public void write(int b) throws IOException {
+            writer.write(b);
+          }
+        };
+    boolean injected =
+        ServletOutputStreamInjectionHelper.handleWrite(
+            originalFirstPartBytes, 0, originalFirstPartBytes.length, obj, sp);
+
+    writer.flush();
+    String result = writer.toString();
+    assertThat(obj.getHeadTagBytesSeen()).isEqualTo(3);
+    assertThat(result).isEqualTo("");
+    assertThat(injected).isEqualTo(false);
+    String originalSecondPart =
+        "ad>\n"
+            + "  <meta charset=\"UTF-8\">\n"
+            + "  <title>Title</title>\n"
+            + "</head>\n"
+            + "<body>\n"
+            + "\n"
+            + "</body>\n"
+            + "</html>";
+    byte[] originalSecondPartBytes = originalSecondPart.getBytes(StandardCharsets.UTF_8);
+    injected =
+        ServletOutputStreamInjectionHelper.handleWrite(
+            originalSecondPartBytes, 0, originalSecondPartBytes.length, obj, sp);
+    assertThat(obj.getHeadTagBytesSeen()).isEqualTo(-1);
+    assertThat(injected).isEqualTo(true);
+    String correctSecondPart =
+        "ad>\n"
+            + "  <script type=\"text/javascript\"> Test </script>\n"
+            + "  <meta charset=\"UTF-8\">\n"
+            + "  <title>Title</title>\n"
+            + "</head>\n"
+            + "<body>\n"
+            + "\n"
+            + "</body>\n"
+            + "</html>";
+    writer.flush();
+    result = writer.toString();
+    assertThat(result).isEqualTo(correctSecondPart);
+  }
+}

instrumentation/servlet/servlet-3.0/javaagent-unit-tests/src/test/java/io/opentelemetry/javaagent/instrumentation/servlet/v3_0/snippet/SnippetInjectingResponseWrapperTest.java

@@ -0,0 +1,189 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.servlet.v3_0.snippet;
+
+import static io.opentelemetry.javaagent.instrumentation.servlet.v3_0.snippet.TestUtil.readFile;
+import static org.assertj.core.api.AssertionsForClassTypes.assertThat;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import io.opentelemetry.javaagent.bootstrap.servlet.ExperimentalSnippetHolder;
+import java.io.IOException;
+import java.io.PrintWriter;
+import java.io.StringWriter;
+import java.nio.charset.Charset;
+import javax.servlet.http.HttpServletResponse;
+import org.junit.jupiter.api.Disabled;
+import org.junit.jupiter.api.Test;
+
+class SnippetInjectingResponseWrapperTest {
+
+  @Test
+  void testInjectToTextHtml() throws IOException {
+
+    // read the originalFile
+    String original = readFile("staticHtmlOrigin.html");
+    String correct = readFile("staticHtmlAfter.html");
+    HttpServletResponse response = mock(HttpServletResponse.class);
+    when(response.getContentType()).thenReturn("text/html");
+    when(response.getStatus()).thenReturn(200);
+    when(response.containsHeader("content-type")).thenReturn(true);
+    StringWriter writer = new StringWriter();
+    when(response.getWriter()).thenReturn(new PrintWriter(writer));
+    ExperimentalSnippetHolder.setSnippet("\n  <script type=\"text/javascript\"> Test </script>");
+    SnippetInjectingResponseWrapper responseWrapper = new SnippetInjectingResponseWrapper(response);
+    responseWrapper.getWriter().write(original);
+    responseWrapper.getWriter().flush();
+    responseWrapper.getWriter().close();
+
+    // read file get result
+    String result = writer.toString();
+    writer.close();
+    // check whether new response == correct answer
+    assertThat(result).isEqualTo(correct);
+  }
+
+  @Test
+  @Disabled
+  void testInjectToChineseTextHtml() throws IOException {
+
+    // read the originalFile
+    String original = readFile("staticHtmlChineseOrigin.html");
+    String correct = readFile("staticHtmlChineseAfter.html");
+    HttpServletResponse response = mock(HttpServletResponse.class);
+    when(response.getContentType()).thenReturn("text/html");
+
+    StringWriter writer = new StringWriter();
+    when(response.getWriter()).thenReturn(new PrintWriter(writer));
+    ExperimentalSnippetHolder.setSnippet("\n  <script type=\"text/javascript\"> Test </script>");
+    SnippetInjectingResponseWrapper responseWrapper = new SnippetInjectingResponseWrapper(response);
+    responseWrapper.getWriter().write(original);
+    responseWrapper.getWriter().flush();
+    responseWrapper.getWriter().close();
+
+    // read file get result
+    String result = writer.toString();
+    writer.close();
+    // check whether new response == correct answer
+    assertThat(result).isEqualTo(correct);
+  }
+
+  @Test
+  void shouldNotInjectToTextHtml() throws IOException {
+
+    // read the originalFile
+    String original = readFile("staticHtmlOrigin.html");
+
+    StringWriter writer = new StringWriter();
+    HttpServletResponse response = mock(HttpServletResponse.class);
+    when(response.getContentType()).thenReturn("not/text");
+    when(response.getStatus()).thenReturn(200);
+    when(response.containsHeader("content-type")).thenReturn(true);
+
+    when(response.getWriter()).thenReturn(new PrintWriter(writer, true));
+    ExperimentalSnippetHolder.setSnippet("\n  <script type=\"text/javascript\"> Test </script>");
+
+    SnippetInjectingResponseWrapper responseWrapper = new SnippetInjectingResponseWrapper(response);
+    responseWrapper.getWriter().write(original);
+    responseWrapper.getWriter().flush();
+    responseWrapper.getWriter().close();
+
+    // read file get result
+    String result = writer.toString();
+    writer.close();
+    // check whether new response == correct answer
+    assertThat(result).isEqualTo(original);
+  }
+
+  @Test
+  void testWriteInt() throws IOException {
+
+    // read the originalFile
+    String original = readFile("staticHtmlOrigin.html");
+    String correct = readFile("staticHtmlAfter.html");
+    HttpServletResponse response = mock(HttpServletResponse.class);
+    when(response.getContentType()).thenReturn("text/html");
+
+    StringWriter writer = new StringWriter();
+    //    StringWriter correctWriter = new StringWriter();
+    when(response.getWriter()).thenReturn(new PrintWriter(writer));
+    ExperimentalSnippetHolder.setSnippet("\n  <script type=\"text/javascript\"> Test </script>");
+    SnippetInjectingResponseWrapper responseWrapper = new SnippetInjectingResponseWrapper(response);
+    byte[] originalBytes = original.getBytes(Charset.defaultCharset().name());
+    //    byte[] correctBytes = correct.getBytes(UTF_8);
+    //    PrintWriter correctPw = new PrintWriter(correctWriter);
+    for (int i = 0; i < originalBytes.length; i++) {
+      responseWrapper.getWriter().write(originalBytes[i]);
+    }
+
+    responseWrapper.getWriter().flush();
+    responseWrapper.getWriter().close();
+
+    // read file get result
+    String result = writer.toString();
+    writer.close();
+    // check whether new response == correct answer
+    assertThat(result).isEqualTo(correct);
+  }
+
+  @Test
+  void testWriteCharArray() throws IOException {
+
+    // read the originalFile
+    String original = readFile("staticHtmlChineseOrigin.html");
+    String correct = readFile("staticHtmlChineseAfter.html");
+    HttpServletResponse response = mock(HttpServletResponse.class);
+    when(response.getContentType()).thenReturn("text/html");
+    when(response.getStatus()).thenReturn(200);
+    when(response.containsHeader("content-type")).thenReturn(true);
+
+    StringWriter writer = new StringWriter();
+    when(response.getWriter()).thenReturn(new PrintWriter(writer));
+    ExperimentalSnippetHolder.setSnippet("\n  <script type=\"text/javascript\"> Test </script>");
+    SnippetInjectingResponseWrapper responseWrapper = new SnippetInjectingResponseWrapper(response);
+    char[] originalChars = original.toCharArray();
+    responseWrapper.getWriter().write(originalChars, 0, originalChars.length);
+    responseWrapper.getWriter().flush();
+    responseWrapper.getWriter().close();
+
+    // read file get result
+    String result = writer.toString();
+    writer.close();
+    // check whether new response == correct answer
+    assertThat(result).isEqualTo(correct);
+  }
+
+  @Test
+  void testWriteWithOffset() throws IOException {
+
+    // read the originalFile
+    String original = readFile("staticHtmlChineseOrigin.html");
+    String correct = readFile("staticHtmlChineseAfter.html");
+    String extraBuffer = "this buffer should not be print out";
+    original = extraBuffer + original;
+    HttpServletResponse response = mock(HttpServletResponse.class);
+    when(response.getContentType()).thenReturn("text/html");
+    when(response.getStatus()).thenReturn(200);
+    when(response.containsHeader("content-type")).thenReturn(true);
+
+    StringWriter writer = new StringWriter();
+    when(response.getWriter()).thenReturn(new PrintWriter(writer));
+    ExperimentalSnippetHolder.setSnippet("\n  <script type=\"text/javascript\"> Test </script>");
+    SnippetInjectingResponseWrapper responseWrapper = new SnippetInjectingResponseWrapper(response);
+
+    responseWrapper
+        .getWriter()
+        .write(original, extraBuffer.length(), original.length() - extraBuffer.length());
+    responseWrapper.getWriter().flush();
+    responseWrapper.getWriter().close();
+
+    // read file get result
+    String result = writer.toString();
+    writer.close();
+    // check whether new response == correct answer
+    assertThat(result).isEqualTo(correct);
+  }
+}