Skip to content

Commit ca3b9ed

Browse files
authored
Merge branch 'master' into alejandro.gonzalez/add-http-route-play-in-iast
2 parents 39bf456 + 16eee52 commit ca3b9ed

61 files changed

Lines changed: 2118 additions & 28 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

.github/workflows/analyze-changes.yaml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -109,7 +109,7 @@ jobs:
109109
ls -laR "./workspace/.trivy"
110110
111111
- name: Run Trivy security scanner
112-
uses: aquasecurity/trivy-action@76071ef0d7ec797419534a183b498b4d6366cf37 # v0.31.0
112+
uses: aquasecurity/trivy-action@dc5a429b52fcf669ce959baa2c2dd26090d2a6c4 # v0.32.0
113113
with:
114114
scan-type: rootfs
115115
scan-ref: './workspace/.trivy/'

.gitignore

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,3 +68,11 @@ replay_pid*
6868
benchmark/reports
6969
benchmark/tracer
7070
benchmark/dacapo/scratch
71+
72+
# JDK provisioning tools #
73+
# mise
74+
mise*.local.toml
75+
.mise*.local.toml
76+
.config/mise*.toml
77+
# asdf
78+
.tool-versions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,65 @@
1+
package datadog.trace.bootstrap.instrumentation.buffer;
2+
3+
import static java.util.concurrent.TimeUnit.MICROSECONDS;
4+
import static java.util.concurrent.TimeUnit.SECONDS;
5+
6+
import java.io.ByteArrayOutputStream;
7+
import java.io.IOException;
8+
import java.io.InputStream;
9+
import java.io.PrintWriter;
10+
import java.net.URL;
11+
import java.nio.charset.StandardCharsets;
12+
import java.util.List;
13+
import org.apache.commons.io.IOUtils;
14+
import org.openjdk.jmh.annotations.Benchmark;
15+
import org.openjdk.jmh.annotations.BenchmarkMode;
16+
import org.openjdk.jmh.annotations.Fork;
17+
import org.openjdk.jmh.annotations.Measurement;
18+
import org.openjdk.jmh.annotations.Mode;
19+
import org.openjdk.jmh.annotations.OutputTimeUnit;
20+
import org.openjdk.jmh.annotations.Scope;
21+
import org.openjdk.jmh.annotations.State;
22+
import org.openjdk.jmh.annotations.Warmup;
23+
24+
/*
25+
* Benchmark Mode Cnt Score Error Units
26+
* InjectingPipeOutputStreamBenchmark.withPipe avgt 2 15.515 us/op
27+
* InjectingPipeOutputStreamBenchmark.withoutPipe avgt 2 12.861 us/op
28+
*/
29+
@State(Scope.Benchmark)
30+
@Warmup(iterations = 1, time = 30, timeUnit = SECONDS)
31+
@Measurement(iterations = 2, time = 30, timeUnit = SECONDS)
32+
@BenchmarkMode(Mode.AverageTime)
33+
@OutputTimeUnit(MICROSECONDS)
34+
@Fork(value = 1)
35+
public class InjectingPipeOutputStreamBenchmark {
36+
private static final List<String> htmlContent;
37+
private static final byte[] marker;
38+
private static final byte[] content;
39+
40+
static {
41+
try (InputStream is = new URL("https://www.google.com").openStream()) {
42+
htmlContent = IOUtils.readLines(is, StandardCharsets.UTF_8);
43+
} catch (IOException ioe) {
44+
throw new RuntimeException(ioe);
45+
}
46+
marker = "</head>".getBytes(StandardCharsets.UTF_8);
47+
content = "<script/>".getBytes(StandardCharsets.UTF_8);
48+
}
49+
50+
@Benchmark
51+
public void withPipe() throws Exception {
52+
try (final PrintWriter out =
53+
new PrintWriter(
54+
new InjectingPipeOutputStream(new ByteArrayOutputStream(), marker, content, null))) {
55+
htmlContent.forEach(out::println);
56+
}
57+
}
58+
59+
@Benchmark
60+
public void withoutPipe() throws Exception {
61+
try (final PrintWriter out = new PrintWriter(new ByteArrayOutputStream())) {
62+
htmlContent.forEach(out::println);
63+
}
64+
}
65+
}
Lines changed: 153 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,153 @@
1+
package datadog.trace.bootstrap.instrumentation.buffer;
2+
3+
import java.io.FilterOutputStream;
4+
import java.io.IOException;
5+
import java.io.OutputStream;
6+
7+
/**
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.
10+
*/
11+
public class InjectingPipeOutputStream extends FilterOutputStream {
12+
private final byte[] lookbehind;
13+
private int pos;
14+
private boolean bufferFilled;
15+
private final byte[] marker;
16+
private final byte[] contentToInject;
17+
private boolean found = false;
18+
private int matchingPos = 0;
19+
private final Runnable onContentInjected;
20+
private final int bulkWriteThreshold;
21+
22+
/**
23+
* @param downstream the delegate output stream
24+
* @param marker the marker to find in the stream
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 InjectingPipeOutputStream(
29+
final OutputStream downstream,
30+
final byte[] marker,
31+
final byte[] contentToInject,
32+
final Runnable onContentInjected) {
33+
super(downstream);
34+
this.marker = marker;
35+
this.lookbehind = new byte[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+
out.write(b);
46+
return;
47+
}
48+
49+
if (bufferFilled) {
50+
out.write(lookbehind[pos]);
51+
}
52+
53+
lookbehind[pos] = (byte) 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+
out.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 write(byte[] b, int off, int len) throws IOException {
76+
if (found) {
77+
out.write(b, off, len);
78+
return;
79+
}
80+
if (len > bulkWriteThreshold) {
81+
// if the content is large enough, we can bulk write everything but the N trail and tail.
82+
// This because the buffer can already contain some byte from a previous single write.
83+
// 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);
85+
if (idx >= 0) {
86+
// we have a full match. just write everything
87+
found = true;
88+
drain();
89+
out.write(b, off, idx);
90+
out.write(contentToInject);
91+
if (onContentInjected != null) {
92+
onContentInjected.run();
93+
}
94+
out.write(b, off + idx, len - idx);
95+
} else {
96+
// we don't have a full match. write everything in a bulk except the lookbehind buffer
97+
// sequentially
98+
for (int i = off; i < off + marker.length - 1; i++) {
99+
write(b[i]);
100+
}
101+
drain();
102+
out.write(b, off + marker.length - 1, len - bulkWriteThreshold);
103+
for (int i = len - marker.length + 1; i < len; i++) {
104+
write(b[i]);
105+
}
106+
}
107+
} else {
108+
// use slow path because the length to write is small and within the lookbehind buffer size
109+
super.write(b, off, len);
110+
}
111+
}
112+
113+
private int arrayContains(byte[] array, byte[] search) {
114+
for (int i = 0; i < array.length - search.length; i++) {
115+
if (array[i] == search[0]) {
116+
boolean found = true;
117+
int k = i;
118+
for (int j = 1; j < search.length; j++) {
119+
k++;
120+
if (array[k] != search[j]) {
121+
found = false;
122+
break;
123+
}
124+
}
125+
if (found) {
126+
return i;
127+
}
128+
}
129+
}
130+
return -1;
131+
}
132+
133+
private void drain() throws IOException {
134+
if (bufferFilled) {
135+
for (int i = 0; i < lookbehind.length; i++) {
136+
out.write(lookbehind[(pos + i) % lookbehind.length]);
137+
}
138+
} else {
139+
out.write(this.lookbehind, 0, pos);
140+
}
141+
pos = 0;
142+
matchingPos = 0;
143+
bufferFilled = false;
144+
}
145+
146+
@Override
147+
public void close() throws IOException {
148+
if (!found) {
149+
drain();
150+
}
151+
super.close();
152+
}
153+
}

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,7 @@ public abstract class HttpServerDecorator<REQUEST, CONNECTION, RESPONSE, REQUEST
5757

5858
public static final String DD_SPAN_ATTRIBUTE = "datadog.span";
5959
public static final String DD_DISPATCH_SPAN_ATTRIBUTE = "datadog.span.dispatch";
60+
public static final String DD_RUM_INJECTED = "datadog.rum.injected";
6061
public static final String DD_FIN_DISP_LIST_SPAN_ATTRIBUTE =
6162
"datadog.span.finish_dispatch_listener";
6263
public static final String DD_RESPONSE_ATTRIBUTE = "datadog.response";
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
package datadog.trace.bootstrap.instrumentation.buffer
2+
3+
import datadog.trace.test.util.DDSpecification
4+
5+
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+
24+
def 'should filter a buffer and inject if found #found'() {
25+
setup:
26+
def downstream = new ByteArrayOutputStream()
27+
def piped = new OutputStreamWriter(new InjectingPipeOutputStream(downstream, marker.getBytes("UTF-8"), contentToInject.getBytes("UTF-8"), null),
28+
"UTF-8")
29+
when:
30+
try (def closeme = piped) {
31+
piped.write(body)
32+
}
33+
then:
34+
assert downstream.toByteArray() == expected.getBytes("UTF-8")
35+
where:
36+
body | marker | contentToInject | found | expected
37+
"<html><head><foo/></head><body/></html>" | "</head>" | "<script>true</script>" | true | "<html><head><foo/><script>true</script></head><body/></html>"
38+
"<html><body/></html>" | "</head>" | "<something/>" | false | "<html><body/></html>"
39+
"<foo/>" | "<longerThanFoo>" | "<nothing>" | false | "<foo/>"
40+
}
41+
}

dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/agent/ClassesToRetransformFinder.java

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -111,4 +111,8 @@ private static boolean lookupClass(Trie changedClasses, Class<?> clazz) {
111111
String simpleName = extractSimpleName(clazz);
112112
return changedClasses.contains(reverseStr(simpleName));
113113
}
114+
115+
ConcurrentMap<String, String> getClassNamesBySourceFile() {
116+
return classNamesBySourceFile;
117+
}
114118
}

dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/agent/DebuggerAgent.java

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -324,6 +324,10 @@ private static String getDiagnosticEndpoint(
324324

325325
private static void setupSourceFileTracking(
326326
Instrumentation instrumentation, ClassesToRetransformFinder finder) {
327+
if (!Config.get().isDebuggerSourceFileTrackingEnabled()) {
328+
LOGGER.debug("Source file tracking is disabled");
329+
return;
330+
}
327331
SourceFileTrackingTransformer sourceFileTrackingTransformer =
328332
new SourceFileTrackingTransformer(finder);
329333
sourceFileTrackingTransformer.start();

dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/agent/SourceFileTrackingTransformer.java

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,9 @@ private void registerSourceFile(String className, byte[] classfileBuffer) {
7575
if (sourceFile == null) {
7676
return;
7777
}
78+
if (!isExtensionAllowed(sourceFile)) {
79+
return;
80+
}
7881
String simpleClassName = stripPackagePath(className);
7982
String simpleSourceFile = removeExtension(sourceFile);
8083
if (simpleClassName.equals(simpleSourceFile)) {
@@ -83,6 +86,13 @@ private void registerSourceFile(String className, byte[] classfileBuffer) {
8386
finder.register(sourceFile, className);
8487
}
8588

89+
private boolean isExtensionAllowed(String sourceFile) {
90+
return sourceFile.endsWith(".java")
91+
|| sourceFile.endsWith(".kt")
92+
|| sourceFile.endsWith(".scala")
93+
|| sourceFile.endsWith(".groovy");
94+
}
95+
8696
private static class SourceFileItem {
8797
final String className;
8898
final byte[] classfileBuffer;

dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/agent/SourceFileTrackingTransformerTest.java

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,38 @@ void transformInner() throws IllegalClassFormatException {
7979
assertEquals(InnerHelper.MySecondInner.class, changedClasses.get(2));
8080
}
8181

82+
@Test
83+
void transformNotAllowed() throws IllegalClassFormatException {
84+
ClassesToRetransformFinder finder = new ClassesToRetransformFinder();
85+
SourceFileTrackingTransformer sourceFileTrackingTransformer =
86+
new SourceFileTrackingTransformer(finder);
87+
ConfigurationComparer comparer = createComparer("TopLevelHelper.java");
88+
byte[] classFileBytes = getClassFileBytes(TopLevelHelper.class);
89+
replaceInByteArray(
90+
classFileBytes, "TopLevelHelper.java".getBytes(), "TopLevelHelper.cloj".getBytes());
91+
sourceFileTrackingTransformer.transform(null, "", null, null, classFileBytes);
92+
sourceFileTrackingTransformer.flush();
93+
List<Class<?>> changedClasses =
94+
finder.getAllLoadedChangedClasses(new Class[] {InnerHelper.class}, comparer);
95+
assertEquals(0, finder.getClassNamesBySourceFile().size());
96+
}
97+
98+
private static void replaceInByteArray(byte[] buffer, byte[] oldBytes, byte[] newBytes) {
99+
int oldIdx = 0;
100+
for (int i = 0; i < buffer.length; i++) {
101+
if (buffer[i] == oldBytes[oldIdx]) {
102+
oldIdx++;
103+
if (oldIdx == oldBytes.length) {
104+
// Found the oldBytes, replace with newBytes
105+
System.arraycopy(newBytes, 0, buffer, i - oldIdx + 1, newBytes.length);
106+
oldIdx = 0; // Reset for next search
107+
}
108+
} else {
109+
oldIdx = 0; // Reset if current byte does not match
110+
}
111+
}
112+
}
113+
82114
private ConfigurationComparer createComparer(String sourceFile) {
83115
Configuration emptyConfig = Configuration.builder().setService("service-name").build();
84116
Configuration newConfig =

0 commit comments

Comments
 (0)