Skip to content

Commit 47394fd

Browse files
amarzialiamarziali
authored
Avoid multiple injections on dispatch on jboss (#9392)
1 parent 7b8238e commit 47394fd

File tree

24 files changed

+558
-33
lines changed

24 files changed

+558
-33
lines changed

‎dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/rum/RumControllableResponse.java‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
package datadog.trace.bootstrap.instrumentation.rum;
2+
3+
public interface RumControllableResponse {
4+
/** Drain the held buffer. */
5+
void commit();
6+
7+
/** Stops filtering the response. */
8+
void stopFiltering();
9+
}

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

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
import datadog.trace.agent.tooling.Instrumenter;
1111
import datadog.trace.agent.tooling.InstrumenterModule;
1212
import datadog.trace.api.InstrumenterConfig;
13+
import datadog.trace.bootstrap.instrumentation.rum.RumControllableResponse;
1314
import javax.servlet.AsyncContext;
1415
import net.bytebuddy.asm.Advice;
1516
import net.bytebuddy.description.type.TypeDescription;
@@ -56,8 +57,8 @@ public static class CommitAdvice {
5657
public static void commitRumBuffer(@Advice.This final AsyncContext asyncContext) {
5758
final Object maybeRumWrappedResponse =
5859
asyncContext.getRequest().getAttribute(DD_RUM_INJECTED);
59-
if (maybeRumWrappedResponse instanceof RumHttpServletResponseWrapper) {
60-
((RumHttpServletResponseWrapper) maybeRumWrappedResponse).commit();
60+
if (maybeRumWrappedResponse instanceof RumControllableResponse) {
61+
((RumControllableResponse) maybeRumWrappedResponse).commit();
6162
}
6263
}
6364
}

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

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_RUM_INJECTED;
44

5+
import datadog.trace.bootstrap.instrumentation.rum.RumControllableResponse;
56
import javax.servlet.AsyncContext;
67
import javax.servlet.ServletRequest;
78
import javax.servlet.ServletResponse;
@@ -11,7 +12,7 @@
1112

1213
public class RumHttpServletRequestWrapper extends HttpServletRequestWrapper {
1314

14-
private final HttpServletResponse response;
15+
private HttpServletResponse response;
1516

1617
public RumHttpServletRequestWrapper(
1718
final HttpServletRequest request, final HttpServletResponse response) {
@@ -30,17 +31,16 @@ public AsyncContext startAsync(ServletRequest servletRequest, ServletResponse se
3031
throws IllegalStateException {
3132
// deactivate the previous wrapper
3233
final Object maybeRumWrappedResponse = (servletRequest.getAttribute(DD_RUM_INJECTED));
33-
if (maybeRumWrappedResponse instanceof RumHttpServletResponseWrapper) {
34-
((RumHttpServletResponseWrapper) maybeRumWrappedResponse).commit();
35-
((RumHttpServletResponseWrapper) maybeRumWrappedResponse).stopFiltering();
34+
if (maybeRumWrappedResponse instanceof RumControllableResponse) {
35+
((RumControllableResponse) maybeRumWrappedResponse).commit();
36+
((RumControllableResponse) maybeRumWrappedResponse).stopFiltering();
3637
}
37-
ServletResponse actualResponse = servletResponse;
3838
// rewrap it
3939
if (servletResponse instanceof HttpServletResponse) {
40-
actualResponse =
40+
this.response =
4141
new RumHttpServletResponseWrapper(this, (HttpServletResponse) servletResponse);
42-
servletRequest.setAttribute(DD_RUM_INJECTED, actualResponse);
42+
servletRequest.setAttribute(DD_RUM_INJECTED, this.response);
4343
}
44-
return super.startAsync(servletRequest, actualResponse);
44+
return super.startAsync(servletRequest, this.response);
4545
}
4646
}

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

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import datadog.trace.api.rum.RumInjector;
44
import datadog.trace.bootstrap.instrumentation.buffer.InjectingPipeWriter;
5+
import datadog.trace.bootstrap.instrumentation.rum.RumControllableResponse;
56
import datadog.trace.util.MethodHandles;
67
import java.io.IOException;
78
import java.io.PrintWriter;
@@ -14,7 +15,8 @@
1415
import javax.servlet.http.HttpServletResponse;
1516
import javax.servlet.http.HttpServletResponseWrapper;
1617

17-
public class RumHttpServletResponseWrapper extends HttpServletResponseWrapper {
18+
public class RumHttpServletResponseWrapper extends HttpServletResponseWrapper
19+
implements RumControllableResponse {
1820
private final RumInjector rumInjector;
1921
private final String servletVersion;
2022
private WrappedServletOutputStream outputStream;
@@ -203,6 +205,7 @@ public void setContentType(String type) {
203205
super.setContentType(type);
204206
}
205207

208+
@Override
206209
public void commit() {
207210
if (wrappedPipeWriter != null) {
208211
try {
@@ -218,6 +221,7 @@ public void commit() {
218221
}
219222
}
220223

224+
@Override
221225
public void stopFiltering() {
222226
shouldInject = false;
223227
if (wrappedPipeWriter != null) {

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

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
import datadog.trace.api.gateway.Flow;
1818
import datadog.trace.api.rum.RumInjector;
1919
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
20+
import datadog.trace.bootstrap.instrumentation.rum.RumControllableResponse;
2021
import datadog.trace.instrumentation.servlet.ServletBlockingHelper;
2122
import java.security.Principal;
2223
import java.util.concurrent.atomic.AtomicBoolean;
@@ -35,7 +36,7 @@ public static boolean onEnter(
3536
@Advice.Local("isDispatch") boolean isDispatch,
3637
@Advice.Local("finishSpan") boolean finishSpan,
3738
@Advice.Local("contextScope") ContextScope scope,
38-
@Advice.Local("rumServletWrapper") RumHttpServletResponseWrapper rumServletWrapper) {
39+
@Advice.Local("rumServletWrapper") RumControllableResponse rumServletWrapper) {
3940
final boolean invalidRequest =
4041
!(request instanceof HttpServletRequest) || !(response instanceof HttpServletResponse);
4142
if (invalidRequest) {
@@ -47,14 +48,16 @@ public static boolean onEnter(
4748

4849
if (RumInjector.get().isEnabled()) {
4950
final Object maybeRumWrapper = httpServletRequest.getAttribute(DD_RUM_INJECTED);
50-
if (maybeRumWrapper instanceof RumHttpServletResponseWrapper) {
51-
rumServletWrapper = (RumHttpServletResponseWrapper) maybeRumWrapper;
51+
if (maybeRumWrapper instanceof RumControllableResponse) {
52+
rumServletWrapper = (RumControllableResponse) maybeRumWrapper;
5253
} else {
5354
rumServletWrapper =
5455
new RumHttpServletResponseWrapper(httpServletRequest, (HttpServletResponse) response);
5556
httpServletRequest.setAttribute(DD_RUM_INJECTED, rumServletWrapper);
56-
response = rumServletWrapper;
57-
request = new RumHttpServletRequestWrapper(httpServletRequest, rumServletWrapper);
57+
response = (ServletResponse) rumServletWrapper;
58+
request =
59+
new RumHttpServletRequestWrapper(
60+
httpServletRequest, (HttpServletResponse) rumServletWrapper);
5861
}
5962
}
6063

@@ -116,7 +119,7 @@ public static void stopSpan(
116119
@Advice.Local("contextScope") final ContextScope scope,
117120
@Advice.Local("isDispatch") boolean isDispatch,
118121
@Advice.Local("finishSpan") boolean finishSpan,
119-
@Advice.Local("rumServletWrapper") RumHttpServletResponseWrapper rumServletWrapper,
122+
@Advice.Local("rumServletWrapper") RumControllableResponse rumServletWrapper,
120123
@Advice.Thrown final Throwable throwable) {
121124
if (rumServletWrapper != null) {
122125
rumServletWrapper.commit();

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

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import datadog.trace.api.rum.RumInjector;
2020
import datadog.trace.bootstrap.CallDepthThreadLocalMap;
2121
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
22+
import datadog.trace.bootstrap.instrumentation.rum.RumControllableResponse;
2223
import jakarta.servlet.ServletRequest;
2324
import jakarta.servlet.ServletResponse;
2425
import jakarta.servlet.http.HttpServletRequest;
@@ -71,7 +72,7 @@ public static class JakartaServletAdvice {
7172
public static AgentSpan before(
7273
@Advice.Argument(value = 0, readOnly = false) ServletRequest request,
7374
@Advice.Argument(value = 1, readOnly = false) ServletResponse response,
74-
@Advice.Local("rumServletWrapper") RumHttpServletResponseWrapper rumServletWrapper) {
75+
@Advice.Local("rumServletWrapper") RumControllableResponse rumServletWrapper) {
7576
if (!(request instanceof HttpServletRequest)) {
7677
return null;
7778
}
@@ -81,15 +82,17 @@ public static AgentSpan before(
8182

8283
if (RumInjector.get().isEnabled()) {
8384
final Object maybeRumWrapper = httpServletRequest.getAttribute(DD_RUM_INJECTED);
84-
if (maybeRumWrapper instanceof RumHttpServletResponseWrapper) {
85-
rumServletWrapper = (RumHttpServletResponseWrapper) maybeRumWrapper;
85+
if (maybeRumWrapper instanceof RumControllableResponse) {
86+
rumServletWrapper = (RumControllableResponse) maybeRumWrapper;
8687
} else {
8788
rumServletWrapper =
8889
new RumHttpServletResponseWrapper(
8990
httpServletRequest, (HttpServletResponse) response);
9091
httpServletRequest.setAttribute(DD_RUM_INJECTED, rumServletWrapper);
91-
response = rumServletWrapper;
92-
request = new RumHttpServletRequestWrapper(httpServletRequest, rumServletWrapper);
92+
response = (ServletResponse) rumServletWrapper;
93+
request =
94+
new RumHttpServletRequestWrapper(
95+
httpServletRequest, (HttpServletResponse) rumServletWrapper);
9396
}
9497
}
9598
}
@@ -108,7 +111,7 @@ public static AgentSpan before(
108111
public static void after(
109112
@Advice.Enter final AgentSpan span,
110113
@Advice.Argument(0) final ServletRequest request,
111-
@Advice.Local("rumServletWrapper") RumHttpServletResponseWrapper rumServletWrapper) {
114+
@Advice.Local("rumServletWrapper") RumControllableResponse rumServletWrapper) {
112115
if (span == null) {
113116
return;
114117
}

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

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
import datadog.trace.agent.tooling.Instrumenter;
1111
import datadog.trace.agent.tooling.InstrumenterModule;
1212
import datadog.trace.api.InstrumenterConfig;
13+
import datadog.trace.bootstrap.instrumentation.rum.RumControllableResponse;
1314
import jakarta.servlet.AsyncContext;
1415
import net.bytebuddy.asm.Advice;
1516
import net.bytebuddy.description.type.TypeDescription;
@@ -56,8 +57,8 @@ public static class CommitAdvice {
5657
public static void commitRumBuffer(@Advice.This final AsyncContext asyncContext) {
5758
final Object maybeRumWrappedResponse =
5859
asyncContext.getRequest().getAttribute(DD_RUM_INJECTED);
59-
if (maybeRumWrappedResponse instanceof RumHttpServletResponseWrapper) {
60-
((RumHttpServletResponseWrapper) maybeRumWrappedResponse).commit();
60+
if (maybeRumWrappedResponse instanceof RumControllableResponse) {
61+
((RumControllableResponse) maybeRumWrappedResponse).commit();
6162
}
6263
}
6364
}

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

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_RUM_INJECTED;
44

5+
import datadog.trace.bootstrap.instrumentation.rum.RumControllableResponse;
56
import jakarta.servlet.AsyncContext;
67
import jakarta.servlet.ServletRequest;
78
import jakarta.servlet.ServletResponse;
@@ -30,9 +31,9 @@ public AsyncContext startAsync(ServletRequest servletRequest, ServletResponse se
3031
throws IllegalStateException {
3132
// deactivate the previous wrapper
3233
final Object maybeRumWrappedResponse = (servletRequest.getAttribute(DD_RUM_INJECTED));
33-
if (maybeRumWrappedResponse instanceof RumHttpServletResponseWrapper) {
34-
((RumHttpServletResponseWrapper) maybeRumWrappedResponse).commit();
35-
((RumHttpServletResponseWrapper) maybeRumWrappedResponse).stopFiltering();
34+
if (maybeRumWrappedResponse instanceof RumControllableResponse) {
35+
((RumControllableResponse) maybeRumWrappedResponse).commit();
36+
((RumControllableResponse) maybeRumWrappedResponse).stopFiltering();
3637
}
3738
ServletResponse actualResponse = servletResponse;
3839
// rewrap it

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

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import datadog.trace.api.rum.RumInjector;
44
import datadog.trace.bootstrap.instrumentation.buffer.InjectingPipeWriter;
5+
import datadog.trace.bootstrap.instrumentation.rum.RumControllableResponse;
56
import jakarta.servlet.ServletContext;
67
import jakarta.servlet.ServletOutputStream;
78
import jakarta.servlet.http.HttpServletRequest;
@@ -11,7 +12,8 @@
1112
import java.io.PrintWriter;
1213
import java.nio.charset.Charset;
1314

14-
public class RumHttpServletResponseWrapper extends HttpServletResponseWrapper {
15+
public class RumHttpServletResponseWrapper extends HttpServletResponseWrapper
16+
implements RumControllableResponse {
1517
private final RumInjector rumInjector;
1618
private final String servletVersion;
1719
private WrappedServletOutputStream outputStream;
@@ -180,6 +182,7 @@ public void setContentType(String type) {
180182
super.setContentType(type);
181183
}
182184

185+
@Override
183186
public void commit() {
184187
if (wrappedPipeWriter != null) {
185188
try {
@@ -195,6 +198,7 @@ public void commit() {
195198
}
196199
}
197200

201+
@Override
198202
public void stopFiltering() {
199203
shouldInject = false;
200204
if (wrappedPipeWriter != null) {

‎dd-smoke-tests/rum/src/main/groovy/datadog/smoketest/rum/AbstractRumServerSmokeTest.groovy‎

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -14,9 +14,13 @@ class AbstractRumServerSmokeTest extends AbstractServerSmokeTest {
1414
"-Ddd.rum.remote.configuration.id=12345",
1515
]
1616

17+
String mountPoint() {
18+
""
19+
}
20+
1721
void 'test RUM SDK injection on html for path #servletPath'() {
1822
given:
19-
def url = "http://localhost:${httpPort}/${servletPath}"
23+
def url = "http://localhost:${httpPort}${mountPoint()}/${servletPath}"
2024
def request = new Request.Builder()
2125
.url(url)
2226
.get()
@@ -34,7 +38,7 @@ class AbstractRumServerSmokeTest extends AbstractServerSmokeTest {
3438

3539
void 'test RUM SDK injection skip on unsupported mime type'() {
3640
given:
37-
def url = "http://localhost:${httpPort}/xml"
41+
def url = "http://localhost:${httpPort}${mountPoint()}/xml"
3842
def request = new Request.Builder()
3943
.url(url)
4044
.get()
@@ -52,13 +56,14 @@ class AbstractRumServerSmokeTest extends AbstractServerSmokeTest {
5256
assert response.header('x-datadog-rum-injected') == '1': 'RUM injected header missing'
5357
def content = response.body().string()
5458
assert content.contains('https://www.datadoghq-browser-agent.com'): 'RUM script not injected'
55-
assert content.endsWith('</html>'): 'Response not fully flushed'
59+
assert content.trim().endsWith('</html>'): 'Response not fully flushed'
60+
assert content.indexOf("DD_RUM.init(") == content.lastIndexOf("DD_RUM.init("): 'script injected more than once'
5661
}
5762

5863
static void assertRumNotInjected(Response response) {
5964
assert response.header('x-datadog-rum-injected') == null: 'RUM header unexpectedly injected'
6065
def content = response.body().string()
6166
assert !content.contains('https://www.datadoghq-browser-agent.com'): 'RUM script unexpectedly injected'
62-
assert content.endsWith('</response>'): 'Response not fully flushed'
67+
assert content.trim().endsWith('</response>'): 'Response not fully flushed'
6368
}
6469
}

0 commit comments

Comments
 (0)