change based on comments

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

@@ -5,7 +5,7 @@
 
 package io.opentelemetry.javaagent.instrumentation.servlet.v3_0.snippet;
 
-import static io.opentelemetry.javaagent.instrumentation.servlet.v3_0.snippet.TestUtil.readFile;
+import static io.opentelemetry.javaagent.instrumentation.servlet.v3_0.snippet.TestUtil.readFileAsString;
 import static org.assertj.core.api.AssertionsForClassTypes.assertThat;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
@@ -23,7 +23,7 @@ class SnippetPrintWriterTest {
  @Test
  void testInjectToTextHtml() throws IOException {
  String snippet = "\n <script type=\"text/javascript\"> Test </script>";
- String html = readFile("beforeSnippetInjection.html");
+ String html = readFileAsString("beforeSnippetInjection.html");
 
  InMemoryHttpServletResponse response = createInMemoryHttpServletResponse("text/html");
  SnippetInjectingResponseWrapper responseWrapper =
@@ -32,14 +32,14 @@ void testInjectToTextHtml() throws IOException {
  responseWrapper.getWriter().write(html);
  responseWrapper.getWriter().flush();
 
- String expectedHtml = readFile("afterSnippetInjection.html");
+ String expectedHtml = readFileAsString("afterSnippetInjection.html");
  assertThat(response.getStringContent()).isEqualTo(expectedHtml);
  }
 
  @Test
  void testInjectToChineseTextHtml() throws IOException {
  String snippet = "\n <script type=\"text/javascript\"> Test </script>";
- String html = readFile("beforeSnippetInjectionChinese.html");
+ String html = readFileAsString("beforeSnippetInjectionChinese.html");
 
  InMemoryHttpServletResponse response = createInMemoryHttpServletResponse("text/html");
  SnippetInjectingResponseWrapper responseWrapper =
@@ -48,14 +48,14 @@ void testInjectToChineseTextHtml() throws IOException {
  responseWrapper.getWriter().write(html);
  responseWrapper.getWriter().flush();
 
- String expectedHtml = readFile("afterSnippetInjectionChinese.html");
+ String expectedHtml = readFileAsString("afterSnippetInjectionChinese.html");
  assertThat(response.getStringContent()).isEqualTo(expectedHtml);
  }
 
  @Test
  void shouldNotInjectToTextHtml() throws IOException {
  String snippet = "\n <script type=\"text/javascript\"> Test </script>";
- String html = readFile("beforeSnippetInjection.html");
+ String html = readFileAsString("beforeSnippetInjection.html");
 
  InMemoryHttpServletResponse response = createInMemoryHttpServletResponse("not/text");
 
@@ -71,7 +71,7 @@ void shouldNotInjectToTextHtml() throws IOException {
  @Test
  void testWriteInt() throws IOException {
  String snippet = "\n <script type=\"text/javascript\"> Test </script>";
- String html = readFile("beforeSnippetInjection.html");
+ String html = readFileAsString("beforeSnippetInjection.html");
 
  InMemoryHttpServletResponse response = createInMemoryHttpServletResponse("text/html");
  SnippetInjectingResponseWrapper responseWrapper =
@@ -83,14 +83,14 @@ void testWriteInt() throws IOException {
  }
  responseWrapper.getWriter().flush();
 
- String expectedHtml = readFile("afterSnippetInjection.html");
+ String expectedHtml = readFileAsString("afterSnippetInjection.html");
  assertThat(response.getStringContent()).isEqualTo(expectedHtml);
  }
 
  @Test
  void testWriteCharArray() throws IOException {
  String snippet = "\n <script type=\"text/javascript\"> Test </script>";
- String html = readFile("beforeSnippetInjectionChinese.html");
+ String html = readFileAsString("beforeSnippetInjectionChinese.html");
 
  InMemoryHttpServletResponse response = createInMemoryHttpServletResponse("text/html");
  SnippetInjectingResponseWrapper responseWrapper =
@@ -100,14 +100,14 @@ void testWriteCharArray() throws IOException {
  responseWrapper.getWriter().write(originalChars, 0, originalChars.length);
  responseWrapper.getWriter().flush();
 
- String expectedHtml = readFile("afterSnippetInjectionChinese.html");
+ String expectedHtml = readFileAsString("afterSnippetInjectionChinese.html");
  assertThat(response.getStringContent()).isEqualTo(expectedHtml);
  }
 
  @Test
  void testWriteWithOffset() throws IOException {
  String snippet = "\n <script type=\"text/javascript\"> Test </script>";
- String html = readFile("beforeSnippetInjectionChinese.html");
+ String html = readFileAsString("beforeSnippetInjectionChinese.html");
  String extraBuffer = "this buffer should not be print out";
  html = extraBuffer + html;
 
@@ -120,7 +120,7 @@ void testWriteWithOffset() throws IOException {
  .write(html, extraBuffer.length(), html.length() - extraBuffer.length());
  responseWrapper.getWriter().flush();
 
- String expectedHtml = readFile("afterSnippetInjectionChinese.html");
+ String expectedHtml = readFileAsString("afterSnippetInjectionChinese.html");
  assertThat(response.getStringContent()).isEqualTo(expectedHtml);
  }
 

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

@@ -5,7 +5,7 @@
 
 package io.opentelemetry.javaagent.instrumentation.servlet.v3_0.snippet;
 
-import static io.opentelemetry.javaagent.instrumentation.servlet.v3_0.snippet.TestUtil.readFileBytes;
+import static io.opentelemetry.javaagent.instrumentation.servlet.v3_0.snippet.TestUtil.readFileAsBytes;
 import static java.nio.charset.StandardCharsets.UTF_8;
 import static org.assertj.core.api.AssertionsForClassTypes.assertThat;
 import static org.mockito.Mockito.mock;
@@ -23,7 +23,7 @@ class SnippetServletOutputStreamTest {
  @Test
  void testInjectionForStringContainHeadTag() throws IOException {
  String snippet = "\n <script type=\"text/javascript\"> Test </script>";
- byte[] html = readFileBytes("beforeSnippetInjection.html");
+ byte[] html = readFileAsBytes("beforeSnippetInjection.html");
 
  InjectionState obj = createInjectionStateForTesting(snippet, UTF_8);
  InMemoryServletOutputStream out = new InMemoryServletOutputStream();
@@ -33,22 +33,22 @@ void testInjectionForStringContainHeadTag() throws IOException {
  assertThat(obj.getHeadTagBytesSeen()).isEqualTo(-1);
  assertThat(injected).isEqualTo(true);
 
- byte[] expectedHtml = readFileBytes("afterSnippetInjection.html");
+ byte[] expectedHtml = readFileAsBytes("afterSnippetInjection.html");
  assertThat(out.getBytes()).isEqualTo(expectedHtml);
  }
 
  @Test
  void testInjectionForChinese() throws IOException {
  String snippet = "\n <script type=\"text/javascript\"> Test </script>";
- byte[] html = readFileBytes("beforeSnippetInjectionChinese.html");
+ byte[] html = readFileAsBytes("beforeSnippetInjectionChinese.html");
 
  InjectionState obj = createInjectionStateForTesting(snippet, UTF_8);
  InMemoryServletOutputStream out = new InMemoryServletOutputStream();
 
  OutputStreamSnippetInjectionHelper helper = new OutputStreamSnippetInjectionHelper(snippet);
  boolean injected = helper.handleWrite(obj, out, html, 0, html.length);
 
- byte[] expectedHtml = readFileBytes("afterSnippetInjectionChinese.html");
+ byte[] expectedHtml = readFileAsBytes("afterSnippetInjectionChinese.html");
  assertThat(injected).isTrue();
  assertThat(obj.getHeadTagBytesSeen()).isEqualTo(-1);
  assertThat(out.getBytes()).isEqualTo(expectedHtml);
@@ -57,7 +57,7 @@ void testInjectionForChinese() throws IOException {
  @Test
  void testInjectionForStringWithoutHeadTag() throws IOException {
  String snippet = "\n <script type=\"text/javascript\"> Test </script>";
- byte[] html = readFileBytes("htmlWithoutHeadTag.html");
+ byte[] html = readFileAsBytes("htmlWithoutHeadTag.html");
 
  InjectionState obj = createInjectionStateForTesting(snippet, UTF_8);
  InMemoryServletOutputStream out = new InMemoryServletOutputStream();

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

@@ -13,7 +13,7 @@
 
 public class TestUtil {
 
- public static byte[] readFileBytes(String resourceName) throws IOException {
+ protected static byte[] readFileAsBytes(String resourceName) throws IOException {
  InputStream in =
  SnippetPrintWriterTest.class.getClassLoader().getResourceAsStream(resourceName);
  ByteArrayOutputStream result = new ByteArrayOutputStream();
@@ -25,8 +25,8 @@ public static byte[] readFileBytes(String resourceName) throws IOException {
  return result.toByteArray();
  }
 
- public static String readFile(String resourceName) throws IOException {
- return new String(readFileBytes(resourceName), UTF_8);
+ protected static String readFileAsString(String resourceName) throws IOException {
+ return new String(readFileAsBytes(resourceName), UTF_8);
  }
 
  private TestUtil() {}

instrumentation/servlet/servlet-3.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/servlet/v3_0/snippet/InjectionState.java

@@ -24,7 +24,7 @@ public String getCharacterEncoding() {
  return wrapper.getCharacterEncoding();
  }
 
- public void setHeadTagWritten() {
+ private void setHeadTagWritten() {
  headTagBytesSeen = HEAD_TAG_WRITTEN_FAKE_VALUE;
  }
 
@@ -45,7 +45,12 @@ public boolean processByte(int b) {
  } else {
  headTagBytesSeen = 0;
  }
- return headTagBytesSeen == HEAD_TAG_LENGTH;
+ if (headTagBytesSeen == HEAD_TAG_LENGTH) {
+ setHeadTagWritten();
+ return true;
+ } else {
+ return false;
+ }
  }
 
  private boolean inHeadTag(int b) {

instrumentation/servlet/servlet-3.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/servlet/v3_0/snippet/OutputStreamSnippetInjectionHelper.java

@@ -34,19 +34,20 @@ public boolean handleWrite(
  if (state.isHeadTagWritten()) {
  return false;
  }
- int i;
+ int endOfHeadTagPosition;
  boolean endOfHeadTagFound = false;
- for (i = off; i < length && i - off < length; i++) {
- if (state.processByte(original[i])) {
+ for (endOfHeadTagPosition = off;
+ endOfHeadTagPosition < length && endOfHeadTagPosition - off < length;
+ endOfHeadTagPosition++) {
+ if (state.processByte(original[endOfHeadTagPosition])) {
  endOfHeadTagFound = true;
  break;
  }
  }
  if (!endOfHeadTagFound) {
  return false;
  }
- // set before write to avoid recursive loop
- state.setHeadTagWritten();
+
  if (state.getWrapper().isNotSafeToInject()) {
  return false;
  }
@@ -59,9 +60,9 @@ public boolean handleWrite(
  }
  // updating Content-Length before any further writing in case that writing triggers a flush
  state.getWrapper().updateContentLengthIfPreviouslySet();
- out.write(original, off, i + 1);
+ out.write(original, off, endOfHeadTagPosition + 1);
  out.write(snippetBytes);
- out.write(original, i + 1, length - i - 1);
+ out.write(original, endOfHeadTagPosition + 1, length - endOfHeadTagPosition - 1);
  return true;
  }
 
@@ -72,8 +73,6 @@ public boolean handleWrite(InjectionState state, OutputStream out, int b) throws
  if (!state.processByte(b)) {
  return false;
  }
- // set before write to avoid recursive loop
- state.setHeadTagWritten();
 
  if (state.getWrapper().isNotSafeToInject()) {
  return false;

instrumentation/servlet/servlet-3.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/servlet/v3_0/snippet/ServletOutputStreamInjectionState.java

@@ -6,6 +6,7 @@
 package io.opentelemetry.javaagent.instrumentation.servlet.v3_0.snippet;
 
 import io.opentelemetry.instrumentation.api.util.VirtualField;
+import javax.annotation.Nullable;
 import javax.servlet.ServletOutputStream;
 
 public class ServletOutputStreamInjectionState {
@@ -26,6 +27,7 @@ public static void initializeInjectionStateIfNeeded(
  }
  }
 
+ @Nullable
  public static InjectionState getInjectionState(ServletOutputStream servletOutputStream) {
  return virtualField.get(servletOutputStream);
  }

instrumentation/servlet/servlet-3.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/servlet/v3_0/snippet/SnippetInjectingPrintWriter.java

@@ -39,8 +39,7 @@ public void write(int b) {
  if (!endOfHeadTagFound) {
  return;
  }
- // set before write to avoid recursive loop
- state.setHeadTagWritten();
+
  if (state.getWrapper().isNotSafeToInject()) {
  return;
  }

instrumentation/servlet/servlet-3.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/servlet/v3_0/snippet/SnippetInjectingResponseWrapper.java

@@ -38,7 +38,8 @@ public class SnippetInjectingResponseWrapper extends HttpServletResponseWrapper
  private static final int UNSET = -1;
 
  // this is for Servlet 3.1 support
- @Nullable private static final MethodHandle setContentLengthLongHandler = getMethodHandle();
+ @Nullable
+ private static final MethodHandle setContentLengthLongHandler = findSetContentLengthLongMethod();
 
  private final String snippet;
  private final int snippetLength;
@@ -119,7 +120,7 @@ public void setContentLength(int len) {
  }
 
  @Nullable
- private static MethodHandle getMethodHandle() {
+ private static MethodHandle findSetContentLengthLongMethod() {
  try {
  return MethodHandles.lookup()
  .findSpecial(
@@ -143,7 +144,7 @@ public void setContentLengthLong(long length) throws Throwable {
  }
  }
 
- public boolean isContentTypeTextHtml() {
+ boolean isContentTypeTextHtml() {
  String contentType = super.getContentType();
  if (contentType == null) {
  contentType = super.getHeader("content-type");
@@ -170,13 +171,13 @@ public PrintWriter getWriter() throws IOException {
  return snippetInjectingPrintWriter;
  }
 
- public void updateContentLengthIfPreviouslySet() {
+ void updateContentLengthIfPreviouslySet() {
  if (contentLength != UNSET) {
  setContentLength((int) contentLength + snippetLength);
  }
  }
 
- public boolean isNotSafeToInject() {
+ boolean isNotSafeToInject() {
  // if content-length was set and response was already committed (headers sent to the client),
  // then it is not safe to inject because the content-length header cannot be updated to account
  // for the snippet length

instrumentation/servlet/servlet-3.0/javaagent/src/test/groovy/TestServlet3.groovy

@@ -228,7 +228,8 @@ class TestServlet3 {
  resp.writer.print(endpoint.body)
  throw new ServletException(endpoint.body)
  case HTML_PRINT_WRITER:
- // intentionally testing setting status before contentType here to cover that case somewhere resp.status = endpoint.status
+ // intentionally testing setting status before contentType here to cover that case somewhere
+ resp.status = endpoint.status
  resp.contentType = "text/html"
  resp.writer.print(endpoint.body)
  break

instrumentation/servlet/servlet-common/bootstrap/src/main/java/io/opentelemetry/javaagent/bootstrap/servlet/ExperimentalSnippetHolder.java

@@ -9,8 +9,14 @@ public class ExperimentalSnippetHolder {
 
  private static volatile String snippet = "";
 
+ private static boolean isSet = false;
+
  public static void setSnippet(String snippet) {
+ if (isSet) {
+ return;
+ }
  ExperimentalSnippetHolder.snippet = snippet;
+ isSet = true;
  }
 
  public static String getSnippet() {