Skip to content

Commit 7bd7a48

Browse files
sarahchen6sarahchen6
committed
Fix content-length retrieval and add test for injection timing
1 parent 87c7127 commit 7bd7a48

File tree

7 files changed

+96
-53
lines changed

7 files changed

+96
-53
lines changed

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

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

33
import java.io.IOException;
44
import java.io.OutputStream;
5+
import java.util.function.LongConsumer;
56

67
/**
78
* An OutputStream containing a circular buffer with a lookbehind buffer of n bytes. The first time
@@ -18,29 +19,36 @@ public class InjectingPipeOutputStream extends OutputStream {
1819
private final Runnable onContentInjected;
1920
private final int bulkWriteThreshold;
2021
private final OutputStream downstream;
22+
private final LongConsumer onBytesWritten;
23+
private long bytesWritten = 0;
2124

2225
/**
2326
* @param downstream the delegate output stream
2427
* @param marker the marker to find in the stream. Must at least be one byte.
2528
* @param contentToInject the content to inject once before the marker if found.
2629
* @param onContentInjected callback called when and if the content is injected.
30+
* @param onBytesWritten callback called when stream is closed to report total bytes written.
2731
*/
2832
public InjectingPipeOutputStream(
2933
final OutputStream downstream,
3034
final byte[] marker,
3135
final byte[] contentToInject,
32-
final Runnable onContentInjected) {
36+
final Runnable onContentInjected,
37+
final LongConsumer onBytesWritten) {
3338
this.downstream = downstream;
3439
this.marker = marker;
3540
this.lookbehind = new byte[marker.length];
3641
this.pos = 0;
3742
this.contentToInject = contentToInject;
3843
this.onContentInjected = onContentInjected;
44+
this.onBytesWritten = onBytesWritten;
3945
this.bulkWriteThreshold = marker.length * 2 - 2;
4046
}
4147

