From bb035b01743a84ce47e474367cbdc9a80832b328 Mon Sep 17 00:00:00 2001 From: Andrea Marziali Date: Tue, 15 Jul 2025 14:44:04 +0200 Subject: [PATCH] Improve InjectingPipe outputstream and writer --- .../buffer/InjectingPipeOutputStream.java | 5 +++++ .../buffer/InjectingPipeWriter.java | 10 +++++----- .../RumHttpServletResponseWrapper.java | 9 ++++++--- .../RumHttpServletResponseWrapper.java | 9 ++++++--- .../agent/test/base/HttpServerTest.groovy | 2 +- .../datadog/trace/api/rum/RumInjector.java | 8 ++++---- .../trace/api/rum/RumInjectorTest.groovy | 20 +++++++++---------- 7 files changed, 37 insertions(+), 26 deletions(-) diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/buffer/InjectingPipeOutputStream.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/buffer/InjectingPipeOutputStream.java index 6e20308aef5..633db0c8d4c 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/buffer/InjectingPipeOutputStream.java +++ b/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) { diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/buffer/InjectingPipeWriter.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/buffer/InjectingPipeWriter.java index 00746a77ecd..f012c04cae4 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/buffer/InjectingPipeWriter.java +++ b/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,9 +40,9 @@ 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; } @@ -50,14 +50,14 @@ public void write(int b) throws IOException { 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); diff --git a/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/RumHttpServletResponseWrapper.java b/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/RumHttpServletResponseWrapper.java index c1c7bcf5e3a..d7b173c6827 100644 --- a/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/RumHttpServletResponseWrapper.java +++ b/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; } diff --git a/dd-java-agent/instrumentation/servlet/request-5/src/main/java/datadog/trace/instrumentation/servlet5/RumHttpServletResponseWrapper.java b/dd-java-agent/instrumentation/servlet/request-5/src/main/java/datadog/trace/instrumentation/servlet5/RumHttpServletResponseWrapper.java index c7da4982ff2..ce86c11e863 100644 --- a/dd-java-agent/instrumentation/servlet/request-5/src/main/java/datadog/trace/instrumentation/servlet5/RumHttpServletResponseWrapper.java +++ b/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; } diff --git a/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/base/HttpServerTest.groovy b/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/base/HttpServerTest.groovy index 452373a6d91..69760e48eb9 100644 --- a/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/base/HttpServerTest.groovy +++ b/dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/base/HttpServerTest.groovy @@ -2216,7 +2216,7 @@ abstract class HttpServerTest extends WithHttpServer { 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 diff --git a/internal-api/src/main/java/datadog/trace/api/rum/RumInjector.java b/internal-api/src/main/java/datadog/trace/api/rum/RumInjector.java index 3b59173de7a..03917bc0cfc 100644 --- a/internal-api/src/main/java/datadog/trace/api/rum/RumInjector.java +++ b/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; } diff --git a/internal-api/src/test/groovy/datadog/trace/api/rum/RumInjectorTest.groovy b/internal-api/src/test/groovy/datadog/trace/api/rum/RumInjectorTest.groovy index a72b8fc0c8b..c64f756cac0 100644 --- a/internal-api/src/test/groovy/datadog/trace/api/rum/RumInjectorTest.groovy +++ b/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 } }