Skip to content

Commit 23e0455

Browse files
sarahchen6sarahchen6
committed
Fix content-length retrieval and add test for injection timing
1 parent feb40be commit 23e0455

File tree

7 files changed

+86
-51
lines changed

7 files changed

+86
-51
lines changed

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

Lines changed: 17 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
import javax.annotation.concurrent.NotThreadSafe;
67

78
/**
@@ -23,18 +24,22 @@ public class InjectingPipeOutputStream extends OutputStream {
2324
private final Runnable onContentInjected;
2425
private final int bulkWriteThreshold;
2526
private final OutputStream downstream;
27+
private final LongConsumer onBytesWritten;
28+
private long bytesWritten = 0;
2629

2730
/**
2831
* @param downstream the delegate output stream
2932
* @param marker the marker to find in the stream. Must at least be one byte.
3033
* @param contentToInject the content to inject once before the marker if found.
3134
* @param onContentInjected callback called when and if the content is injected.
35+
* @param onBytesWritten callback called when stream is closed to report total bytes written.
3236
*/
3337
public InjectingPipeOutputStream(
3438
final OutputStream downstream,
3539
final byte[] marker,
3640
final byte[] contentToInject,
37-
final Runnable onContentInjected) {
41+
final Runnable onContentInjected,
42+
final LongConsumer onBytesWritten) {
3843
this.downstream = downstream;
3944
this.marker = marker;
4045
this.lookbehind = new byte[marker.length];
@@ -46,11 +51,13 @@ public InjectingPipeOutputStream(
4651
this.filter = true;
4752
this.contentToInject = contentToInject;
4853
this.onContentInjected = onContentInjected;
54+
this.onBytesWritten = onBytesWritten;
4955
this.bulkWriteThreshold = marker.length * 2 - 2;
5056
}
5157

5258
@Override
5359
public void write(int b) throws IOException {
60+
bytesWritten++;
5461
if (!filter) {
5562
if (wasDraining) {
5663
// continue draining
@@ -85,6 +92,7 @@ public void write(int b) throws IOException {
8592

8693
@Override
8794
public void write(byte[] array, int off, int len) throws IOException {
95+
bytesWritten += len;
8896
if (!filter) {
8997
if (wasDraining) {
9098
// needs drain
@@ -113,6 +121,7 @@ public void write(byte[] array, int off, int len) throws IOException {
113121
// we don't have a full match. write everything in a bulk except the lookbehind buffer
114122
// sequentially
115123
for (int i = off; i < off + marker.length - 1; i++) {
124+
bytesWritten--; // avoid double counting
116125
write(array[i]);
117126
}
118127
drain();
@@ -123,12 +132,14 @@ public void write(byte[] array, int off, int len) throws IOException {
123132
downstream.write(array, off + marker.length - 1, len - bulkWriteThreshold);
124133
filter = wasFiltering;
125134
for (int i = len - marker.length + 1; i < len; i++) {
135+
bytesWritten--; // avoid double counting
126136
write(array[i]);
127137
}
128138
}
129139
} else {
130140
// use slow path because the length to write is small and within the lookbehind buffer size
131141
for (int i = off; i < off + len; i++) {
142+
bytesWritten--; // avoid double counting
132143
write(array[i]);
133144
}
134145
}
@@ -185,6 +196,11 @@ public void flush() throws IOException {
185196
public void close() throws IOException {
186197
try {
187198
commit();
199+
// report the size of the original HTTP response before injecting via callback
200+
if (onBytesWritten != null) {
201+
onBytesWritten.accept(bytesWritten);
202+
}
203+
bytesWritten = 0;
188204
} finally {
189205
downstream.close();
190206
}

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

Lines changed: 2 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -104,16 +104,8 @@ public PrintWriter getWriter() throws IOException {
104104
@Override
105105
public void setHeader(String name, String value) {
106106
if (name != null) {
107-
String lowerName = name.toLowerCase();
108-
if (lowerName.startsWith("content-security-policy")) {
107+
if (name.toLowerCase().startsWith("content-security-policy")) {
109108
RumInjector.getTelemetryCollector().onContentSecurityPolicyDetected();
110-
} else if (lowerName.equals("content-length") && value != null) {
111-
try {
112-
long contentLength = Long.parseLong(value);
113-
RumInjector.getTelemetryCollector().onInjectionResponseSize(contentLength);
114-
} catch (NumberFormatException ignored) {
115-
// ignore?
116-
}
117109
}
118110
}
119111
super.setHeader(name, value);
@@ -122,8 +114,7 @@ public void setHeader(String name, String value) {
122114
@Override
123115
public void addHeader(String name, String value) {
124116
if (name != null) {
125-
String lowerName = name.toLowerCase();
126-
if (lowerName.startsWith("content-security-policy")) {
117+
if (name.toLowerCase().startsWith("content-security-policy")) {
127118
RumInjector.getTelemetryCollector().onContentSecurityPolicyDetected();
128119
}
129120
}
@@ -132,9 +123,6 @@ public void addHeader(String name, String value) {
132123

133124
@Override
134125
public void setContentLength(int len) {
135-
if (len >= 0) {
136-
RumInjector.getTelemetryCollector().onInjectionResponseSize(len);
137-
}
138126
// don't set it since we don't know if we will inject
139127
if (!shouldInject) {
140128
super.setContentLength(len);

‎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
@@ -4,6 +4,7 @@
44
import datadog.trace.util.MethodHandles;
55
import java.io.IOException;
66
import java.lang.invoke.MethodHandle;
7+
import java.util.function.LongConsumer;
78
import javax.servlet.ServletOutputStream;
89
import javax.servlet.WriteListener;
910

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

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

‎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: 2 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -85,16 +85,8 @@ public PrintWriter getWriter() throws IOException {
8585
@Override
8686
public void setHeader(String name, String value) {
8787
if (name != null) {
88-
String lowerName = name.toLowerCase();
89-
if (lowerName.startsWith("content-security-policy")) {
88+
if (name.toLowerCase().startsWith("content-security-policy")) {
9089
RumInjector.getTelemetryCollector().onContentSecurityPolicyDetected();
91-
} else if (lowerName.equals("content-length") && value != null) {
92-
try {
93-
long contentLength = Long.parseLong(value);
94-
RumInjector.getTelemetryCollector().onInjectionResponseSize(contentLength);
95-
} catch (NumberFormatException ignored) {
96-
// ignore?
97-
}
9890
}
9991
}
10092
super.setHeader(name, value);
@@ -103,8 +95,7 @@ public void setHeader(String name, String value) {
10395
@Override
10496
public void addHeader(String name, String value) {
10597
if (name != null) {
106-
String lowerName = name.toLowerCase();
107-
if (lowerName.startsWith("content-security-policy")) {
98+
if (name.toLowerCase().startsWith("content-security-policy")) {
10899
RumInjector.getTelemetryCollector().onContentSecurityPolicyDetected();
109100
}
110101
}
@@ -113,9 +104,6 @@ public void addHeader(String name, String value) {
113104

114105
@Override
115106
public void setContentLength(int len) {
116-
if (len >= 0) {
117-
RumInjector.getTelemetryCollector().onInjectionResponseSize(len);
118-
}
119107
// don't set it since we don't know if we will inject
120108
if (!shouldInject) {
121109
super.setContentLength(len);

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

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,14 +4,22 @@
44
import jakarta.servlet.ServletOutputStream;
55
import jakarta.servlet.WriteListener;
66
import java.io.IOException;
7+
import java.io.OutputStream;
8+
import java.util.function.LongConsumer;
79

810
public class WrappedServletOutputStream extends ServletOutputStream {
911
private final InjectingPipeOutputStream filtered;
1012
private final ServletOutputStream delegate;
1113

1214
public WrappedServletOutputStream(
13-
ServletOutputStream delegate, byte[] marker, byte[] contentToInject, Runnable onInjected) {
14-
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);
1523
this.delegate = delegate;
1624
}
1725

‎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)