Implement Rum injection for servlet 3

dd-java-agent/agent-bootstrap/src/jmh/java/datadog/trace/bootstrap/instrumentation/buffer/InjectingPipeOutputStreamBenchmark.java

@@ -0,0 +1,60 @@
+package datadog.trace.bootstrap.instrumentation.buffer;
+
+import static java.util.concurrent.TimeUnit.MICROSECONDS;
+import static java.util.concurrent.TimeUnit.SECONDS;
+
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.PrintWriter;
+import java.net.URL;
+import java.nio.charset.StandardCharsets;
+import java.util.List;
+import org.apache.commons.io.IOUtils;
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Fork;
+import org.openjdk.jmh.annotations.Measurement;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.OutputTimeUnit;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.Warmup;
+
+@State(Scope.Benchmark)
+@Warmup(iterations = 1, time = 30, timeUnit = SECONDS)
+@Measurement(iterations = 2, time = 30, timeUnit = SECONDS)
+@BenchmarkMode(Mode.AverageTime)
+@OutputTimeUnit(MICROSECONDS)
+@Fork(value = 1)
+public class InjectingPipeOutputStreamBenchmark {
+  private static final List<String> htmlContent;
+  private static final byte[] marker;
+  private static final byte[] content;
+
+  static {
+    try (InputStream is = new URL("https://www.google.com").openStream()) {
+      htmlContent = IOUtils.readLines(is, StandardCharsets.UTF_8);
+    } catch (IOException ioe) {
+      throw new RuntimeException(ioe);
+    }
+    marker = "</head>".getBytes(StandardCharsets.UTF_8);
+    content = "<script/>".getBytes(StandardCharsets.UTF_8);
+  }
+
+  @Benchmark
+  public void withPipe() throws Exception {
+    try (final PrintWriter out =
+        new PrintWriter(
+            new InjectingPipeOutputStream(new ByteArrayOutputStream(), marker, content, null))) {
+      htmlContent.forEach(out::println);
+    }
+  }
+
+  @Benchmark
+  public void withoutPipe() throws Exception {
+    try (final PrintWriter out = new PrintWriter(new ByteArrayOutputStream())) {
+      htmlContent.forEach(out::println);
+    }
+  }
+}

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/buffer/InjectingPipeOutputStream.java

@@ -0,0 +1,91 @@
+package datadog.trace.bootstrap.instrumentation.buffer;
+
+import java.io.FilterOutputStream;
+import java.io.IOException;
+import java.io.OutputStream;
+import java.util.function.Consumer;
+
+/**
+ * A circular buffer that holds n+1 bytes and with a lookbehind buffer of n bytes. The first time
+ * that the latest n bytes matches the marker, a content is injected before.
+ */
+public class InjectingPipeOutputStream extends FilterOutputStream {
+  private final byte[] lookbehind;
+  private int pos;
+  private boolean bufferFilled;
+  private final byte[] marker;
+  private final byte[] contentToInject;
+  private boolean found = false;
+  private int matchingPos = 0;
+  private final Consumer<Void> onContentInjected;
+
+  /**
+   * @param downstream the delegate output stream
+   * @param marker the marker to find in the stream
+   * @param contentToInject the content to inject once before the marker if found.
+   * @param onContentInjected callback called when and if the content is injected.
+   */
+  public InjectingPipeOutputStream(
+      final OutputStream downstream,
+      final byte[] marker,
+      final byte[] contentToInject,
+      final Consumer<Void> onContentInjected) {
+    super(downstream);
+    this.marker = marker;
+    this.lookbehind = new byte[marker.length + 1];
+    this.pos = 0;
+    this.contentToInject = contentToInject;
+    this.onContentInjected = onContentInjected;
+  }
+
+  @Override
+  public void write(int b) throws IOException {
+    if (found) {
+      out.write(b);
+      return;
+    }
+    lookbehind[pos] = (byte) b;
+    pos = (pos + 1) % lookbehind.length;
+
+    if (marker[matchingPos++] == b) {
+      if (matchingPos == marker.length) {
+        found = true;
+        out.write(contentToInject);
+        if (onContentInjected != null) {
+          onContentInjected.accept(null);
+        }
+        drain((pos + 1) % lookbehind.length, marker.length);
+        return;
+      }
+    } else {
+      matchingPos = 0;
+    }
+
+    if (!bufferFilled) {
+      bufferFilled = pos == lookbehind.length - 1;
+    }
+
+    if (bufferFilled) {
+      super.write(lookbehind[(pos + 1) % lookbehind.length]);
+    }
+  }
+
+  private void drain(int from, int size) throws IOException {
+    while (size-- > 0) {
+      super.write(Character.valueOf((char) lookbehind[from]));
+      from = (from + 1) % lookbehind.length;
+    }
+  }
+
+  @Override
+  public void close() throws IOException {
+    if (!found) {
+      if (bufferFilled) {
+        drain((pos + 2) % lookbehind.length, marker.length - 1);
+      } else {
+        drain(0, pos);
+      }
+    }
+    super.close();
+  }
+}

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java

