Skip to content

Commit b3d3879

Browse files
cataphractcataphract
committed
IAST: support webflux
1 parent 61ecb0f commit b3d3879

File tree

39 files changed

+1928
-23
lines changed

39 files changed

+1928
-23
lines changed

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

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

33
import com.datadog.iast.model.json.SourceIndex;
44
import java.util.Objects;
5+
import java.util.StringJoiner;
56
import javax.annotation.Nonnegative;
67
import javax.annotation.Nonnull;
78

@@ -36,6 +37,15 @@ public boolean equals(Object o) {
3637
return start == range.start && length == range.length && Objects.equals(source, range.source);
3738
}
3839

40+
@Override
41+
public String toString() {
42+
return new StringJoiner(", ", Range.class.getSimpleName() + "[", "]")
43+
.add("start=" + start)
44+
.add("length=" + length)
45+
.add("source=" + source)
46+
.toString();
47+
}
48+
3949
@Override
4050
public int hashCode() {
4151
return Objects.hash(start, length, source);

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

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,10 @@
11
package com.datadog.iast.model;
22

33
import com.datadog.iast.model.json.SourceTypeString;
4+
import datadog.trace.api.iast.SourceTypes;
45
import datadog.trace.api.iast.Taintable;
56
import java.util.Objects;
7+
import java.util.StringJoiner;
68

79
public final class Source implements Taintable.Source {
810
private final @SourceTypeString byte origin;
@@ -30,6 +32,15 @@ public String getValue() {
3032
return value;
3133
}
3234

35+
@Override
36+
public String toString() {
37+
return new StringJoiner(", ", Source.class.getSimpleName() + "[", "]")
38+
.add("origin=" + SourceTypes.toString(origin))
39+
.add("name='" + name + "'")
40+
.add("value='" + value + "'")
41+
.toString();
42+
}
43+
3344
@Override
3445
public boolean equals(Object o) {
3546
if (this == o) return true;

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

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -96,6 +96,19 @@ public void taint(final byte origin, @Nullable final Object... toTaintArray) {
9696
}
9797
}
9898

99+
@Override
100+
public boolean isTainted(@Nullable Object obj) {
101+
final TaintedObjects taintedObjects = getTaintedObjects();
102+
if (obj == null) {
103+
return false;
104+
}
105+
106+
if (taintedObjects == null) {
107+
return false;
108+
}
109+
return taintedObjects.get(obj) != null;
110+
}
111+
99112
private static Source getFirstSource(final TaintedObjects taintedObjects, final Object object) {
100113
if (object instanceof Taintable) {
101114
return (Source) ((Taintable) object).$$DD$getSource();

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ import datadog.trace.api.gateway.RequestContextSlot
99
import datadog.trace.api.iast.sink.SqlInjectionModule
1010
import datadog.trace.bootstrap.instrumentation.api.AgentSpan
1111

12-
import static com.datadog.iast.IastAgentTestRunner.EMPTY_SOURCE
12+
import static com.datadog.iast.test.IastAgentTestRunner.EMPTY_SOURCE
1313

1414
class SqlInjectionModuleTest extends IastModuleImplTestBase {
1515

dd-java-agent/agent-iast/src/testFixtures/groovy/com/datadog/iast/IastAgentTestRunner.groovy renamed to dd-java-agent/agent-iast/src/testFixtures/groovy/com/datadog/iast/test/IastAgentTestRunner.groovy

Lines changed: 12 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,13 @@
1-
package com.datadog.iast
1+
package com.datadog.iast.test
22

3+
import com.datadog.iast.IastRequestContext
34
import com.datadog.iast.model.Source
45
import com.datadog.iast.taint.TaintedObjects
5-
import com.datadog.iast.telemetry.NoOpTelemetry
66
import datadog.trace.agent.test.AgentTestRunner
77
import datadog.trace.api.gateway.CallbackProvider
88
import datadog.trace.api.gateway.Events
99
import datadog.trace.api.gateway.Flow
1010
import datadog.trace.api.gateway.RequestContextSlot
11-
import datadog.trace.api.iast.InstrumentationBridge
1211
import datadog.trace.api.iast.SourceTypes
1312
import datadog.trace.bootstrap.instrumentation.api.AgentSpan
1413
import datadog.trace.bootstrap.instrumentation.api.AgentTracer
@@ -17,31 +16,32 @@ import datadog.trace.core.DDSpan
1716

1817
import java.util.function.Supplier
1918

20-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.get
21-
22-
class IastAgentTestRunner extends AgentTestRunner {
19+
class IastAgentTestRunner extends AgentTestRunner implements IastRequestContextPreparationTrait {
2320
public static final EMPTY_SOURCE = new Source(SourceTypes.NONE, '', '')
2421

2522
void configurePreAgent() {
2623
super.configurePreAgent()
2724
injectSysConfig('dd.iast.enabled', 'true')
2825
}
2926

27+
protected Closure getRequestEndAction() { }
28+
3029
void setupSpec() {
31-
// Register the Instrumentation Gateway callbacks
32-
def ss = AgentTracer.get().getSubscriptionService(RequestContextSlot.IAST)
33-
IastSystem.start(ss, new NoopOverheadController(), new NoOpTelemetry())
30+
iastSystemSetup(requestEndAction)
3431
}
3532

3633
void cleanupSpec() {
37-
get().getSubscriptionService(RequestContextSlot.IAST).reset()
38-
InstrumentationBridge.clearIastModules()
34+
iastSystemCleanup()
3935
}
4036

41-
protected TaintedObjects getTaintedObjects() {
37+
protected TaintedObjects getLocalTaintedObjects() {
4238
IastRequestContext.get().taintedObjects
4339
}
4440

41+
protected TaintedObjectCollection getLocalTaintedObjectCollection() {
42+
new TaintedObjectCollection(localTaintedObjects)
43+
}
44+
4545
protected DDSpan runUnderIastTrace(Closure cl) {
4646
CallbackProvider iastCbp = TEST_TRACER.getCallbackProvider(RequestContextSlot.IAST)
4747
Supplier<Flow<Object>> reqStartCb = iastCbp.getCallback(Events.EVENTS.requestStarted())

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

Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,102 @@
1+
package com.datadog.iast.test
2+
3+
import com.datadog.iast.IastRequestContext
4+
import com.datadog.iast.IastSystem
5+
import com.datadog.iast.model.Range
6+
import com.datadog.iast.model.Source
7+
import com.datadog.iast.taint.TaintedMap
8+
import com.datadog.iast.taint.TaintedObject
9+
import com.datadog.iast.taint.TaintedObjects
10+
import com.datadog.iast.telemetry.NoOpTelemetry
11+
import datadog.trace.api.gateway.CallbackProvider
12+
import datadog.trace.api.gateway.EventType
13+
import datadog.trace.api.gateway.Events
14+
import datadog.trace.api.gateway.Flow
15+
import datadog.trace.api.gateway.IGSpanInfo
16+
import datadog.trace.api.gateway.RequestContext
17+
import datadog.trace.api.gateway.RequestContextSlot
18+
import datadog.trace.api.iast.InstrumentationBridge
19+
import datadog.trace.bootstrap.instrumentation.api.AgentTracer
20+
21+
import java.util.function.BiFunction
22+
import java.util.function.Supplier
23+
24+
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.get
25+
26+
trait IastRequestContextPreparationTrait {
27+
28+
static void iastSystemSetup(Closure reqEndAction = null) {
29+
def ss = AgentTracer.get().getSubscriptionService(RequestContextSlot.IAST)
30+
IastSystem.start(ss, new NoopOverheadController(), new NoOpTelemetry())
31+
32+
EventType<Supplier<Flow<Object>>> requestStarted = Events.get().requestStarted()
33+
EventType<BiFunction<RequestContext, IGSpanInfo, Flow<Void>>> requestEnded =
34+
Events.get().requestEnded()
35+
36+
// get original callbacks
37+
CallbackProvider provider = AgentTracer.get().getCallbackProvider(RequestContextSlot.IAST)
38+
def origRequestStarted = provider.getCallback(requestStarted)
39+
def origRequestEnded = provider.getCallback(requestEnded)
40+
41+
// wrap the original IG callbacks
42+
ss.reset()
43+
ss.registerCallback(requestStarted, new TaintedMapSaveStrongRefsRequestStarted(orig: origRequestStarted))
44+
if (reqEndAction != null) {
45+
ss.registerCallback(requestEnded, new TaintedMapSavingRequestEnded(
46+
original: origRequestEnded, beforeAction: reqEndAction))
47+
}
48+
}
49+
50+
static void iastSystemCleanup() {
51+
get().getSubscriptionService(RequestContextSlot.IAST).reset()
52+
InstrumentationBridge.clearIastModules()
53+
}
54+
55+
static class TaintedMapSaveStrongRefsRequestStarted implements Supplier<Flow<Object>> {
56+
Supplier<Flow<Object>> orig
57+
58+
@Override
59+
Flow<Object> get() {
60+
IastRequestContext reqCtx = orig.get().result
61+
def taintedObjectWrapper = new TaintedObjectSaveStrongReference(delegate: reqCtx.taintedObjects)
62+
new Flow.ResultFlow<>(new IastRequestContext(taintedObjectWrapper))
63+
}
64+
65+
static class TaintedObjectSaveStrongReference implements TaintedObjects {
66+
@Delegate
67+
TaintedObjects delegate
68+
69+
List<Object> objects = []
70+
71+
TaintedMap getTaintedMap() {
72+
delegate.map
73+
}
74+
75+
TaintedObject taintInputString(String obj, Source source) {
76+
objects << obj
77+
this.delegate.taintInputString(obj, source)
78+
}
79+
80+
TaintedObject taintInputObject(Object obj, Source source) {
81+
objects << obj
82+
this.delegate.taintInputObject(obj, source)
83+
}
84+
85+
TaintedObject taint(Object obj, Range[] ranges) {
86+
objects << obj
87+
this.delegate.taintInputString(obj, ranges)
88+
}
89+
}
90+
}
91+
92+
static class TaintedMapSavingRequestEnded implements BiFunction<RequestContext, IGSpanInfo, Flow<Void>> {
93+
BiFunction<RequestContext, IGSpanInfo, Flow<Void>> original
94+
Closure beforeAction
95+
96+
@Override
97+
Flow<Void> apply(RequestContext requestContext, IGSpanInfo igSpanInfo) {
98+
beforeAction.call(requestContext, igSpanInfo)
99+
original.apply(requestContext, igSpanInfo)
100+
}
101+
}
102+
}

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

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
package com.datadog.iast.test
2+
3+
import com.datadog.iast.IastRequestContext
4+
import com.datadog.iast.taint.TaintedObjects
5+
import datadog.trace.api.gateway.IGSpanInfo
6+
import datadog.trace.api.gateway.RequestContext
7+
8+
import java.util.concurrent.LinkedBlockingQueue
9+
import java.util.concurrent.TimeUnit
10+
11+
class IastRequestTestRunner extends IastAgentTestRunner implements IastRequestContextPreparationTrait {
12+
13+
private static final LinkedBlockingQueue<TaintedObjectCollection> TAINTED_OBJECTS = new LinkedBlockingQueue<>()
14+
15+
@Override
16+
protected Closure getRequestEndAction() {
17+
{ RequestContext requestContext, IGSpanInfo igSpanInfo ->
18+
// request end action
19+
IastRequestContext iastRequestContext = IastRequestContext.get(requestContext)
20+
if (iastRequestContext) {
21+
TaintedObjects taintedObjects = iastRequestContext.getTaintedObjects()
22+
TAINTED_OBJECTS.offer(new TaintedObjectCollection(taintedObjects))
23+
}
24+
}
25+
}
26+
27+
void cleanup() {
28+
TAINTED_OBJECTS.clear()
29+
}
30+
31+
protected TaintedObjectCollection getFinReqTaintedObjects() {
32+
TAINTED_OBJECTS.poll(15, TimeUnit.SECONDS)
33+
}
34+
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
package com.datadog.iast
1+
package com.datadog.iast.test
22

33
import com.datadog.iast.overhead.Operation
44
import com.datadog.iast.overhead.OverheadController

0 commit comments

Comments
 (0)