Skip to content

Commit 416dbae

Browse files
jandro996jandro996
committed
wip
1 parent 5069ac3 commit 416dbae

File tree

12 files changed

+393
-17
lines changed

12 files changed

+393
-17
lines changed

‎dd-java-agent/agent-iast/src/jmh/java/com/datadog/iast/overhead/OverheadControllerBenchmark.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,6 @@ public void acquireReleaseRequestNoSampling() {
4141

4242
@Benchmark
4343
public void consumeQuota() {
44-
overheadController.consumeQuota(Operations.REPORT_VULNERABILITY, null);
44+
overheadController.consumeQuota(Operations.REPORT_VULNERABILITY, null, null);
4545
}
4646
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,7 @@ public IastContext resolve() {
4646

4747
@Override
4848
public IastContext buildRequestContext() {
49-
return new IastRequestContext(globalContext.getTaintedObjects());
49+
return new IastRequestContext((TaintedObjects) globalContext.getTaintedObjects());
5050
}
5151

5252
@Override

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@ public IastContext resolve() {
3232

3333
@Override
3434
public IastContext buildRequestContext() {
35-
return new IastRequestContext(optOutContext.getTaintedObjects());
35+
return new IastRequestContext((TaintedObjects) optOutContext.getTaintedObjects());
3636
}
3737

3838
@Override

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

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,18 @@ public IastRequestContext(final TaintedObjects taintedObjects) {
5353
this.taintedObjects = taintedObjects;
5454
}
5555

56+
/**
57+
* Use this constructor only when you want to create a new context with a fresh overhead context
58+
* (e.g. for testing purposes).
59+
*
60+
* @param overheadContext the overhead context to use
61+
*/
62+
public IastRequestContext(final OverheadContext overheadContext) {
63+
this.vulnerabilityBatch = new VulnerabilityBatch();
64+
this.overheadContext = overheadContext;
65+
this.taintedObjects = TaintedObjects.build(TaintedMap.build(MAP_SIZE));
66+
}
67+
5668
public VulnerabilityBatch getVulnerabilityBatch() {
5769
return vulnerabilityBatch;
5870
}
@@ -188,6 +200,7 @@ public void releaseRequestContext(@Nonnull final IastContext context) {
188200
pool.offer(unwrapped);
189201
iastCtx.setTaintedObjects(TaintedObjects.NoOp.INSTANCE);
190202
}
203+
iastCtx.overheadContext.resetMaps();
191204
}
192205
}
193206
}

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

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,17 +2,53 @@
22

33
import static datadog.trace.api.iast.IastDetectionMode.UNLIMITED;
44

5+
import com.datadog.iast.model.VulnerabilityType;
56
import com.datadog.iast.util.NonBlockingSemaphore;
7+
import java.util.HashMap;
8+
import java.util.LinkedHashMap;
9+
import java.util.Map;
10+
import java.util.Set;
11+
import javax.annotation.Nullable;
612

