Improve InjectingPipe outputstream and writer (#9176)

Author: amarziali

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

@@ -145,6 +145,11 @@ private void drain() throws IOException {
     bufferFilled = false;
   }
 
+  @Override
+  public void flush() throws IOException {
+    downstream.flush();
+  }
+
   @Override
   public void close() throws IOException {
     if (!found) {

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

@@ -4,7 +4,7 @@
 import java.io.Writer;
 
 /**
- * An Writer containing a circular buffer with a lookbehind buffer of n bytes. The first time that
+ * A Writer containing a circular buffer 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 InjectingPipeWriter extends Writer {
@@ -40,24 +40,24 @@ public InjectingPipeWriter(
   }
 
   @Override
-  public void write(int b) throws IOException {
+  public void write(int c) throws IOException {
     if (found) {
-      downstream.write(b);
+      downstream.write(c);
       return;
     }
 
     if (bufferFilled) {
       downstream.write(lookbehind[pos]);
     }
 
-    lookbehind[pos] = (char) b;
+    lookbehind[pos] = (char) c;
     pos = (pos + 1) % lookbehind.length;
 
     if (!bufferFilled) {
       bufferFilled = pos == 0;
     }
 
-    if (marker[matchingPos++] == b) {
+    if (marker[matchingPos++] == c) {
       if (matchingPos == marker.length) {
         found = true;
         downstream.write(contentToInject);

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

@@ -33,8 +33,8 @@ public ServletOutputStream getOutputStream() throws IOException {
       outputStream =
           new WrappedServletOutputStream(
               super.getOutputStream(),
-              rumInjector.getMarker(encoding),
-              rumInjector.getSnippet(encoding),
+              rumInjector.getMarkerBytes(encoding),
+              rumInjector.getSnippetBytes(encoding),
               this::onInjected);
     }
     return outputStream;
@@ -50,7 +50,10 @@ public PrintWriter getWriter() throws IOException {
       printWriter =
           new PrintWriter(
               new InjectingPipeWriter(
-                  delegate, rumInjector.getMarker(), rumInjector.getSnippet(), this::onInjected));
+                  delegate,
+                  rumInjector.getMarkerChars(),
+                  rumInjector.getSnippetChars(),
+                  this::onInjected));
     }
     return printWriter;
   }

dd-java-agent/instrumentation/servlet/request-5/src/main/java/datadog/trace/instrumentation/servlet5/RumHttpServletResponseWrapper.java

@@ -33,8 +33,8 @@ public ServletOutputStream getOutputStream() throws IOException {
       outputStream =
           new WrappedServletOutputStream(
               super.getOutputStream(),
-              rumInjector.getMarker(encoding),
-              rumInjector.getSnippet(encoding),
+              rumInjector.getMarkerBytes(encoding),
+              rumInjector.getSnippetBytes(encoding),
               this::onInjected);
     }
     return outputStream;
@@ -50,7 +50,10 @@ public PrintWriter getWriter() throws IOException {
       printWriter =
           new PrintWriter(
               new InjectingPipeWriter(
-                  delegate, rumInjector.getMarker(), rumInjector.getSnippet(), this::onInjected));
+                  delegate,
+                  rumInjector.getMarkerChars(),
+                  rumInjector.getSnippetChars(),
+                  this::onInjected));
     }
     return printWriter;
   }

dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/base/HttpServerTest.groovy

@@ -2216,7 +2216,7 @@ abstract class HttpServerTest<SERVER> extends WithHttpServer<SERVER> {
     def response = client.newCall(request).execute()
     then:
     assert response.code() == 200
-    assert response.body().string().contains(new String(RumInjector.get().getSnippet("UTF-8"), "UTF-8")) == expected
+    assert response.body().string().contains(new String(RumInjector.get().getSnippetBytes("UTF-8"), "UTF-8")) == expected
     assert response.header("x-datadog-rum-injected") == (expected ? "1" : null)
     where:
     mime   | expected

internal-api/src/main/java/datadog/trace/api/rum/RumInjector.java

@@ -74,7 +74,7 @@ public boolean isEnabled() {
    * @return The HTML snippet chars to inject, {@code null} if RUM injection is disabled.
    */
   @Nullable
-  public char[] getSnippet() {
+  public char[] getSnippetChars() {
     if (!this.enabled) {
       return null;
     }
@@ -87,7 +87,7 @@ public char[] getSnippet() {
    * @return The HTML snippet to inject, {@code null} if RUM injection is disabled.
    */
   @Nullable
-  public byte[] getSnippet(String encoding) {
+  public byte[] getSnippetBytes(String encoding) {
     if (!this.enabled) {
       return null;
     }
@@ -100,7 +100,7 @@ public byte[] getSnippet(String encoding) {
    * @return The marker chars, {@code null} if RUM injection is disabled.
    */
   @Nullable
-  public char[] getMarker() {
+  public char[] getMarkerChars() {
     if (!this.enabled) {
       return null;
     }
@@ -114,7 +114,7 @@ public char[] getMarker() {
    * @return The marker bytes, {@code null} if RUM injection is disabled.
    */
   @Nullable
-  public byte[] getMarker(String encoding) {
+  public byte[] getMarkerBytes(String encoding) {
     if (!this.enabled) {
       return null;
     }

internal-api/src/test/groovy/datadog/trace/api/rum/RumInjectorTest.groovy

@@ -20,8 +20,8 @@ class RumInjectorTest extends DDSpecification {
 
     then:
     !injector.isEnabled()
-    injector.getMarker(UTF8) == null
-    injector.getSnippet(UTF8) == null
+    injector.getMarkerBytes(UTF8) == null
+    injector.getSnippetBytes(UTF8) == null
   }
 
   void 'invalid config injector'() {
@@ -36,10 +36,10 @@ class RumInjectorTest extends DDSpecification {
 
     then:
     !injector.isEnabled()
-    injector.getMarker(UTF8) == null
-    injector.getSnippet(UTF8) == null
-    injector.getSnippet() == null
-    injector.getMarker() == null
+    injector.getMarkerBytes(UTF8) == null
+    injector.getSnippetBytes(UTF8) == null
+    injector.getSnippetChars() == null
+    injector.getMarkerChars() == null
   }
 
   void 'enabled injector'() {
@@ -56,9 +56,9 @@ class RumInjectorTest extends DDSpecification {
 
     then:
     injector.isEnabled()
-    injector.getMarker(UTF8) != null
-    injector.getSnippet(UTF8) != null
-    injector.getSnippet() != null
-    injector.getMarker() != null
+    injector.getMarkerBytes(UTF8) != null
+    injector.getSnippetBytes(UTF8) != null
+    injector.getSnippetChars() != null
+    injector.getMarkerChars() != null
   }
 }