Skip to content

Commit 8f323b5

Browse files
pawelchckidevflow.devflow-routing-intake
pawelchcki
and
devflow.devflow-routing-intake
authored
Fix RUM injection for offset response writes (#11949)
Add RUM offset write regression coverage Fix RUM injection for offset writes Merge branch 'master' into codex/fix-rum-injection-offset-writes Preserve boundary-spanning RUM injection matches Refine RUM injection offset writes Consolidate Spring Boot WebMVC smoke tests Co-authored-by: devflow.devflow-routing-intake <devflow.devflow-routing-intake@kubernetes.us1.ddbuild.io>
1 parent d3d2e12 commit 8f323b5

8 files changed

Lines changed: 267 additions & 27 deletions

File tree

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

Lines changed: 32 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -118,6 +118,7 @@ public void write(int b) throws IOException {
118118

119119
@Override
120120
public void write(byte[] array, int off, int len) throws IOException {
121+
final int end = off + len;
121122
if (!filter) {
122123
if (wasDraining) {
123124
// needs drain
@@ -129,15 +130,25 @@ public void write(byte[] array, int off, int len) throws IOException {
129130
}
130131

131132
if (len > bulkWriteThreshold) {
133+
// A match that started in the previous write precedes every match wholly in this array.
134+
if (matchingPos > 0 && arrayCompletesPendingMatch(array, off, end)) {
135+
int pendingMatchLength = marker.length - matchingPos;
136+
for (int i = off; i < off + pendingMatchLength; i++) {
137+
write(array[i]);
138+
}
139+
write(array, off + pendingMatchLength, len - pendingMatchLength);
140+
return;
141+
}
142+
132143
// if the content is large enough, we can bulk write everything but the N trail and tail.
133144
// This because the buffer can already contain some byte from a previous single write.
134145
// Also we need to fill the buffer with the tail since we don't know about the next write.
135-
int idx = arrayContains(array, off, len, marker);
146+
int idx = arrayContains(array, off, end, marker);
136147
if (idx >= 0) {
137148
// we have a full match. just write everything
138149
filter = false;
139150
drain();
140-
int bytesToWrite = idx;
151+
int bytesToWrite = idx - off;
141152
downstream.write(array, off, bytesToWrite);
142153
bytesWritten += bytesToWrite;
143154
long injectionStart = System.nanoTime();
@@ -149,8 +160,8 @@ public void write(byte[] array, int off, int len) throws IOException {
149160
if (onContentInjected != null) {
150161
onContentInjected.run();
151162
}
152-
bytesToWrite = len - idx;
153-
downstream.write(array, off + idx, bytesToWrite);
163+
bytesToWrite = end - idx;
164+
downstream.write(array, idx, bytesToWrite);
154165
bytesWritten += bytesToWrite;
155166
} else {
156167
// we don't have a full match. write everything in a bulk except the lookbehind buffer
@@ -167,20 +178,33 @@ public void write(byte[] array, int off, int len) throws IOException {
167178
downstream.write(array, off + marker.length - 1, bytesToWrite);
168179
bytesWritten += bytesToWrite;
169180
filter = wasFiltering;
170-
for (int i = len - marker.length + 1; i < len; i++) {
181+
for (int i = end - marker.length + 1; i < end; i++) {
171182
write(array[i]);
172183
}
173184
}
174185
} else {
175186
// use slow path because the length to write is small and within the lookbehind buffer size
176-
for (int i = off; i < off + len; i++) {
187+
for (int i = off; i < end; i++) {
177188
write(array[i]);
178189
}
179190
}
180191
}
181192

182-
private int arrayContains(byte[] array, int off, int len, byte[] search) {
183-
for (int i = off; i < len - search.length; i++) {
193+
private boolean arrayCompletesPendingMatch(byte[] array, int off, int end) {
194+
int pendingMatchLength = marker.length - matchingPos;
195+
if (off + pendingMatchLength > end) {
196+
return false;
197+
}
198+
for (int i = 0; i < pendingMatchLength; i++) {
199+
if (array[off + i] != marker[matchingPos + i]) {
200+
return false;
201+
}
202+
}
203+
return true;
204+
}
205+
206+
private int arrayContains(byte[] array, int off, int end, byte[] search) {
207+
for (int i = off; i <= end - search.length; i++) {
184208
if (array[i] == search[0]) {
185209
boolean found = true;
186210
int k = i;

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

Lines changed: 32 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -118,6 +118,7 @@ public void write(int c) throws IOException {
118118

119119
@Override
120120
public void write(char[] array, int off, int len) throws IOException {
121+
final int end = off + len;
121122
if (!filter) {
122123
if (wasDraining) {
123124
// needs drain
@@ -129,15 +130,25 @@ public void write(char[] array, int off, int len) throws IOException {
129130
}
130131

131132
if (len > bulkWriteThreshold) {
133+
// A match that started in the previous write precedes every match wholly in this array.
134+
if (matchingPos > 0 && arrayCompletesPendingMatch(array, off, end)) {
135+
int pendingMatchLength = marker.length - matchingPos;
136+
for (int i = off; i < off + pendingMatchLength; i++) {
137+
write(array[i]);
138+
}
139+
write(array, off + pendingMatchLength, len - pendingMatchLength);
140+
return;
141+
}
142+
132143
// if the content is large enough, we can bulk write everything but the N trail and tail.
133144
// This because the buffer can already contain some byte from a previous single write.
134145
// Also we need to fill the buffer with the tail since we don't know about the next write.
135-
int idx = arrayContains(array, off, len, marker);
146+
int idx = arrayContains(array, off, end, marker);
136147
if (idx >= 0) {
137148
// we have a full match. just write everything
138149
filter = false;
139150
drain();
140-
int bytesToWrite = idx;
151+
int bytesToWrite = idx - off;
141152
downstream.write(array, off, bytesToWrite);
142153
bytesWritten += bytesToWrite;
143154
long injectionStart = System.nanoTime();
@@ -149,8 +160,8 @@ public void write(char[] array, int off, int len) throws IOException {
149160
if (onContentInjected != null) {
150161
onContentInjected.run();
151162
}
152-
bytesToWrite = len - idx;
153-
downstream.write(array, off + idx, bytesToWrite);
163+
bytesToWrite = end - idx;
164+
downstream.write(array, idx, bytesToWrite);
154165
bytesWritten += bytesToWrite;
155166
} else {
156167
// we don't have a full match. write everything in a bulk except the lookbehind buffer
@@ -168,20 +179,33 @@ public void write(char[] array, int off, int len) throws IOException {
168179
bytesWritten += bytesToWrite;
169180
filter = wasFiltering;
170181

171-
for (int i = len - marker.length + 1; i < len; i++) {
182+
for (int i = end - marker.length + 1; i < end; i++) {
172183
write(array[i]);
173184
}
174185
}
175186
} else {
176187
// use slow path because the length to write is small and within the lookbehind buffer size
177-
for (int i = off; i < off + len; i++) {
188+
for (int i = off; i < end; i++) {
178189
write(array[i]);
179190
}
180191
}
181192
}
182193

183-
private int arrayContains(char[] array, int off, int len, char[] search) {
184-
for (int i = off; i < len - search.length; i++) {
194+
private boolean arrayCompletesPendingMatch(char[] array, int off, int end) {
195+
int pendingMatchLength = marker.length - matchingPos;
196+
if (off + pendingMatchLength > end) {
197+
return false;
198+
}
199+
for (int i = 0; i < pendingMatchLength; i++) {
200+
if (array[off + i] != marker[matchingPos + i]) {
201+
return false;
202+
}
203+
}
204+
return true;
205+
}
206+
207+
private int arrayContains(char[] array, int off, int end, char[] search) {
208+
for (int i = off; i <= end - search.length; i++) {
185209
if (array[i] == search[0]) {
186210
boolean found = true;
187211
int k = i;

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

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package datadog.trace.bootstrap.instrumentation.buffer
22

33
import datadog.trace.test.util.DDSpecification
4+
import java.util.concurrent.atomic.AtomicInteger
45
import java.util.function.LongConsumer
56

67
class InjectingPipeOutputStreamTest extends DDSpecification {
@@ -152,6 +153,49 @@ class InjectingPipeOutputStreamTest extends DDSpecification {
152153
downstream.toByteArray() == testBytes.getBytes("UTF-8")
153154
}
154155

156+
def 'should honor non-zero offsets in bulk writes: #scenario'() {
157+
setup:
158+
def prefix = "ignored-prefix"
159+
def payloadBytes = payload.getBytes("UTF-8")
160+
def source = (prefix + payload + "ignored-suffix").getBytes("UTF-8")
161+
def downstream = new ByteArrayOutputStream()
162+
def counter = new Counter()
163+
def injections = new AtomicInteger()
164+
def piped = new InjectingPipeOutputStream(downstream, MARKER_BYTES, CONTEXT_BYTES, injections.&incrementAndGet, { long bytes -> counter.incr(bytes) }, null)
165+
166+
when:
167+
piped.write(source, prefix.getBytes("UTF-8").length, payloadBytes.length)
168+
piped.close()
169+
170+
then:
171+
downstream.toByteArray() == expected.getBytes("UTF-8")
172+
injections.get() == expectedInjections
173+
counter.value == payloadBytes.length
174+
175+
where:
176+
scenario | payload | expected | expectedInjections
177+
"without a marker" | "<html><body>safe</body></html>" | "<html><body>safe</body></html>" | 0
178+
"with a marker" | "<html><head>dynamic</head><body>safe</body></html>" | "<html><head>dynamic<script></script></head><body>safe</body></html>" | 1
179+
"with a marker at the end" | "<html><head>dynamic-content</head>" | "<html><head>dynamic-content<script></script></head>" | 1
180+
}
181+
182+
def 'should prioritize a marker spanning a bulk write boundary'() {
183+
setup:
184+
def prefix = "ignored-prefix"
185+
def payload = ">0123456789</head>"
186+
def source = (prefix + payload + "ignored-suffix").getBytes("UTF-8")
187+
def downstream = new ByteArrayOutputStream()
188+
def piped = new InjectingPipeOutputStream(downstream, MARKER_BYTES, CONTEXT_BYTES)
189+
190+
when:
191+
piped.write("abc</head".getBytes("UTF-8"))
192+
piped.write(source, prefix.getBytes("UTF-8").length, payload.getBytes("UTF-8").length)
193+
piped.close()
194+
195+
then:
196+
downstream.toByteArray() == "abc<script></script></head>0123456789</head>".getBytes("UTF-8")
197+
}
198+
155199
def 'should be resilient to exceptions when onBytesWritten callback is null'() {
156200
setup:
157201
def testBytes = "test content".getBytes("UTF-8")

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

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package datadog.trace.bootstrap.instrumentation.buffer
22

33
import datadog.trace.test.util.DDSpecification
4+
import java.util.concurrent.atomic.AtomicInteger
45
import java.util.function.LongConsumer
56

67
class InjectingPipeWriterTest extends DDSpecification {
@@ -168,6 +169,48 @@ class InjectingPipeWriterTest extends DDSpecification {
168169
downstream.toString() == testBytes
169170
}
170171

172+
def 'should honor non-zero offsets in bulk writes: #scenario'() {
173+
setup:
174+
def prefix = "ignored-prefix"
175+
def source = (prefix + payload + "ignored-suffix").toCharArray()
176+
def downstream = new StringWriter()
177+
def counter = new Counter()
178+
def injections = new AtomicInteger()
179+
def piped = new InjectingPipeWriter(downstream, MARKER_CHARS, CONTEXT_CHARS, injections.&incrementAndGet, { long bytes -> counter.incr(bytes) }, null)
180+
181+
when:
182+
piped.write(source, prefix.length(), payload.length())
183+
piped.close()
184+
185+
then:
186+
downstream.toString() == expected
187+
injections.get() == expectedInjections
188+
counter.value == payload.length()
189+
190+
where:
191+
scenario | payload | expected | expectedInjections
192+
"without a marker" | "<html><body>safe</body></html>" | "<html><body>safe</body></html>" | 0
193+
"with a marker" | "<html><head>dynamic</head><body>safe</body></html>" | "<html><head>dynamic<script></script></head><body>safe</body></html>" | 1
194+
"with a marker at the end" | "<html><head>dynamic-content</head>" | "<html><head>dynamic-content<script></script></head>" | 1
195+
}
196+
197+
def 'should prioritize a marker spanning a bulk write boundary'() {
198+
setup:
199+
def prefix = "ignored-prefix"
200+
def payload = ">0123456789</head>"
201+
def source = (prefix + payload + "ignored-suffix").toCharArray()
202+
def downstream = new StringWriter()
203+
def piped = new InjectingPipeWriter(downstream, MARKER_CHARS, CONTEXT_CHARS)
204+
205+
when:
206+
piped.write("abc</head".toCharArray())
207+
piped.write(source, prefix.length(), payload.length())
208+
piped.close()
209+
210+
then:
211+
downstream.toString() == "abc<script></script></head>0123456789</head>"
212+
}
213+
171214
def 'should be resilient to exceptions when onBytesWritten callback is null'() {
172215
setup:
173216
def downstream = new StringWriter()

‎dd-smoke-tests/spring-boot-3.3-webmvc/application/build.gradle‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ java {
2424

2525
dependencies {
2626
implementation 'org.springframework.boot:spring-boot-starter-web'
27+
implementation 'org.springframework.boot:spring-boot-starter-thymeleaf'
2728
implementation 'org.springframework.boot:spring-boot-starter-data-jpa'
2829
implementation group: 'com.h2database', name: 'h2', version: '2.1.214'
2930
compileOnly group:"com.google.code.findbugs", name:"jsr305", version:"3.0.2"

‎dd-smoke-tests/spring-boot-3.3-webmvc/application/src/main/java/datadog/smoketest/springboot/controller/FormResponseController.java‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
package datadog.smoketest.springboot.controller;
2+
3+
import jakarta.servlet.http.HttpServletResponse;
4+
import java.io.IOException;
5+
import java.util.Locale;
6+
import org.springframework.stereotype.Controller;
7+
import org.springframework.web.bind.annotation.GetMapping;
8+
import org.thymeleaf.context.Context;
9+
import org.thymeleaf.spring6.SpringTemplateEngine;
10+
11+
@Controller
12+
public class FormResponseController {
13+
private static final int RESPONSE_OFFSET = 6;
14+
15+
private final SpringTemplateEngine templateEngine;
16+
17+
public FormResponseController(SpringTemplateEngine templateEngine) {
18+
this.templateEngine = templateEngine;
19+
}
20+
21+
@GetMapping("/form-response")
22+
public void formResponse(HttpServletResponse response) throws IOException {
23+
Context context = new Context(Locale.ROOT);
24+
context.setVariable("returnUrl", "https://app.example.test/flows/complete?request=request-123");
25+
context.setVariable("formAction", "https://provider.example.test/flows/continue");
26+
context.setVariable("requestId", "request-123");
27+
28+
String content = templateEngine.process("form-response", context);
29+
// Exercise response writers that receive a slice of a reusable buffer.
30+
char[] responseBuffer = new char[RESPONSE_OFFSET + content.length()];
31+
content.getChars(0, content.length(), responseBuffer, RESPONSE_OFFSET);
32+
33+
response.setContentType("text/html");
34+
response.getWriter().write(responseBuffer, RESPONSE_OFFSET, content.length());
35+
}
36+
}

‎dd-smoke-tests/spring-boot-3.3-webmvc/application/src/main/resources/templates/form-response.html‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
<!DOCTYPE html>
2+
<html lang="en" xmlns:th="http://www.thymeleaf.org">
3+
<head>
4+
<meta charset="UTF-8">
5+
<title>Continue request</title>
6+
<script th:inline="javascript">
7+
const returnUrl = /*[[${returnUrl}]]*/ "https://app.example.test/flows/complete";
8+
window.formResponse = { returnUrl: returnUrl };
9+
</script>
10+
</head>
11+
<body onload="document.getElementById('response-form').submit()">
12+
<form id="response-form" method="post" th:action="${formAction}">
13+
<input type="hidden" name="requestId" th:value="${requestId}">
14+
</form>
15+
</body>
16+
</html>

0 commit comments

Comments
 (0)