@@ -57,6 +57,7 @@ public abstract class HttpServerDecorator<REQUEST, CONNECTION, RESPONSE, REQUEST
 
   public static final String DD_SPAN_ATTRIBUTE = "datadog.span";
   public static final String DD_DISPATCH_SPAN_ATTRIBUTE = "datadog.span.dispatch";
+  public static final String DD_RUM_INJECTED = "datadog.rum.injected";
   public static final String DD_FIN_DISP_LIST_SPAN_ATTRIBUTE =
       "datadog.span.finish_dispatch_listener";
   public static final String DD_RESPONSE_ATTRIBUTE = "datadog.response";

dd-java-agent/agent-bootstrap/src/test/groovy/datadog/trace/bootstrap/instrumentation/buffer/InjectingPipeOutputStreamTest.groovy

@@ -0,0 +1,23 @@
+package datadog.trace.bootstrap.instrumentation.buffer
+
+import spock.lang.Specification
+
+class InjectingPipeOutputStreamTest extends Specification {
+  def 'should filter a buffer and inject if found #found'() {
+    setup:
+    def downstream = new ByteArrayOutputStream()
+    def piped = new OutputStreamWriter(new InjectingPipeOutputStream(downstream, marker.getBytes("UTF-8"), contentToInject.getBytes("UTF-8"), null),
+    "UTF-8")
+    when:
+    try (def closeme = piped) {
+      piped.write(body)
+    }
+    then:
+    assert downstream.toByteArray() == expected.getBytes("UTF-8")
+    where:
+    body                                      | marker            | contentToInject         | found | expected
+    "<html><head><foo/></head><body/></html>" | "</head>"         | "<script>true</script>" | true  | "<html><head><foo/><script>true</script></head><body/></html>"
+    "<html><body/></html>"                    | "</head>"         | "<something/>"          | false | "<html><body/></html>"
+    "<foo/>"                                  | "<longerThanFoo>" | "<nothing>"             | false | "<foo/>"
+  }
+}

dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/RumHttpServletResponseWrapper.java

