Skip to content

Commit

Permalink
Collect and report RASP events including Stack traces
Browse files Browse the repository at this point in the history
  • Loading branch information
ValentinZakharov committed Jun 11, 2024
1 parent 54461ee commit 324d272
Show file tree
Hide file tree
Showing 9 changed files with 296 additions and 11 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,10 @@
import com.datadog.appsec.event.data.Address;
import com.datadog.appsec.event.data.DataBundle;
import com.datadog.appsec.report.AppSecEvent;
import com.datadog.appsec.stack_trace.StackTraceCollection;
import com.datadog.appsec.stack_trace.StackTraceEvent;
import com.datadog.appsec.util.StandardizedLogging;
import datadog.trace.api.Config;
import datadog.trace.api.http.StoredBodySupplier;
import datadog.trace.api.internal.TraceSegment;
import io.sqreen.powerwaf.Additive;
Expand Down Expand Up @@ -72,7 +75,8 @@ public class AppSecRequestContext implements DataBundle, Closeable {
}

private final ConcurrentHashMap<Address<?>, Object> persistentData = new ConcurrentHashMap<>();
private Collection<AppSecEvent> collectedEvents; // guarded by this
private Collection<AppSecEvent> appSecEvents; // guarded by this
private Collection<StackTraceEvent> stackTraceEvents; // guarded by this

// assume these will always be written and read by the same thread
private String scheme;
Expand Down Expand Up @@ -405,27 +409,42 @@ public CharSequence getStoredRequestBody() {
return storedRequestBodySupplier.get();
}

