JavaScript Snippet Injection Support Servlet 5

instrumentation/servlet/servlet-5.0/javaagent/build.gradle.kts

@@ -17,6 +17,8 @@ dependencies {
 
   compileOnly("jakarta.servlet:jakarta.servlet-api:5.0.0")
 
+  testImplementation(project(":instrumentation:servlet:servlet-common:bootstrap"))
+
   testLibrary("org.eclipse.jetty:jetty-server:11.0.0")
   testLibrary("org.eclipse.jetty:jetty-servlet:11.0.0")
   testLibrary("org.apache.tomcat.embed:tomcat-embed-core:10.0.0")

instrumentation/servlet/servlet-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/servlet/v5_0/JakartaServletInstrumentationModule.java

@@ -13,6 +13,7 @@
 import io.opentelemetry.javaagent.instrumentation.servlet.common.async.AsyncStartInstrumentation;
 import io.opentelemetry.javaagent.instrumentation.servlet.common.response.HttpServletResponseInstrumentation;
 import io.opentelemetry.javaagent.instrumentation.servlet.common.service.ServletAndFilterInstrumentation;
+import io.opentelemetry.javaagent.instrumentation.servlet.common.service.ServletOutputStreamInstrumentation;
 import java.util.Arrays;
 import java.util.List;
 
@@ -37,6 +38,11 @@ BASE_PACKAGE, adviceClassName(".async.AsyncContextStartAdvice")),
             adviceClassName(".service.JakartaServletServiceAdvice"),
             adviceClassName(".service.JakartaServletInitAdvice"),
             adviceClassName(".service.JakartaServletFilterInitAdvice")),
+        new ServletOutputStreamInstrumentation(
+            BASE_PACKAGE,
+            adviceClassName(".Servlet5OutputStreamWriteBytesAndOffsetAdvice"),
+            adviceClassName(".Servlet5OutputStreamWriteBytesAdvice"),
+            adviceClassName(".Servlet5OutputStreamWriteIntAdvice")),
         new HttpServletResponseInstrumentation(
             BASE_PACKAGE, adviceClassName(".response.ResponseSendAdvice")));
   }

instrumentation/servlet/servlet-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/servlet/v5_0/Servlet5OutputStreamWriteBytesAdvice.java

@@ -0,0 +1,34 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.servlet.v5_0;
+
+import static io.opentelemetry.javaagent.instrumentation.servlet.v5_0.Servlet5Singletons.getSnippetInjectionHelper;
+
+import io.opentelemetry.javaagent.instrumentation.servlet.snippet.InjectionState;
+import io.opentelemetry.javaagent.instrumentation.servlet.v5_0.snippet.ServletOutputStreamInjectionState;
+import jakarta.servlet.ServletOutputStream;
+import java.io.IOException;
+import net.bytebuddy.asm.Advice;
+
+public class Servlet5OutputStreamWriteBytesAdvice {
+
+  @Advice.OnMethodEnter(skipOn = Advice.OnDefaultValue.class, suppress = Throwable.class)
+  public static boolean methodEnter(
+      @Advice.This ServletOutputStream servletOutputStream, @Advice.Argument(0) byte[] write)
+      throws IOException {
+    InjectionState state = ServletOutputStreamInjectionState.getInjectionState(servletOutputStream);
+    if (state == null) {
+      return true;
+    }
+    // if handleWrite returns true, then it means the original bytes + the snippet were written
+    // to the servletOutputStream, and so we no longer need to execute the original method
+    // call (see skipOn above)
+    // if it returns false, then it means nothing was written to the servletOutputStream and the
+    // original method call should be executed
+    return !getSnippetInjectionHelper()
+        .handleWrite(state, servletOutputStream, write, 0, write.length);
+  }
+}

instrumentation/servlet/servlet-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/servlet/v5_0/Servlet5OutputStreamWriteBytesAndOffsetAdvice.java