713
public class OverheadContext {
814

15+
/**
16+
* Maximum number of distinct endpoints to remember in the global cache (LRU eviction beyond this
17+
* size).
18+
*/
19+
private static final int GLOBAL_MAP_MAX_SIZE = 4096;
20+
21+
/**
22+
* Global LRU cache mapping each “method + path” key to its historical vulnerabilityCounts map.
23+
* Key: HTTP_METHOD + " " + HTTP_PATH Value: Map<vulnerabilityType, count>
24+
*/
25+
static final Map<String, Map<VulnerabilityType, Integer>> globalMap =
26+
new LinkedHashMap<String, Map<VulnerabilityType, Integer>>(GLOBAL_MAP_MAX_SIZE, 0.75f, true) {
27+
@Override
28+
protected boolean removeEldestEntry(
29+
Map.Entry<String, Map<VulnerabilityType, Integer>> eldest) {
30+
return size() > GLOBAL_MAP_MAX_SIZE;
31+
}
32+
};
33+
34+
@Nullable final Map<String, Map<VulnerabilityType, Integer>> copyMap;
35+
@Nullable final Map<String, Map<VulnerabilityType, Integer>> requestMap;
36+
937
private final NonBlockingSemaphore availableVulnerabilities;
38+
private final boolean isGlobal;
1039

1140
public OverheadContext(final int vulnerabilitiesPerRequest) {
41+
this(vulnerabilitiesPerRequest, false);
42+
}
43+
44+
public OverheadContext(final int vulnerabilitiesPerRequest, final boolean isGlobal) {
1245
availableVulnerabilities =
1346
vulnerabilitiesPerRequest == UNLIMITED
1447
? NonBlockingSemaphore.unlimited()
1548
: NonBlockingSemaphore.withPermitCount(vulnerabilitiesPerRequest);
49+
this.isGlobal = isGlobal;
50+
this.requestMap = isGlobal ? null : new HashMap<>();
51+
this.copyMap = isGlobal ? null : new HashMap<>();
1652
}
1753

1854
public int getAvailableQuota() {
@@ -26,4 +62,48 @@ public boolean consumeQuota(final int delta) {
2662
public void reset() {
2763
availableVulnerabilities.reset();
2864
}
65+
66+
public void resetMaps() {
67+
if (isGlobal || requestMap == null || copyMap == null) {
68+
return;
69+
}
70+
// If the budget is not consumed, we can reset the maps
71+
Set<String> keys = requestMap.keySet();
72+
if (getAvailableQuota() > 0) {
73+
keys.forEach(globalMap::remove);
74+
keys.clear();
75+
requestMap.clear();
76+
copyMap.clear();
77+
return;
78+
}
79+
keys.forEach(
80+
key -> {
81+
Map<VulnerabilityType, Integer> countMap = requestMap.get(key);
82+
if (countMap == null || countMap.isEmpty()) {
83+
globalMap.remove(key);
84+
return;
85+
}
86+
countMap.forEach(
87+
(key1, counter) -> {
88+
Map<VulnerabilityType, Integer> globalCountMap = globalMap.get(key);
89+
if (globalCountMap != null) {
90+
Integer globalCounter = globalCountMap.getOrDefault(key1, 0);
91+
if (counter > globalCounter) {
92+
globalCountMap.put(key1, counter);
93+
}
94+
} else {
95+
globalCountMap = new HashMap<>();
96+
globalCountMap.put(key1, counter);
97+
globalMap.put(key, globalCountMap);
98+
}
99+
});
100+
});
101+
keys.clear();
102+
requestMap.clear();
103+
copyMap.clear();
104+
}
105+
106+
public boolean isGlobal() {
107+
return isGlobal;
108+
}
29109
}

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

Lines changed: 82 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,11 @@
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;
@@ -12,7 +14,11 @@
1214
import datadog.trace.api.telemetry.LogCollector;
1315
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
1416
import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
17+
import datadog.trace.bootstrap.instrumentation.api.Tags;
1518
import datadog.trace.util.AgentTaskScheduler;
19+
import java.util.HashMap;
20+
import java.util.Map;
21+
import java.util.Set;
1622
import java.util.concurrent.TimeUnit;
1723
import java.util.concurrent.atomic.AtomicLong;
1824
import javax.annotation.Nullable;
@@ -29,7 +35,10 @@ public interface OverheadController {
2935

3036
boolean hasQuota(final Operation operation, @Nullable final AgentSpan span);
3137

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

3443
static OverheadController build(final Config config, final AgentTaskScheduler scheduler) {
3544
return build(
@@ -99,15 +108,19 @@ public boolean hasQuota(final Operation operation, @Nullable final AgentSpan spa
99108
}
100109

101110
@Override
102-
public boolean consumeQuota(final Operation operation, @Nullable final AgentSpan span) {
103-
final boolean result = delegate.consumeQuota(operation, span);
111+
public boolean consumeQuota(
112+
final Operation operation,
113+
@Nullable final AgentSpan span,
114+
@Nullable final VulnerabilityType type) {
115+
final boolean result = delegate.consumeQuota(operation, span, type);
104116
if (LOGGER.isDebugEnabled()) {
105117
LOGGER.debug(
106-
"consumeQuota: operation={}, result={}, availableQuota={}, span={}",
118+
"consumeQuota: operation={}, result={}, availableQuota={}, span={}, type={}",
107119
operation,
108120
result,
109121
getAvailableQuote(span),
110-
span);
122+
span,
123+
type);
111124
}
112125
return result;
113126
}
@@ -147,7 +160,7 @@ class OverheadControllerImpl implements OverheadController {
147160
private volatile long lastAcquiredTimestamp = Long.MAX_VALUE;
148161

149162
final OverheadContext globalContext =
150-
new OverheadContext(Config.get().getIastVulnerabilitiesPerRequest());
163+
new OverheadContext(Config.get().getIastVulnerabilitiesPerRequest(), true);
151164

152165
public OverheadControllerImpl(
153166
final float requestSampling,
@@ -191,8 +204,69 @@ public boolean hasQuota(final Operation operation, @Nullable final AgentSpan spa
191204
}
192205

193206
@Override
194-
public boolean consumeQuota(final Operation operation, @Nullable final AgentSpan span) {
195-
return operation.consumeQuota(getContext(span));
207+
public boolean consumeQuota(
208+
final Operation operation,
209+
@Nullable final AgentSpan span,
210+
@Nullable final VulnerabilityType type) {
211+
212+
OverheadContext ctx = getContext(span);
213+
if (ctx == null) {
214+
return false;
215+
}
216+
if (ctx.isGlobal()) {
217+
return operation.consumeQuota(ctx);
218+
}
219+
if (operation.hasQuota(ctx)) {
220+
String method = null;
221+
String path = null;
222+
if (span != null) {
223+
Object methodTag = span.getLocalRootSpan().getTag(Tags.HTTP_METHOD);
224+
method = (methodTag == null) ? "" : methodTag.toString();
225+
Object routeTag = span.getLocalRootSpan().getTag(Tags.HTTP_ROUTE);
226+
path = (routeTag == null) ? "" : routeTag.toString();
227+
}
228+
if (!maybeSkipVulnerability(ctx, type, method, path)) {
229+
return operation.consumeQuota(ctx);
230+
}
231+
}
232+
return false;
233+
}
234+
235+
/**
236+
* Method to be called when a vulnerability of a certain type is detected. Implements the
237+
* RFC-1029 algorithm.
238+
*
239+
* @param type the type of vulnerability detected
240+
*/
241+
private boolean maybeSkipVulnerability(
242+
@Nullable final OverheadContext ctx,
243+
@Nullable final VulnerabilityType type,
244+
@Nullable final String httpMethod,
245+
@Nullable final String httpPath) {
246+
247+
if (ctx == null || type == null || ctx.requestMap == null || ctx.copyMap == null) {
248+
return false;
249+
}
250+
251+
String currentKey = httpMethod + " " + httpPath;
252+
Set<String> keys = ctx.requestMap.keySet();
253+
254+
if (!keys.contains(currentKey)) {
255+
ctx.copyMap.put(currentKey, globalMap.getOrDefault(currentKey, new HashMap<>()));
256+
}
257+
258+
ctx.requestMap.computeIfAbsent(currentKey, k -> new HashMap<>());
259+
260+
Integer counter = ctx.requestMap.get(currentKey).getOrDefault(type, 0);
261+
ctx.requestMap.get(currentKey).put(type, counter + 1);
262+
263+
Integer storedCounter = 0;
264+
Map<VulnerabilityType, Integer> copyCountMap = ctx.copyMap.get(currentKey);
265+
if (copyCountMap != null) {
266+
storedCounter = copyCountMap.getOrDefault(type, 0);
267+
}
268+
269+
return counter < storedCounter;
196270
}
197271

198272
@Nullable

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

Lines changed: 6 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,11 @@ 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+
// TODO decide if we remove this one quota for all vulnerabilities as new IAST sampling
70+
// algorithm is able to report all endpoint vulnerabilities
71+
if (!overheadController.consumeQuota(
72+
Operations.REPORT_VULNERABILITY, span, INSECURE_COOKIE // we need a type to check quota
73+
)) {
6974
return;
7075
}
7176
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/IastRequestContextTest.groovy‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package com.datadog.iast
22

33
import com.datadog.iast.model.Range
4+
import com.datadog.iast.overhead.OverheadContext
45
import com.datadog.iast.taint.TaintedObjects
56
import datadog.trace.api.Config
67
import datadog.trace.api.gateway.RequestContext
@@ -120,4 +121,16 @@ class IastRequestContextTest extends DDSpecification {
120121
then:
121122
ctx.taintedObjects.count() == 0
122123
}
124+
125+
void 'on release context overheadContext reset is called'() {
126+
setup:
127+
final overheadCtx = Mock(OverheadContext)
128+
final ctx = new IastRequestContext(overheadCtx)
129+
130+
when:
131+
provider.releaseRequestContext(ctx)
132+
133+
then:
134+
1 * overheadCtx.resetMaps()
135+
}
123136
}

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

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

3+
import com.datadog.iast.model.VulnerabilityType
34
import com.datadog.iast.overhead.Operation
45
import com.datadog.iast.overhead.OverheadController
56
import com.github.javaparser.quality.Nullable
@@ -24,7 +25,7 @@ class NoopOverheadController implements OverheadController {
2425
}
2526

2627
@Override
27-
boolean consumeQuota(Operation operation, @Nullable AgentSpan span) {
28+
boolean consumeQuota(Operation operation, @Nullable AgentSpan span, @Nullable VulnerabilityType type) {
2829
true
2930
}
3031

0 commit comments

Comments
 (0)