public void reportEvents(Collection<AppSecEvent> events) {
for (AppSecEvent event : events) {
public void reportEvents(Collection<AppSecEvent> appSecEvents) {
for (AppSecEvent event : appSecEvents) {
StandardizedLogging.attackDetected(log, event);
}
synchronized (this) {
if (this.collectedEvents == null) {
this.collectedEvents = new ArrayList<>();
if (this.appSecEvents == null) {
this.appSecEvents = new ArrayList<>();
}
try {
this.collectedEvents.addAll(events);
this.appSecEvents.addAll(appSecEvents);
} catch (UnsupportedOperationException e) {
throw new IllegalStateException("Events cannot be added anymore");
}
}
}

public void reportStackTrace(StackTraceEvent stackTraceEvent) {
synchronized (this) {
if (this.stackTraceEvents == null) {
this.stackTraceEvents = new ArrayList<>();
}
try {
if (stackTraceEvents.size() <= Config.get().getAppSecMaxStackTraces()) {
this.stackTraceEvents.add(stackTraceEvent);
}
} catch (UnsupportedOperationException e) {
throw new IllegalStateException("Stacktrace cannot be added anymore");
}
}
}

Collection<AppSecEvent> transferCollectedEvents() {
Collection<AppSecEvent> events;
synchronized (this) {
events = this.collectedEvents;
this.collectedEvents = Collections.emptyList();
events = this.appSecEvents;
this.appSecEvents = Collections.emptyList();
}
if (events != null) {
return events;
Expand All @@ -434,6 +453,19 @@ Collection<AppSecEvent> transferCollectedEvents() {
}
}

StackTraceCollection transferStackTracesCollection() {
Collection<StackTraceEvent> stackTraces;
synchronized (this) {
stackTraces = this.stackTraceEvents;
this.stackTraceEvents = Collections.emptyList();
}
if (stackTraces != null) {
return new StackTraceCollection(stackTraces);
} else {
return null;
}
}

public void reportApiSchemas(Map<String, String> schemas) {
if (schemas == null || schemas.isEmpty()) return;

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,8 @@
import com.datadog.appsec.event.data.SingletonDataBundle;
import com.datadog.appsec.report.AppSecEvent;
import com.datadog.appsec.report.AppSecEventWrapper;
import com.datadog.appsec.stack_trace.StackTraceCollection;
import com.datadog.appsec.util.ObjectFlattener;
import datadog.trace.api.Config;
import datadog.trace.api.DDTags;
import datadog.trace.api.function.TriConsumer;
Expand Down Expand Up @@ -162,6 +164,14 @@ public void init() {
writeRequestHeaders(traceSeg, REQUEST_HEADERS_ALLOW_LIST, ctx.getRequestHeaders());
writeResponseHeaders(
traceSeg, RESPONSE_HEADERS_ALLOW_LIST, ctx.getResponseHeaders());

// Report collected stack traces
StackTraceCollection stackTraceCollection = ctx.transferStackTracesCollection();
if (stackTraceCollection != null) {
Object flatStruct = ObjectFlattener.flatten(stackTraceCollection);
traceSeg.setMetaStructTop("_dd.stack", flatStruct);
}

} else if (hasUserTrackingEvent(traceSeg)) {
// Report all collected request headers on user tracking event
writeRequestHeaders(traceSeg, REQUEST_HEADERS_ALLOW_LIST, ctx.getRequestHeaders());
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,8 @@
import com.datadog.appsec.event.data.KnownAddresses;
import com.datadog.appsec.gateway.AppSecRequestContext;
import com.datadog.appsec.report.AppSecEvent;
import com.datadog.appsec.stack_trace.StackTraceEvent;
import com.datadog.appsec.stack_trace.StackTraceEvent.Frame;
import com.datadog.appsec.util.StandardizedLogging;
import com.squareup.moshi.JsonAdapter;
import com.squareup.moshi.Moshi;
Expand All @@ -26,6 +28,7 @@
import datadog.trace.api.telemetry.WafMetricCollector;
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
import datadog.trace.util.stacktrace.StackWalkerFactory;
import io.sqreen.powerwaf.Additive;
import io.sqreen.powerwaf.Powerwaf;
import io.sqreen.powerwaf.PowerwafConfig;
Expand Down Expand Up @@ -57,6 +60,7 @@
import java.util.UUID;
import java.util.concurrent.atomic.AtomicReference;
import java.util.stream.Collectors;
import java.util.stream.IntStream;
import javax.annotation.Nonnull;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
Expand All @@ -76,6 +80,8 @@ public class PowerWAFModule implements AppSecModule {

private static final Map<String, ActionInfo> DEFAULT_ACTIONS;

private static final String EXPLOIT_DETECTED_MSG = "Exploit detected";

private static class ActionInfo {
final String type;
final Map<String, Object> parameters;
Expand Down Expand Up @@ -430,6 +436,11 @@ public void onDataAvailable(
Flow.Action.RequestBlockingAction rba = createRedirectRequestAction(actionInfo);
flow.setAction(rba);
break;
} else if ("generate_stack".equals(actionInfo.type)
&& Config.get().isAppSecStackTraceEnabled()) {
String stackId = (String) actionInfo.parameters.get("stack_id");
StackTraceEvent stackTraceEvent = createExploitStackTraceEvent(stackId);
reqCtx.reportStackTrace(stackTraceEvent);
} else {
log.info("Ignoring action with type {}", actionInfo.type);
}
Expand Down Expand Up @@ -497,6 +508,29 @@ private Flow.Action.RequestBlockingAction createRedirectRequestAction(ActionInfo
}
}

private StackTraceEvent createExploitStackTraceEvent(String stackId) {
if (stackId == null || stackId.isEmpty()) {
return null;
}
List<Frame> result = generateUserCodeStackTrace();
return new StackTraceEvent(stackId, EXPLOIT_DETECTED_MSG, result);
}

/** Function generates stack trace of the user code (excluding datadog classes) */
private List<Frame> generateUserCodeStackTrace() {
int stackCapacity = Config.get().getAppSecMaxStackTraceDepth();
List<StackTraceElement> elements =
StackWalkerFactory.INSTANCE.walk(
stream ->
stream
.filter(elem -> !elem.getClassName().startsWith("com.datadog"))
.limit(stackCapacity)
.collect(Collectors.toList()));
return IntStream.range(0, elements.size())
.mapToObj(idx -> new Frame(elements.get(idx), idx))
.collect(Collectors.toList());
}

private Powerwaf.ResultWithData doRunPowerwaf(
AppSecRequestContext reqCtx,
DataBundle newData,
Expand Down Expand Up @@ -563,6 +597,7 @@ private AppSecEvent buildEvent(PowerWAFResultData wafResult) {
.withRule(wafResult.rule)
.withRuleMatches(wafResult.rule_matches)
.withSpanId(spanId)
.withStackId(wafResult.stack_id)
.build();
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
public class PowerWAFResultData {
Rule rule;
List<RuleMatch> rule_matches;
String stack_id;

public static class RuleMatch {
String operator;
Expand All @@ -19,10 +20,16 @@ public static class Rule {
Map<String, String> tags;
}

public static class Parameter {
public static class Parameter extends MatchInfo {
MatchInfo resources;
MatchInfo params;
MatchInfo db_type;
List<String> highlight;
}

public static class MatchInfo {
String address;
List<Object> key_path;
String value;
List<String> highlight;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,9 @@ public class AppSecEvent {
@com.squareup.moshi.Json(name = "span_id")
private Long spanId;

@com.squareup.moshi.Json(name = "stack_id")
private String stackId;

public Rule getRule() {
return rule;
}
Expand All @@ -30,6 +33,10 @@ public Long getSpanId() {
return spanId;
}

public String getStackId() {
return stackId;
}

@Override
public String toString() {
StringBuilder sb = new StringBuilder();
Expand All @@ -48,6 +55,9 @@ public String toString() {
sb.append("spanId");
sb.append('=');
sb.append(((this.spanId == null) ? "<null>" : this.spanId));
sb.append("stackId");
sb.append('=');
sb.append(((this.stackId == null) ? "<null>" : this.stackId));
if (sb.charAt((sb.length() - 1)) == ',') {
sb.setCharAt((sb.length() - 1), ']');
} else {
Expand All @@ -62,6 +72,7 @@ public int hashCode() {
result = ((result * 31) + ((this.rule == null) ? 0 : this.rule.hashCode()));
result = ((result * 31) + ((this.ruleMatches == null) ? 0 : this.ruleMatches.hashCode()));
result = ((result * 31) + ((this.spanId == null) ? 0 : this.spanId.hashCode()));
result = ((result * 31) + ((this.stackId == null) ? 0 : this.stackId.hashCode()));
return result;
}

Expand All @@ -76,7 +87,8 @@ public boolean equals(Object other) {
AppSecEvent rhs = ((AppSecEvent) other);
return ((Objects.equals(this.rule, rhs.rule))
&& (Objects.equals(this.ruleMatches, rhs.ruleMatches))
&& (Objects.equals(this.spanId, rhs.spanId)));
&& (Objects.equals(this.spanId, rhs.spanId))
&& (Objects.equals(this.stackId, rhs.stackId)));
}

public static class Builder {
Expand Down Expand Up @@ -108,5 +120,10 @@ public Builder withSpanId(Long spanId) {
this.instance.spanId = spanId;
return this;
}

public Builder withStackId(String stackId) {
this.instance.stackId = stackId;
return this;
}
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,12 @@
package com.datadog.appsec.stack_trace;

import java.util.Collection;

public class StackTraceCollection {
public final Collection<StackTraceEvent> exploit;
// TODO: Add vulnerability and exception collections for future use in IAST and APM

public StackTraceCollection(Collection<StackTraceEvent> exploit) {
this.exploit = exploit;
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,35 @@
package com.datadog.appsec.stack_trace;

import java.util.List;

public class StackTraceEvent {

public final String id;
public final String language = "java";
public final String message;
public final List<Frame> frames;

public StackTraceEvent(String id, String message, List<Frame> frames) {
this.id = id;
this.message = message;
this.frames = frames;
}

public static class Frame {
public final int id;
public final String text;
public final String file;
public final int line;
public final String class_name;
public final String function;

public Frame(StackTraceElement element, int id) {
this.id = id;
this.text = element.toString();
this.file = element.getFileName();
this.line = element.getLineNumber();
this.class_name = element.getClassName();
this.function = element.getMethodName();
}
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,26 @@
package com.datadog.appsec.util;

import com.squareup.moshi.JsonAdapter;
import com.squareup.moshi.Moshi;

/**
* This class provides a utility method to flatten an object into a Map. The object's fields are
* used as keys in the Map, and the field values are used as values. The field values are
* recursively flattened if they are custom objects.
*/
public class ObjectFlattener {

private static final JsonAdapter<Object> JSON_ADAPTER =
new Moshi.Builder().build().adapter(Object.class);

/**
* Flattens an object into a Map.
*
* @param obj the object to flatten
* @return the flattened object as a Map, or the original object if it's a primitive type or a
* Collection or a Map. Returns null if the input object is null.
*/
public static Object flatten(Object obj) {
return JSON_ADAPTER.toJsonValue(obj);
}
}
Loading

0 comments on commit 324d272

Please sign in to comment.