@@ -0,0 +1,35 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.servlet.v5_0;
+
+import static io.opentelemetry.javaagent.instrumentation.servlet.v5_0.Servlet5Singletons.getSnippetInjectionHelper;
+
+import io.opentelemetry.javaagent.instrumentation.servlet.snippet.InjectionState;
+import io.opentelemetry.javaagent.instrumentation.servlet.v5_0.snippet.ServletOutputStreamInjectionState;
+import jakarta.servlet.ServletOutputStream;
+import java.io.IOException;
+import net.bytebuddy.asm.Advice;
+
+public class Servlet5OutputStreamWriteBytesAndOffsetAdvice {
+  @Advice.OnMethodEnter(skipOn = Advice.OnDefaultValue.class, suppress = Throwable.class)
+  public static boolean methodEnter(
+      @Advice.This ServletOutputStream servletOutputStream,
+      @Advice.Argument(value = 0) byte[] write,
+      @Advice.Argument(value = 1) int off,
+      @Advice.Argument(value = 2) int len)
+      throws IOException {
+    InjectionState state = ServletOutputStreamInjectionState.getInjectionState(servletOutputStream);
+    if (state == null) {
+      return true;
+    }
+    // if handleWrite returns true, then it means the original bytes + the snippet were written
+    // to the servletOutputStream, and so we no longer need to execute the original method
+    // call (see skipOn above)
+    // if it returns false, then it means nothing was written to the servletOutputStream and the
+    // original method call should be executed
+    return !getSnippetInjectionHelper().handleWrite(state, servletOutputStream, write, off, len);
+  }
+}

instrumentation/servlet/servlet-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/servlet/v5_0/Servlet5OutputStreamWriteIntAdvice.java

@@ -0,0 +1,33 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.servlet.v5_0;
+
+import static io.opentelemetry.javaagent.instrumentation.servlet.v5_0.Servlet5Singletons.getSnippetInjectionHelper;
+
+import io.opentelemetry.javaagent.instrumentation.servlet.snippet.InjectionState;
+import io.opentelemetry.javaagent.instrumentation.servlet.v5_0.snippet.ServletOutputStreamInjectionState;
+import jakarta.servlet.ServletOutputStream;
+import java.io.IOException;
+import net.bytebuddy.asm.Advice;
+
+public class Servlet5OutputStreamWriteIntAdvice {
+
+  @Advice.OnMethodEnter(skipOn = Advice.OnDefaultValue.class, suppress = Throwable.class)
+  public static boolean methodEnter(
+      @Advice.This ServletOutputStream servletOutputStream, @Advice.Argument(0) int write)
+      throws IOException {
+    InjectionState state = ServletOutputStreamInjectionState.getInjectionState(servletOutputStream);
+    if (state == null) {
+      return true;
+    }
+    // if handleWrite returns true, then it means the original bytes + the snippet were written
+    // to the servletOutputStream, and so we no longer need to execute the original method
+    // call (see skipOn above)
+    // if it returns false, then it means nothing was written to the servletOutputStream and the
+    // original method call should be executed
+    return !getSnippetInjectionHelper().handleWrite(state, servletOutputStream, write);
+  }
+}

instrumentation/servlet/servlet-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/servlet/v5_0/Servlet5Singletons.java

@@ -8,12 +8,14 @@
 import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
 import io.opentelemetry.instrumentation.api.instrumenter.util.ClassAndMethod;
 import io.opentelemetry.instrumentation.api.util.VirtualField;
+import io.opentelemetry.javaagent.bootstrap.servlet.ExperimentalSnippetHolder;
 import io.opentelemetry.javaagent.bootstrap.servlet.MappingResolver;
 import io.opentelemetry.javaagent.instrumentation.servlet.ServletHelper;
 import io.opentelemetry.javaagent.instrumentation.servlet.ServletInstrumenterBuilder;
 import io.opentelemetry.javaagent.instrumentation.servlet.ServletRequestContext;
 import io.opentelemetry.javaagent.instrumentation.servlet.ServletResponseContext;
 import io.opentelemetry.javaagent.instrumentation.servlet.common.response.ResponseInstrumenterFactory;
+import io.opentelemetry.javaagent.instrumentation.servlet.snippet.OutputStreamSnippetInjectionHelper;
 import jakarta.servlet.Filter;
 import jakarta.servlet.Servlet;
 import jakarta.servlet.http.HttpServletRequest;