4248
@Override
4349
public void write(int b) throws IOException {
50+
bytesWritten++;
51+
4452
if (found) {
4553
downstream.write(b);
4654
return;
@@ -73,6 +81,8 @@ public void write(int b) throws IOException {
7381

7482
@Override
7583
public void write(byte[] array, int off, int len) throws IOException {
84+
bytesWritten += len;
85+
7686
if (found) {
7787
downstream.write(array, off, len);
7888
return;
@@ -96,17 +106,20 @@ public void write(byte[] array, int off, int len) throws IOException {
96106
// we don't have a full match. write everything in a bulk except the lookbehind buffer
97107
// sequentially
98108
for (int i = off; i < off + marker.length - 1; i++) {
109+
bytesWritten--; // avoid double counting
99110
write(array[i]);
100111
}
101112
drain();
102113
downstream.write(array, off + marker.length - 1, len - bulkWriteThreshold);
103114
for (int i = len - marker.length + 1; i < len; i++) {
115+
bytesWritten--; // avoid double counting
104116
write(array[i]);
105117
}
106118
}
107119
} else {
108120
// use slow path because the length to write is small and within the lookbehind buffer size
109121
for (int i = off; i < off + len; i++) {
122+
bytesWritten--; // avoid double counting
110123
write(array[i]);
111124
}
112125
}
@@ -155,6 +168,16 @@ public void close() throws IOException {
155168
if (!found) {
156169
drain();
157170
}
171+
172+
// report the size of the original HTTP response before injecting via callback
173+
if (onBytesWritten != null) {
174+
try {
175+
onBytesWritten.accept(bytesWritten);
176+
} catch (Exception e) {
177+
}
178+
}
179+
bytesWritten = 0;
180+
158181
downstream.close();
159182
}
160183
}

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

Lines changed: 4 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,8 @@ public ServletOutputStream getOutputStream() throws IOException {
4242
super.getOutputStream(),
4343
rumInjector.getMarkerBytes(encoding),
4444
rumInjector.getSnippetBytes(encoding),
45-
this::onInjected);
45+
this::onInjected,
46+
bytes -> RumInjector.getTelemetryCollector().onInjectionResponseSize(bytes));
4647
} catch (Exception e) {
4748
RumInjector.getTelemetryCollector().onInjectionFailed();
4849
throw e;
@@ -82,16 +83,8 @@ public PrintWriter getWriter() throws IOException {
8283
@Override
8384
public void setHeader(String name, String value) {
8485
if (name != null) {
85-
String lowerName = name.toLowerCase();
86-
if (lowerName.startsWith("content-security-policy")) {
86+
if (name.toLowerCase().startsWith("content-security-policy")) {
8787
RumInjector.getTelemetryCollector().onContentSecurityPolicyDetected();
88-
} else if (lowerName.equals("content-length") && value != null) {
89-
try {
90-
long contentLength = Long.parseLong(value);
91-
RumInjector.getTelemetryCollector().onInjectionResponseSize(contentLength);
92-
} catch (NumberFormatException ignored) {
93-
// ignore?
94-
}
9588
}
9689
}
9790
super.setHeader(name, value);
@@ -100,8 +93,7 @@ public void setHeader(String name, String value) {
10093
@Override
10194
public void addHeader(String name, String value) {
10295
if (name != null) {
103-
String lowerName = name.toLowerCase();
104-
if (lowerName.startsWith("content-security-policy")) {
96+
if (name.toLowerCase().startsWith("content-security-policy")) {
10597
RumInjector.getTelemetryCollector().onContentSecurityPolicyDetected();
10698
}
10799
}
@@ -110,9 +102,6 @@ public void addHeader(String name, String value) {
110102

111103
@Override
112104
public void setContentLength(int len) {
113-
if (len >= 0) {
114-
RumInjector.getTelemetryCollector().onInjectionResponseSize(len);
115-
}
116105
// don't set it since we don't know if we will inject
117106
}
118107

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

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
import java.io.IOException;
66
import java.io.OutputStream;
77
import java.lang.invoke.MethodHandle;
8+
import java.util.function.LongConsumer;
89
import javax.servlet.ServletOutputStream;
910
import javax.servlet.WriteListener;
1011

@@ -30,8 +31,14 @@ private static <E extends Throwable> void sneakyThrow(Throwable e) throws E {
3031
}
3132

3233
public WrappedServletOutputStream(
33-
ServletOutputStream delegate, byte[] marker, byte[] contentToInject, Runnable onInjected) {
34-
this.filtered = new InjectingPipeOutputStream(delegate, marker, contentToInject, onInjected);
34+
ServletOutputStream delegate,
35+
byte[] marker,
36+
byte[] contentToInject,
37+
Runnable onInjected,
38+
LongConsumer onBytesWritten) {
39+
this.filtered =
40+
new InjectingPipeOutputStream(
41+
delegate, marker, contentToInject, onInjected, onBytesWritten);
3542
this.delegate = delegate;
3643
}
3744

‎dd-java-agent/instrumentation/servlet/request-3/src/test/groovy/RumHttpServletResponseWrapperTest.groovy

Lines changed: 23 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -119,22 +119,36 @@ class RumHttpServletResponseWrapperTest extends AgentTestRunner {
119119
1 * mockResponse.addHeader("X-Content-Security-Policy", "test")
120120
}
121121

122-
void 'setHeader with Content-Length reports response size'() {
122+
void 'response sizes are reported correctly'() {
123+
setup:
124+
def downstream = Mock(java.io.OutputStream)
125+
def marker = "</head>".getBytes("UTF-8")
126+
def contentToInject = "<script></script>".getBytes("UTF-8")
127+
def onBytesWritten = Mock(java.util.function.LongConsumer)
128+
def stream = new datadog.trace.bootstrap.instrumentation.buffer.InjectingPipeOutputStream(
129+
downstream, marker, contentToInject, null, onBytesWritten)
130+
123131
when:
124-
wrapper.setHeader("Content-Length", "1024")
132+
stream.write("test".getBytes("UTF-8"))
133+
stream.write("content".getBytes("UTF-8"))
134+
stream.close()
125135

126136
then:
127-
1 * mockTelemetryCollector.onInjectionResponseSize(1024)
128-
1 * mockResponse.setHeader("Content-Length", "1024")
137+
1 * onBytesWritten.accept(11)
129138
}
130139

131-
void 'setContentLength method reports response size'() {
140+
void 'injection timing is reported when injection is successful'() {
141+
setup:
142+
wrapper.setContentType("text/html")
143+
132144
when:
133-
wrapper.setContentLength(1024)
145+
try {
146+
wrapper.getOutputStream() // set injectionStartTime
147+
} catch (Exception ignored) {} // expect failure due to improper setup
148+
Thread.sleep(1) // ensure some time passes
149+
wrapper.onInjected() // report timing when injection "is successful"
134150

135151
then:
136-
1 * mockTelemetryCollector.onInjectionResponseSize(1024)
152+
1 * mockTelemetryCollector.onInjectionTime({ it > 0 })
137153
}
138-
139-
// test injection timing?
140154
}

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

Lines changed: 4 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,8 @@ public ServletOutputStream getOutputStream() throws IOException {
4242
super.getOutputStream(),
4343
rumInjector.getMarkerBytes(encoding),
4444
rumInjector.getSnippetBytes(encoding),
45-
this::onInjected);
45+
this::onInjected,
46+
bytes -> RumInjector.getTelemetryCollector().onInjectionResponseSize(bytes));
4647
} catch (Exception e) {
4748
RumInjector.getTelemetryCollector().onInjectionFailed();
4849
throw e;
@@ -82,16 +83,8 @@ public PrintWriter getWriter() throws IOException {
8283
@Override
8384
public void setHeader(String name, String value) {
8485
if (name != null) {
85-
String lowerName = name.toLowerCase();
86-
if (lowerName.startsWith("content-security-policy")) {
86+
if (name.toLowerCase().startsWith("content-security-policy")) {
8787
RumInjector.getTelemetryCollector().onContentSecurityPolicyDetected();
88-
} else if (lowerName.equals("content-length") && value != null) {
89-
try {
90-
long contentLength = Long.parseLong(value);
91-
RumInjector.getTelemetryCollector().onInjectionResponseSize(contentLength);
92-
} catch (NumberFormatException ignored) {
93-
// ignore?
94-
}
9588
}
9689
}
9790
super.setHeader(name, value);
@@ -100,8 +93,7 @@ public void setHeader(String name, String value) {
10093
@Override
10194
public void addHeader(String name, String value) {
10295
if (name != null) {
103-
String lowerName = name.toLowerCase();
104-
if (lowerName.startsWith("content-security-policy")) {
96+
if (name.toLowerCase().startsWith("content-security-policy")) {
10597
RumInjector.getTelemetryCollector().onContentSecurityPolicyDetected();
10698
}
10799
}
@@ -110,9 +102,6 @@ public void addHeader(String name, String value) {
110102

111103
@Override
112104
public void setContentLength(int len) {
113-
if (len >= 0) {
114-
RumInjector.getTelemetryCollector().onInjectionResponseSize(len);
115-
}
116105
// don't set it since we don't know if we will inject
117106
}
118107

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

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,14 +5,21 @@
55
import jakarta.servlet.WriteListener;
66
import java.io.IOException;
77
import java.io.OutputStream;
8+
import java.util.function.LongConsumer;
89

910
public class WrappedServletOutputStream extends ServletOutputStream {
1011
private final OutputStream filtered;
1112
private final ServletOutputStream delegate;
1213

1314
public WrappedServletOutputStream(
14-
ServletOutputStream delegate, byte[] marker, byte[] contentToInject, Runnable onInjected) {
15-
this.filtered = new InjectingPipeOutputStream(delegate, marker, contentToInject, onInjected);
15+
ServletOutputStream delegate,
16+
byte[] marker,
17+
byte[] contentToInject,
18+
Runnable onInjected,
19+
LongConsumer onBytesWritten) {
20+
this.filtered =
21+
new InjectingPipeOutputStream(
22+
delegate, marker, contentToInject, onInjected, onBytesWritten);
1623
this.delegate = delegate;
1724
}
1825

‎dd-java-agent/instrumentation/servlet/request-5/src/test/groovy/RumHttpServletResponseWrapperTest.groovy

Lines changed: 23 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -119,22 +119,36 @@ class RumHttpServletResponseWrapperTest extends AgentTestRunner {
119119
1 * mockResponse.addHeader("X-Content-Security-Policy", "test")
120120
}
121121

122-
void 'setHeader with Content-Length reports response size'() {
122+
void 'response sizes are reported correctly'() {
123+
setup:
124+
def downstream = Mock(java.io.OutputStream)
125+
def marker = "</head>".getBytes("UTF-8")
126+
def contentToInject = "<script></script>".getBytes("UTF-8")
127+
def onBytesWritten = Mock(java.util.function.LongConsumer)
128+
def stream = new datadog.trace.bootstrap.instrumentation.buffer.InjectingPipeOutputStream(
129+
downstream, marker, contentToInject, null, onBytesWritten)
130+
123131
when:
124-
wrapper.setHeader("Content-Length", "1024")
132+
stream.write("test".getBytes("UTF-8"))
133+
stream.write("content".getBytes("UTF-8"))
134+
stream.close()
125135

126136
then:
127-
1 * mockTelemetryCollector.onInjectionResponseSize(1024)
128-
1 * mockResponse.setHeader("Content-Length", "1024")
137+
1 * onBytesWritten.accept(11)
129138
}
130139

131-
void 'setContentLength method reports response size'() {
140+
void 'injection timing is reported when injection is successful'() {
141+
setup:
142+
wrapper.setContentType("text/html")
143+
132144
when:
133-
wrapper.setContentLength(1024)
145+
try {
146+
wrapper.getOutputStream() // set injectionStartTime
147+
} catch (Exception ignored) {} // expect failure due to improper setup
148+
Thread.sleep(1) // ensure some time passes
149+
wrapper.onInjected() // report timing when injection "is successful"
134150

135151
then:
136-
1 * mockTelemetryCollector.onInjectionResponseSize(1024)
152+
1 * mockTelemetryCollector.onInjectionTime({ it > 0 })
137153
}
138-
139-
// test injection timing?
140154
}

0 commit comments

Comments
 (0)