Skip to content

Commit 1ae2692

Browse files
amarzialiamarziali
authored
Wrap servlet original PrintWriter on rum injector (#9146)
1 parent 51532a8 commit 1ae2692

File tree

8 files changed

+273
-46
lines changed

8 files changed

+273
-46
lines changed

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

Lines changed: 26 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,13 @@
11
package datadog.trace.bootstrap.instrumentation.buffer;
22

3-
import java.io.FilterOutputStream;
43
import java.io.IOException;
54
import java.io.OutputStream;
65

76
/**
8-
* A circular buffer with a lookbehind buffer of n bytes. The first time that the latest n bytes
9-
* matches the marker, a content is injected before.
7+
* An OutputStream containing a circular buffer with a lookbehind buffer of n bytes. The first time
8+
* that the latest n bytes matches the marker, a content is injected before.
109
*/
11-
public class InjectingPipeOutputStream extends FilterOutputStream {
10+
public class InjectingPipeOutputStream extends OutputStream {
1211
private final byte[] lookbehind;
1312
private int pos;
1413
private boolean bufferFilled;
@@ -18,10 +17,11 @@ public class InjectingPipeOutputStream extends FilterOutputStream {
1817
private int matchingPos = 0;
1918
private final Runnable onContentInjected;
2019
private final int bulkWriteThreshold;
20+
private final OutputStream downstream;
2121

2222
/**
2323
* @param downstream the delegate output stream
24-
* @param marker the marker to find in the stream
24+
* @param marker the marker to find in the stream. Must at least be one byte.
2525
* @param contentToInject the content to inject once before the marker if found.
2626
* @param onContentInjected callback called when and if the content is injected.
2727
*/
@@ -30,7 +30,7 @@ public InjectingPipeOutputStream(
3030
final byte[] marker,
3131
final byte[] contentToInject,
3232
final Runnable onContentInjected) {
33-
super(downstream);
33+
this.downstream = downstream;
3434
this.marker = marker;
3535
this.lookbehind = new byte[marker.length];
3636
this.pos = 0;
@@ -42,12 +42,12 @@ public InjectingPipeOutputStream(
4242
@Override
4343
public void write(int b) throws IOException {
4444
if (found) {
45-
out.write(b);
45+
downstream.write(b);
4646
return;
4747
}
4848

4949
if (bufferFilled) {
50-
out.write(lookbehind[pos]);
50+
downstream.write(lookbehind[pos]);
5151
}
5252

5353
lookbehind[pos] = (byte) b;
@@ -60,7 +60,7 @@ public void write(int b) throws IOException {
6060
if (marker[matchingPos++] == b) {
6161
if (matchingPos == marker.length) {
6262
found = true;
63-
out.write(contentToInject);
63+
downstream.write(contentToInject);
6464
if (onContentInjected != null) {
6565
onContentInjected.run();
6666
}
@@ -72,46 +72,48 @@ public void write(int b) throws IOException {
7272
}
7373

7474
@Override
75-
public void write(byte[] b, int off, int len) throws IOException {
75+
public void write(byte[] array, int off, int len) throws IOException {
7676
if (found) {
77-
out.write(b, off, len);
77+
downstream.write(array, off, len);
7878
return;
7979
}
8080
if (len > bulkWriteThreshold) {
8181
// if the content is large enough, we can bulk write everything but the N trail and tail.
8282
// This because the buffer can already contain some byte from a previous single write.
8383
// Also we need to fill the buffer with the tail since we don't know about the next write.
84-
int idx = arrayContains(b, marker);
84+
int idx = arrayContains(array, off, len, marker);
8585
if (idx >= 0) {
8686
// we have a full match. just write everything
8787
found = true;
8888
drain();
89-
out.write(b, off, idx);
90-
out.write(contentToInject);
89+
downstream.write(array, off, idx);
90+
downstream.write(contentToInject);
9191
if (onContentInjected != null) {
9292
onContentInjected.run();
9393
}
94-
out.write(b, off + idx, len - idx);
94+
downstream.write(array, off + idx, len - idx);
9595
} else {
9696
// we don't have a full match. write everything in a bulk except the lookbehind buffer
9797
// sequentially
9898
for (int i = off; i < off + marker.length - 1; i++) {
99-
write(b[i]);
99+
write(array[i]);
100100
}
101101
drain();
102-
out.write(b, off + marker.length - 1, len - bulkWriteThreshold);
102+
downstream.write(array, off + marker.length - 1, len - bulkWriteThreshold);
103103
for (int i = len - marker.length + 1; i < len; i++) {
104-
write(b[i]);
104+
write(array[i]);
105105
}
106106
}
107107
} else {
108108
// use slow path because the length to write is small and within the lookbehind buffer size
109-
super.write(b, off, len);
109+
for (int i = off; i < off + len; i++) {
110+
write(array[i]);
111+
}
110112
}
111113
}
112114

113-
private int arrayContains(byte[] array, byte[] search) {
114-
for (int i = 0; i < array.length - search.length; i++) {
115+
private int arrayContains(byte[] array, int off, int len, byte[] search) {
116+
for (int i = off; i < len - search.length; i++) {
115117
if (array[i] == search[0]) {
116118
boolean found = true;
117119
int k = i;
@@ -133,10 +135,10 @@ private int arrayContains(byte[] array, byte[] search) {
133135
private void drain() throws IOException {
134136
if (bufferFilled) {
135137
for (int i = 0; i < lookbehind.length; i++) {
136-
out.write(lookbehind[(pos + i) % lookbehind.length]);
138+
downstream.write(lookbehind[(pos + i) % lookbehind.length]);
137139
}
138140
} else {
139-
out.write(this.lookbehind, 0, pos);
141+
downstream.write(this.lookbehind, 0, pos);
140142
}
141143
pos = 0;
142144
matchingPos = 0;
@@ -148,6 +150,6 @@ public void close() throws IOException {
148150
if (!found) {
149151
drain();
150152
}
151-
super.close();
153+
downstream.close();
152154
}
153155
}

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

Lines changed: 160 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,160 @@
1+
package datadog.trace.bootstrap.instrumentation.buffer;
2+
3+
import java.io.IOException;
4+
import java.io.Writer;
5+
6+
/**
7+
* An Writer containing a circular buffer with a lookbehind buffer of n bytes. The first time that
8+
* the latest n bytes matches the marker, a content is injected before.
9+
*/
10+
public class InjectingPipeWriter extends Writer {
11+
private final char[] lookbehind;
12+
private int pos;
13+
private boolean bufferFilled;
14+
private final char[] marker;
15+
private final char[] contentToInject;
16+
private boolean found = false;
17+
private int matchingPos = 0;
18+
private final Runnable onContentInjected;
19+
private final int bulkWriteThreshold;
20+
private final Writer downstream;
21+
22+
/**
23+
* @param downstream the delegate writer
24+
* @param marker the marker to find in the stream. Must at least be one char.
25+
* @param contentToInject the content to inject once before the marker if found.
26+
* @param onContentInjected callback called when and if the content is injected.
27+
*/
28+
public InjectingPipeWriter(
29+
final Writer downstream,
30+
final char[] marker,
31+
final char[] contentToInject,
32+
final Runnable onContentInjected) {
33+
this.downstream = downstream;
34+
this.marker = marker;
35+
this.lookbehind = new char[marker.length];
36+
this.pos = 0;
37+
this.contentToInject = contentToInject;
38+
this.onContentInjected = onContentInjected;
39+
this.bulkWriteThreshold = marker.length * 2 - 2;
40+
}
41+
42+
@Override
43+
public void write(int b) throws IOException {
44+
if (found) {
45+
downstream.write(b);
46+
return;
47+
}
48+
49+
if (bufferFilled) {
50+
downstream.write(lookbehind[pos]);
51+
}
52+
53+
lookbehind[pos] = (char) b;
54+
pos = (pos + 1) % lookbehind.length;
55+
56+
if (!bufferFilled) {
57+
bufferFilled = pos == 0;
58+
}
59+
60+
if (marker[matchingPos++] == b) {
61+
if (matchingPos == marker.length) {
62+
found = true;
63+
downstream.write(contentToInject);
64+
if (onContentInjected != null) {
65+
onContentInjected.run();
66+
}
67+
drain();
68+
}
69+
} else {
70+
matchingPos = 0;
71+
}
72+
}
73+
74+
@Override
75+
public void flush() throws IOException {
76+
downstream.flush();
77+
}
78+
79+
@Override
80+
public void write(char[] array, int off, int len) throws IOException {
81+
if (found) {
82+
downstream.write(array, off, len);
83+
return;
84+
}
85+
if (len > bulkWriteThreshold) {
86+
// if the content is large enough, we can bulk write everything but the N trail and tail.
87+
// This because the buffer can already contain some byte from a previous single write.
88+
// Also we need to fill the buffer with the tail since we don't know about the next write.
89+
int idx = arrayContains(array, off, len, marker);
90+
if (idx >= 0) {
91+
// we have a full match. just write everything
92+
found = true;
93+
drain();
94+
downstream.write(array, off, idx);
95+
downstream.write(contentToInject);
96+
if (onContentInjected != null) {
97+
onContentInjected.run();
98+
}
99+
downstream.write(array, off + idx, len - idx);
100+
} else {
101+
// we don't have a full match. write everything in a bulk except the lookbehind buffer
102+
// sequentially
103+
for (int i = off; i < off + marker.length - 1; i++) {
104+
write(array[i]);
105+
}
106+
drain();
107+
downstream.write(array, off + marker.length - 1, len - bulkWriteThreshold);
108+
for (int i = len - marker.length + 1; i < len; i++) {
109+
write(array[i]);
110+
}
111+
}
112+
} else {
113+
// use slow path because the length to write is small and within the lookbehind buffer size
114+
for (int i = off; i < off + len; i++) {
115+
write(array[i]);
116+
}
117+
}
118+
}
119+
120+
private int arrayContains(char[] array, int off, int len, char[] search) {
121+
for (int i = off; i < len - search.length; i++) {
122+
if (array[i] == search[0]) {
123+
boolean found = true;
124+
int k = i;
125+
for (int j = 1; j < search.length; j++) {
126+
k++;
127+
if (array[k] != search[j]) {
128+
found = false;
129+
break;
130+
}
131+
}
132+
if (found) {
133+
return i;
134+
}
135+
}
136+
}
137+
return -1;
138+
}
139+
140+
private void drain() throws IOException {
141+
if (bufferFilled) {
142+
for (int i = 0; i < lookbehind.length; i++) {
143+
downstream.write(lookbehind[(pos + i) % lookbehind.length]);
144+
}
145+
} else {
146+
downstream.write(this.lookbehind, 0, pos);
147+
}
148+
pos = 0;
149+
matchingPos = 0;
150+
bufferFilled = false;
151+
}
152+
153+
@Override
154+
public void close() throws IOException {
155+
if (!found) {
156+
drain();
157+
}
158+
downstream.close();
159+
}
160+
}

‎dd-java-agent/agent-bootstrap/src/test/groovy/datadog/trace/bootstrap/instrumentation/buffer/InjectingPipeOutputStreamTest.groovy

Lines changed: 0 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -3,24 +3,6 @@ package datadog.trace.bootstrap.instrumentation.buffer
33
import datadog.trace.test.util.DDSpecification
44

55
class InjectingPipeOutputStreamTest extends DDSpecification {
6-
7-
static class ExceptionControlledOutputStream extends FilterOutputStream {
8-
9-
boolean failWrite = false
10-
11-
ExceptionControlledOutputStream(OutputStream out) {
12-
super(out)
13-
}
14-
15-
@Override
16-
void write(int b) throws IOException {
17-
if (failWrite) {
18-
throw new IOException("Failed")
19-
}
20-
super.write(b)
21-
}
22-
}
23-
246
def 'should filter a buffer and inject if found #found'() {
257
setup:
268
def downstream = new ByteArrayOutputStream()

‎dd-java-agent/agent-bootstrap/src/test/groovy/datadog/trace/bootstrap/instrumentation/buffer/InjectingPipeWriterTest.groovy

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
package datadog.trace.bootstrap.instrumentation.buffer
2+
3+
import datadog.trace.test.util.DDSpecification
4+
5+
class InjectingPipeWriterTest extends DDSpecification {
6+
def 'should filter a buffer and inject if found #found using write'() {
7+
setup:
8+
def downstream = new StringWriter()
9+
def piped = new PrintWriter(new InjectingPipeWriter(downstream, marker.toCharArray(), contentToInject.toCharArray(), null))
10+
when:
11+
try (def closeme = piped) {
12+
piped.write(body)
13+
}
14+
then:
15+
assert downstream.toString() == expected
16+
where:
17+
body | marker | contentToInject | found | expected
18+
"<html><head><foo/></head><body/></html>" | "</head>" | "<script>true</script>" | true | "<html><head><foo/><script>true</script></head><body/></html>"
19+
"<html><body/></html>" | "</head>" | "<something/>" | false | "<html><body/></html>"
20+
"<foo/>" | "<longerThanFoo>" | "<nothing>" | false | "<foo/>"
21+
}
22+
23+
def 'should filter a buffer and inject if found #found using append'() {
24+
setup:
25+
def downstream = new StringWriter()
26+
def piped = new PrintWriter(new InjectingPipeWriter(downstream, marker.toCharArray(), contentToInject.toCharArray(), null))
27+
when:
28+
try (def closeme = piped) {
29+
piped.append(body)
30+
}
31+
then:
32+
assert downstream.toString() == expected
33+
where:
34+
body | marker | contentToInject | found | expected
35+
"<html><head><foo/></head><body/></html>" | "</head>" | "<script>true</script>" | true | "<html><head><foo/><script>true</script></head><body/></html>"
36+
"<html><body/></html>" | "</head>" | "<something/>" | false | "<html><body/></html>"
37+
"<foo/>" | "<longerThanFoo>" | "<nothing>" | false | "<foo/>"
38+
}
39+
}

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

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package datadog.trace.instrumentation.servlet3;
22

33
import datadog.trace.api.rum.RumInjector;
4+
import datadog.trace.bootstrap.instrumentation.buffer.InjectingPipeWriter;
45
import java.io.IOException;
56
import java.io.PrintWriter;
67
import java.nio.charset.Charset;
@@ -41,11 +42,15 @@ public ServletOutputStream getOutputStream() throws IOException {
4142

4243
@Override
4344
public PrintWriter getWriter() throws IOException {
45+
final PrintWriter delegate = super.getWriter();
4446
if (!shouldInject) {
45-
return super.getWriter();
47+
return delegate;
4648
}
4749
if (printWriter == null) {
48-
printWriter = new PrintWriter(getOutputStream());
50+
printWriter =
51+
new PrintWriter(
52+
new InjectingPipeWriter(
53+
delegate, rumInjector.getMarker(), rumInjector.getSnippet(), this::onInjected));
4954
}
5055
return printWriter;
5156
}

0 commit comments

Comments
 (0)