Skip to content

Commit 99ecab7

Browse files
jandro996jandro996
authored
Optimize IAST Vulnerability Detection (#8885)
What Does This Do Implements the new algorithm for detecting IAST vulnerabilities, where vulnerabilities that were already explored in previous runs for a given endpoint are skipped, ensuring that all remaining ones are eventually explored. This addresses the current limitation where only the first matching vulnerabilities are consistently reported, causing others to remain hidden. Changes to OverheadContext The OverheadContext class has been extended to support three separate tracking maps: globalMap Used to track vulnerability detection counts per endpoint across all requests. Keys are strings combining the request method and route (GET /login, POST /submit, etc.). Values are maps from vulnerabilityType → int (count of occurrences). Capped at 4,096 entries using a clear‐on‐overflow strategy, to ensure bounded memory usage. Oldest entries are cleared once the limit is reached. copyMap Created per request to copy the global counts at the start of the request, ensuring a consistent baseline to compare against throughout the lifecycle of the request. requestMap Tracks vulnerability type counts within the request. An additional field, isGlobal, has been added to indicate whether the context is global or request-scoped. If isGlobal is true, the maps are not used, and quota checks proceed using the global strategy only. A new method, resetMaps(), has been added to update globalMap when the request ends and vulnerability data has been reported. Two scenarios are supported: Case 1: Budget not fully used → The entry for the endpoint in globalMap is cleared, since the request stayed within budget. Case 2: Budget fully used → The counts from requestMap are compared to those in copyMap. For each vulnerability type, if the value in requestMap is greater, it is used to update the corresponding entry in globalMap. Changes to OverheadController The method consumeQuota() has been extended to receive a vulnerabilityType and modified to support the new logic: If an OverheadContext is present and not global, and there is remaining quota and a valid span, the controller now invokes a new method maybeSkipVulnerability() to determine whether quota should actually be consumed or not, based on endpoint-specific history. It's better to check the Algorithm execution example flow diagram to understand how this should work Changes to IastRequestContext In releaseRequestContext(), the request now calls resetMaps() on the associated OverheadContext, ensuring globalMap is updated at the end of each request. Motivation [RFC-1029] Optimizing IAST Vulnerability Detection implementation Additional Notes java tracer needs to implement also [RFC-1029-A1] Solution for dynamic http routes
1 parent b3e2ecd commit 99ecab7

File tree

16 files changed

+857
-17
lines changed

16 files changed

+857
-17
lines changed

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

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,8 +48,27 @@ public IastRequestContext() {
4848
}
4949

5050
public IastRequestContext(final TaintedObjects taintedObjects) {
51+
this(taintedObjects, false);
52+
}
53+
54+
public IastRequestContext(final TaintedObjects taintedObjects, boolean isGlobal) {
55+
this.vulnerabilityBatch = new VulnerabilityBatch();
56+
this.overheadContext =
57+
new OverheadContext(Config.get().getIastVulnerabilitiesPerRequest(), isGlobal);
58+
this.taintedObjects = taintedObjects;
59+
}
60+
61+
/**
62+
* Use this constructor only when you want to create a new context with a fresh overhead context
63+
* (e.g. for testing purposes).
64+
*
65+
* @param taintedObjects the tainted objects to use
66+
* @param overheadContext the overhead context to use
67+
*/
68+
public IastRequestContext(
69+
final TaintedObjects taintedObjects, final OverheadContext overheadContext) {
5170
this.vulnerabilityBatch = new VulnerabilityBatch();
52-
this.overheadContext = new OverheadContext(Config.get().getIastVulnerabilitiesPerRequest());
71+
this.overheadContext = overheadContext;
5372
this.taintedObjects = taintedObjects;
5473
}
5574

@@ -188,6 +207,7 @@ public void releaseRequestContext(@Nonnull final IastContext context) {
188207
pool.offer(unwrapped);
189208
iastCtx.setTaintedObjects(TaintedObjects.NoOp.INSTANCE);
190209
}
210+
iastCtx.overheadContext.resetMaps();
191211
}
192212
}
193213
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -140,7 +140,7 @@ private VulnerabilityBatch getOrCreateVulnerabilityBatch(final AgentSpan span) {
140140
private AgentSpan startNewSpan() {
141141
final AgentSpanContext tagContext =
142142
new TagContext()
143-
.withRequestContextDataIast(new IastRequestContext(TaintedObjects.NoOp.INSTANCE));
143+
.withRequestContextDataIast(new IastRequestContext(TaintedObjects.NoOp.INSTANCE, true));
144144
final AgentSpan span =
145145
tracer()
146146
.startSpan("iast", VULNERABILITY_SPAN_NAME, tagContext)

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

Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,16 +3,60 @@
33
import static datadog.trace.api.iast.IastDetectionMode.UNLIMITED;
44

55
import com.datadog.iast.util.NonBlockingSemaphore;
6+
import datadog.trace.api.iast.VulnerabilityTypes;
7+
import java.util.Map;
8+
import java.util.Set;
9+
import java.util.concurrent.ConcurrentHashMap;
10+
import java.util.concurrent.ConcurrentMap;
11+
import java.util.concurrent.atomic.AtomicIntegerArray;
12+
import java.util.function.Function;
13+
import javax.annotation.Nullable;
14+
import org.jetbrains.annotations.NotNull;
615

716
public class OverheadContext {
817

18+
/** Maximum number of distinct endpoints to remember in the global cache. */
19+
private static final int GLOBAL_MAP_MAX_SIZE = 4096;
20+
21+
/**
22+
* Global concurrent cache mapping each “method + path” key to its historical vulnerabilityCounts
23+
* map. As soon as size() > GLOBAL_MAP_MAX_SIZE, we clear() the whole map.
24+
*/
25+
static final ConcurrentMap<String, AtomicIntegerArray> globalMap =
26+
new ConcurrentHashMap<String, AtomicIntegerArray>() {
27+
28+
@Override
29+
public AtomicIntegerArray computeIfAbsent(
30+
String key,
31+
@NotNull Function<? super String, ? extends AtomicIntegerArray> mappingFunction) {
32+
if (this.size() >= GLOBAL_MAP_MAX_SIZE) {
33+
super.clear();
34+
}
35+
return super.computeIfAbsent(key, mappingFunction);
36+
}
37+
};
38+
39+
// Snapshot of the globalMap for the current request
40+
private @Nullable final Map<String, int[]> copyMap;
41+
// Map of vulnerabilities per endpoint for the current request, needs to use AtomicIntegerArray
42+
// because it's possible to have concurrent updates in the same request
43+
private @Nullable final Map<String, AtomicIntegerArray> requestMap;
44+
945
private final NonBlockingSemaphore availableVulnerabilities;
46+
private final boolean isGlobal;
1047

1148
public OverheadContext(final int vulnerabilitiesPerRequest) {
49+
this(vulnerabilitiesPerRequest, false);
50+
}
51+
52+
public OverheadContext(final int vulnerabilitiesPerRequest, final boolean isGlobal) {
1253
availableVulnerabilities =
1354
vulnerabilitiesPerRequest == UNLIMITED
1455
? NonBlockingSemaphore.unlimited()
1556
: NonBlockingSemaphore.withPermitCount(vulnerabilitiesPerRequest);
57+
this.isGlobal = isGlobal;
58+
this.requestMap = isGlobal ? null : new ConcurrentHashMap<>();
59+
this.copyMap = isGlobal ? null : new ConcurrentHashMap<>();
1660
}
1761

1862
public int getAvailableQuota() {
@@ -26,4 +70,52 @@ public boolean consumeQuota(final int delta) {
2670
public void reset() {
2771
availableVulnerabilities.reset();
2872
}
73+
74+
public void resetMaps() {
75+
// If this is a global context, we do not reset the maps
76+
if (isGlobal || requestMap == null || copyMap == null) {
77+
return;
78+
}
79+
Set<String> endpoints = requestMap.keySet();
80+
// If the budget is not consumed, we can reset the maps
81+
if (getAvailableQuota() > 0) {
82+
// clean endpoints from globalMap
83+
endpoints.forEach(globalMap::remove);
84+
return;
85+
}
86+
// If the budget is consumed, we need to merge the requestMap into the globalMap
87+
endpoints.forEach(
88+
endpoint -> {
89+
AtomicIntegerArray countMap = requestMap.get(endpoint);
90+
// should not happen, but just in case
91+
if (countMap == null) {
92+
globalMap.remove(endpoint);
93+
return;
94+
}
95+
// Iterate over the vulnerabilities and update the globalMap
96+
int numberOfVulnerabilities = VulnerabilityTypes.STRINGS.length;
97+
for (int i = 0; i < numberOfVulnerabilities; i++) {
98+
int counter = countMap.get(i);
99+
if (counter > 0) {
100+
AtomicIntegerArray globalCountMap =
101+
globalMap.computeIfAbsent(
102+
endpoint, value -> new AtomicIntegerArray(numberOfVulnerabilities));
103+
104+
globalCountMap.accumulateAndGet(i, counter, Math::max);
105+
}
106+
}
107+
});
108+
}
109+
110+
public boolean isGlobal() {
111+
return isGlobal;
112+
}
113+
114+
public @Nullable Map<String, int[]> getCopyMap() {
115+
return copyMap;
116+
}
117+
118+
public @Nullable Map<String, AtomicIntegerArray> getRequestMap() {
119+
return requestMap;
120+
}
29121
}

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

Lines changed: 113 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,24 @@
11
package com.datadog.iast.overhead;
22

3+
import static com.datadog.iast.overhead.OverheadContext.globalMap;
34
import static datadog.trace.api.iast.IastDetectionMode.UNLIMITED;
45

56
import com.datadog.iast.IastRequestContext;
67
import com.datadog.iast.IastSystem;
8+
import com.datadog.iast.model.VulnerabilityType;
79
import com.datadog.iast.util.NonBlockingSemaphore;
810
import datadog.trace.api.Config;
911
import datadog.trace.api.gateway.RequestContext;
1012
import datadog.trace.api.gateway.RequestContextSlot;
1113
import datadog.trace.api.iast.IastContext;
14+
import datadog.trace.api.iast.VulnerabilityTypes;
1215
import datadog.trace.api.telemetry.LogCollector;
1316
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
1417
import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
18+
import datadog.trace.bootstrap.instrumentation.api.Tags;
1519
import datadog.trace.util.AgentTaskScheduler;
1620
import java.util.concurrent.TimeUnit;
21+
import java.util.concurrent.atomic.AtomicIntegerArray;
1722
import java.util.concurrent.atomic.AtomicLong;
1823
import javax.annotation.Nullable;
1924
import org.slf4j.Logger;
@@ -27,9 +32,12 @@ public interface OverheadController {
2732

2833
int releaseRequest();
2934

30-
boolean hasQuota(final Operation operation, @Nullable final AgentSpan span);
35+
boolean hasQuota(Operation operation, @Nullable AgentSpan span);
3136

32-
boolean consumeQuota(final Operation operation, @Nullable final AgentSpan span);
37+
boolean consumeQuota(Operation operation, @Nullable AgentSpan span);
38+
39+
boolean consumeQuota(
40+
Operation operation, @Nullable AgentSpan span, @Nullable VulnerabilityType type);
3341

3442
static OverheadController build(final Config config, final AgentTaskScheduler scheduler) {
3543
return build(
@@ -100,14 +108,23 @@ public boolean hasQuota(final Operation operation, @Nullable final AgentSpan spa
100108

101109
@Override
102110
public boolean consumeQuota(final Operation operation, @Nullable final AgentSpan span) {
103-
final boolean result = delegate.consumeQuota(operation, span);
111+
return consumeQuota(operation, span, null);
112+
}
113+
114+
@Override
115+
public boolean consumeQuota(
116+
final Operation operation,
117+
@Nullable final AgentSpan span,
118+
@Nullable final VulnerabilityType type) {
119+
final boolean result = delegate.consumeQuota(operation, span, type);
104120
if (LOGGER.isDebugEnabled()) {
105121
LOGGER.debug(
106-
"consumeQuota: operation={}, result={}, availableQuota={}, span={}",
122+
"consumeQuota: operation={}, result={}, availableQuota={}, span={}, type={}",
107123
operation,
108124
result,
109125
getAvailableQuote(span),
110-
span);
126+
span,
127+
type);
111128
}
112129
return result;
113130
}
@@ -147,7 +164,7 @@ class OverheadControllerImpl implements OverheadController {
147164
private volatile long lastAcquiredTimestamp = Long.MAX_VALUE;
148165

149166
final OverheadContext globalContext =
150-
new OverheadContext(Config.get().getIastVulnerabilitiesPerRequest());
167+
new OverheadContext(Config.get().getIastVulnerabilitiesPerRequest(), true);
151168

152169
public OverheadControllerImpl(
153170
final float requestSampling,
@@ -192,7 +209,96 @@ public boolean hasQuota(final Operation operation, @Nullable final AgentSpan spa
192209

193210
@Override
194211
public boolean consumeQuota(final Operation operation, @Nullable final AgentSpan span) {
195-
return operation.consumeQuota(getContext(span));
212+
return consumeQuota(operation, span, null);
213+
}
214+
215+
@Override
216+
public boolean consumeQuota(
217+
final Operation operation,
218+
@Nullable final AgentSpan span,
219+
@Nullable final VulnerabilityType type) {
220+
221+
OverheadContext ctx = getContext(span);
222+
if (ctx == null) {
223+
return false;
224+
}
225+
if (ctx.isGlobal()) {
226+
return operation.consumeQuota(ctx);
227+
}
228+
if (operation.hasQuota(ctx)) {
229+
String method = null;
230+
String path = null;
231+
if (span != null) {
232+
AgentSpan rootSpan = span.getLocalRootSpan();
233+
Object methodTag = rootSpan.getTag(Tags.HTTP_METHOD);
234+
method = (methodTag == null) ? "" : methodTag.toString();
235+
Object routeTag = rootSpan.getTag(Tags.HTTP_ROUTE);
236+
path = (routeTag == null) ? "" : routeTag.toString();
237+
}
238+
if (!maybeSkipVulnerability(ctx, type, method, path)) {
239+
return operation.consumeQuota(ctx);
240+
}
241+
}
242+
return false;
243+
}
244+
245+
/**
246+
* Method to be called when a vulnerability of a certain type is detected. Implements the
247+
* RFC-1029 algorithm.
248+
*
249+
* @param ctx the overhead context for the current request
250+
* @param type the type of vulnerability detected
251+
* @param httpMethod the HTTP method of the request (e.g., GET, POST)
252+
* @param httpPath the HTTP path of the request
253+
* @return true if the vulnerability should be skipped, false otherwise
254+
*/
255+
private boolean maybeSkipVulnerability(
256+
@Nullable final OverheadContext ctx,
257+
@Nullable final VulnerabilityType type,
258+
@Nullable final String httpMethod,
259+
@Nullable final String httpPath) {
260+
261+
if (ctx == null || type == null || ctx.getRequestMap() == null || ctx.getCopyMap() == null) {
262+
return false;
263+
}
264+
265+
int numberOfVulnerabilities = VulnerabilityTypes.STRINGS.length;
266+
267+
String currentEndpoint = httpMethod + " " + httpPath;
268+
269+
AtomicIntegerArray requestArray = ctx.getRequestMap().get(currentEndpoint);
270+
int[] copyArray;
271+
272+
if (requestArray == null) {
273+
AtomicIntegerArray globalArray =
274+
globalMap.computeIfAbsent(
275+
currentEndpoint, k -> new AtomicIntegerArray(numberOfVulnerabilities));
276+
copyArray = toIntArray(globalArray);
277+
ctx.getCopyMap().put(currentEndpoint, copyArray);
278+
requestArray =
279+
ctx.getRequestMap()
280+
.computeIfAbsent(
281+
currentEndpoint, k -> new AtomicIntegerArray(numberOfVulnerabilities));
282+
} else {
283+
copyArray = ctx.getCopyMap().get(currentEndpoint);
284+
}
285+
286+
int counter = requestArray.getAndIncrement(type.type());
287+
int storedCounter = 0;
288+
if (copyArray != null) {
289+
storedCounter = copyArray[type.type()];
290+
}
291+
292+
return counter < storedCounter;
293+
}
294+
295+
private static int[] toIntArray(AtomicIntegerArray atomic) {
296+
int length = atomic.length();
297+
int[] result = new int[length];
298+
for (int i = 0; i < length; i++) {
299+
result[i] = atomic.get(i);
300+
}
301+
return result;
196302
}
197303

198304
@Nullable

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

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package com.datadog.iast.sink;
22

3+
import static com.datadog.iast.model.VulnerabilityType.INSECURE_COOKIE;
34
import static com.datadog.iast.util.HttpHeader.SET_COOKIE;
45
import static com.datadog.iast.util.HttpHeader.SET_COOKIE2;
56
import static java.util.Collections.singletonList;
@@ -65,7 +66,9 @@ private void onCookies(final List<Cookie> cookies) {
6566
return;
6667
}
6768
final AgentSpan span = AgentTracer.activeSpan();
68-
if (!overheadController.consumeQuota(Operations.REPORT_VULNERABILITY, span)) {
69+
if (!overheadController.consumeQuota(
70+
Operations.REPORT_VULNERABILITY, span, INSECURE_COOKIE // we need a type to check quota
71+
)) {
6972
return;
7073
}
7174
final Location location = Location.forSpanAndStack(span, getCurrentStackTrace());

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

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,8 @@ protected void report(final Vulnerability vulnerability) {
5858
}
5959

6060
protected void report(@Nullable final AgentSpan span, final Vulnerability vulnerability) {
61-
if (!overheadController.consumeQuota(Operations.REPORT_VULNERABILITY, span)) {
61+
if (!overheadController.consumeQuota(
62+
Operations.REPORT_VULNERABILITY, span, vulnerability.getType())) {
6263
return;
6364
}
6465
reporter.report(span, vulnerability);
@@ -70,7 +71,7 @@ protected void report(final VulnerabilityType type, final Evidence evidence) {
7071

7172
protected void report(
7273
@Nullable final AgentSpan span, final VulnerabilityType type, final Evidence evidence) {
73-
if (!overheadController.consumeQuota(Operations.REPORT_VULNERABILITY, span)) {
74+
if (!overheadController.consumeQuota(Operations.REPORT_VULNERABILITY, span, type)) {
7475
return;
7576
}
7677
final Vulnerability vulnerability =
@@ -170,7 +171,7 @@ protected final Evidence checkInjection(
170171
}
171172

172173
final AgentSpan span = AgentTracer.activeSpan();
173-
if (!overheadController.consumeQuota(Operations.REPORT_VULNERABILITY, span)) {
174+
if (!overheadController.consumeQuota(Operations.REPORT_VULNERABILITY, span, type)) {
174175
return null;
175176
}
176177

@@ -251,7 +252,7 @@ protected final Evidence checkInjection(
251252
if (!spanFetched && valueRanges != null && valueRanges.length > 0) {
252253
span = AgentTracer.activeSpan();
253254
spanFetched = true;
254-
if (!overheadController.consumeQuota(Operations.REPORT_VULNERABILITY, span)) {
255+
if (!overheadController.consumeQuota(Operations.REPORT_VULNERABILITY, span, type)) {
255256
return null;
256257
}
257258
}

‎dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/IastModuleImplTestBase.groovy

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -114,6 +114,7 @@ class IastModuleImplTestBase extends DDSpecification {
114114
return Stub(OverheadController) {
115115
acquireRequest() >> true
116116
consumeQuota(_ as Operation, _) >> true
117+
consumeQuota(_ as Operation, _, _) >> true
117118
}
118119
}
119120
}

0 commit comments

Comments
 (0)