Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit f46c8f1

Browse files
amarzialiamarziali
committed
Improve RUM injection matching and avoid truncating responses
1 parent 993bf57 commit f46c8f1

File tree

15 files changed

+223
-124
lines changed

15 files changed

+223
-124
lines changed

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

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,13 +2,15 @@
22

33
import java.io.IOException;
44
import java.io.OutputStream;
5+
import javax.annotation.concurrent.NotThreadSafe;
56

67
/**
78
* An OutputStream containing a circular buffer with a lookbehind buffer of n bytes. The first time
89
* that the latest n bytes matches the marker, a content is injected before. In case of IOException
910
* thrown by the downstream, the buffer will be lost unless the error occurred when draining it. In
1011
* this case the draining will be resumed.
1112
*/
13+
@NotThreadSafe
1214
public class InjectingPipeOutputStream extends OutputStream {
1315
private final byte[] lookbehind;
1416
private int pos;
@@ -168,6 +170,13 @@ private void drain() throws IOException {
168170
}
169171
}
170172

173+
public void commit() throws IOException {
174+
if (filter || wasDraining) {
175+
filter = false;
176+
drain();
177+
}
178+
}
179+
171180
@Override
172181
public void flush() throws IOException {
173182
downstream.flush();
@@ -176,9 +185,7 @@ public void flush() throws IOException {
176185
@Override
177186
public void close() throws IOException {
178187
try {
179-
if (filter || wasDraining) {
180-
drain();
181-
}
188+
commit();
182189
} finally {
183190
downstream.close();
184191
}

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

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,13 +2,15 @@
22

33
import java.io.IOException;
44
import java.io.Writer;
5+
import javax.annotation.concurrent.NotThreadSafe;
56

67
/**
78
* A Writer containing a circular buffer with a lookbehind buffer of n bytes. The first time that
89
* the latest n bytes matches the marker, a content is injected before. In case of IOException
910
* thrown by the downstream, the buffer will be lost unless the error occurred when draining it. In
1011
* this case the draining will be resumed.
1112
*/
13+
@NotThreadSafe
1214
public class InjectingPipeWriter extends Writer {
1315
private final char[] lookbehind;
1416
private int pos;
@@ -169,6 +171,13 @@ private void drain() throws IOException {
169171
}
170172
}
171173

174+
public void commit() throws IOException {
175+
if (filter || wasDraining) {
176+
filter = false;
177+
drain();
178+
}
179+
}
180+
172181
@Override
173182
public void flush() throws IOException {
174183
downstream.flush();
@@ -177,9 +186,7 @@ public void flush() throws IOException {
177186
@Override
178187
public void close() throws IOException {
179188
try {
180-
if (filter || wasDraining) {
181-
drain();
182-
}
189+
commit();
183190
} finally {
184191
downstream.close();
185192
}

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

Lines changed: 61 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -11,9 +11,10 @@
1111

1212
public class RumHttpServletResponseWrapper extends HttpServletResponseWrapper {
1313
private final RumInjector rumInjector;
14-
private ServletOutputStream outputStream;
14+
private WrappedServletOutputStream outputStream;
1515
private PrintWriter printWriter;
16-
private boolean shouldInject = false;
16+
private InjectingPipeWriter wrappedPipeWriter;
17+
private boolean shouldInject = true;
1718

1819
public RumHttpServletResponseWrapper(HttpServletResponse response) {
1920
super(response);
@@ -22,50 +23,64 @@ public RumHttpServletResponseWrapper(HttpServletResponse response) {
2223

2324
@Override
2425
public ServletOutputStream getOutputStream() throws IOException {
26+
if (outputStream != null) {
27+
return outputStream;
28+
}
2529
if (!shouldInject) {
2630
return super.getOutputStream();
2731
}
28-
if (outputStream == null) {
29-
String encoding = getCharacterEncoding();
30-
if (encoding == null) {
31-
encoding = Charset.defaultCharset().name();
32-
}
33-
outputStream =
34-
new WrappedServletOutputStream(
35-
super.getOutputStream(),
36-
rumInjector.getMarkerBytes(encoding),
37-
rumInjector.getSnippetBytes(encoding),
38-
this::onInjected);
32+
String encoding = getCharacterEncoding();
33+
if (encoding == null) {
34+
encoding = Charset.defaultCharset().name();
3935
}
36+
outputStream =
37+
new WrappedServletOutputStream(
38+
super.getOutputStream(),
39+
rumInjector.getMarkerBytes(encoding),
40+
rumInjector.getSnippetBytes(encoding),
41+
this::onInjected);
42+
4043
return outputStream;
4144
}
4245

4346
@Override
4447
public PrintWriter getWriter() throws IOException {
45-
final PrintWriter delegate = super.getWriter();
46-
if (!shouldInject) {
47-
return delegate;
48+
if (printWriter != null) {
49+
return printWriter;
4850
}
49-
if (printWriter == null) {
50-
printWriter =
51-
new PrintWriter(
52-
new InjectingPipeWriter(
53-
delegate,
54-
rumInjector.getMarkerChars(),
55-
rumInjector.getSnippetChars(),
56-
this::onInjected));
51+
if (!shouldInject) {
52+
return super.getWriter();
5753
}
54+
wrappedPipeWriter =
55+
new InjectingPipeWriter(
56+
super.getWriter(),
57+
rumInjector.getMarkerChars(),
58+
rumInjector.getSnippetChars(),
59+
this::onInjected);
60+
printWriter = new PrintWriter(wrappedPipeWriter);
61+
5862
return printWriter;
5963
}
6064

6165
@Override
6266
public void setContentLength(int len) {
6367
// don't set it since we don't know if we will inject
68+
if (!shouldInject) {
69+
super.setContentLength(len);
70+
}
71+
}
72+
73+
@Override
74+
public void setContentLengthLong(long len) {
75+
if (!shouldInject) {
76+
super.setContentLengthLong(len);
77+
}
6478
}
6579

6680
@Override
6781
public void reset() {
6882
this.outputStream = null;
83+
this.wrappedPipeWriter = null;
6984
this.printWriter = null;
7085
this.shouldInject = false;
7186
super.reset();
@@ -74,8 +89,8 @@ public void reset() {
7489
@Override
7590
public void resetBuffer() {
7691
this.outputStream = null;
92+
this.wrappedPipeWriter = null;
7793
this.printWriter = null;
78-
this.shouldInject = false;
7994
super.resetBuffer();
8095
}
8196

@@ -89,7 +104,27 @@ public void onInjected() {
89104

90105
@Override
91106
public void setContentType(String type) {
92-
shouldInject = type != null && type.contains("text/html");
107+
if (shouldInject) {
108+
shouldInject = type != null && type.contains("text/html");
109+
}
110+
if (!shouldInject) {
111+
commit();
112+
}
93113
super.setContentType(type);
94114
}
115+
116+
public void commit() {
117+
if (wrappedPipeWriter != null) {
118+
try {
119+
wrappedPipeWriter.commit();
120+
} catch (Throwable ignored) {
121+
}
122+
}
123+
if (outputStream != null) {
124+
try {
125+
outputStream.commit();
126+
} catch (Throwable ignored) {
127+
}
128+
}
129+
}
95130
}

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

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,8 @@ public static boolean onEnter(
3535
@Advice.Argument(value = 1, readOnly = false) ServletResponse response,
3636
@Advice.Local("isDispatch") boolean isDispatch,
3737
@Advice.Local("finishSpan") boolean finishSpan,
38-
@Advice.Local("contextScope") ContextScope scope) {
38+
@Advice.Local("contextScope") ContextScope scope,
39+
@Advice.Local("rumServletWrapper") RumHttpServletResponseWrapper rumServletWrapper) {
3940
final boolean invalidRequest =
4041
!(request instanceof HttpServletRequest) || !(response instanceof HttpServletResponse);
4142
if (invalidRequest) {
@@ -47,7 +48,8 @@ public static boolean onEnter(
4748

4849
if (RumInjector.get().isEnabled() && httpServletRequest.getAttribute(DD_RUM_INJECTED) == null) {
4950
httpServletRequest.setAttribute(DD_RUM_INJECTED, Boolean.TRUE);
50-
httpServletResponse = new RumHttpServletResponseWrapper(httpServletResponse);
51+
rumServletWrapper = new RumHttpServletResponseWrapper(httpServletResponse);
52+
httpServletResponse = rumServletWrapper;
5153
response = httpServletResponse;
5254
}
5355

@@ -108,7 +110,11 @@ public static void stopSpan(
108110
@Advice.Local("contextScope") final ContextScope scope,
109111
@Advice.Local("isDispatch") boolean isDispatch,
110112
@Advice.Local("finishSpan") boolean finishSpan,
113+
@Advice.Local("rumServletWrapper") RumHttpServletResponseWrapper rumServletWrapper,
111114
@Advice.Thrown final Throwable throwable) {
115+
if (rumServletWrapper != null) {
116+
rumServletWrapper.commit();
117+
}
112118
// Set user.principal regardless of who created this span.
113119
final Object spanAttr = request.getAttribute(DD_SPAN_ATTRIBUTE);
114120
if (Config.get().isServletPrincipalEnabled()

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

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,13 +3,12 @@
33
import datadog.trace.bootstrap.instrumentation.buffer.InjectingPipeOutputStream;
44
import datadog.trace.util.MethodHandles;
55
import java.io.IOException;
6-
import java.io.OutputStream;
76
import java.lang.invoke.MethodHandle;
87
import javax.servlet.ServletOutputStream;
98
import javax.servlet.WriteListener;
109

1110
public class WrappedServletOutputStream extends ServletOutputStream {
12-
private final OutputStream filtered;
11+
private final InjectingPipeOutputStream filtered;
1312
private final ServletOutputStream delegate;
1413

1514
private static final MethodHandle IS_READY_MH = getMh("isReady");
@@ -83,4 +82,8 @@ public void setWriteListener(WriteListener writeListener) {
8382
sneakyThrow(e);
8483
}
8584
}
85+
86+
public void commit() throws IOException {
87+
filtered.commit();
88+
}
8689
}

‎dd-java-agent/instrumentation/servlet/request-3/src/testFixtures/groovy/datadog/trace/instrumentation/servlet3/RumServlet.groovy‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,8 +14,8 @@ class RumServlet extends HttpServlet {
1414

1515
@Override
1616
protected void service(HttpServletRequest req, HttpServletResponse resp) throws ServletException, IOException {
17-
resp.setContentType(mimeType)
1817
try (def writer = resp.getWriter()) {
18+
resp.setContentType(mimeType)
1919
writer.println("\n" +
2020
"<!doctype html>\n" +
2121
"<html>\n" +

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

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -68,7 +68,8 @@ public static class JakartaServletAdvice {
6868
@Advice.OnMethodEnter(suppress = Throwable.class)
6969
public static AgentSpan before(
7070
@Advice.Argument(0) final ServletRequest request,
71-
@Advice.Argument(value = 1, readOnly = false) ServletResponse response) {
71+
@Advice.Argument(value = 1, readOnly = false) ServletResponse response,
72+
@Advice.Local("rumServletWrapper") RumHttpServletResponseWrapper rumServletWrapper) {
7273
if (!(request instanceof HttpServletRequest)) {
7374
return null;
7475
}
@@ -79,7 +80,8 @@ public static AgentSpan before(
7980
if (RumInjector.get().isEnabled()
8081
&& httpServletRequest.getAttribute(DD_RUM_INJECTED) == null) {
8182
httpServletRequest.setAttribute(DD_RUM_INJECTED, Boolean.TRUE);
82-
response = new RumHttpServletResponseWrapper((HttpServletResponse) response);
83+
rumServletWrapper = new RumHttpServletResponseWrapper((HttpServletResponse) response);
84+
response = rumServletWrapper;
8385
}
8486
}
8587

@@ -95,10 +97,15 @@ public static AgentSpan before(
9597

9698
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
9799
public static void after(
98-
@Advice.Enter final AgentSpan span, @Advice.Argument(0) final ServletRequest request) {
100+
@Advice.Enter final AgentSpan span,
101+
@Advice.Argument(0) final ServletRequest request,
102+
@Advice.Local("rumServletWrapper") RumHttpServletResponseWrapper rumServletWrapper) {
99103
if (span == null) {
100104
return;
101105
}
106+
if (rumServletWrapper != null) {
107+
rumServletWrapper.commit();
108+
}
102109

103110
CallDepthThreadLocalMap.reset(HttpServletRequest.class);
104111
final HttpServletRequest httpServletRequest =

0 commit comments

Comments
 (0)