Skip to content

Commit 8d0b1ce

Browse files
manuel-alvarez-alvarezmanuel-alvarez-alvarez
committed
Extract interface for tainted objects
1 parent 2b24697 commit 8d0b1ce

File tree

320 files changed

+3168
-4730
lines changed

Some content is hidden

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

320 files changed

+3168
-4730
lines changed

‎dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/iast/NamedContext.java‎

Lines changed: 10 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,9 @@
22

33
import datadog.trace.api.iast.IastContext;
44
import datadog.trace.api.iast.InstrumentationBridge;
5-
import datadog.trace.api.iast.Taintable.Source;
65
import datadog.trace.api.iast.propagation.PropagationModule;
6+
import datadog.trace.api.iast.taint.Source;
7+
import datadog.trace.api.iast.taint.TaintedObjects;
78
import datadog.trace.bootstrap.ContextStore;
89
import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
910
import javax.annotation.Nonnull;
@@ -30,7 +31,8 @@ public static <E> NamedContext getOrCreate(
3031
}
3132
final PropagationModule module = InstrumentationBridge.PROPAGATION;
3233
if (module != null) {
33-
final Source source = module.findSource(target);
34+
final TaintedObjects to = IastContext.Provider.taintedObjects();
35+
final Source source = module.findSource(to, target);
3436
if (source != null) {
3537
result = new NamedContextImpl(module, source);
3638
}
@@ -60,7 +62,7 @@ private static class NamedContextImpl extends NamedContext {
6062
@Nullable private String currentName;
6163

6264
private boolean fetched;
63-
@Nullable private IastContext context;
65+
@Nullable private TaintedObjects to;
6466

6567
public NamedContextImpl(@Nonnull final PropagationModule module, @Nonnull final Source source) {
6668
this.module = module;
@@ -69,7 +71,7 @@ public NamedContextImpl(@Nonnull final PropagationModule module, @Nonnull final
6971

7072
@Override
7173
public void taintValue(@Nullable final String value) {
72-
module.taintString(iastCtx(), value, source.getOrigin(), currentName, source.getValue());
74+
module.taintObject(to(), value, source.getOrigin(), currentName, source.getValue());
7375
}
7476

7577
@Override
@@ -79,7 +81,7 @@ public void taintName(@Nullable final String name) {
7981
// prevent tainting the same name more than once
8082
if (currentName != name) {
8183
currentName = name;
82-
module.taintString(iastCtx(), name, source.getOrigin(), name, source.getValue());
84+
module.taintObject(to(), name, source.getOrigin(), name, source.getValue());
8385
}
8486
}
8587

@@ -88,12 +90,12 @@ public void setCurrentName(@Nullable final String name) {
8890
currentName = name;
8991
}
9092

91-
private IastContext iastCtx() {
93+
private TaintedObjects to() {
9294
if (!fetched) {
9395
fetched = true;
94-
context = IastContext.Provider.get();
96+
to = IastContext.Provider.taintedObjects();
9597
}
96-
return context;
98+
return to;
9799
}
98100
}
99101
}

‎dd-java-agent/agent-bootstrap/src/test/groovy/datadog/trace/bootstrap/instrumentation/iast/NamedContextTest.groovy‎

Lines changed: 18 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,8 @@ package datadog.trace.bootstrap.instrumentation.iast
33

44
import datadog.trace.api.iast.InstrumentationBridge
55
import datadog.trace.api.iast.SourceTypes
6-
import datadog.trace.api.iast.Taintable.Source
76
import datadog.trace.api.iast.propagation.PropagationModule
7+
import datadog.trace.api.iast.taint.Source
88
import datadog.trace.bootstrap.ContextStore
99
import datadog.trace.test.util.DDSpecification
1010

@@ -38,7 +38,7 @@ class NamedContextTest extends DDSpecification {
3838
context.taintName(name)
3939

4040
then:
41-
1 * module.taintString(_, name, source.origin, name, source.value)
41+
1 * module.taintObject(_, name, source.origin, name, source.value)
4242

4343
when:
4444
context.taintName(name)
@@ -50,7 +50,7 @@ class NamedContextTest extends DDSpecification {
5050
context.taintValue(value)
5151

5252
then:
53-
1 * module.taintString(_, value, source.origin, name, source.value)
53+
1 * module.taintObject(_, value, source.origin, name, source.value)
5454
0 * _
5555
}
5656

@@ -82,5 +82,20 @@ class NamedContextTest extends DDSpecification {
8282
byte origin
8383
String name
8484
String value
85+
86+
@Override
87+
Source attachValue(Object newValue) {
88+
return new SourceImpl(origin: origin, name: name, value: newValue as String)
89+
}
90+
91+
@Override
92+
boolean isReference() {
93+
return false
94+
}
95+
96+
@Override
97+
Object getRawValue() {
98+
return value
99+
}
85100
}
86101
}

‎dd-java-agent/agent-builder/src/main/java/datadog/trace/agent/tooling/AgentInstaller.java‎

Lines changed: 0 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,6 @@
66
import static net.bytebuddy.matcher.ElementMatchers.isDefaultFinalizer;
77

88
import datadog.trace.agent.tooling.bytebuddy.SharedTypePools;
9-
import datadog.trace.agent.tooling.bytebuddy.iast.TaintableRedefinitionStrategyListener;
109
import datadog.trace.agent.tooling.bytebuddy.matcher.DDElementMatchers;
1110
import datadog.trace.agent.tooling.bytebuddy.memoize.MemoizedMatchers;
1211
import datadog.trace.agent.tooling.bytebuddy.outline.TypePoolFacade;
@@ -146,7 +145,6 @@ public static ClassFileTransformer installBytebuddyAgent(
146145
.with(AgentStrategies.transformerDecorator())
147146
.with(AgentBuilder.RedefinitionStrategy.RETRANSFORMATION)
148147
.with(AgentStrategies.rediscoveryStrategy())
149-
.with(redefinitionStrategyListener(enabledSystems))
150148
.with(AgentStrategies.locationStrategy())
151149
.with(AgentStrategies.poolStrategy())
152150
.with(AgentBuilder.DescriptionStrategy.Default.POOL_ONLY)
@@ -163,7 +161,6 @@ public static ClassFileTransformer installBytebuddyAgent(
163161
agentBuilder
164162
.with(AgentBuilder.RedefinitionStrategy.RETRANSFORMATION)
165163
.with(AgentStrategies.rediscoveryStrategy())
166-
.with(redefinitionStrategyListener(enabledSystems))
167164
.with(new RedefinitionLoggingListener())
168165
.with(new TransformLoggingListener());
169166
}
@@ -331,15 +328,6 @@ private static void addByteBuddyRawSetting() {
331328
}
332329
}
333330

334-
private static AgentBuilder.RedefinitionStrategy.Listener redefinitionStrategyListener(
335-
final Set<InstrumenterModule.TargetSystem> enabledSystems) {
336-
if (enabledSystems.contains(InstrumenterModule.TargetSystem.IAST)) {
337-
return TaintableRedefinitionStrategyListener.INSTANCE;
338-
} else {
339-
return AgentBuilder.RedefinitionStrategy.Listener.NoOp.INSTANCE;
340-
}
341-
}
342-
343331
static class RedefinitionLoggingListener implements AgentBuilder.RedefinitionStrategy.Listener {
344332

345333
private static final Logger log = LoggerFactory.getLogger(RedefinitionLoggingListener.class);

‎dd-java-agent/agent-iast/src/jmh/java/com/datadog/iast/propagation/AbstractBenchmark.java‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -3,14 +3,15 @@
33
import static java.util.concurrent.TimeUnit.NANOSECONDS;
44

55
import com.datadog.iast.IastSystem;
6-
import com.datadog.iast.model.Range;
7-
import com.datadog.iast.model.Source;
8-
import com.datadog.iast.taint.TaintedObjects;
6+
import com.datadog.iast.model.SourceImpl;
97
import datadog.trace.api.Config;
108
import datadog.trace.api.ProductActivation;
119
import datadog.trace.api.gateway.InstrumentationGateway;
1210
import datadog.trace.api.gateway.RequestContextSlot;
1311
import datadog.trace.api.iast.IastContext;
12+
import datadog.trace.api.iast.taint.Range;
13+
import datadog.trace.api.iast.taint.Source;
14+
import datadog.trace.api.iast.taint.TaintedObjects;
1415
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
1516
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
1617
import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
@@ -95,7 +96,7 @@ protected <E> E notTainted(final E value) {
9596
}
9697

9798
protected Source source() {
98-
return new Source((byte) 0, "key", "value");
99+
return new SourceImpl((byte) 0, "key", "value");
99100
}
100101

101102
private static long computeHash(final Object value) {

‎dd-java-agent/agent-iast/src/jmh/java/com/datadog/iast/propagation/StringBuilderAppendBenchmark.java‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
import static datadog.trace.api.iast.VulnerabilityMarks.NOT_MARKED;
44

55
import com.datadog.iast.IastRequestContext;
6-
import com.datadog.iast.model.Range;
6+
import com.datadog.iast.model.RangeImpl;
77
import datadog.trace.api.iast.IastContext;
88
import datadog.trace.instrumentation.java.lang.StringBuilderCallSite;
99
import org.openjdk.jmh.annotations.Benchmark;
@@ -17,14 +17,14 @@ protected Context initializeContext() {
1717
final IastRequestContext context = new IastRequestContext();
1818
final String notTainted = notTainted("I am not a tainted string");
1919
final String tainted =
20-
tainted(context, "I am a tainted string", new Range(5, 6, source(), NOT_MARKED));
20+
tainted(context, "I am a tainted string", new RangeImpl(5, 6, source(), NOT_MARKED));
2121
final StringBuilder notTaintedBuilder =
2222
notTainted(new StringBuilder("I am not a tainted string builder"));
2323
final StringBuilder taintedBuilder =
2424
tainted(
2525
context,
2626
new StringBuilder("I am a tainted string builder"),
27-
new Range(5, 6, source(), NOT_MARKED));
27+
new RangeImpl(5, 6, source(), NOT_MARKED));
2828
return new Context(context, notTainted, tainted, notTaintedBuilder, taintedBuilder);
2929
}
3030

‎dd-java-agent/agent-iast/src/jmh/java/com/datadog/iast/propagation/StringBuilderBatchBenchmark.java‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
import static java.util.concurrent.TimeUnit.MICROSECONDS;
55

66
import com.datadog.iast.IastRequestContext;
7-
import com.datadog.iast.model.Range;
7+
import com.datadog.iast.model.RangeImpl;
88
import datadog.trace.api.iast.IastContext;
99
import datadog.trace.instrumentation.java.lang.StringBuilderCallSite;
1010
import java.util.ArrayList;
@@ -36,7 +36,8 @@ protected StringBuilderBatchBenchmark.Context initializeContext() {
3636
final String value;
3737
if (current < limit) {
3838
value =
39-
tainted(context, UUID.randomUUID().toString(), new Range(3, 6, source(), NOT_MARKED));
39+
tainted(
40+
context, UUID.randomUUID().toString(), new RangeImpl(3, 6, source(), NOT_MARKED));
4041
} else {
4142
value = notTainted(UUID.randomUUID().toString());
4243
}

‎dd-java-agent/agent-iast/src/jmh/java/com/datadog/iast/propagation/StringBuilderInitBenchmark.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
import static datadog.trace.api.iast.VulnerabilityMarks.NOT_MARKED;
44

55
import com.datadog.iast.IastRequestContext;
6-
import com.datadog.iast.model.Range;
6+
import com.datadog.iast.model.RangeImpl;
77
import datadog.trace.api.iast.IastContext;
88
import datadog.trace.instrumentation.java.lang.StringBuilderCallSite;
99
import org.openjdk.jmh.annotations.Benchmark;
@@ -17,7 +17,7 @@ protected Context initializeContext() {
1717
final IastRequestContext context = new IastRequestContext();
1818
final String notTainted = notTainted("I am not a tainted string");
1919
final String tainted =
20-
tainted(context, "I am a tainted string", new Range(3, 6, source(), NOT_MARKED));
20+
tainted(context, "I am a tainted string", new RangeImpl(3, 6, source(), NOT_MARKED));
2121
return new Context(context, notTainted, tainted);
2222
}
2323

‎dd-java-agent/agent-iast/src/jmh/java/com/datadog/iast/propagation/StringBuilderToStringBenchmark.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
import static datadog.trace.api.iast.VulnerabilityMarks.NOT_MARKED;
44

55
import com.datadog.iast.IastRequestContext;
6-
import com.datadog.iast.model.Range;
6+
import com.datadog.iast.model.RangeImpl;
77
import datadog.trace.api.iast.IastContext;
88
import datadog.trace.instrumentation.java.lang.StringBuilderCallSite;
99
import org.openjdk.jmh.annotations.Benchmark;
@@ -21,7 +21,7 @@ protected Context initializeContext() {
2121
tainted(
2222
context,
2323
new StringBuilder("I am a tainted string builder"),
24-
new Range(5, 7, source(), NOT_MARKED));
24+
new RangeImpl(5, 7, source(), NOT_MARKED));
2525
return new Context(context, notTaintedBuilder, taintedBuilder);
2626
}
2727

‎dd-java-agent/agent-iast/src/jmh/java/com/datadog/iast/propagation/StringConcatBenchmark.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
import static datadog.trace.api.iast.VulnerabilityMarks.NOT_MARKED;
44

55
import com.datadog.iast.IastRequestContext;
6-
import com.datadog.iast.model.Range;
6+
import com.datadog.iast.model.RangeImpl;
77
import datadog.trace.api.iast.IastContext;
88
import datadog.trace.instrumentation.java.lang.StringCallSite;
99
import org.openjdk.jmh.annotations.Benchmark;
@@ -16,7 +16,7 @@ protected StringConcatBenchmark.Context initializeContext() {
1616
final IastRequestContext context = new IastRequestContext();
1717
final String notTainted = notTainted("I am not a tainted string");
1818
final String tainted =
19-
tainted(context, "I am a tainted string", new Range(3, 5, source(), NOT_MARKED));
19+
tainted(context, "I am a tainted string", new RangeImpl(3, 5, source(), NOT_MARKED));
2020
return new StringConcatBenchmark.Context(context, notTainted, tainted);
2121
}
2222

‎dd-java-agent/agent-iast/src/jmh/java/com/datadog/iast/propagation/StringConcatFactoryBatchBenchmark.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
import static java.util.concurrent.TimeUnit.MICROSECONDS;
55

66
import com.datadog.iast.IastRequestContext;
7-
import com.datadog.iast.model.Range;
7+
import com.datadog.iast.model.RangeImpl;
88
import datadog.trace.api.iast.IastContext;
99
import datadog.trace.api.iast.InstrumentationBridge;
1010
import java.lang.invoke.MethodHandle;
@@ -56,7 +56,7 @@ protected StringConcatFactoryBatchBenchmark.Context initializeContext() {
5656
double current = i / (double) stringCount;
5757
final String value;
5858
if (current < limit) {
59-
value = tainted(context, "Yep, tainted", new Range(3, 5, source(), NOT_MARKED));
59+
value = tainted(context, "Yep, tainted", new RangeImpl(3, 5, source(), NOT_MARKED));
6060
} else {
6161
value = notTainted("Nop, tainted");
6262
}

0 commit comments

Comments
 (0)