@@ -0,0 +1,84 @@
+package datadog.trace.instrumentation.servlet3;
+
+import datadog.trace.api.rum.RumInjector;
+import java.io.IOException;
+import java.io.PrintWriter;
+import javax.servlet.ServletOutputStream;
+import javax.servlet.http.HttpServletResponse;
+import javax.servlet.http.HttpServletResponseWrapper;
+
+public class RumHttpServletResponseWrapper extends HttpServletResponseWrapper {
+  private ServletOutputStream outputStream;
+  private PrintWriter printWriter;
+  private boolean shouldInject;
+
+  public RumHttpServletResponseWrapper(HttpServletResponse response) {
+    super(response);
+  }
+
+  @Override
+  public ServletOutputStream getOutputStream() throws IOException {
+    if (!shouldInject) {
+      return super.getOutputStream();
+    }
+    if (outputStream == null) {
+      String encoding = getCharacterEncoding();
+      if (encoding == null) {
+        encoding = "UTF-8";
+      }
+      outputStream =
+          new WrappedServletOutputStream(
+              super.getOutputStream(),
+              RumInjector.getMarker(encoding),
+              RumInjector.getSnippet(encoding),
+              this::onInjected);
+    }
+    return outputStream;
+  }
+
+  @Override
+  public PrintWriter getWriter() throws IOException {
+    if (!shouldInject) {
+      return super.getWriter();
+    }
+    if (printWriter == null) {
+      printWriter = new PrintWriter(getOutputStream());
+    }
+    return printWriter;
+  }
+
+  @Override
+  public void setContentLength(int len) {
+    // don't set it since we don't know if we will inject
+  }
+
+  @Override
+  public void reset() {
+    this.outputStream = null;
+    this.printWriter = null;
+    this.shouldInject = false;
+    super.reset();
+  }
+
+  @Override
+  public void resetBuffer() {
+    this.outputStream = null;
+    this.printWriter = null;
+    this.shouldInject = false;
+    super.resetBuffer();
+  }
+
+  public void onInjected(Void ignored) {
+    try {
+      setHeader("x-datadog-rum-injected", "1");
+    } catch (Throwable ignored2) {
+    }
+  }
+
+  @Override
+  public void setContentType(String type) {
+    if (type != null && type.contains("html")) {
+      shouldInject = true;
+    }
+  }
+}

dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet3Advice.java

@@ -4,6 +4,7 @@
 import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.spanFromContext;
 import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_DISPATCH_SPAN_ATTRIBUTE;
 import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_FIN_DISP_LIST_SPAN_ATTRIBUTE;
+import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_RUM_INJECTED;
 import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE;
 import static datadog.trace.instrumentation.servlet3.Servlet3Decorator.DECORATE;
 
@@ -15,6 +16,7 @@
 import datadog.trace.api.DDTags;
 import datadog.trace.api.GlobalTracer;
 import datadog.trace.api.gateway.Flow;
+import datadog.trace.api.rum.RumInjector;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.instrumentation.servlet.ServletBlockingHelper;
 import java.security.Principal;
@@ -30,7 +32,7 @@ public class Servlet3Advice {
   @Advice.OnMethodEnter(suppress = Throwable.class, skipOn = Advice.OnNonDefaultValue.class)
   public static boolean onEnter(
       @Advice.Argument(value = 0, readOnly = false) ServletRequest request,
-      @Advice.Argument(value = 1) ServletResponse response,
+      @Advice.Argument(value = 1, readOnly = false) ServletResponse response,
       @Advice.Local("isDispatch") boolean isDispatch,
       @Advice.Local("finishSpan") boolean finishSpan,
       @Advice.Local("contextScope") ContextScope scope) {
@@ -41,7 +43,13 @@ public static boolean onEnter(
     }
 
     final HttpServletRequest httpServletRequest = (HttpServletRequest) request;
-    final HttpServletResponse httpServletResponse = (HttpServletResponse) response;
+    HttpServletResponse httpServletResponse = (HttpServletResponse) response;
+
+    if (RumInjector.isEnabled() && httpServletRequest.getAttribute(DD_RUM_INJECTED) == null) {
+      httpServletRequest.setAttribute(DD_RUM_INJECTED, Boolean.TRUE);
+      httpServletResponse = new RumHttpServletResponseWrapper(httpServletResponse);
+      response = httpServletResponse;
+    }
 
     Object dispatchSpan = request.getAttribute(DD_DISPATCH_SPAN_ATTRIBUTE);
     if (dispatchSpan instanceof AgentSpan) {

dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet3Instrumentation.java

@@ -52,6 +52,8 @@ public String[] helperClassNames() {
       packageName + ".Servlet3Decorator",
       packageName + ".ServletRequestURIAdapter",
       packageName + ".FinishAsyncDispatchListener",
+      packageName + ".RumHttpServletResponseWrapper",
+      packageName + ".WrappedServletOutputStream",
       "datadog.trace.instrumentation.servlet.ServletBlockingHelper",
     };
   }