@@ -39,6 +41,8 @@ public final class Servlet5Singletons {
 
   private static final Instrumenter<ClassAndMethod, Void> RESPONSE_INSTRUMENTER =
       ResponseInstrumenterFactory.createInstrumenter(INSTRUMENTATION_NAME);
+  private static final OutputStreamSnippetInjectionHelper SNIPPET_INJECTION_HELPER =
+      new OutputStreamSnippetInjectionHelper(() -> ExperimentalSnippetHolder.getSnippet());
 
   public static ServletHelper<HttpServletRequest, HttpServletResponse> helper() {
     return HELPER;
@@ -48,6 +52,10 @@ public static Instrumenter<ClassAndMethod, Void> responseInstrumenter() {
     return RESPONSE_INSTRUMENTER;
   }
 
+  public static OutputStreamSnippetInjectionHelper getSnippetInjectionHelper() {
+    return SNIPPET_INJECTION_HELPER;
+  }
+
   public static MappingResolver getMappingResolver(Object servletOrFilter) {
     MappingResolver.Factory factory = getMappingResolverFactory(servletOrFilter);
     if (factory != null) {

instrumentation/servlet/servlet-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/servlet/v5_0/service/JakartaServletServiceAdvice.java

@@ -5,6 +5,7 @@
 
 package io.opentelemetry.javaagent.instrumentation.servlet.v5_0.service;
 
+import static io.opentelemetry.javaagent.instrumentation.servlet.v5_0.Servlet5Singletons.getSnippetInjectionHelper;
 import static io.opentelemetry.javaagent.instrumentation.servlet.v5_0.Servlet5Singletons.helper;
 
 import io.opentelemetry.context.Context;
@@ -17,6 +18,7 @@
 import io.opentelemetry.javaagent.instrumentation.servlet.ServletRequestContext;
 import io.opentelemetry.javaagent.instrumentation.servlet.v5_0.Servlet5Accessor;
 import io.opentelemetry.javaagent.instrumentation.servlet.v5_0.Servlet5Singletons;
+import io.opentelemetry.javaagent.instrumentation.servlet.v5_0.snippet.Servlet5SnippetInjectingResponseWrapper;
 import jakarta.servlet.Servlet;
 import jakarta.servlet.ServletRequest;
 import jakarta.servlet.ServletResponse;
@@ -43,6 +45,13 @@ public static void onEnter(
     }
     HttpServletRequest httpServletRequest = (HttpServletRequest) request;
 
+    String snippet = getSnippetInjectionHelper().getSnippet();
+    if (!snippet.isEmpty()
+        && !((HttpServletResponse) response)
+            .containsHeader(Servlet5SnippetInjectingResponseWrapper.FAKE_SNIPPET_HEADER)) {
+      response = new Servlet5SnippetInjectingResponseWrapper((HttpServletResponse) response, snippet);
+    }
+
     callDepth = CallDepth.forClass(AppServerBridge.getCallDepthKey());
     callDepth.getAndIncrement();
 

instrumentation/servlet/servlet-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/servlet/v5_0/snippet/Servlet5SnippetInjectingResponseWrapper.java

@@ -0,0 +1,154 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.servlet.v5_0.snippet;
+
+import static io.opentelemetry.javaagent.instrumentation.servlet.v5_0.snippet.ServletOutputStreamInjectionState.initializeInjectionStateIfNeeded;
+import static java.util.logging.Level.FINE;
+
+import io.opentelemetry.javaagent.instrumentation.servlet.snippet.SnippetInjectingPrintWriter;
+import io.opentelemetry.javaagent.instrumentation.servlet.snippet.SnippetInjectingResponseWrapper;
+import jakarta.servlet.ServletOutputStream;
+import jakarta.servlet.http.HttpServletResponse;
+import jakarta.servlet.http.HttpServletResponseWrapper;
+import java.io.IOException;
+import java.io.PrintWriter;
+import java.util.logging.Logger;
+
+public class Servlet5SnippetInjectingResponseWrapper extends HttpServletResponseWrapper
+    implements SnippetInjectingResponseWrapper {
+
+  private static final Logger logger = Logger.getLogger(HttpServletResponseWrapper.class.getName());
+
+  public static final String FAKE_SNIPPET_HEADER = "FAKE_SNIPPET_HEADER";
+
+  private static final int UNSET = -1;
+  private final String snippet;
+  private final int snippetLength;
+
+  private long contentLength = UNSET;
+
+  private SnippetInjectingPrintWriter snippetInjectingPrintWriter = null;
+
+  public Servlet5SnippetInjectingResponseWrapper(HttpServletResponse response, String snippet) {
+    super(response);
+    this.snippet = snippet;
+    snippetLength = snippet.length();
+  }
+
+  @Override
+  public boolean containsHeader(String name) {
+    // this function is overridden in order to make sure the response is wrapped
+    // but not wrapped twice
+    // we don't use req.setAttribute
+    // because async requests pass down their attributes, but don't pass down our wrapped response
+    // and so we would see the presence of the attribute and think the response was already wrapped
+    // when it really is not
+    // see also https://docs.oracle.com/javaee/7/api/javax/servlet/AsyncContext.html
+    if (name.equals(FAKE_SNIPPET_HEADER)) {
+      return true;
+    }
+    return super.containsHeader(name);
+  }
+
+  @Override
+  public void setHeader(String name, String value) {
+    // checking content-type is just an optimization to avoid unnecessary parsing
+    if ("Content-Length".equalsIgnoreCase(name) && isContentTypeTextHtml()) {
+      try {
+        contentLength = Long.parseLong(value);
+      } catch (NumberFormatException ex) {
+        logger.log(FINE, "NumberFormatException", ex);
+      }
+    }
+    super.setHeader(name, value);
+  }
+
+  @Override
+  public void addHeader(String name, String value) {
+    // checking content-type is just an optimization to avoid unnecessary parsing
+    if ("Content-Length".equalsIgnoreCase(name) && isContentTypeTextHtml()) {
+      try {
+        contentLength = Long.parseLong(value);
+      } catch (NumberFormatException ex) {
+        logger.log(FINE, "NumberFormatException", ex);
+      }
+    }
+    super.addHeader(name, value);
+  }
+
+  @Override
+  public void setIntHeader(String name, int value) {
+    // checking content-type is just an optimization to avoid unnecessary parsing
+    if ("Content-Length".equalsIgnoreCase(name) && isContentTypeTextHtml()) {
+      contentLength = value;
+    }
+    super.setIntHeader(name, value);
+  }
+
+  @Override
+  public void addIntHeader(String name, int value) {
+    // checking content-type is just an optimization to avoid unnecessary parsing
+    if ("Content-Length".equalsIgnoreCase(name) && isContentTypeTextHtml()) {
+      contentLength = value;
+    }
+    super.addIntHeader(name, value);
+  }
+
+  @Override
+  public void setContentLength(int len) {
+    contentLength = len;
+    super.setContentLength(len);
+  }
+
+  @Override
+  public void setContentLengthLong(long length) {
+    contentLength = length;
+    super.setContentLengthLong(length);
+  }
+
+  @Override
+  public ServletOutputStream getOutputStream() throws IOException {
+    ServletOutputStream output = super.getOutputStream();
+    initializeInjectionStateIfNeeded(output, this);
+    return output;
+  }
+
+  @Override
+  public PrintWriter getWriter() throws IOException {
+    if (!isContentTypeTextHtml()) {
+      return super.getWriter();
+    }
+    if (snippetInjectingPrintWriter == null) {
+      snippetInjectingPrintWriter =
+          new SnippetInjectingPrintWriter(super.getWriter(), snippet, this);
+    }
+    return snippetInjectingPrintWriter;
+  }
+
+  @Override
+  public boolean isContentTypeTextHtml() {
+    String contentType = super.getContentType();
+    if (contentType == null) {
+      contentType = super.getHeader("content-type");
+    }
+    return contentType != null && contentType.startsWith("text/html");
+  }
+
+  @Override
+  public void updateContentLengthIfPreviouslySet() {
+    if (contentLength != UNSET) {
+      setContentLength((int) contentLength + snippetLength);
+    }
+  }
+
+  @Override
+  public 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
+    return contentLength != UNSET && isCommitted();
+  }
+}

instrumentation/servlet/servlet-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/servlet/v5_0/snippet/ServletOutputStreamInjectionState.java

@@ -0,0 +1,36 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.servlet.v5_0.snippet;
+
+import io.opentelemetry.instrumentation.api.util.VirtualField;
+import io.opentelemetry.javaagent.instrumentation.servlet.snippet.InjectionState;
+import jakarta.servlet.ServletOutputStream;
+import javax.annotation.Nullable;
+
+public class ServletOutputStreamInjectionState {
+  private static final VirtualField<ServletOutputStream, InjectionState> virtualField =
+      VirtualField.find(ServletOutputStream.class, InjectionState.class);
+
+  public static void initializeInjectionStateIfNeeded(
+      ServletOutputStream servletOutputStream, Servlet5SnippetInjectingResponseWrapper wrapper) {
+    InjectionState state = virtualField.get(servletOutputStream);
+    if (!wrapper.isContentTypeTextHtml()) {
+      virtualField.set(servletOutputStream, null);
+      return;
+    }
+    if (state == null || state.getWrapper() != wrapper) {
+      state = new InjectionState(wrapper);
+      virtualField.set(servletOutputStream, state);
+    }
+  }
+
+  @Nullable
+  public static InjectionState getInjectionState(ServletOutputStream servletOutputStream) {
+    return virtualField.get(servletOutputStream);
+  }
+
+  private ServletOutputStreamInjectionState() {}
+}