Skip to content

Commit 7010cb8

Browse files
jandro996jandro996
authored
Attach stacktrace to IAST vulnerabilities (#7757)
What Does This Do Rework current stack trace utilities for RASP to: - Be accesible for IAST ( and other product in the future) - Add missing fields and highlight if they are optional - Replace the FlattenerObject strategy for the MsgPackWriter to avoid double serialisation - Be able to retrieve the metastruct map to share it between different products simultaneously Add stack trace to IAST vulnerabilities Add more tests Motivation Give more context to the source code vulnerabilities by adding the stack trace
1 parent ec5eedc commit 7010cb8

File tree

32 files changed

+729
-276
lines changed

32 files changed

+729
-276
lines changed

‎communication/src/main/java/datadog/communication/serialization/Codec.java‎

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,28 @@
11
package datadog.communication.serialization;
22

3+
import datadog.communication.serialization.custom.stacktrace.StackTraceEventFrameWriter;
4+
import datadog.communication.serialization.custom.stacktrace.StackTraceEventWriter;
5+
import datadog.trace.util.stacktrace.StackTraceEvent;
6+
import datadog.trace.util.stacktrace.StackTraceFrame;
37
import java.nio.ByteBuffer;
48
import java.nio.CharBuffer;
59
import java.util.Collection;
610
import java.util.Collections;
711
import java.util.Map;
12+
import java.util.stream.Collectors;
13+
import java.util.stream.Stream;
814

915
public final class Codec extends ClassValue<ValueWriter<?>> {
1016

11-
public static final Codec INSTANCE = new Codec();
17+
private static final Map<Class<?>, ValueWriter<?>> defaultConfig =
18+
Stream.of(
19+
new Object[][] {
20+
{StackTraceEvent.class, new StackTraceEventWriter()},
21+
{StackTraceFrame.class, new StackTraceEventFrameWriter()},
22+
})
23+
.collect(Collectors.toMap(data -> (Class<?>) data[0], data -> (ValueWriter<?>) data[1]));
24+
25+
public static final Codec INSTANCE = new Codec(defaultConfig);
1226

1327
private final Map<Class<?>, ValueWriter<?>> config;
1428

‎communication/src/main/java/datadog/communication/serialization/custom/stacktrace/StackTraceEventFrameWriter.java‎

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
1+
package datadog.communication.serialization.custom.stacktrace;
2+
3+
import datadog.communication.serialization.EncodingCache;
4+
import datadog.communication.serialization.ValueWriter;
5+
import datadog.communication.serialization.Writable;
6+
import datadog.trace.util.Strings;
7+
import datadog.trace.util.stacktrace.StackTraceFrame;
8+
9+
public class StackTraceEventFrameWriter implements ValueWriter<StackTraceFrame> {
10+
11+
@Override
12+
public void write(StackTraceFrame value, Writable writable, EncodingCache encodingCache) {
13+
int mapSize = 1; // id always present
14+
boolean hasText = Strings.isNotBlank(value.getText());
15+
boolean hasFile = Strings.isNotBlank(value.getFile());
16+
boolean hasLine = value.getLine() != null;
17+
boolean hasClass = Strings.isNotBlank(value.getClass_name());
18+
boolean hasFunction = Strings.isNotBlank(value.getFunction());
19+
if (hasText) {
20+
mapSize++;
21+
}
22+
if (hasFile) {
23+
mapSize++;
24+
}
25+
if (hasLine) {
26+
mapSize++;
27+
}
28+
if (hasClass) {
29+
mapSize++;
30+
}
31+
if (hasFunction) {
32+
mapSize++;
33+
}
34+
writable.startMap(mapSize);
35+
writable.writeString("id", encodingCache);
36+
writable.writeInt(value.getId());
37+
if (hasText) {
38+
writable.writeString("text", encodingCache);
39+
writable.writeString(value.getText(), encodingCache);
40+
}
41+
if (hasFile) {
42+
writable.writeString("file", encodingCache);
43+
writable.writeString(value.getFile(), encodingCache);
44+
}
45+
if (hasLine) {
46+
writable.writeString("line", encodingCache);
47+
writable.writeInt(value.getLine());
48+
}
49+
if (hasClass) {
50+
writable.writeString("class_name", encodingCache);
51+
writable.writeString(value.getClass_name(), encodingCache);
52+
}
53+
if (hasFunction) {
54+
writable.writeString("function", encodingCache);
55+
writable.writeString(value.getFunction(), encodingCache);
56+
}
57+
}
58+
}

‎communication/src/main/java/datadog/communication/serialization/custom/stacktrace/StackTraceEventWriter.java‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
package datadog.communication.serialization.custom.stacktrace;
2+
3+
import datadog.communication.serialization.EncodingCache;
4+
import datadog.communication.serialization.ValueWriter;
5+
import datadog.communication.serialization.Writable;
6+
import datadog.trace.util.Strings;
7+
import datadog.trace.util.stacktrace.StackTraceEvent;
8+
9+
public class StackTraceEventWriter implements ValueWriter<StackTraceEvent> {
10+
11+
@Override
12+
public void write(StackTraceEvent value, Writable writable, EncodingCache encodingCache) {
13+
int mapSize = 1; // frames always present
14+
boolean hasId = Strings.isNotBlank(value.getId());
15+
boolean hasLanguage = Strings.isNotBlank(value.getLanguage());
16+
boolean hasMessage = Strings.isNotBlank(value.getMessage());
17+
if (hasId) {
18+
mapSize++;
19+
}
20+
if (hasLanguage) {
21+
mapSize++;
22+
}
23+
if (hasMessage) {
24+
mapSize++;
25+
}
26+
writable.startMap(mapSize);
27+
if (hasId) {
28+
writable.writeString("id", encodingCache);
29+
writable.writeString(value.getId(), encodingCache);
30+
}
31+
if (hasLanguage) {
32+
writable.writeString("language", encodingCache);
33+
writable.writeString(value.getLanguage(), encodingCache);
34+
}
35+
if (hasMessage) {
36+
writable.writeString("message", encodingCache);
37+
writable.writeString(value.getMessage(), encodingCache);
38+
}
39+
writable.writeString("frames", encodingCache);
40+
writable.writeObject(value.getFrames(), encodingCache);
41+
}
42+
}

‎communication/src/test/java/datadog/communication/serialization/msgpack/MsgPackWriterTest.java‎

Lines changed: 176 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package datadog.communication.serialization.msgpack;
22

3+
import static datadog.trace.util.stacktrace.StackTraceEvent.DEFAULT_LANGUAGE;
34
import static org.junit.jupiter.api.Assertions.assertArrayEquals;
45
import static org.junit.jupiter.api.Assertions.assertEquals;
56
import static org.junit.jupiter.api.Assertions.assertFalse;
@@ -12,12 +13,16 @@
1213
import datadog.communication.serialization.MessageFormatter;
1314
import datadog.communication.serialization.StreamingBuffer;
1415
import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString;
16+
import datadog.trace.util.stacktrace.StackTraceEvent;
17+
import datadog.trace.util.stacktrace.StackTraceFrame;
1518
import java.io.IOException;
1619
import java.math.BigDecimal;
1720
import java.nio.ByteBuffer;
1821
import java.nio.CharBuffer;
1922
import java.nio.charset.StandardCharsets;
2023
import java.util.ArrayList;
24+
import java.util.Arrays;
25+
import java.util.Collections;
2126
import java.util.HashMap;
2227
import java.util.List;
2328
import java.util.Map;
@@ -786,6 +791,177 @@ public void testStartStringHeader() {
786791
writer.writeStringHeader(0x10000);
787792
}
788793

794+
@Test
795+
public void testStackTraceBatch() {
796+
Map<String, List<StackTraceEvent>> batch = new HashMap<>();
797+
batch.put("product_key", buildStacktraceEvents());
798+
testStackTraceBatch(batch);
799+
}
800+
801+
@Test
802+
public void tesEmptyContentInStackTraceBatch() {
803+
Map<String, List<StackTraceEvent>> batch = new HashMap<>();
804+
testStackTraceBatch(batch);
805+
}
806+
807+
@Test
808+
public void testSimpleStackTraceEvent() {
809+
final StackTraceFrame frame =
810+
new StackTraceFrame(1, new StackTraceElement("class1", "method1", "file1", 1));
811+
StackTraceEvent data = new StackTraceEvent(Collections.singletonList(frame), null, null, null);
812+
MessageFormatter messageFormatter =
813+
new MsgPackWriter(
814+
newBuffer(
815+
1000,
816+
(messageCount, buffer) -> {
817+
MessageUnpacker unpacker = MessagePack.newDefaultUnpacker(buffer);
818+
try {
819+
checkStackTraceEvent(data, unpacker);
820+
} catch (IOException e) {
821+
Assertions.fail(e.getMessage());
822+
}
823+
}));
824+
messageFormatter.format(data, (x, w) -> w.writeObject(x, null));
825+
messageFormatter.flush();
826+
}
827+
828+
@Test
829+
public void testSimpleStackTraceFrame() {
830+
final StackTraceFrame frame = new StackTraceFrame(1, null, "null", -1, null, null);
831+
MessageFormatter messageFormatter =
832+
new MsgPackWriter(
833+
newBuffer(
834+
1000,
835+
(messageCount, buffer) -> {
836+
MessageUnpacker unpacker = MessagePack.newDefaultUnpacker(buffer);
837+
try {
838+
checkStackTraceFrame(frame, unpacker);
839+
} catch (IOException e) {
840+
Assertions.fail(e.getMessage());
841+
}
842+
}));
843+
messageFormatter.format(frame, (x, w) -> w.writeObject(x, null));
844+
messageFormatter.flush();
845+
}
846+
847+
private void testStackTraceBatch(final Map<String, List<StackTraceEvent>> batch) {
848+
MessageFormatter messageFormatter =
849+
new MsgPackWriter(
850+
newBuffer(
851+
100000,
852+
(messageCount, buffer) -> {
853+
MessageUnpacker unpacker = MessagePack.newDefaultUnpacker(buffer);
854+
try {
855+
checkStackTraceBatch(batch, unpacker);
856+
} catch (IOException e) {
857+
Assertions.fail(e.getMessage());
858+
}
859+
}));
860+
messageFormatter.format(batch, (x, w) -> w.writeObject(x, null));
861+
messageFormatter.flush();
862+
}
863+
864+
private void checkStackTraceBatch(
865+
Map<String, List<StackTraceEvent>> batch, MessageUnpacker unpacker) throws IOException {
866+
867+
assertEquals(batch.size(), unpacker.unpackMapHeader());
868+
if (batch.isEmpty()) {
869+
return;
870+
}
871+
assertEquals("product_key", unpacker.unpackString());
872+
List<StackTraceEvent> events = batch.get("product_key");
873+
assertEquals(events.size(), unpacker.unpackArrayHeader());
874+
for (StackTraceEvent event : events) {
875+
checkStackTraceEvent(event, unpacker);
876+
}
877+
}
878+
879+
private void checkStackTraceEvent(StackTraceEvent event, MessageUnpacker unpacker)
880+
throws IOException {
881+
boolean hasId = event.getId() != null && !event.getId().isEmpty();
882+
boolean hasLanguage = event.getLanguage() != null && !event.getLanguage().isEmpty();
883+
boolean hasMessage = event.getMessage() != null && !event.getMessage().isEmpty();
884+
int mapSize = 1 + (hasId ? 1 : 0) + (hasLanguage ? 1 : 0) + (hasMessage ? 1 : 0);
885+
assertEquals(mapSize, unpacker.unpackMapHeader());
886+
if (hasId) {
887+
assertEquals("id", unpacker.unpackString());
888+
assertEquals(event.getId(), unpacker.unpackString());
889+
}
890+
if (hasLanguage) {
891+
assertEquals("language", unpacker.unpackString());
892+
assertEquals(event.getLanguage(), unpacker.unpackString());
893+
}
894+
if (hasMessage) {
895+
assertEquals("message", unpacker.unpackString());
896+
assertEquals(event.getMessage(), unpacker.unpackString());
897+
}
898+
assertEquals("frames", unpacker.unpackString());
899+
assertEquals(event.getFrames().size(), unpacker.unpackArrayHeader());
900+
for (StackTraceFrame frame : event.getFrames()) {
901+
checkStackTraceFrame(frame, unpacker);
902+
}
903+
}
904+
905+
private void checkStackTraceFrame(StackTraceFrame frame, MessageUnpacker unpacker)
906+
throws IOException {
907+
boolean hasText = frame.getText() != null && !frame.getText().isEmpty();
908+
boolean hasFile = frame.getFile() != null && !frame.getFile().isEmpty();
909+
boolean hasLine = frame.getLine() != null;
910+
boolean hasClass = frame.getClass_name() != null && !frame.getClass_name().isEmpty();
911+
boolean hasFunction = frame.getFunction() != null && !frame.getFunction().isEmpty();
912+
913+
int mapSize = 1; // id is always present
914+
if (hasText) {
915+
mapSize++;
916+
}
917+
if (hasFile) {
918+
mapSize++;
919+
}
920+
if (hasLine) {
921+
mapSize++;
922+
}
923+
if (hasClass) {
924+
mapSize++;
925+
}
926+
if (hasFunction) {
927+
mapSize++;
928+
}
929+
assertEquals(mapSize, unpacker.unpackMapHeader());
930+
assertEquals("id", unpacker.unpackString());
931+
assertEquals(frame.getId(), unpacker.unpackInt());
932+
if (hasText) {
933+
assertEquals("text", unpacker.unpackString());
934+
assertEquals(frame.getText(), unpacker.unpackString());
935+
}
936+
if (hasFile) {
937+
assertEquals("file", unpacker.unpackString());
938+
assertEquals(frame.getFile(), unpacker.unpackString());
939+
}
940+
if (hasLine) {
941+
assertEquals("line", unpacker.unpackString());
942+
assertEquals(frame.getLine(), unpacker.unpackInt());
943+
}
944+
if (hasClass) {
945+
assertEquals("class_name", unpacker.unpackString());
946+
assertEquals(frame.getClass_name(), unpacker.unpackString());
947+
}
948+
if (hasFunction) {
949+
assertEquals("function", unpacker.unpackString());
950+
assertEquals(frame.getFunction(), unpacker.unpackString());
951+
}
952+
}
953+
954+
private List<StackTraceEvent> buildStacktraceEvents() {
955+
final StackTraceFrame frame1 =
956+
new StackTraceFrame(1, new StackTraceElement("class1", "method1", "file1", 1));
957+
final StackTraceFrame frame2 =
958+
new StackTraceFrame(2, new StackTraceElement("class2", "method2", "file2", 2));
959+
960+
return Arrays.asList(
961+
new StackTraceEvent(Arrays.asList(frame1, frame2), DEFAULT_LANGUAGE, "id1", "event1"),
962+
new StackTraceEvent(Arrays.asList(frame1, frame2), DEFAULT_LANGUAGE, "id2", "event2"));
963+
}
964+
789965
private StreamingBuffer newBuffer(int capacity, ByteBufferConsumer consumer) {
790966
return new FlushingBuffer(capacity, consumer);
791967
}

‎dd-java-agent/agent-iast/src/main/java/com/datadog/iast/Reporter.java‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import static com.datadog.iast.IastTag.Enabled.ANALYZED;
44
import static datadog.trace.api.telemetry.LogCollector.SEND_TELEMETRY;
5+
import static datadog.trace.util.stacktrace.StackTraceEvent.DEFAULT_LANGUAGE;
56

67
import com.datadog.iast.model.Vulnerability;
78
import com.datadog.iast.model.VulnerabilityBatch;
@@ -12,6 +13,11 @@
1213
import datadog.trace.api.internal.TraceSegment;
1314
import datadog.trace.bootstrap.instrumentation.api.*;
1415
import datadog.trace.util.AgentTaskScheduler;
16+
import datadog.trace.util.stacktrace.StackTraceEvent;
17+
import datadog.trace.util.stacktrace.StackTraceFrame;
18+
import datadog.trace.util.stacktrace.StackUtils;
19+
import java.util.Collections;
20+
import java.util.List;
1521
import java.util.Set;
1622
import java.util.concurrent.ConcurrentHashMap;
1723
import java.util.concurrent.TimeUnit;
@@ -29,6 +35,7 @@ public class Reporter {
2935
private static final String IAST_TAG = "iast";
3036

3137
private static final String VULNERABILITY_SPAN_NAME = "vulnerability";
38+
public static final String METASTRUCT_VULNERABILITY = "vulnerability";
3239

3340
private final Predicate<Vulnerability> duplicated;
3441

@@ -69,7 +76,27 @@ private void reportVulnerability(
6976
final VulnerabilityBatch batch = getOrCreateVulnerabilityBatch(span);
7077
if (batch != null) {
7178
batch.add(vulnerability);
79+
if (Config.get().isIastStackTraceEnabled() && batch.getVulnerabilities() != null) {
80+
String stackId =
81+
addVulnerabilityStackTrace(span, String.valueOf(batch.getVulnerabilities().size()));
82+
if (stackId != null) {
83+
vulnerability.setStackId(stackId);
84+
}
85+
}
86+
}
87+
}
88+
89+
@Nullable
90+
private static String addVulnerabilityStackTrace(@Nonnull AgentSpan span, String index) {
91+
final RequestContext reqCtx = span.getRequestContext();
92+
if (reqCtx == null) {
93+
return null;
7294
}
95+
List<StackTraceFrame> frames = StackUtils.generateUserCodeStackTrace();
96+
StackTraceEvent stackTraceEvent = new StackTraceEvent(frames, DEFAULT_LANGUAGE, index, null);
97+
StackUtils.addStacktraceEventsToMetaStruct(
98+
reqCtx, METASTRUCT_VULNERABILITY, Collections.singletonList(stackTraceEvent));
99+
return stackTraceEvent.getId();
73100
}
74101

75102
@Nullable

0 commit comments

Comments
 (0)