From 2b187de53b35ed71d7a22c5ed1c62792c45b153c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Manuel=20=C3=81lvarez=20=C3=81lvarez?= Date: Mon, 25 Mar 2024 09:54:40 +0100 Subject: [PATCH] Refactor propagation module to include two different APIs for strings and objects --- .../instrumentation/iast/NamedContext.java | 29 +- .../iast/NamedContextTest.groovy | 36 +- .../iast/propagation/FastCodecModule.java | 26 +- .../propagation/PropagationModuleImpl.java | 456 ++++++++---------- .../iast/GrpcRequestMessageHandlerTest.groovy | 4 - .../propagation/BaseCodecModuleTest.groovy | 21 - .../propagation/PropagationModuleTest.groovy | 96 ++-- .../iast/sink/AbstractSinkModuleTest.groovy | 2 +- .../iast/test/TaintMarkerHelpers.groovy | 3 +- .../iast/PathMatcherInstrumentation.java | 15 +- .../helpers/TaintSingleParameterFunction.java | 4 +- .../iast/helpers/TaintUnmarshaller.java | 8 +- .../iast/helpers/TaintUriFunction.java | 2 +- .../iast/helpers/TaintParametersFunction.java | 4 +- .../apache-httpcore-4/build.gradle | 1 + .../IastHttpHostInstrumentation.java | 8 +- .../IastHttpHostInstrumentationTest.groovy | 20 +- .../IastHttpMethodBaseInstrumentation.java | 9 +- .../commons-lang-2/build.gradle | 1 + .../StringEscapeUtilsCallSite.java | 15 +- .../StringEscapeUtilsCallSiteTest.groovy | 19 +- .../commons-lang-3/build.gradle | 1 + .../StringEscapeUtilsCallSite.java | 15 +- .../StringEscapeUtilsCallSiteTest.groovy | 26 +- .../instrumentation/commons-text/build.gradle | 1 + .../StringEscapeUtilsCallSite.java | 15 +- .../StringEscapeUtilsCallSiteTest.groovy | 26 +- .../instrumentation/freemarker/build.gradle | 1 + .../freemarker/StringUtilCallSite.java | 8 +- .../freemarker/StringUtilCallSiteTest.groovy | 21 +- .../instrumentation/gson-1.6/build.gradle | 1 + .../gson/JsonReaderInstrumentation.java | 11 +- .../gson/JsonReaderInstrumentationTest.groovy | 21 +- .../instrumentation/jackson-core/build.gradle | 1 + .../jackson-core/jackson-core-1/build.gradle | 1 + .../core/Json1FactoryInstrumentation.java | 15 +- .../Json1FactoryInstrumentationTest.groovy | 40 +- .../Json1ParserInstrumentationTest.groovy | 45 +- .../core/Json2FactoryInstrumentation.java | 15 +- .../core/TokenBufferInstrumentation.java | 8 +- .../instrumentation/java-io/build.gradle | 1 + .../java/lang/ByteBufferCallSite.java | 11 +- .../java/lang/InputStreamInstrumentation.java | 15 +- .../java/lang/InputStreamReaderCallSite.java | 8 +- .../java/lang/StringReaderCallSite.java | 12 +- .../java/io/BaseIoCallSiteTest.groovy | 9 +- .../java/io/ByteBufferTest.groovy | 19 +- .../io/InputStreamInstrumentationTest.groovy | 15 +- .../io/InputStreamReaderCallSiteTest.groovy | 5 +- .../java/io/StringReaderCallSiteTest.groovy | 9 +- .../java/lang/StringCallSite.java | 11 +- .../instrumentation/java/net/URICallSite.java | 25 +- .../instrumentation/java/net/URLCallSite.java | 20 +- .../java/net/URLEncoderCallSite.java | 6 +- ...actParamValueExtractorInstrumentation.java | 2 +- .../jersey/AbstractStringReaderAdvice.java | 2 +- .../JSONObjectUtilsInstrumentation.java | 2 +- .../kafka_clients/KafkaIastHelper.java | 13 +- .../kafka_clients/UtilsInstrumentation.java | 13 +- .../ByteBufInputStreamInstrumentation.java | 12 +- .../okhttp2/IastHttpUrlInstrumentation.java | 15 +- .../okhttp3/IastHttpUrlInstrumentation.java | 8 +- .../instrumentation/org-json/build.gradle | 1 + .../json/JSONArrayInstrumentation.java | 44 +- .../json/JSONCookieInstrumentation.java | 6 +- .../json/JSONObjectInstrumentation.java | 44 +- .../json/JSONTokenerInstrumentation.java | 6 +- .../JSONArrayInstrumentationTest.groovy | 43 +- .../JSONCookieInstrumentationTest.groovy | 13 +- .../JSONObjectInstrumentationTest.groovy | 62 ++- .../JSONTokenerInstrumentationTest.groovy | 11 +- .../owasp/esapi/EncoderCallSite.java | 51 +- .../iast/PathMatcherInstrumentation.java | 15 +- .../iast/helpers/TaintParametersFunction.java | 4 +- .../helpers/TaintRequestContextFunction.java | 2 +- .../iast/helpers/TaintRequestFunction.java | 2 +- .../helpers/TaintSingleParameterFunction.java | 4 +- .../iast/helpers/TaintUnmarshaller.java | 8 +- .../iast/helpers/TaintUriFunction.java | 2 +- .../resteasy/CookieParamInjectorAdvice.java | 4 +- .../resteasy/FormParamInjectorAdvice.java | 4 +- .../resteasy/HeaderParamInjectorAdvice.java | 20 +- .../resteasy/PathParamInjectorAdvice.java | 4 +- .../resteasy/QueryParamInjectorAdvice.java | 4 +- .../HttpServletResponseInstrumentation.java | 8 +- ...pServletResponseInstrumentationTest.groovy | 21 +- ...rtaHttpServletResponseInstrumentation.java | 8 +- .../StreamUtilsInstrumentation.java | 8 +- .../iast/DataBufferAsInputStreamAdvice.java | 6 +- .../springweb/EscapeUtilsCallSite.java | 15 +- .../unbescape/EscapeUtilsCallSite.java | 16 +- .../vertx_3_4/core/BufferInstrumentation.java | 34 +- .../HttpServerRequestInstrumentation.java | 2 +- .../vertx_4_0/core/BufferInstrumentation.java | 34 +- .../iast/propagation/PropagationModule.java | 204 ++++---- 95 files changed, 1127 insertions(+), 867 deletions(-) diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/iast/NamedContext.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/iast/NamedContext.java index 9d843cd68d6..beac5518297 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/iast/NamedContext.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/iast/NamedContext.java @@ -28,9 +28,10 @@ public static NamedContext getOrCreate( } final PropagationModule module = InstrumentationBridge.PROPAGATION; if (module != null) { - final Source source = module.findSource(target); - if (source != null) { - result = new NamedContextImpl(module, source); + final IastContext ctx = IastContext.Provider.get(); + final Source source; + if (ctx != null && (source = module.findSource(ctx, target)) != null) { + result = new NamedContextImpl(module, ctx, source); } } result = result == null ? NoOp.INSTANCE : result; @@ -51,20 +52,22 @@ public void taintName(@Nullable final String name) {} private static class NamedContextImpl extends NamedContext { @Nonnull private final PropagationModule module; + @Nonnull private final IastContext ctx; @Nonnull private final Source source; @Nullable private String currentName; - private boolean fetched; - @Nullable private IastContext context; - - public NamedContextImpl(@Nonnull final PropagationModule module, @Nonnull final Source source) { + public NamedContextImpl( + @Nonnull final PropagationModule module, + @Nonnull final IastContext ctx, + @Nonnull final Source source) { this.module = module; + this.ctx = ctx; this.source = source; } @Override public void taintValue(@Nullable final String value) { - module.taint(iastCtx(), value, source.getOrigin(), currentName, source.getValue()); + module.taint(ctx, value, source.getOrigin(), currentName, source.getValue()); } @Override @@ -74,16 +77,8 @@ public void taintName(@Nullable final String name) { // prevent tainting the same name more than once if (currentName != name) { currentName = name; - module.taint(iastCtx(), name, source.getOrigin(), name, source.getValue()); - } - } - - private IastContext iastCtx() { - if (!fetched) { - fetched = true; - context = IastContext.Provider.get(); + module.taint(ctx, name, source.getOrigin(), name, source.getValue()); } - return context; } } } diff --git a/dd-java-agent/agent-bootstrap/src/test/groovy/datadog/trace/bootstrap/instrumentation/iast/NamedContextTest.groovy b/dd-java-agent/agent-bootstrap/src/test/groovy/datadog/trace/bootstrap/instrumentation/iast/NamedContextTest.groovy index 8d1f7591fd2..dc976ad5983 100644 --- a/dd-java-agent/agent-bootstrap/src/test/groovy/datadog/trace/bootstrap/instrumentation/iast/NamedContextTest.groovy +++ b/dd-java-agent/agent-bootstrap/src/test/groovy/datadog/trace/bootstrap/instrumentation/iast/NamedContextTest.groovy @@ -1,24 +1,50 @@ package datadog.trace.bootstrap.instrumentation.iast - +import datadog.trace.api.gateway.RequestContext +import datadog.trace.api.gateway.RequestContextSlot +import datadog.trace.api.iast.IastContext import datadog.trace.api.iast.InstrumentationBridge import datadog.trace.api.iast.SourceTypes import datadog.trace.api.iast.Taintable.Source import datadog.trace.api.iast.propagation.PropagationModule import datadog.trace.bootstrap.ContextStore +import datadog.trace.bootstrap.instrumentation.api.AgentSpan +import datadog.trace.bootstrap.instrumentation.api.AgentTracer +import datadog.trace.bootstrap.instrumentation.api.AgentTracer.TracerAPI import datadog.trace.test.util.DDSpecification +import spock.lang.Shared class NamedContextTest extends DDSpecification { + @Shared + protected static final TracerAPI ORIGINAL_TRACER = AgentTracer.get() + protected PropagationModule module protected ContextStore store + protected TracerAPI tracer + protected IastContext ctx void setup() { + ctx = Stub(IastContext) + final reqCtx = Stub(RequestContext) { + getData(RequestContextSlot.IAST) >> ctx + } + final span = Stub(AgentSpan) { + getRequestContext() >> reqCtx + } + tracer = Stub(TracerAPI) { + activeSpan() >> span + } + AgentTracer.forceRegister(tracer) module = Mock(PropagationModule) InstrumentationBridge.registerIastModule(module) store = Mock(ContextStore) } + void cleanup() { + AgentTracer.forceRegister(ORIGINAL_TRACER) + } + void 'test that the context taints names and values'() { given: final source = new SourceImpl(origin: SourceTypes.REQUEST_PARAMETER_NAME) @@ -31,14 +57,14 @@ class NamedContextTest extends DDSpecification { then: 1 * store.get(target) >> null - 1 * module.findSource(target) >> source + 1 * module.findSource(ctx, target) >> source 1 * store.put(target, _) when: context.taintName(name) then: - 1 * module.taint(_, name, source.origin, name, source.value) + 1 * module.taint(ctx, name, source.origin, name, source.value) when: context.taintName(name) @@ -50,7 +76,7 @@ class NamedContextTest extends DDSpecification { context.taintValue(value) then: - 1 * module.taint(_, value, source.origin, name, source.value) + 1 * module.taint(ctx, value, source.origin, name, source.value) 0 * _ } @@ -62,7 +88,7 @@ class NamedContextTest extends DDSpecification { final ctx = NamedContext.getOrCreate(store, target) then: - 1 * module.findSource(target) >> null + 1 * module.findSource(ctx, target) >> null 1 * store.put(target, _) when: diff --git a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/propagation/FastCodecModule.java b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/propagation/FastCodecModule.java index a0fad39c309..9bf5a172d76 100644 --- a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/propagation/FastCodecModule.java +++ b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/propagation/FastCodecModule.java @@ -2,6 +2,7 @@ import static datadog.trace.api.iast.VulnerabilityMarks.NOT_MARKED; +import datadog.trace.api.iast.IastContext; import datadog.trace.api.iast.propagation.CodecModule; import javax.annotation.Nonnull; import javax.annotation.Nullable; @@ -11,7 +12,10 @@ public class FastCodecModule extends PropagationModuleImpl implements CodecModul @Override public void onUrlDecode( @Nonnull final String value, @Nullable final String encoding, @Nonnull final String result) { - taintIfTainted(result, value); + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + taintIfTainted(ctx, result, value); + } } @Override @@ -22,22 +26,34 @@ public void onStringFromBytes( @Nullable final String charset, @Nonnull final String result) { // create a new range shifted to the result string coordinates - taintIfTainted(result, value, offset, length, false, NOT_MARKED); + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + taintIfTainted(ctx, result, value, offset, length, false, NOT_MARKED); + } } @Override public void onStringGetBytes( @Nonnull final String value, @Nullable final String charset, @Nonnull final byte[] result) { - taintIfTainted(result, value); + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + taintIfTainted(ctx, result, value); + } } @Override public void onBase64Encode(@Nullable byte[] value, @Nullable byte[] result) { - taintIfTainted(result, value); + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + taintIfTainted(ctx, result, value); + } } @Override public void onBase64Decode(@Nullable byte[] value, @Nullable byte[] result) { - taintIfTainted(result, value); + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + taintIfTainted(ctx, result, value); + } } } diff --git a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/propagation/PropagationModuleImpl.java b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/propagation/PropagationModuleImpl.java index ec07035c38c..ea10ee4afbf 100644 --- a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/propagation/PropagationModuleImpl.java +++ b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/propagation/PropagationModuleImpl.java @@ -10,7 +10,6 @@ import com.datadog.iast.taint.Ranges; import com.datadog.iast.taint.TaintedObject; import com.datadog.iast.taint.TaintedObjects; -import com.datadog.iast.taint.Tainteds; import com.datadog.iast.util.ObjectVisitor; import com.datadog.iast.util.Ranged; import datadog.trace.api.Config; @@ -18,152 +17,155 @@ import datadog.trace.api.iast.Taintable; import datadog.trace.api.iast.propagation.PropagationModule; import java.lang.ref.WeakReference; -import java.lang.reflect.Array; import java.util.function.Predicate; import javax.annotation.Nonnull; import javax.annotation.Nullable; -import org.jetbrains.annotations.Contract; +@SuppressWarnings("DuplicatedCode") public class PropagationModuleImpl implements PropagationModule { /** Prevent copy of values bigger than this threshold */ private static final int MAX_VALUE_LENGTH = Config.get().getIastTruncationMaxValueLength(); - @Override - public void taint(@Nullable final Object target, final byte origin) { - taint(target, origin, null); - } - - @Override - public void taint( - @Nullable final Object target, final byte origin, @Nullable final CharSequence name) { - taint(target, origin, name, target); - } - @Override public void taint( + @Nonnull final IastContext ctx, @Nullable final Object target, final byte origin, - @Nullable final CharSequence name, - @Nullable final Object value) { - if (!canBeTainted(target)) { + final int start, + final int length) { + if (target == null || length == 0) { return; } - taint(LazyContext.build(), target, origin, name, value); - } - - @Override - public void taint( - @Nullable final IastContext ctx, @Nullable final Object target, final byte origin) { - taint(ctx, target, origin, null); - } - - @Override - public void taint( - @Nullable final Object target, final byte origin, final int start, final int length) { - taint(LazyContext.build(), target, origin, start, length); + final TaintedObjects to = ctx.getTaintedObjects(); + final Range range = + new Range(start, length, newSource(target, origin, null, target), NOT_MARKED); + internalTaint(to, target, new Range[] {range}, NOT_MARKED, true); } @Override public void taint( - @Nullable final IastContext ctx, - @Nullable final Object target, + @Nonnull final IastContext ctx, + @Nullable final String target, final byte origin, final int start, final int length) { - if (!canBeTainted(target) || length == 0) { + if (target == null || length == 0) { return; } + final TaintedObjects to = ctx.getTaintedObjects(); final Range range = new Range(start, length, newSource(target, origin, null, target), NOT_MARKED); - internalTaint(ctx, target, new Range[] {range}, NOT_MARKED); + internalTaint(to, target, new Range[] {range}, NOT_MARKED, true); } @Override public void taint( - @Nullable final IastContext ctx, - @Nullable final Object target, - final byte origin, - @Nullable final CharSequence name) { - taint(ctx, target, origin, name, target); - } - - @Override - public void taint( - @Nullable final IastContext ctx, + @Nonnull final IastContext ctx, @Nullable final Object target, final byte origin, @Nullable final CharSequence name, @Nullable final Object value) { - if (!canBeTainted(target)) { + if (target == null) { return; } - internalTaint(ctx, target, newSource(target, origin, name, value), NOT_MARKED); - } - - @Override - public void taintIfTainted(@Nullable final Object target, @Nullable final Object input) { - taintIfTainted(target, input, false, NOT_MARKED); + final TaintedObjects to = ctx.getTaintedObjects(); + internalTaint(to, target, newSource(target, origin, name, value), NOT_MARKED); } @Override - public void taintIfTainted( - @Nullable final Object target, @Nullable final Object input, boolean keepRanges, int mark) { - if (!canBeTainted(target) || !canBeTainted(input)) { + public void taint( + @Nonnull final IastContext ctx, + @Nullable final String target, + final byte origin, + @Nullable final CharSequence name, + @Nullable final CharSequence value) { + if (target == null) { return; } - taintIfTainted(LazyContext.build(), target, input, keepRanges, mark); + final TaintedObjects to = ctx.getTaintedObjects(); + internalTaint(to, target, newSource(target, origin, name, value), NOT_MARKED); } @Override public void taintIfTainted( - @Nullable final IastContext ctx, + @Nonnull final IastContext ctx, @Nullable final Object target, - @Nullable final Object input) { - taintIfTainted(ctx, target, input, false, NOT_MARKED); + @Nullable final Object input, + boolean keepRanges, + int mark) { + if (target == null || input == null) { + return; + } + final TaintedObjects to = ctx.getTaintedObjects(); + if (keepRanges) { + internalTaint(to, target, getRanges(to, input), mark); + } else { + internalTaint(to, target, highestPrioritySource(to, input), mark); + } } @Override public void taintIfTainted( - @Nullable final IastContext ctx, - @Nullable final Object target, + @Nonnull final IastContext ctx, + @Nullable final String target, @Nullable final Object input, boolean keepRanges, int mark) { - if (!canBeTainted(target) || !canBeTainted(input)) { + if (target == null || input == null) { return; } + final TaintedObjects to = ctx.getTaintedObjects(); if (keepRanges) { - internalTaint(ctx, target, getRanges(ctx, input), mark); + internalTaint(to, target, getRanges(to, input), mark); } else { - internalTaint(ctx, target, highestPrioritySource(ctx, input), mark); + internalTaint(to, target, highestPrioritySource(to, input), mark); } } @Override public void taintIfTainted( + @Nonnull final IastContext ctx, @Nullable final Object target, @Nullable final Object input, final int start, final int length, boolean keepRanges, int mark) { - taintIfTainted(LazyContext.build(), target, input, start, length, keepRanges, mark); + if (target == null || input == null || length == 0) { + return; + } + final TaintedObjects to = ctx.getTaintedObjects(); + final Range[] ranges = getRanges(to, input); + if (ranges == null || ranges.length == 0) { + return; + } + final Range[] intersection = Ranges.intersection(Ranged.build(start, length), ranges); + if (intersection == null || intersection.length == 0) { + return; + } + if (keepRanges) { + internalTaint(to, target, intersection, mark); + } else { + final Range range = highestPriorityRange(intersection); + internalTaint(to, target, range.getSource(), mark); + } } @Override public void taintIfTainted( - @Nullable final IastContext ctx, - @Nullable final Object target, + @Nonnull final IastContext ctx, + @Nullable final String target, @Nullable final Object input, final int start, final int length, boolean keepRanges, int mark) { - if (!canBeTainted(target) || !canBeTainted(input) || length == 0) { + if (target == null || input == null || length == 0) { return; } - final Range[] ranges = getRanges(ctx, input); + final TaintedObjects to = ctx.getTaintedObjects(); + final Range[] ranges = getRanges(to, input); if (ranges == null || ranges.length == 0) { return; } @@ -172,148 +174,107 @@ public void taintIfTainted( return; } if (keepRanges) { - internalTaint(ctx, target, intersection, mark); + internalTaint(to, target, intersection, mark); } else { final Range range = highestPriorityRange(intersection); - internalTaint(ctx, target, range.getSource(), mark); + internalTaint(to, target, range.getSource(), mark); } } @Override public void taintIfTainted( - @Nullable final Object target, @Nullable final Object input, final byte origin) { - taintIfTainted(target, input, origin, null); - } - - @Override - public void taintIfTainted( - @Nullable final Object target, - @Nullable final Object input, - final byte origin, - @Nullable final CharSequence name) { - taintIfTainted(target, input, origin, name, target); - } - - @Override - public void taintIfTainted( + @Nonnull final IastContext ctx, @Nullable final Object target, @Nullable final Object input, final byte origin, @Nullable final CharSequence name, @Nullable final Object value) { - if (!canBeTainted(target) || !canBeTainted(input)) { + if (target == null || input == null) { return; } - taintIfTainted(LazyContext.build(), target, input, origin, name, value); - } - - @Override - public void taintIfTainted( - @Nullable final IastContext ctx, - @Nullable final Object target, - @Nullable final Object input, - final byte origin) { - taintIfTainted(ctx, target, input, origin, null); - } - - @Override - public void taintIfTainted( - @Nullable final IastContext ctx, - @Nullable final Object target, - @Nullable final Object input, - final byte origin, - @Nullable final CharSequence name) { - taintIfTainted(ctx, target, input, origin, name, target); + final TaintedObjects to = ctx.getTaintedObjects(); + if (to.get(input) != null) { + internalTaint(to, target, newSource(target, origin, name, value), NOT_MARKED); + } } @Override public void taintIfTainted( - @Nullable final IastContext ctx, - @Nullable final Object target, + @Nonnull final IastContext ctx, + @Nullable final String target, @Nullable final Object input, final byte origin, @Nullable final CharSequence name, @Nullable final Object value) { - if (!canBeTainted(target) || !canBeTainted(input)) { + if (target == null || input == null) { return; } - if (isTainted(ctx, input)) { - internalTaint(ctx, target, newSource(target, origin, name, value), NOT_MARKED); + final TaintedObjects to = ctx.getTaintedObjects(); + if (to.get(input) != null) { + internalTaint(to, target, newSource(target, origin, name, value), NOT_MARKED); } } - @Override - public void taintIfAnyTainted(@Nullable final Object target, @Nullable final Object[] inputs) { - taintIfAnyTainted(target, inputs, false, NOT_MARKED); - } - @Override public void taintIfAnyTainted( + @Nonnull final IastContext ctx, @Nullable final Object target, @Nullable final Object[] inputs, final boolean keepRanges, final int mark) { - if (!canBeTainted(target) || !canBeTainted(inputs)) { + if (target == null || inputs == null || inputs.length == 0) { return; } - taintIfAnyTainted(LazyContext.build(), target, inputs, keepRanges, mark); - } - - @Override - public void taintIfAnyTainted( - @Nullable final IastContext ctx, - @Nullable final Object target, - @Nullable final Object[] inputs) { - taintIfAnyTainted(ctx, target, inputs, false, NOT_MARKED); + final TaintedObjects to = ctx.getTaintedObjects(); + if (keepRanges) { + final Range[] ranges = getRangesInArray(to, inputs); + if (ranges != null) { + internalTaint(to, target, ranges, mark); + } + } else { + final Source source = highestPrioritySourceInArray(to, inputs); + if (source != null) { + internalTaint(to, target, source, mark); + } + } } @Override public void taintIfAnyTainted( - @Nullable final IastContext ctx, - @Nullable final Object target, + @Nonnull final IastContext ctx, + @Nullable final String target, @Nullable final Object[] inputs, final boolean keepRanges, final int mark) { - if (!canBeTainted(target) || !canBeTainted(inputs)) { + if (target == null || inputs == null || inputs.length == 0) { return; } + final TaintedObjects to = ctx.getTaintedObjects(); if (keepRanges) { - final Range[] ranges = getRangesInArray(ctx, inputs); + final Range[] ranges = getRangesInArray(to, inputs); if (ranges != null) { - internalTaint(ctx, target, ranges, mark); + internalTaint(to, target, ranges, mark); } } else { - final Source source = highestPrioritySourceInArray(ctx, inputs); + final Source source = highestPrioritySourceInArray(to, inputs); if (source != null) { - internalTaint(ctx, target, source, mark); + internalTaint(to, target, source, mark); } } } @Override public int taintDeeply( - @Nullable final Object target, final byte origin, final Predicate> classFilter) { - if (!canBeTainted(target)) { - return 0; - } - return taintDeeply(LazyContext.build(), target, origin, classFilter); - } - - @Override - public int taintDeeply( - @Nullable final IastContext ctx, + @Nonnull final IastContext ctx, @Nullable final Object target, final byte origin, final Predicate> classFilter) { - if (!canBeTainted(target)) { - return 0; - } - final TaintedObjects to = getTaintedObjects(ctx); - if (to == null) { + if (target == null) { return 0; } + final TaintedObjects to = ctx.getTaintedObjects(); if (target instanceof CharSequence) { - internalTaint(ctx, target, newSource(target, origin, null, target), NOT_MARKED); + internalTaint(to, target, newSource(target, origin, null, target), NOT_MARKED); return 1; } else { final TaintingVisitor visitor = new TaintingVisitor(to, origin); @@ -322,29 +283,24 @@ public int taintDeeply( } } - @Nullable - @Override - public Taintable.Source findSource(@Nullable final Object target) { - return target == null ? null : findSource(LazyContext.build(), target); - } - @Nullable @Override public Taintable.Source findSource( - @Nullable final IastContext ctx, @Nullable final Object target) { + @Nonnull final IastContext ctx, @Nullable final Object target) { if (target == null) { return null; } - return highestPrioritySource(ctx, target); + final TaintedObjects to = ctx.getTaintedObjects(); + return highestPrioritySource(to, target); } @Override - public boolean isTainted(@Nullable final Object target) { - return target != null && isTainted(LazyContext.build(), target); + public boolean isTainted(@Nonnull final IastContext ctx, @Nullable final Object target) { + return target != null && findSource(ctx, target) != null; } @Override - public boolean isTainted(@Nullable final IastContext ctx, @Nullable final Object target) { + public boolean isTainted(@Nonnull final IastContext ctx, @Nullable final String target) { return target != null && findSource(ctx, target) != null; } @@ -353,10 +309,24 @@ public boolean isTainted(@Nullable final IastContext ctx, @Nullable final Object * properties */ private static Source newSource( - @Nonnull final Object tainted, + @Nonnull final Object target, final byte origin, @Nullable final CharSequence name, @Nullable final Object value) { + final Object sourceValue = sourceReference(target, value, true); + final Object sourceName = name == value ? sourceValue : sourceReference(target, name, false); + return new Source(origin, sourceName, sourceValue); + } + + /** + * Ensures that the reference is not kept due to a strong reference via the name or value + * properties + */ + private static Source newSource( + @Nonnull final String tainted, + final byte origin, + @Nullable final CharSequence name, + @Nullable final CharSequence value) { final Object sourceValue = sourceReference(tainted, value, true); final Object sourceName = name == value ? sourceValue : sourceReference(tainted, name, false); return new Source(origin, sourceName, sourceValue); @@ -384,38 +354,11 @@ private static Object sourceReference( return value ? PROPAGATION_PLACEHOLDER : null; } - @Contract("null -> false") - private static boolean canBeTainted(@Nullable final Object target) { - if (target == null) { - return false; - } - if (target instanceof CharSequence) { - return Tainteds.canBeTainted((CharSequence) target); - } - if (target.getClass().isArray()) { - return Array.getLength(target) > 0; - } - return true; - } - - @Contract("null -> false") - private static boolean canBeTainted(@Nullable final Object[] target) { - if (target == null || target.length == 0) { - return false; - } - return true; - } - - @Nullable - private static TaintedObjects getTaintedObjects(@Nullable final IastContext ctx) { - return ctx == null ? null : ctx.getTaintedObjects(); - } - @Nullable private static Range[] getRangesInArray( - final @Nullable IastContext ctx, final @Nonnull Object[] objects) { + final @Nonnull TaintedObjects to, final @Nonnull Object[] objects) { for (final Object object : objects) { - final Range[] ranges = getRanges(ctx, object); + final Range[] ranges = getRanges(to, object); if (ranges != null) { return ranges; } @@ -424,28 +367,24 @@ private static Range[] getRangesInArray( } @Nullable - private static Range[] getRanges(final @Nullable IastContext ctx, final @Nonnull Object object) { + private static Range[] getRanges(final @Nonnull TaintedObjects to, final @Nonnull Object object) { if (object instanceof Taintable) { - final Source source = highestPrioritySource(ctx, object); + final Source source = highestPrioritySource(to, object); if (source == null) { return null; } else { return new Range[] {new Range(0, Integer.MAX_VALUE, source, NOT_MARKED)}; } } - final TaintedObjects to = getTaintedObjects(ctx); - if (to == null) { - return null; - } final TaintedObject tainted = to.get(object); return tainted == null ? null : tainted.getRanges(); } @Nullable private static Source highestPrioritySourceInArray( - final @Nullable IastContext ctx, final @Nonnull Object[] objects) { + final @Nonnull TaintedObjects to, final @Nonnull Object[] objects) { for (final Object object : objects) { - final Source source = highestPrioritySource(ctx, object); + final Source source = highestPrioritySource(to, object); if (source != null) { return source; } @@ -455,17 +394,17 @@ private static Source highestPrioritySourceInArray( @Nullable private static Source highestPrioritySource( - final @Nullable IastContext ctx, final @Nonnull Object object) { + final @Nonnull TaintedObjects to, final @Nonnull Object object) { if (object instanceof Taintable) { return (Source) ((Taintable) object).$$DD$getSource(); } else { - final Range[] ranges = getRanges(ctx, object); + final Range[] ranges = getRanges(to, object); return ranges != null && ranges.length > 0 ? highestPriorityRange(ranges).getSource() : null; } } private static void internalTaint( - @Nullable final IastContext ctx, + @Nonnull final TaintedObjects to, @Nonnull final Object value, @Nullable Source source, int mark) { @@ -475,10 +414,6 @@ private static void internalTaint( if (value instanceof Taintable) { ((Taintable) value).$$DD$setSource(source); } else { - final TaintedObjects to = getTaintedObjects(ctx); - if (to == null) { - return; - } if (value instanceof CharSequence) { source = source.attachValue((CharSequence) value); to.taint(value, Ranges.forCharSequence((CharSequence) value, source, mark)); @@ -489,10 +424,39 @@ private static void internalTaint( } private static void internalTaint( - @Nullable final IastContext ctx, + @Nonnull final TaintedObjects to, + @Nonnull final String value, + @Nullable Source source, + int mark) { + if (source == null) { + return; + } + source = source.attachValue(value); + to.taint(value, Ranges.forCharSequence(value, source, mark)); + } + + private static void internalTaint( + @Nonnull final TaintedObjects to, @Nonnull final Object value, @Nullable Range[] ranges, final int mark) { + internalTaint(to, value, ranges, mark, false); + } + + private static void internalTaint( + @Nonnull final TaintedObjects to, + @Nonnull final String value, + @Nullable Range[] ranges, + final int mark) { + internalTaint(to, value, ranges, mark, false); + } + + private static void internalTaint( + @Nonnull final TaintedObjects to, + @Nonnull final Object value, + @Nullable Range[] ranges, + final int mark, + final boolean appendRanges) { if (ranges == null || ranges.length == 0) { return; } @@ -502,25 +466,48 @@ private static void internalTaint( taintable.$$DD$setSource(ranges[0].getSource()); } } else { - final TaintedObjects to = getTaintedObjects(ctx); - if (to != null) { - if (value instanceof CharSequence) { - ranges = attachSourceValue(ranges, (CharSequence) value); - } - ranges = markRanges(ranges, mark); + if (value instanceof CharSequence) { + ranges = attachSourceValue(ranges, (CharSequence) value); + } + ranges = markRanges(ranges, mark); + if (appendRanges) { final TaintedObject tainted = to.get(value); - if (tainted == null) { - // taint new value - to.taint(value, ranges); - } else { + if (tainted != null) { // append ranges final Range[] newRanges = Ranges.mergeRangesSorted(tainted.getRanges(), ranges); tainted.setRanges(newRanges); + return; } } + // taint new value + to.taint(value, ranges); } } + private static void internalTaint( + @Nonnull final TaintedObjects to, + @Nonnull final String value, + @Nullable Range[] ranges, + final int mark, + final boolean appendRanges) { + if (ranges == null || ranges.length == 0) { + return; + } + ranges = attachSourceValue(ranges, value); + ranges = markRanges(ranges, mark); + if (appendRanges) { + final TaintedObject tainted = to.get(value); + if (tainted != null) { + // append ranges + final Range[] newRanges = Ranges.mergeRangesSorted(tainted.getRanges(), ranges); + tainted.setRanges(newRanges); + return; + } + } + // taint new value + to.taint(value, ranges); + } + @Nonnull private static Range[] markRanges(@Nonnull final Range[] ranges, final int mark) { if (mark == NOT_MARKED) { @@ -549,33 +536,6 @@ private static Range[] attachSourceValue( : Ranges.forCharSequence(value, newSource, range.getMarks()); } - private static class LazyContext implements IastContext { - - private boolean fetched; - @Nullable private IastContext delegate; - - @Nullable - private IastContext getDelegate() { - if (!fetched) { - fetched = true; - delegate = IastContext.Provider.get(); - } - return delegate; - } - - public static IastContext build() { - return new LazyContext(); - } - - @SuppressWarnings("unchecked") - @Nonnull - @Override - public TaintedObjects getTaintedObjects() { - final IastContext delegate = getDelegate(); - return delegate == null ? TaintedObjects.NoOp.INSTANCE : delegate.getTaintedObjects(); - } - } - private static class TaintingVisitor implements ObjectVisitor.Visitor { private final TaintedObjects taintedObjects; @@ -592,12 +552,10 @@ private TaintingVisitor(@Nonnull final TaintedObjects taintedObjects, final byte public ObjectVisitor.State visit(@Nonnull final String path, @Nonnull final Object value) { if (value instanceof CharSequence) { final CharSequence charSequence = (CharSequence) value; - if (canBeTainted(charSequence)) { - final Source source = newSource(charSequence, origin, path, charSequence); - count++; - taintedObjects.taint( - charSequence, Ranges.forCharSequence(charSequence, source, NOT_MARKED)); - } + final Source source = newSource(charSequence, origin, path, charSequence); + count++; + taintedObjects.taint( + charSequence, Ranges.forCharSequence(charSequence, source, NOT_MARKED)); } return CONTINUE; } diff --git a/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/GrpcRequestMessageHandlerTest.groovy b/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/GrpcRequestMessageHandlerTest.groovy index c5c565a681b..76d823e16f0 100644 --- a/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/GrpcRequestMessageHandlerTest.groovy +++ b/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/GrpcRequestMessageHandlerTest.groovy @@ -3,16 +3,12 @@ package com.datadog.iast import com.datadog.iast.propagation.PropagationModuleImpl import com.datadog.iast.protobuf.Test2 import com.datadog.iast.protobuf.Test3 -import com.datadog.iast.taint.TaintedObjects import com.datadog.iast.util.ObjectVisitor -import datadog.trace.api.gateway.RequestContext -import datadog.trace.api.gateway.RequestContextSlot import datadog.trace.api.iast.InstrumentationBridge import datadog.trace.api.iast.SourceTypes import datadog.trace.api.iast.propagation.PropagationModule import datadog.trace.api.iast.telemetry.IastMetric import datadog.trace.api.iast.telemetry.IastMetricCollector -import datadog.trace.test.util.DDSpecification import foo.bar.VisitableClass import java.util.function.Predicate diff --git a/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/propagation/BaseCodecModuleTest.groovy b/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/propagation/BaseCodecModuleTest.groovy index b555ee6b87c..5201d9b64be 100644 --- a/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/propagation/BaseCodecModuleTest.groovy +++ b/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/propagation/BaseCodecModuleTest.groovy @@ -25,27 +25,6 @@ abstract class BaseCodecModuleTest extends IastModuleImplTestBase { } } - void '#method null or empty'() { - when: - module.&"$method".call(args.toArray()) - - then: - 0 * _ - - where: - method | args - 'onUrlDecode' | ['test', 'utf-8', null] - 'onUrlDecode' | ['test', 'utf-8', ''] - 'onStringGetBytes' | ['test', 'utf-8', null] - 'onStringGetBytes' | ['test', 'utf-8', [] as byte[]] - 'onStringFromBytes' | ['test'.bytes, 0, 2, 'utf-8', null] - 'onStringFromBytes' | ['test'.bytes, 0, 2, 'utf-8', ''] - 'onBase64Encode' | ['test'.bytes, null] - 'onBase64Encode' | ['test'.bytes, [] as byte[]] - 'onBase64Decode' | ['test'.bytes, null] - 'onBase64Decode' | ['test'.bytes, [] as byte[]] - } - void '#method no context'() { when: module.&"$method".call(args.toArray()) diff --git a/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/propagation/PropagationModuleTest.groovy b/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/propagation/PropagationModuleTest.groovy index bc134c31e4f..d61b5ab2fba 100644 --- a/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/propagation/PropagationModuleTest.groovy +++ b/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/propagation/PropagationModuleTest.groovy @@ -40,12 +40,6 @@ class PropagationModuleTest extends IastModuleImplTestBase { } void '#method(#args) not taintable'() { - when: 'there is no context by default' - module.&"$method".call(args.toArray()) - - then: 'no mock calls should happen' - 0 * _ - when: 'there is a context' args.add(0, ctx) module.&"$method".call(args.toArray()) @@ -87,38 +81,6 @@ class PropagationModuleTest extends IastModuleImplTestBase { 'isTainted' | [null] } - void '#method without span'() { - when: - module.&"$method".call(args.toArray()) - - then: - 1 * tracer.activeSpan() >> null - 0 * _ - - where: - method | args - 'taint' | ['test', SourceTypes.REQUEST_PARAMETER_VALUE] - 'taint' | ['test', SourceTypes.REQUEST_PARAMETER_VALUE, 'name'] - 'taint' | ['test', SourceTypes.REQUEST_PARAMETER_VALUE, 'name', 'value'] - 'taint' | ['test', SourceTypes.REQUEST_PARAMETER_VALUE, 0, 10] - 'taintIfTainted' | ['test', 'test'] - 'taintIfTainted' | ['test', 'test', 0, 4, false, NOT_MARKED] - 'taintIfTainted' | ['test', 'test', false, NOT_MARKED] - 'taintIfTainted' | ['test', 'test', SourceTypes.REQUEST_PARAMETER_VALUE] - 'taintIfTainted' | ['test', 'test', SourceTypes.REQUEST_PARAMETER_VALUE, 'name'] - 'taintIfTainted' | ['test', 'test', SourceTypes.REQUEST_PARAMETER_VALUE, 'name', 'value'] - 'taintIfAnyTainted' | ['test', ['test']] - 'taintDeeply' | [ - 'test', - SourceTypes.REQUEST_PARAMETER_VALUE, - { - true - } - ] - 'findSource' | ['test'] - 'isTainted' | ['test'] - } - void 'test taint'() { given: final value = (target instanceof CharSequence) ? target.toString() : null @@ -126,7 +88,7 @@ class PropagationModuleTest extends IastModuleImplTestBase { final ranges = Ranges.forObject(source) when: - module.taint(target, source.origin, source.name, source.value) + module.taint(ctx, target, source.origin, source.name, source.value) then: final tainted = getTaintedObject(target) @@ -151,7 +113,7 @@ class PropagationModuleTest extends IastModuleImplTestBase { final ranges = [new Range(start, length, source, NOT_MARKED)] as Range[] when: - module.taint(target, source.origin, start, length) + module.taint(ctx, target, source.origin, start, length) then: final tainted = getTaintedObject(target) @@ -172,14 +134,14 @@ class PropagationModuleTest extends IastModuleImplTestBase { final ranges = [new Range(0, 1, source, NOT_MARKED), new Range(1, 1, source, NOT_MARKED)] as Range[] when: 'input is not tainted' - module.taintIfTainted(target, input, true, NOT_MARKED) + module.taintIfTainted(ctx, target, input, true, NOT_MARKED) then: assert getTaintedObject(target) == null when: 'input is tainted' final taintedFrom = taintObject(input, ranges) - module.taintIfTainted(target, input, true, NOT_MARKED) + module.taintIfTainted(ctx, target, input, true, NOT_MARKED) then: final tainted = getTaintedObject(target) @@ -201,14 +163,14 @@ class PropagationModuleTest extends IastModuleImplTestBase { final ranges = [new Range(0, 2, source, NOT_MARKED)] as Range[] when: 'input is not tainted' - module.taintIfTainted(target, input, 0, 2, false, NOT_MARKED) + module.taintIfTainted(ctx, target, input, 0, 2, false, NOT_MARKED) then: assert getTaintedObject(target) == null when: 'input tainted but range does not overlap' final firstTaintedForm = taintObject(input, ranges) - module.taintIfTainted(target, input, 4, 3, false, NOT_MARKED) + module.taintIfTainted(ctx, target, input, 4, 3, false, NOT_MARKED) then: final firstTainted = getTaintedObject(target) @@ -222,7 +184,7 @@ class PropagationModuleTest extends IastModuleImplTestBase { when: 'input is tainted and range overlaps' ctx.taintedObjects.clear() final secondTaintedFrom = taintObject(input, ranges) - module.taintIfTainted(target, input, 0, 2, false, NOT_MARKED) + module.taintIfTainted(ctx, target, input, 0, 2, false, NOT_MARKED) then: final secondTainted = getTaintedObject(target) @@ -246,14 +208,14 @@ class PropagationModuleTest extends IastModuleImplTestBase { final mark = VulnerabilityMarks.UNVALIDATED_REDIRECT_MARK when: 'input is not tainted' - module.taintIfTainted(target, input, true, mark) + module.taintIfTainted(ctx, target, input, true, mark) then: assert getTaintedObject(target) == null when: 'input is tainted' final taintedFrom = taintObject(input, ranges) - module.taintIfTainted(target, input, true, mark) + module.taintIfTainted(ctx, target, input, true, mark) then: final tainted = getTaintedObject(target) @@ -270,14 +232,14 @@ class PropagationModuleTest extends IastModuleImplTestBase { final ranges = [new Range(0, 1, source, NOT_MARKED), new Range(1, 1, source, NOT_MARKED)] as Range[] when: 'input is not tainted' - module.taintIfTainted(target, input, false, NOT_MARKED) + module.taintIfTainted(ctx, target, input, false, NOT_MARKED) then: assert getTaintedObject(target) == null when: 'input is tainted' final taintedFrom = taintObject(input, ranges) - module.taintIfTainted(target, input, false, NOT_MARKED) + module.taintIfTainted(ctx, target, input, false, NOT_MARKED) then: final tainted = getTaintedObject(target) @@ -296,14 +258,14 @@ class PropagationModuleTest extends IastModuleImplTestBase { final mark = VulnerabilityMarks.LDAP_INJECTION_MARK when: 'input is not tainted' - module.taintIfTainted(target, input, false, mark) + module.taintIfTainted(ctx, target, input, false, mark) then: assert getTaintedObject(target) == null when: 'input is tainted' final taintedFrom = taintObject(input, ranges) - module.taintIfTainted(target, input, false, mark) + module.taintIfTainted(ctx, target, input, false, mark) then: final tainted = getTaintedObject(target) @@ -321,14 +283,14 @@ class PropagationModuleTest extends IastModuleImplTestBase { final ranges = [new Range(0, 1, source, NOT_MARKED), new Range(1, 1, source, NOT_MARKED)] as Range[] when: 'input is not tainted' - module.taintIfAnyTainted(target, inputs, true, NOT_MARKED) + module.taintIfAnyTainted(ctx, target, inputs, true, NOT_MARKED) then: assert getTaintedObject(target) == null when: 'input is tainted' final taintedFrom = taintObject(input, ranges) - module.taintIfAnyTainted(target, inputs, true, NOT_MARKED) + module.taintIfAnyTainted(ctx, target, inputs, true, NOT_MARKED) then: final tainted = getTaintedObject(target) @@ -353,14 +315,14 @@ class PropagationModuleTest extends IastModuleImplTestBase { final mark = VulnerabilityMarks.UNVALIDATED_REDIRECT_MARK when: 'input is not tainted' - module.taintIfAnyTainted(target, inputs, true, mark) + module.taintIfAnyTainted(ctx, target, inputs, true, mark) then: assert getTaintedObject(target) == null when: 'input is tainted' final taintedFrom = taintObject(input, ranges) - module.taintIfAnyTainted(target, inputs, true, mark) + module.taintIfAnyTainted(ctx, target, inputs, true, mark) then: final tainted = getTaintedObject(target) @@ -378,14 +340,14 @@ class PropagationModuleTest extends IastModuleImplTestBase { final ranges = [new Range(0, 1, source, NOT_MARKED), new Range(1, 1, source, NOT_MARKED)] as Range[] when: 'input is not tainted' - module.taintIfAnyTainted(target, inputs, false, NOT_MARKED) + module.taintIfAnyTainted(ctx, target, inputs, false, NOT_MARKED) then: assert getTaintedObject(target) == null when: 'input is tainted' final taintedFrom = taintObject(input, ranges) - module.taintIfAnyTainted(target, inputs, false, NOT_MARKED) + module.taintIfAnyTainted(ctx, target, inputs, false, NOT_MARKED) then: final tainted = getTaintedObject(target) @@ -405,14 +367,14 @@ class PropagationModuleTest extends IastModuleImplTestBase { final mark = VulnerabilityMarks.LDAP_INJECTION_MARK when: 'input is not tainted' - module.taintIfAnyTainted(target, inputs, false, mark) + module.taintIfAnyTainted(ctx, target, inputs, false, mark) then: assert getTaintedObject(target) == null when: 'input is tainted' final taintedFrom = taintObject(input, ranges) - module.taintIfAnyTainted(target, inputs, false, mark) + module.taintIfAnyTainted(ctx, target, inputs, false, mark) then: final tainted = getTaintedObject(target) @@ -427,7 +389,7 @@ class PropagationModuleTest extends IastModuleImplTestBase { final target = [Hello: " World!", Age: 25] when: - module.taintDeeply(target, SourceTypes.GRPC_BODY, { true }) + module.taintDeeply(ctx, target, SourceTypes.GRPC_BODY, { true }) then: final taintedObjects = ctx.taintedObjects @@ -443,7 +405,7 @@ class PropagationModuleTest extends IastModuleImplTestBase { final target = stringBuilder('taint me') when: - module.taintDeeply(target, SourceTypes.GRPC_BODY, { true }) + module.taintDeeply(ctx, target, SourceTypes.GRPC_BODY, { true }) then: final taintedObjects = ctx.taintedObjects @@ -462,13 +424,13 @@ class PropagationModuleTest extends IastModuleImplTestBase { } when: - final tainted = module.isTainted(target) + final tainted = module.isTainted(ctx, target) then: tainted == (source != null) when: - final foundSource = module.findSource(target) + final foundSource = module.findSource(ctx, target) then: foundSource == source @@ -490,7 +452,7 @@ class PropagationModuleTest extends IastModuleImplTestBase { assert target.length() > maxValueLength when: - module.taint(target, SourceTypes.REQUEST_PARAMETER_VALUE) + module.taint(ctx, target, SourceTypes.REQUEST_PARAMETER_VALUE) then: final tainted = ctx.getTaintedObjects().get(target) @@ -510,7 +472,7 @@ class PropagationModuleTest extends IastModuleImplTestBase { void 'test that source names/values should not make a strong reference over the value'() { when: - module.taint(toTaint, SourceTypes.REQUEST_PARAMETER_NAME, name, value) + module.taint(ctx, toTaint, SourceTypes.REQUEST_PARAMETER_NAME, name, value) then: final tainted = ctx.getTaintedObjects().get(toTaint) @@ -557,7 +519,7 @@ class PropagationModuleTest extends IastModuleImplTestBase { final baos = toTaint.bytes when: 'tainting a non char sequence object' - module.taint(baos, SourceTypes.KAFKA_MESSAGE_KEY) + module.taint(ctx, baos, SourceTypes.KAFKA_MESSAGE_KEY) then: with(ctx.taintedObjects.get(baos)) { @@ -569,7 +531,7 @@ class PropagationModuleTest extends IastModuleImplTestBase { } when: 'the object is propagated' - module.taintIfTainted(toTaint, baos) + module.taintIfTainted(ctx, toTaint, baos) then: with(ctx.taintedObjects.get(toTaint)) { diff --git a/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/sink/AbstractSinkModuleTest.groovy b/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/sink/AbstractSinkModuleTest.groovy index 2fc3f897998..ecaed48fce3 100644 --- a/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/sink/AbstractSinkModuleTest.groovy +++ b/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/sink/AbstractSinkModuleTest.groovy @@ -58,7 +58,7 @@ class AbstractSinkModuleTest extends IastModuleImplTestBase { ctx.getTaintedObjects().taint(input, Ranges.forCharSequence(input, source)) when: - propagation.taintIfTainted(toReport, input) + propagation.taintIfTainted(ctx, toReport, input) final evidence = sink.checkInjection(SSRF, toReport) then: diff --git a/dd-java-agent/agent-iast/src/testFixtures/groovy/com/datadog/iast/test/TaintMarkerHelpers.groovy b/dd-java-agent/agent-iast/src/testFixtures/groovy/com/datadog/iast/test/TaintMarkerHelpers.groovy index 0a3d83586d7..ba5694ef19e 100644 --- a/dd-java-agent/agent-iast/src/testFixtures/groovy/com/datadog/iast/test/TaintMarkerHelpers.groovy +++ b/dd-java-agent/agent-iast/src/testFixtures/groovy/com/datadog/iast/test/TaintMarkerHelpers.groovy @@ -1,10 +1,11 @@ package com.datadog.iast.test +import datadog.trace.api.iast.IastContext import datadog.trace.api.iast.InstrumentationBridge class TaintMarkerHelpers { static t(Object o) { def propagation = InstrumentationBridge.PROPAGATION - propagation.isTainted(o) ? "$o (tainted)" : o + propagation.isTainted(IastContext.Provider.get(), o) ? "$o (tainted)" : o } } diff --git a/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/PathMatcherInstrumentation.java b/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/PathMatcherInstrumentation.java index 30175d95f3c..49e53af9813 100644 --- a/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/PathMatcherInstrumentation.java +++ b/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/PathMatcherInstrumentation.java @@ -56,24 +56,23 @@ static void onExit( return; } + scala.Tuple1 tuple = (scala.Tuple1) extractions; + Object value = tuple._1(); + PropagationModule module = InstrumentationBridge.PROPAGATION; - if (module == null) { + if (module == null || !(value instanceof String)) { return; } - - scala.Tuple1 tuple = (scala.Tuple1) extractions; - Object value = tuple._1(); + final String stringValue = (String) value; final IastContext ctx = reqCtx.getData(RequestContextSlot.IAST); // in the test, 4 instances of PathMatcher$Match are created, all with the same value - if (module.isTainted(ctx, value)) { + if (module.isTainted(ctx, stringValue)) { return; } - if (value instanceof String) { - module.taint(ctx, value, SourceTypes.REQUEST_PATH_PARAMETER); - } + module.taint(ctx, stringValue, SourceTypes.REQUEST_PATH_PARAMETER); } } } diff --git a/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/helpers/TaintSingleParameterFunction.java b/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/helpers/TaintSingleParameterFunction.java index b329facad35..3cbc98c3dbc 100644 --- a/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/helpers/TaintSingleParameterFunction.java +++ b/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/helpers/TaintSingleParameterFunction.java @@ -50,11 +50,11 @@ public Tuple1 apply(Tuple1 v1) { while (iterator.hasNext()) { Object o = iterator.next(); if (o instanceof String) { - mod.taint(ctx, o, SourceTypes.REQUEST_PARAMETER_VALUE, paramName); + mod.taint(ctx, (String) o, SourceTypes.REQUEST_PARAMETER_VALUE, paramName); } } } else if (value instanceof String) { - mod.taint(ctx, value, SourceTypes.REQUEST_PARAMETER_VALUE, paramName); + mod.taint(ctx, (String) value, SourceTypes.REQUEST_PARAMETER_VALUE, paramName); } return v1; diff --git a/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/helpers/TaintUnmarshaller.java b/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/helpers/TaintUnmarshaller.java index 8f74ef298c8..498d7b9504b 100644 --- a/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/helpers/TaintUnmarshaller.java +++ b/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/helpers/TaintUnmarshaller.java @@ -30,9 +30,11 @@ public TaintUnmarshaller(PropagationModule propagationModule, Unmarshaller @Override public Future apply(A value, ExecutionContext ec, Materializer materializer) { - IastContext ctx = IastContext.Provider.get(AgentTracer.activeSpan()); - if (ctx != null) { - propagationModule.taint(ctx, value, SourceTypes.REQUEST_BODY); + if (value != null) { + IastContext ctx = IastContext.Provider.get(AgentTracer.activeSpan()); + if (ctx != null) { + propagationModule.taint(ctx, value, SourceTypes.REQUEST_BODY); + } } return delegate.apply(value, ec, materializer); } diff --git a/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/helpers/TaintUriFunction.java b/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/helpers/TaintUriFunction.java index de1c42a979a..b2748e78b0b 100644 --- a/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/helpers/TaintUriFunction.java +++ b/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/helpers/TaintUriFunction.java @@ -17,7 +17,7 @@ public Tuple1 apply(Tuple1 v1) { Uri uri = v1._1(); PropagationModule mod = InstrumentationBridge.PROPAGATION; - if (mod == null) { + if (mod == null || uri == null) { return v1; } IastContext ctx = IastContext.Provider.get(AgentTracer.activeSpan()); diff --git a/dd-java-agent/instrumentation/akka-http-10.2-iast/src/main/java/datadog/trace/instrumentation/akkahttp102/iast/helpers/TaintParametersFunction.java b/dd-java-agent/instrumentation/akka-http-10.2-iast/src/main/java/datadog/trace/instrumentation/akkahttp102/iast/helpers/TaintParametersFunction.java index 89e328c88bd..887fc36c144 100644 --- a/dd-java-agent/instrumentation/akka-http-10.2-iast/src/main/java/datadog/trace/instrumentation/akkahttp102/iast/helpers/TaintParametersFunction.java +++ b/dd-java-agent/instrumentation/akka-http-10.2-iast/src/main/java/datadog/trace/instrumentation/akkahttp102/iast/helpers/TaintParametersFunction.java @@ -41,11 +41,11 @@ public Tuple1 apply(Tuple1 v1) { while (iterator.hasNext()) { Object o = iterator.next(); if (o instanceof String) { - mod.taint(ctx, o, SourceTypes.REQUEST_PARAMETER_VALUE, paramName); + mod.taint(ctx, (String) o, SourceTypes.REQUEST_PARAMETER_VALUE, paramName); } } } else if (value instanceof String) { - mod.taint(ctx, value, SourceTypes.REQUEST_PARAMETER_VALUE, paramName); + mod.taint(ctx, (String) value, SourceTypes.REQUEST_PARAMETER_VALUE, paramName); } return v1; diff --git a/dd-java-agent/instrumentation/apache-httpcore-4/build.gradle b/dd-java-agent/instrumentation/apache-httpcore-4/build.gradle index 096173402d1..85899d79a98 100644 --- a/dd-java-agent/instrumentation/apache-httpcore-4/build.gradle +++ b/dd-java-agent/instrumentation/apache-httpcore-4/build.gradle @@ -15,6 +15,7 @@ dependencies { compileOnly group: 'org.apache.httpcomponents', name: 'httpcore', version: '4.0' testImplementation group: 'org.apache.httpcomponents', name: 'httpcore', version: '4.0' + testImplementation(testFixtures(project(':dd-java-agent:agent-iast'))) latestDepTestImplementation group: 'org.apache.httpcomponents', name: 'httpcore', version: '+' } diff --git a/dd-java-agent/instrumentation/apache-httpcore-4/src/main/java/datadog/trace/instrumentation/apachehttpcore/IastHttpHostInstrumentation.java b/dd-java-agent/instrumentation/apache-httpcore-4/src/main/java/datadog/trace/instrumentation/apachehttpcore/IastHttpHostInstrumentation.java index 585ec803529..07b40943ef9 100644 --- a/dd-java-agent/instrumentation/apache-httpcore-4/src/main/java/datadog/trace/instrumentation/apachehttpcore/IastHttpHostInstrumentation.java +++ b/dd-java-agent/instrumentation/apache-httpcore-4/src/main/java/datadog/trace/instrumentation/apachehttpcore/IastHttpHostInstrumentation.java @@ -6,6 +6,7 @@ import com.google.auto.service.AutoService; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.api.iast.IastContext; import datadog.trace.api.iast.InstrumentationBridge; import datadog.trace.api.iast.Propagation; import datadog.trace.api.iast.propagation.PropagationModule; @@ -37,8 +38,11 @@ public static class CtorAdvice { public static void afterCtor( @Advice.This final Object self, @Advice.Argument(0) final Object argument) { final PropagationModule module = InstrumentationBridge.PROPAGATION; - if (module != null) { - module.taintIfTainted(self, argument); + if (argument != null && module != null) { + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + module.taintIfTainted(ctx, self, argument); + } } } } diff --git a/dd-java-agent/instrumentation/apache-httpcore-4/src/test/groovy/datadog/trace/instrumentation/apachehttpcore/IastHttpHostInstrumentationTest.groovy b/dd-java-agent/instrumentation/apache-httpcore-4/src/test/groovy/datadog/trace/instrumentation/apachehttpcore/IastHttpHostInstrumentationTest.groovy index 4732411e181..bab1bc605e2 100644 --- a/dd-java-agent/instrumentation/apache-httpcore-4/src/test/groovy/datadog/trace/instrumentation/apachehttpcore/IastHttpHostInstrumentationTest.groovy +++ b/dd-java-agent/instrumentation/apache-httpcore-4/src/test/groovy/datadog/trace/instrumentation/apachehttpcore/IastHttpHostInstrumentationTest.groovy @@ -1,16 +1,12 @@ package datadog.trace.instrumentation.apachehttpcore -import datadog.trace.agent.test.AgentTestRunner +import com.datadog.iast.test.IastAgentTestRunner +import datadog.trace.api.iast.IastContext import datadog.trace.api.iast.InstrumentationBridge import datadog.trace.api.iast.propagation.PropagationModule import org.apache.http.HttpHost -class IastHttpHostInstrumentationTest extends AgentTestRunner { - - @Override - protected void configurePreAgent() { - injectSysConfig('dd.iast.enabled', 'true') - } +class IastHttpHostInstrumentationTest extends IastAgentTestRunner { void 'test'(){ given: @@ -18,15 +14,15 @@ class IastHttpHostInstrumentationTest extends AgentTestRunner { InstrumentationBridge.registerIastModule(module) when: - HttpHost.newInstance(*args) + runUnderIastTrace { HttpHost.newInstance(*args) } then: - 1 * module.taintIfTainted( _ as HttpHost, 'localhost') + 1 * module.taintIfTainted(_ as IastContext, _ as HttpHost, 'localhost') where: - args | _ - ['localhost'] | _ - ['localhost', 8080] | _ + args | _ + ['localhost'] | _ + ['localhost', 8080] | _ ['localhost', 8080, 'http'] | _ } } diff --git a/dd-java-agent/instrumentation/commons-httpclient-2/src/main/java/datadog/trace/instrumentation/commonshttpclient/IastHttpMethodBaseInstrumentation.java b/dd-java-agent/instrumentation/commons-httpclient-2/src/main/java/datadog/trace/instrumentation/commonshttpclient/IastHttpMethodBaseInstrumentation.java index 6cd3f18c101..e9795decdca 100644 --- a/dd-java-agent/instrumentation/commons-httpclient-2/src/main/java/datadog/trace/instrumentation/commonshttpclient/IastHttpMethodBaseInstrumentation.java +++ b/dd-java-agent/instrumentation/commons-httpclient-2/src/main/java/datadog/trace/instrumentation/commonshttpclient/IastHttpMethodBaseInstrumentation.java @@ -8,6 +8,7 @@ import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.agent.tooling.bytebuddy.iast.TaintableVisitor; +import datadog.trace.api.iast.IastContext; import datadog.trace.api.iast.InstrumentationBridge; import datadog.trace.api.iast.Propagation; import datadog.trace.api.iast.propagation.PropagationModule; @@ -51,8 +52,12 @@ public static class CtorAdvice { public static void afterCtor( @Advice.This final Object self, @Advice.Argument(0) final Object argument) { final PropagationModule module = InstrumentationBridge.PROPAGATION; - if (module != null) { - module.taintIfTainted(self, argument); + if (argument != null && module != null) { + final IastContext ctx = IastContext.Provider.get(); + if (ctx == null) { + return; + } + module.taintIfTainted(ctx, self, argument); } } } diff --git a/dd-java-agent/instrumentation/commons-lang-2/build.gradle b/dd-java-agent/instrumentation/commons-lang-2/build.gradle index 9257b3b6f96..f55f8130c4e 100644 --- a/dd-java-agent/instrumentation/commons-lang-2/build.gradle +++ b/dd-java-agent/instrumentation/commons-lang-2/build.gradle @@ -16,6 +16,7 @@ dependencies { compileOnly group: 'commons-lang', name: 'commons-lang', version: '2.1' testImplementation group: 'commons-lang', name: 'commons-lang', version: '2.1' + testImplementation(testFixtures(project(':dd-java-agent:agent-iast'))) testRuntimeOnly project(':dd-java-agent:instrumentation:iast-instrumenter') diff --git a/dd-java-agent/instrumentation/commons-lang-2/src/main/java/datadog/trace/instrumentation/commonslang/StringEscapeUtilsCallSite.java b/dd-java-agent/instrumentation/commons-lang-2/src/main/java/datadog/trace/instrumentation/commonslang/StringEscapeUtilsCallSite.java index d7fdd84417d..8b8d2ec746d 100644 --- a/dd-java-agent/instrumentation/commons-lang-2/src/main/java/datadog/trace/instrumentation/commonslang/StringEscapeUtilsCallSite.java +++ b/dd-java-agent/instrumentation/commons-lang-2/src/main/java/datadog/trace/instrumentation/commonslang/StringEscapeUtilsCallSite.java @@ -2,6 +2,7 @@ import datadog.trace.agent.tooling.csi.CallSite; import datadog.trace.api.iast.IastCallSites; +import datadog.trace.api.iast.IastContext; import datadog.trace.api.iast.InstrumentationBridge; import datadog.trace.api.iast.Propagation; import datadog.trace.api.iast.VulnerabilityMarks; @@ -23,9 +24,12 @@ public class StringEscapeUtilsCallSite { public static String afterEscape( @CallSite.Argument(0) @Nullable final String input, @CallSite.Return final String result) { final PropagationModule module = InstrumentationBridge.PROPAGATION; - if (module != null) { + if (result != null && module != null) { try { - module.taintIfTainted(result, input, false, VulnerabilityMarks.XSS_MARK); + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + module.taintIfTainted(ctx, result, input, false, VulnerabilityMarks.XSS_MARK); + } } catch (final Throwable e) { module.onUnexpectedException("afterEscape threw", e); } @@ -38,9 +42,12 @@ public static String afterEscape( public static String afterEscapeSQL( @CallSite.Argument(0) @Nullable final String input, @CallSite.Return final String result) { final PropagationModule module = InstrumentationBridge.PROPAGATION; - if (module != null) { + if (result != null && module != null) { try { - module.taintIfTainted(result, input, false, VulnerabilityMarks.SQL_INJECTION_MARK); + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + module.taintIfTainted(ctx, result, input, false, VulnerabilityMarks.SQL_INJECTION_MARK); + } } catch (final Throwable e) { module.onUnexpectedException("afterEscapeSQL threw", e); } diff --git a/dd-java-agent/instrumentation/commons-lang-2/src/test/groovy/datadog/trace/instrumentation/commonslang/StringEscapeUtilsCallSiteTest.groovy b/dd-java-agent/instrumentation/commons-lang-2/src/test/groovy/datadog/trace/instrumentation/commonslang/StringEscapeUtilsCallSiteTest.groovy index 348aaa57503..623d861168e 100644 --- a/dd-java-agent/instrumentation/commons-lang-2/src/test/groovy/datadog/trace/instrumentation/commonslang/StringEscapeUtilsCallSiteTest.groovy +++ b/dd-java-agent/instrumentation/commons-lang-2/src/test/groovy/datadog/trace/instrumentation/commonslang/StringEscapeUtilsCallSiteTest.groovy @@ -1,21 +1,15 @@ package datadog.trace.instrumentation.commonslang -import datadog.trace.agent.test.AgentTestRunner +import com.datadog.iast.test.IastAgentTestRunner +import datadog.trace.api.iast.IastContext import datadog.trace.api.iast.InstrumentationBridge import datadog.trace.api.iast.propagation.PropagationModule import foo.bar.TestStringEscapeUtilsSuite -import groovy.transform.CompileDynamic import static datadog.trace.api.iast.VulnerabilityMarks.SQL_INJECTION_MARK import static datadog.trace.api.iast.VulnerabilityMarks.XSS_MARK -@CompileDynamic -class StringEscapeUtilsCallSiteTest extends AgentTestRunner { - - @Override - protected void configurePreAgent() { - injectSysConfig("dd.iast.enabled", "true") - } +class StringEscapeUtilsCallSiteTest extends IastAgentTestRunner { void 'test #method'() { given: @@ -23,11 +17,14 @@ class StringEscapeUtilsCallSiteTest extends AgentTestRunner { InstrumentationBridge.registerIastModule(module) when: - final result = TestStringEscapeUtilsSuite.&"$method".call(args) + def result = null + runUnderIastTrace { + result = TestStringEscapeUtilsSuite.&"$method".call(args) + } then: result == expected - 1 * module.taintIfTainted(_ as String, args[0], false, mark) + 1 * module.taintIfTainted(_ as IastContext, _ as String, args[0], false, mark) 0 * _ where: diff --git a/dd-java-agent/instrumentation/commons-lang-3/build.gradle b/dd-java-agent/instrumentation/commons-lang-3/build.gradle index e6637a18314..a3d55234615 100644 --- a/dd-java-agent/instrumentation/commons-lang-3/build.gradle +++ b/dd-java-agent/instrumentation/commons-lang-3/build.gradle @@ -16,6 +16,7 @@ dependencies { compileOnly group: 'org.apache.commons', name: 'commons-lang3', version: '3.5' testImplementation group: 'org.apache.commons', name: 'commons-lang3', version: '3.5' + testImplementation(testFixtures(project(':dd-java-agent:agent-iast'))) testRuntimeOnly project(':dd-java-agent:instrumentation:iast-instrumenter') diff --git a/dd-java-agent/instrumentation/commons-lang-3/src/main/java/datadog/trace/instrumentation/commonslang3/StringEscapeUtilsCallSite.java b/dd-java-agent/instrumentation/commons-lang-3/src/main/java/datadog/trace/instrumentation/commonslang3/StringEscapeUtilsCallSite.java index b52f57c57e9..1cd10021c92 100644 --- a/dd-java-agent/instrumentation/commons-lang-3/src/main/java/datadog/trace/instrumentation/commonslang3/StringEscapeUtilsCallSite.java +++ b/dd-java-agent/instrumentation/commons-lang-3/src/main/java/datadog/trace/instrumentation/commonslang3/StringEscapeUtilsCallSite.java @@ -2,6 +2,7 @@ import datadog.trace.agent.tooling.csi.CallSite; import datadog.trace.api.iast.IastCallSites; +import datadog.trace.api.iast.IastContext; import datadog.trace.api.iast.InstrumentationBridge; import datadog.trace.api.iast.Propagation; import datadog.trace.api.iast.VulnerabilityMarks; @@ -25,9 +26,12 @@ public class StringEscapeUtilsCallSite { public static String afterEscape( @CallSite.Argument(0) @Nullable final String input, @CallSite.Return final String result) { final PropagationModule module = InstrumentationBridge.PROPAGATION; - if (module != null) { + if (result != null && module != null) { try { - module.taintIfTainted(result, input, false, VulnerabilityMarks.XSS_MARK); + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + module.taintIfTainted(ctx, result, input, false, VulnerabilityMarks.XSS_MARK); + } } catch (final Throwable e) { module.onUnexpectedException("afterEscape threw", e); } @@ -40,9 +44,12 @@ public static String afterEscape( public static String afterEscapeJson( @CallSite.Argument(0) @Nullable final String input, @CallSite.Return final String result) { final PropagationModule module = InstrumentationBridge.PROPAGATION; - if (module != null) { + if (result != null && module != null) { try { - module.taintIfTainted(result, input); + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + module.taintIfTainted(ctx, result, input); + } } catch (final Throwable e) { module.onUnexpectedException("afterEscapeJson threw", e); } diff --git a/dd-java-agent/instrumentation/commons-lang-3/src/test/groovy/datadog/trace/instrumentation/commonslang3/StringEscapeUtilsCallSiteTest.groovy b/dd-java-agent/instrumentation/commons-lang-3/src/test/groovy/datadog/trace/instrumentation/commonslang3/StringEscapeUtilsCallSiteTest.groovy index 18730b70612..9898b6cfaa0 100644 --- a/dd-java-agent/instrumentation/commons-lang-3/src/test/groovy/datadog/trace/instrumentation/commonslang3/StringEscapeUtilsCallSiteTest.groovy +++ b/dd-java-agent/instrumentation/commons-lang-3/src/test/groovy/datadog/trace/instrumentation/commonslang3/StringEscapeUtilsCallSiteTest.groovy @@ -1,19 +1,13 @@ package datadog.trace.instrumentation.commonslang3 -import datadog.trace.agent.test.AgentTestRunner +import com.datadog.iast.test.IastAgentTestRunner +import datadog.trace.api.iast.IastContext import datadog.trace.api.iast.InstrumentationBridge import datadog.trace.api.iast.VulnerabilityMarks import datadog.trace.api.iast.propagation.PropagationModule import foo.bar.TestStringEscapeUtilsSuite -import groovy.transform.CompileDynamic -@CompileDynamic -class StringEscapeUtilsCallSiteTest extends AgentTestRunner { - - @Override - protected void configurePreAgent() { - injectSysConfig("dd.iast.enabled", "true") - } +class StringEscapeUtilsCallSiteTest extends IastAgentTestRunner { void 'test #method with XSS mark'() { given: @@ -21,11 +15,14 @@ class StringEscapeUtilsCallSiteTest extends AgentTestRunner { InstrumentationBridge.registerIastModule(module) when: - final result = TestStringEscapeUtilsSuite.&"$method".call(args) + def result = null + runUnderIastTrace { + result = TestStringEscapeUtilsSuite.&"$method".call(args) + } then: result == expected - 1 * module.taintIfTainted(_ as String, args[0], false, VulnerabilityMarks.XSS_MARK) + 1 * module.taintIfTainted(_ as IastContext, _ as String, args[0], false, VulnerabilityMarks.XSS_MARK) 0 * _ where: @@ -43,11 +40,14 @@ class StringEscapeUtilsCallSiteTest extends AgentTestRunner { InstrumentationBridge.registerIastModule(module) when: - final result = TestStringEscapeUtilsSuite.&"$method".call(args) + def result = null + runUnderIastTrace { + result = TestStringEscapeUtilsSuite.&"$method".call(args) + } then: result == expected - 1 * module.taintIfTainted(_ as String, args[0]) + 1 * module.taintIfTainted(_ as IastContext, _ as String, args[0]) 0 * _ where: diff --git a/dd-java-agent/instrumentation/commons-text/build.gradle b/dd-java-agent/instrumentation/commons-text/build.gradle index 2f7b91eb5ed..21d0dc02d12 100644 --- a/dd-java-agent/instrumentation/commons-text/build.gradle +++ b/dd-java-agent/instrumentation/commons-text/build.gradle @@ -16,6 +16,7 @@ dependencies { compileOnly group: 'org.apache.commons', name: 'commons-text', version: '1.0' testImplementation group: 'org.apache.commons', name: 'commons-text', version: '1.0' + testImplementation(testFixtures(project(':dd-java-agent:agent-iast'))) testRuntimeOnly project(':dd-java-agent:instrumentation:iast-instrumenter') diff --git a/dd-java-agent/instrumentation/commons-text/src/main/java/datadog/trace/instrumentation/commonstext/StringEscapeUtilsCallSite.java b/dd-java-agent/instrumentation/commons-text/src/main/java/datadog/trace/instrumentation/commonstext/StringEscapeUtilsCallSite.java index ebd592fa241..7d7aaf4f4ff 100644 --- a/dd-java-agent/instrumentation/commons-text/src/main/java/datadog/trace/instrumentation/commonstext/StringEscapeUtilsCallSite.java +++ b/dd-java-agent/instrumentation/commons-text/src/main/java/datadog/trace/instrumentation/commonstext/StringEscapeUtilsCallSite.java @@ -2,6 +2,7 @@ import datadog.trace.agent.tooling.csi.CallSite; import datadog.trace.api.iast.IastCallSites; +import datadog.trace.api.iast.IastContext; import datadog.trace.api.iast.InstrumentationBridge; import datadog.trace.api.iast.Propagation; import datadog.trace.api.iast.VulnerabilityMarks; @@ -27,9 +28,12 @@ public class StringEscapeUtilsCallSite { public static String afterEscape( @CallSite.Argument(0) @Nullable final String input, @CallSite.Return final String result) { final PropagationModule module = InstrumentationBridge.PROPAGATION; - if (module != null) { + if (result != null && module != null) { try { - module.taintIfTainted(result, input, false, VulnerabilityMarks.XSS_MARK); + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + module.taintIfTainted(ctx, result, input, false, VulnerabilityMarks.XSS_MARK); + } } catch (final Throwable e) { module.onUnexpectedException("afterEscape threw", e); } @@ -42,9 +46,12 @@ public static String afterEscape( public static String afterEscapeJson( @CallSite.Argument(0) @Nullable final String input, @CallSite.Return final String result) { final PropagationModule module = InstrumentationBridge.PROPAGATION; - if (module != null) { + if (result != null && module != null) { try { - module.taintIfTainted(result, input); + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + module.taintIfTainted(ctx, result, input); + } } catch (final Throwable e) { module.onUnexpectedException("afterEscapeJson threw", e); } diff --git a/dd-java-agent/instrumentation/commons-text/src/test/groovy/datadog/trace/instrumentation/commonstext/StringEscapeUtilsCallSiteTest.groovy b/dd-java-agent/instrumentation/commons-text/src/test/groovy/datadog/trace/instrumentation/commonstext/StringEscapeUtilsCallSiteTest.groovy index 440e31ea3be..c75562f3d5b 100644 --- a/dd-java-agent/instrumentation/commons-text/src/test/groovy/datadog/trace/instrumentation/commonstext/StringEscapeUtilsCallSiteTest.groovy +++ b/dd-java-agent/instrumentation/commons-text/src/test/groovy/datadog/trace/instrumentation/commonstext/StringEscapeUtilsCallSiteTest.groovy @@ -1,19 +1,13 @@ package datadog.trace.instrumentation.commonstext -import datadog.trace.agent.test.AgentTestRunner +import com.datadog.iast.test.IastAgentTestRunner +import datadog.trace.api.iast.IastContext import datadog.trace.api.iast.InstrumentationBridge import datadog.trace.api.iast.VulnerabilityMarks import datadog.trace.api.iast.propagation.PropagationModule import foo.bar.TestStringEscapeUtilsSuite -import groovy.transform.CompileDynamic -@CompileDynamic -class StringEscapeUtilsCallSiteTest extends AgentTestRunner { - - @Override - protected void configurePreAgent() { - injectSysConfig("dd.iast.enabled", "true") - } +class StringEscapeUtilsCallSiteTest extends IastAgentTestRunner { void 'test #method with XSS mark'() { given: @@ -21,11 +15,14 @@ class StringEscapeUtilsCallSiteTest extends AgentTestRunner { InstrumentationBridge.registerIastModule(module) when: - final result = TestStringEscapeUtilsSuite.&"$method".call(args) + def result = null + runUnderIastTrace { + result = TestStringEscapeUtilsSuite.&"$method".call(args) + } then: result == expected - 1 * module.taintIfTainted(_ as String, args[0], false, VulnerabilityMarks.XSS_MARK) + 1 * module.taintIfTainted(_ as IastContext, _ as String, args[0], false, VulnerabilityMarks.XSS_MARK) 0 * _ where: @@ -44,11 +41,14 @@ class StringEscapeUtilsCallSiteTest extends AgentTestRunner { InstrumentationBridge.registerIastModule(module) when: - final result = TestStringEscapeUtilsSuite.&"$method".call(args) + def result = null + runUnderIastTrace { + result = TestStringEscapeUtilsSuite.&"$method".call(args) + } then: result == expected - 1 * module.taintIfTainted(_ as String, args[0]) + 1 * module.taintIfTainted(_ as IastContext, _ as String, args[0]) 0 * _ where: diff --git a/dd-java-agent/instrumentation/freemarker/build.gradle b/dd-java-agent/instrumentation/freemarker/build.gradle index f2d5ef612aa..4417d23102f 100644 --- a/dd-java-agent/instrumentation/freemarker/build.gradle +++ b/dd-java-agent/instrumentation/freemarker/build.gradle @@ -16,6 +16,7 @@ dependencies { compileOnly group: 'org.freemarker', name: 'freemarker', version: '2.3.32' testImplementation group: 'org.freemarker', name: 'freemarker', version: '2.3.32' + testImplementation(testFixtures(project(':dd-java-agent:agent-iast'))) testRuntimeOnly project(':dd-java-agent:instrumentation:iast-instrumenter') diff --git a/dd-java-agent/instrumentation/freemarker/src/main/java/datadog/trace/instrumentation/freemarker/StringUtilCallSite.java b/dd-java-agent/instrumentation/freemarker/src/main/java/datadog/trace/instrumentation/freemarker/StringUtilCallSite.java index 391339f528b..5f20718519b 100644 --- a/dd-java-agent/instrumentation/freemarker/src/main/java/datadog/trace/instrumentation/freemarker/StringUtilCallSite.java +++ b/dd-java-agent/instrumentation/freemarker/src/main/java/datadog/trace/instrumentation/freemarker/StringUtilCallSite.java @@ -2,6 +2,7 @@ import datadog.trace.agent.tooling.csi.CallSite; import datadog.trace.api.iast.IastCallSites; +import datadog.trace.api.iast.IastContext; import datadog.trace.api.iast.InstrumentationBridge; import datadog.trace.api.iast.Propagation; import datadog.trace.api.iast.VulnerabilityMarks; @@ -27,9 +28,12 @@ public class StringUtilCallSite { public static String afterEscape( @CallSite.Argument(0) @Nullable final String input, @CallSite.Return final String result) { final PropagationModule module = InstrumentationBridge.PROPAGATION; - if (module != null) { + if (result != null && module != null) { try { - module.taintIfTainted(result, input, false, VulnerabilityMarks.XSS_MARK); + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + module.taintIfTainted(ctx, result, input, false, VulnerabilityMarks.XSS_MARK); + } } catch (final Throwable e) { module.onUnexpectedException("afterEscape threw", e); } diff --git a/dd-java-agent/instrumentation/freemarker/src/test/groovy/datadog/trace/instrumentation/freemarker/StringUtilCallSiteTest.groovy b/dd-java-agent/instrumentation/freemarker/src/test/groovy/datadog/trace/instrumentation/freemarker/StringUtilCallSiteTest.groovy index 18d86d6aef5..02d4220fdc1 100644 --- a/dd-java-agent/instrumentation/freemarker/src/test/groovy/datadog/trace/instrumentation/freemarker/StringUtilCallSiteTest.groovy +++ b/dd-java-agent/instrumentation/freemarker/src/test/groovy/datadog/trace/instrumentation/freemarker/StringUtilCallSiteTest.groovy @@ -1,17 +1,13 @@ package datadog.trace.instrumentation.freemarker -import datadog.trace.agent.test.AgentTestRunner +import com.datadog.iast.test.IastAgentTestRunner +import datadog.trace.api.iast.IastContext import datadog.trace.api.iast.InstrumentationBridge import datadog.trace.api.iast.VulnerabilityMarks import datadog.trace.api.iast.propagation.PropagationModule import foo.bar.TestStringUtilSuite -class StringUtilCallSiteTest extends AgentTestRunner { - - @Override - protected void configurePreAgent() { - injectSysConfig("dd.iast.enabled", "true") - } +class StringUtilCallSiteTest extends IastAgentTestRunner { void 'test #method'() { given: @@ -19,11 +15,14 @@ class StringUtilCallSiteTest extends AgentTestRunner { InstrumentationBridge.registerIastModule(module) when: - final result = TestStringUtilSuite.&"$method".call(args) + def result = null + runUnderIastTrace { + result = TestStringUtilSuite.&"$method".call(args) + } then: result == expected - 1 * module.taintIfTainted(_ as String, args[0], false, VulnerabilityMarks.XSS_MARK) + 1 * module.taintIfTainted(_ as IastContext, _ as String, args[0], false, VulnerabilityMarks.XSS_MARK) 0 * _ where: @@ -42,10 +41,10 @@ class StringUtilCallSiteTest extends AgentTestRunner { InstrumentationBridge.registerIastModule(module) when: - TestStringUtilSuite.&"$method".call(null) + runUnderIastTrace { TestStringUtilSuite.&"$method".call(null) } then: - def thrownException = thrown (Exception) + def thrownException = thrown(Exception) assert thrownException.stackTrace[0].getClassName().startsWith('freemarker') 0 * _ diff --git a/dd-java-agent/instrumentation/gson-1.6/build.gradle b/dd-java-agent/instrumentation/gson-1.6/build.gradle index 953d6dec9b9..95f0e3c495f 100644 --- a/dd-java-agent/instrumentation/gson-1.6/build.gradle +++ b/dd-java-agent/instrumentation/gson-1.6/build.gradle @@ -16,6 +16,7 @@ dependencies { compileOnly group: 'com.google.code.gson', name: 'gson', version: '1.6' testImplementation group: 'com.google.code.gson', name: 'gson', version: '1.6' + testImplementation(testFixtures(project(':dd-java-agent:agent-iast'))) testRuntimeOnly project(':dd-java-agent:instrumentation:iast-instrumenter') diff --git a/dd-java-agent/instrumentation/gson-1.6/src/main/java/datadog/trace/instrumentation/gson/JsonReaderInstrumentation.java b/dd-java-agent/instrumentation/gson-1.6/src/main/java/datadog/trace/instrumentation/gson/JsonReaderInstrumentation.java index e47b677d19e..8846378f151 100644 --- a/dd-java-agent/instrumentation/gson-1.6/src/main/java/datadog/trace/instrumentation/gson/JsonReaderInstrumentation.java +++ b/dd-java-agent/instrumentation/gson-1.6/src/main/java/datadog/trace/instrumentation/gson/JsonReaderInstrumentation.java @@ -14,6 +14,7 @@ import com.google.auto.service.AutoService; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.api.iast.IastContext; import datadog.trace.api.iast.InstrumentationBridge; import datadog.trace.api.iast.Propagation; import datadog.trace.api.iast.propagation.PropagationModule; @@ -59,7 +60,10 @@ public static void afterInit( @Advice.This Object self, @Advice.Argument(0) final java.io.Reader input) { final PropagationModule iastModule = InstrumentationBridge.PROPAGATION; if (iastModule != null && input != null) { - iastModule.taintIfTainted(self, input); + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + iastModule.taintIfTainted(ctx, self, input); + } } } } @@ -70,7 +74,10 @@ public static class MethodAdvice { public static void afterMethod(@Advice.This Object self, @Advice.Return final String result) { final PropagationModule iastModule = InstrumentationBridge.PROPAGATION; if (iastModule != null && result != null) { - iastModule.taintIfTainted(result, self); + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + iastModule.taintIfTainted(ctx, result, self); + } } } } diff --git a/dd-java-agent/instrumentation/gson-1.6/src/test/groovy/datadog/trace/instrumentation/gson/JsonReaderInstrumentationTest.groovy b/dd-java-agent/instrumentation/gson-1.6/src/test/groovy/datadog/trace/instrumentation/gson/JsonReaderInstrumentationTest.groovy index 58e3621e92e..d3f6d001979 100644 --- a/dd-java-agent/instrumentation/gson-1.6/src/test/groovy/datadog/trace/instrumentation/gson/JsonReaderInstrumentationTest.groovy +++ b/dd-java-agent/instrumentation/gson-1.6/src/test/groovy/datadog/trace/instrumentation/gson/JsonReaderInstrumentationTest.groovy @@ -1,17 +1,13 @@ package datadog.trace.instrumentation.gson +import com.datadog.iast.test.IastAgentTestRunner import com.google.gson.Gson import com.google.gson.stream.JsonReader -import datadog.trace.agent.test.AgentTestRunner +import datadog.trace.api.iast.IastContext import datadog.trace.api.iast.InstrumentationBridge import datadog.trace.api.iast.propagation.PropagationModule -class JsonReaderInstrumentationTest extends AgentTestRunner { - - @Override - protected void configurePreAgent() { - injectSysConfig("dd.iast.enabled", "true") - } +class JsonReaderInstrumentationTest extends IastAgentTestRunner { void 'test'() { given: @@ -20,16 +16,19 @@ class JsonReaderInstrumentationTest extends AgentTestRunner { final gson = new Gson() when: - final reader = new JsonReader(new StringReader(json)) + JsonReader reader = null + runUnderIastTrace { + reader = new JsonReader(new StringReader(json)) + } then: - 1 * module.taintIfTainted(_ as JsonReader, _ as StringReader) + 1 * module.taintIfTainted(_ as IastContext, _ as JsonReader, _ as StringReader) when: - gson.fromJson(reader, clazz) + runUnderIastTrace { gson.fromJson(reader, clazz) } then: - calls * module.taintIfTainted(_ as String, _ as JsonReader) + calls * module.taintIfTainted(_ as IastContext, _ as String, _ as JsonReader) 0 * _ where: diff --git a/dd-java-agent/instrumentation/jackson-core/build.gradle b/dd-java-agent/instrumentation/jackson-core/build.gradle index d542b451c3c..9f2b775e0b5 100644 --- a/dd-java-agent/instrumentation/jackson-core/build.gradle +++ b/dd-java-agent/instrumentation/jackson-core/build.gradle @@ -26,4 +26,5 @@ dependencies { testImplementation(group: 'com.fasterxml.jackson.core', name: 'jackson-core', version: jacksonVersion) testImplementation(group: 'com.fasterxml.jackson.core', name: 'jackson-databind', version: jacksonVersion) + testImplementation(testFixtures(project(':dd-java-agent:agent-iast'))) } diff --git a/dd-java-agent/instrumentation/jackson-core/jackson-core-1/build.gradle b/dd-java-agent/instrumentation/jackson-core/jackson-core-1/build.gradle index cd193f4c5cd..d8edeee7e17 100644 --- a/dd-java-agent/instrumentation/jackson-core/jackson-core-1/build.gradle +++ b/dd-java-agent/instrumentation/jackson-core/jackson-core-1/build.gradle @@ -16,4 +16,5 @@ dependencies { compileOnly(group: 'org.codehaus.jackson', name: 'jackson-core-asl', version: jacksonVersion) testImplementation(group: 'org.codehaus.jackson', name: 'jackson-mapper-asl', version: jacksonVersion) + testImplementation(testFixtures(project(':dd-java-agent:agent-iast'))) } diff --git a/dd-java-agent/instrumentation/jackson-core/jackson-core-1/src/main/java/datadog/trace/instrumentation/jackson/codehouse/core/Json1FactoryInstrumentation.java b/dd-java-agent/instrumentation/jackson-core/jackson-core-1/src/main/java/datadog/trace/instrumentation/jackson/codehouse/core/Json1FactoryInstrumentation.java index 1c16bb4b883..fd2a76aa5a9 100644 --- a/dd-java-agent/instrumentation/jackson-core/jackson-core-1/src/main/java/datadog/trace/instrumentation/jackson/codehouse/core/Json1FactoryInstrumentation.java +++ b/dd-java-agent/instrumentation/jackson-core/jackson-core-1/src/main/java/datadog/trace/instrumentation/jackson/codehouse/core/Json1FactoryInstrumentation.java @@ -9,6 +9,7 @@ import com.google.auto.service.AutoService; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.api.iast.IastContext; import datadog.trace.api.iast.InstrumentationBridge; import datadog.trace.api.iast.Propagation; import datadog.trace.api.iast.Sink; @@ -58,10 +59,13 @@ public static class InstrumenterAdvice { @Sink(VulnerabilityTypes.SSRF) // SSRF takes priority over the propagation one public static void onExit( @Advice.Argument(0) final Object input, @Advice.Return final Object parser) { - if (input != null) { + if (parser != null) { final PropagationModule propagation = InstrumentationBridge.PROPAGATION; if (propagation != null) { - propagation.taintIfTainted(parser, input); + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + propagation.taintIfTainted(ctx, parser, input); + } } if (input instanceof URL) { final SsrfModule ssrf = InstrumentationBridge.SSRF; @@ -82,10 +86,13 @@ public static void onExit( @Advice.Argument(1) final int offset, @Advice.Argument(2) final int length, @Advice.Return final Object parser) { - if (input != null || length <= 0) { + if (parser != null) { final PropagationModule propagation = InstrumentationBridge.PROPAGATION; if (propagation != null) { - propagation.taintIfTainted(parser, input, offset, length, false, NOT_MARKED); + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + propagation.taintIfTainted(ctx, parser, input, offset, length, false, NOT_MARKED); + } } } } diff --git a/dd-java-agent/instrumentation/jackson-core/jackson-core-1/src/test/groovy/Json1FactoryInstrumentationTest.groovy b/dd-java-agent/instrumentation/jackson-core/jackson-core-1/src/test/groovy/Json1FactoryInstrumentationTest.groovy index 839d12baf30..bdff2259d4c 100644 --- a/dd-java-agent/instrumentation/jackson-core/jackson-core-1/src/test/groovy/Json1FactoryInstrumentationTest.groovy +++ b/dd-java-agent/instrumentation/jackson-core/jackson-core-1/src/test/groovy/Json1FactoryInstrumentationTest.groovy @@ -1,5 +1,6 @@ -import datadog.trace.agent.test.AgentTestRunner +import com.datadog.iast.test.IastAgentTestRunner import datadog.trace.agent.test.server.http.TestHttpServer +import datadog.trace.api.iast.IastContext import datadog.trace.api.iast.InstrumentationBridge import datadog.trace.api.iast.VulnerabilityMarks import datadog.trace.api.iast.propagation.PropagationModule @@ -11,7 +12,7 @@ import spock.lang.Shared import static datadog.trace.agent.test.server.http.TestHttpServer.httpServer -class Json1FactoryInstrumentationTest extends AgentTestRunner { +class Json1FactoryInstrumentationTest extends IastAgentTestRunner { @Shared @AutoCleanup @@ -25,11 +26,6 @@ class Json1FactoryInstrumentationTest extends AgentTestRunner { } } - @Override - protected void configurePreAgent() { - injectSysConfig("dd.iast.enabled", "true") - } - void 'test createParser(String)'() { setup: final propagationModule = Mock(PropagationModule) @@ -37,11 +33,12 @@ class Json1FactoryInstrumentationTest extends AgentTestRunner { final content = '{"key":"value"}' when: - final result = new JsonFactory().createJsonParser(content) + JsonParser result = null + runUnderIastTrace { result = new JsonFactory().createJsonParser(content) } then: result != null - 1 * propagationModule.taintIfTainted(_ as JsonParser, content) + 1 * propagationModule.taintIfTainted(_ as IastContext, _ as JsonParser, content) 0 * _ } @@ -52,11 +49,12 @@ class Json1FactoryInstrumentationTest extends AgentTestRunner { final is = Mock(InputStream) when: - final result = new JsonFactory().createJsonParser(is) + JsonParser result = null + runUnderIastTrace { result = new JsonFactory().createJsonParser(is) } then: result != null - 1 * propagationModule.taintIfTainted(_ as JsonParser, is) + 1 * propagationModule.taintIfTainted(_ as IastContext, _ as JsonParser, is) 2 * is.read(_, _, _) 0 * _ } @@ -68,12 +66,13 @@ class Json1FactoryInstrumentationTest extends AgentTestRunner { final reader = Mock(Reader) when: - final result = new JsonFactory().createJsonParser(reader) + JsonParser result = null + runUnderIastTrace { new JsonFactory().createJsonParser(reader) } then: result != null - 1 * propagationModule.taintIfTainted(_ as JsonParser, reader) + 1 * propagationModule.taintIfTainted(_ as IastContext, _ as JsonParser, reader) 0 * _ } @@ -85,11 +84,12 @@ class Json1FactoryInstrumentationTest extends AgentTestRunner { final url = new URL("${clientServer.address}/json") when: - final parser = new JsonFactory().createJsonParser(url) + JsonParser parser = null + runUnderIastTrace { result = new JsonFactory().createJsonParser(url) } then: parser != null - 1 * propagationModule.taintIfTainted(_ as JsonParser, url) + 1 * propagationModule.taintIfTainted(_ as IastContext, _ as JsonParser, url) 1 * ssrfModule.onURLConnection(url) 0 * _ } @@ -102,12 +102,13 @@ class Json1FactoryInstrumentationTest extends AgentTestRunner { final bytes = '{}'.bytes when: - final result = new JsonFactory().createJsonParser(bytes) + JsonParser result = null + runUnderIastTrace { result = new JsonFactory().createJsonParser(bytes) } then: result != null - 1 * propagationModule.taintIfTainted(_ as JsonParser, bytes) + 1 * propagationModule.taintIfTainted(_ as IastContext, _ as JsonParser, bytes) 0 * _ } @@ -118,11 +119,12 @@ class Json1FactoryInstrumentationTest extends AgentTestRunner { final bytes = '{}'.bytes when: - final parser = new JsonFactory().createJsonParser(bytes, 0, 2) + JsonParser parser = null + runUnderIastTrace { parser = new JsonFactory().createJsonParser(bytes, 0, 2) } then: parser != null - 1 * propagationModule.taintIfTainted(_ as JsonParser, bytes, 0, 2, false, VulnerabilityMarks.NOT_MARKED) + 1 * propagationModule.taintIfTainted(_ as IastContext, _ as JsonParser, bytes, 0, 2, false, VulnerabilityMarks.NOT_MARKED) 0 * _ } } diff --git a/dd-java-agent/instrumentation/jackson-core/jackson-core-1/src/test/groovy/Json1ParserInstrumentationTest.groovy b/dd-java-agent/instrumentation/jackson-core/jackson-core-1/src/test/groovy/Json1ParserInstrumentationTest.groovy index d2d92015dce..caf2251c988 100644 --- a/dd-java-agent/instrumentation/jackson-core/jackson-core-1/src/test/groovy/Json1ParserInstrumentationTest.groovy +++ b/dd-java-agent/instrumentation/jackson-core/jackson-core-1/src/test/groovy/Json1ParserInstrumentationTest.groovy @@ -1,22 +1,17 @@ -import datadog.trace.agent.test.AgentTestRunner +import com.datadog.iast.test.IastAgentTestRunner +import datadog.trace.api.iast.IastContext import datadog.trace.api.iast.InstrumentationBridge import datadog.trace.api.iast.SourceTypes import datadog.trace.api.iast.Taintable.Source import datadog.trace.api.iast.propagation.PropagationModule import groovy.json.JsonOutput +import org.codehaus.jackson.JsonParser import org.codehaus.jackson.map.ObjectMapper -import java.nio.charset.Charset - -class Json1ParserInstrumentationTest extends AgentTestRunner { +class Json1ParserInstrumentationTest extends IastAgentTestRunner { private final static String JSON_STRING = '{"root":"root_value","nested":{"nested_array":["array_0","array_1"]}}' - @Override - protected void configurePreAgent() { - injectSysConfig("dd.iast.enabled", "true") - } - void 'test json parsing (tainted)'() { given: final source = new SourceImpl(origin: SourceTypes.REQUEST_BODY, name: 'body', value: JSON_STRING) @@ -27,18 +22,21 @@ class Json1ParserInstrumentationTest extends AgentTestRunner { final reader = new ObjectMapper().reader(Map) when: - final taintedResult = reader.readValue(target) as Map + Map taintedResult = null + runUnderIastTrace { + taintedResult = reader.readValue(target) as Map + } then: JsonOutput.toJson(taintedResult) == JSON_STRING - _ * module.taintIfTainted(_, _) - _ * module.findSource(_) >> source - 1 * module.taint(_, 'root', source.origin, 'root', JSON_STRING) - 1 * module.taint(_, 'root_value', source.origin, 'root', JSON_STRING) - 1 * module.taint(_, 'nested', source.origin, 'nested', JSON_STRING) - 1 * module.taint(_, 'nested_array', source.origin, 'nested_array', JSON_STRING) - 1 * module.taint(_, 'array_0', source.origin, 'nested_array', JSON_STRING) - 1 * module.taint(_, 'array_1', source.origin, 'nested_array', JSON_STRING) + _ * module.taintIfTainted(_ as IastContext, _ as JsonParser, _) + _ * module.findSource(_ as IastContext, _ as JsonParser) >> source + 1 * module.taint(_ as IastContext, _ as JsonParser, 'root', source.origin, 'root', JSON_STRING) + 1 * module.taint(_ as IastContext, _ as JsonParser, 'root_value', source.origin, 'root', JSON_STRING) + 1 * module.taint(_ as IastContext, _ as JsonParser, 'nested', source.origin, 'nested', JSON_STRING) + 1 * module.taint(_ as IastContext, _ as JsonParser, 'nested_array', source.origin, 'nested_array', JSON_STRING) + 1 * module.taint(_ as IastContext, _ as JsonParser, 'array_0', source.origin, 'nested_array', JSON_STRING) + 1 * module.taint(_ as IastContext, _ as JsonParser, 'array_1', source.origin, 'nested_array', JSON_STRING) 0 * _ where: @@ -54,12 +52,15 @@ class Json1ParserInstrumentationTest extends AgentTestRunner { final reader = new ObjectMapper().reader(Map) when: - final taintedResult = reader.readValue(target) as Map + Map taintedResult = null + runUnderIastTrace { + taintedResult = reader.readValue(target) as Map + } then: JsonOutput.toJson(taintedResult) == JSON_STRING - _ * module.taintIfTainted(_, _) - _ * module.findSource(_) >> null + _ * module.taintIfTainted(_ as IastContext, _, _) + _ * module.findSource(_ as IastContext, _) >> null 0 * _ where: @@ -67,7 +68,7 @@ class Json1ParserInstrumentationTest extends AgentTestRunner { } private static List testSuite() { - return [JSON_STRING, new ByteArrayInputStream(JSON_STRING.getBytes(Charset.defaultCharset()))] + return [JSON_STRING] } private static class SourceImpl implements Source { diff --git a/dd-java-agent/instrumentation/jackson-core/src/main/java/datadog/trace/instrumentation/jackson/core/Json2FactoryInstrumentation.java b/dd-java-agent/instrumentation/jackson-core/src/main/java/datadog/trace/instrumentation/jackson/core/Json2FactoryInstrumentation.java index aafb5e7b5b9..f63d96be04c 100644 --- a/dd-java-agent/instrumentation/jackson-core/src/main/java/datadog/trace/instrumentation/jackson/core/Json2FactoryInstrumentation.java +++ b/dd-java-agent/instrumentation/jackson-core/src/main/java/datadog/trace/instrumentation/jackson/core/Json2FactoryInstrumentation.java @@ -9,6 +9,7 @@ import com.google.auto.service.AutoService; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.api.iast.IastContext; import datadog.trace.api.iast.InstrumentationBridge; import datadog.trace.api.iast.Propagation; import datadog.trace.api.iast.Sink; @@ -60,10 +61,13 @@ public static class InstrumenterAdvice { @Sink(VulnerabilityTypes.SSRF) // it's both propagation and Sink but Sink takes priority public static void onExit( @Advice.Argument(0) final Object input, @Advice.Return final Object parser) { - if (input != null) { + if (parser != null) { final PropagationModule propagation = InstrumentationBridge.PROPAGATION; if (propagation != null) { - propagation.taintIfTainted(parser, input); + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + propagation.taintIfTainted(ctx, parser, input); + } } if (input instanceof URL) { final SsrfModule ssrf = InstrumentationBridge.SSRF; @@ -84,10 +88,13 @@ public static void onExit( @Advice.Argument(1) final int offset, @Advice.Argument(2) final int length, @Advice.Return final Object parser) { - if (input != null || length <= 0) { + if (parser != null) { final PropagationModule propagation = InstrumentationBridge.PROPAGATION; if (propagation != null) { - propagation.taintIfTainted(parser, input, offset, length, false, NOT_MARKED); + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + propagation.taintIfTainted(ctx, parser, input, offset, length, false, NOT_MARKED); + } } } } diff --git a/dd-java-agent/instrumentation/jackson-core/src/main/java/datadog/trace/instrumentation/jackson/core/TokenBufferInstrumentation.java b/dd-java-agent/instrumentation/jackson-core/src/main/java/datadog/trace/instrumentation/jackson/core/TokenBufferInstrumentation.java index 85721c1d931..6ffdfa1d040 100644 --- a/dd-java-agent/instrumentation/jackson-core/src/main/java/datadog/trace/instrumentation/jackson/core/TokenBufferInstrumentation.java +++ b/dd-java-agent/instrumentation/jackson-core/src/main/java/datadog/trace/instrumentation/jackson/core/TokenBufferInstrumentation.java @@ -11,6 +11,7 @@ import com.google.auto.service.AutoService; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.api.iast.IastContext; import datadog.trace.api.iast.InstrumentationBridge; import datadog.trace.api.iast.Propagation; import datadog.trace.api.iast.propagation.PropagationModule; @@ -49,8 +50,11 @@ public static class AsParserAdvice { public static void onExit( @Advice.This TokenBuffer tokenBuffer, @Advice.Return JsonParser parser) { final PropagationModule module = InstrumentationBridge.PROPAGATION; - if (module != null) { - module.taintIfTainted(parser, tokenBuffer); + if (parser != null && module != null) { + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + module.taintIfTainted(ctx, parser, tokenBuffer); + } } } } diff --git a/dd-java-agent/instrumentation/java-io/build.gradle b/dd-java-agent/instrumentation/java-io/build.gradle index 00939cbf07e..612e3b18108 100644 --- a/dd-java-agent/instrumentation/java-io/build.gradle +++ b/dd-java-agent/instrumentation/java-io/build.gradle @@ -12,4 +12,5 @@ addTestSuiteForDir('latestDepTest', 'test') dependencies { testRuntimeOnly project(':dd-java-agent:instrumentation:iast-instrumenter') testImplementation group: 'org.apache.tomcat', name: 'tomcat-catalina', version: '9.0.56' + testImplementation(testFixtures(project(':dd-java-agent:agent-iast'))) } diff --git a/dd-java-agent/instrumentation/java-io/src/main/java/datadog/trace/instrumentation/java/lang/ByteBufferCallSite.java b/dd-java-agent/instrumentation/java-io/src/main/java/datadog/trace/instrumentation/java/lang/ByteBufferCallSite.java index 687fc14fa46..3e610e450f7 100644 --- a/dd-java-agent/instrumentation/java-io/src/main/java/datadog/trace/instrumentation/java/lang/ByteBufferCallSite.java +++ b/dd-java-agent/instrumentation/java-io/src/main/java/datadog/trace/instrumentation/java/lang/ByteBufferCallSite.java @@ -4,6 +4,7 @@ import datadog.trace.agent.tooling.csi.CallSite; import datadog.trace.api.iast.IastCallSites; +import datadog.trace.api.iast.IastContext; import datadog.trace.api.iast.InstrumentationBridge; import datadog.trace.api.iast.Propagation; import datadog.trace.api.iast.propagation.PropagationModule; @@ -26,7 +27,10 @@ public static ByteBuffer afterWrap( return result; } try { - module.taintIfTainted(result, bytes, true, NOT_MARKED); // keep ranges + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + module.taintIfTainted(ctx, result, bytes, true, NOT_MARKED); // keep ranges + } } catch (final Throwable e) { module.onUnexpectedException("beforeConstructor threw", e); } @@ -44,7 +48,10 @@ public static byte[] afterArray( return bytes; } try { - module.taintIfTainted(bytes, buffer, true, NOT_MARKED); // keep ranges + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + module.taintIfTainted(ctx, bytes, buffer, true, NOT_MARKED); // keep ranges + } } catch (final Throwable e) { module.onUnexpectedException("afterArray threw", e); } diff --git a/dd-java-agent/instrumentation/java-io/src/main/java/datadog/trace/instrumentation/java/lang/InputStreamInstrumentation.java b/dd-java-agent/instrumentation/java-io/src/main/java/datadog/trace/instrumentation/java/lang/InputStreamInstrumentation.java index c109e6d2944..1984ba912aa 100644 --- a/dd-java-agent/instrumentation/java-io/src/main/java/datadog/trace/instrumentation/java/lang/InputStreamInstrumentation.java +++ b/dd-java-agent/instrumentation/java-io/src/main/java/datadog/trace/instrumentation/java/lang/InputStreamInstrumentation.java @@ -8,6 +8,7 @@ import com.google.auto.service.AutoService; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.api.iast.IastContext; import datadog.trace.api.iast.InstrumentationBridge; import datadog.trace.api.iast.Propagation; import datadog.trace.api.iast.propagation.PropagationModule; @@ -50,17 +51,21 @@ public String[] getClassNamesToBePreloaded() { public static class InputStreamAdvice { + /** This method cannot suppress throwable due to being applied to a constructor */ @Advice.OnMethodExit @Propagation public static void onExit( @Advice.This final InputStream self, @Advice.Argument(0) final InputStream param) { final PropagationModule module = InstrumentationBridge.PROPAGATION; - try { - if (module != null) { - module.taintIfTainted(self, param); + if (param != null && module != null) { + try { + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + module.taintIfTainted(ctx, self, param); + } + } catch (final Throwable e) { + module.onUnexpectedException("InputStreamAdvice onExit threw", e); } - } catch (final Throwable e) { - module.onUnexpectedException("InputStreamAdvice onExit threw", e); } } } diff --git a/dd-java-agent/instrumentation/java-io/src/main/java/datadog/trace/instrumentation/java/lang/InputStreamReaderCallSite.java b/dd-java-agent/instrumentation/java-io/src/main/java/datadog/trace/instrumentation/java/lang/InputStreamReaderCallSite.java index a73115ffb3b..e34915732cd 100644 --- a/dd-java-agent/instrumentation/java-io/src/main/java/datadog/trace/instrumentation/java/lang/InputStreamReaderCallSite.java +++ b/dd-java-agent/instrumentation/java-io/src/main/java/datadog/trace/instrumentation/java/lang/InputStreamReaderCallSite.java @@ -2,6 +2,7 @@ import datadog.trace.agent.tooling.csi.CallSite; import datadog.trace.api.iast.IastCallSites; +import datadog.trace.api.iast.IastContext; import datadog.trace.api.iast.InstrumentationBridge; import datadog.trace.api.iast.Propagation; import datadog.trace.api.iast.propagation.PropagationModule; @@ -20,7 +21,12 @@ public static InputStreamReader afterInit( final PropagationModule module = InstrumentationBridge.PROPAGATION; if (module != null) { try { - module.taintIfTainted(result, params[0]); + if (params[0] != null) { + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + module.taintIfTainted(ctx, result, params[0]); + } + } } catch (final Throwable e) { module.onUnexpectedException("afterInit threw", e); } diff --git a/dd-java-agent/instrumentation/java-io/src/main/java/datadog/trace/instrumentation/java/lang/StringReaderCallSite.java b/dd-java-agent/instrumentation/java-io/src/main/java/datadog/trace/instrumentation/java/lang/StringReaderCallSite.java index 09f203efa60..380c14b2ab5 100644 --- a/dd-java-agent/instrumentation/java-io/src/main/java/datadog/trace/instrumentation/java/lang/StringReaderCallSite.java +++ b/dd-java-agent/instrumentation/java-io/src/main/java/datadog/trace/instrumentation/java/lang/StringReaderCallSite.java @@ -2,6 +2,7 @@ import datadog.trace.agent.tooling.csi.CallSite; import datadog.trace.api.iast.IastCallSites; +import datadog.trace.api.iast.IastContext; import datadog.trace.api.iast.InstrumentationBridge; import datadog.trace.api.iast.Propagation; import datadog.trace.api.iast.propagation.PropagationModule; @@ -18,7 +19,16 @@ public static StringReader afterInit( @CallSite.Return @Nonnull final StringReader result) { final PropagationModule propagationModule = InstrumentationBridge.PROPAGATION; if (propagationModule != null) { - propagationModule.taintIfTainted(result, params[0]); + try { + if (params[0] != null) { + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + propagationModule.taintIfTainted(ctx, result, params[0]); + } + } + } catch (Throwable e) { + propagationModule.onUnexpectedException("afterInit threw", e); + } } return result; } diff --git a/dd-java-agent/instrumentation/java-io/src/test/groovy/datadog/trace/instrumentation/java/io/BaseIoCallSiteTest.groovy b/dd-java-agent/instrumentation/java-io/src/test/groovy/datadog/trace/instrumentation/java/io/BaseIoCallSiteTest.groovy index e9e7e4f1959..0d6cddeab91 100644 --- a/dd-java-agent/instrumentation/java-io/src/test/groovy/datadog/trace/instrumentation/java/io/BaseIoCallSiteTest.groovy +++ b/dd-java-agent/instrumentation/java-io/src/test/groovy/datadog/trace/instrumentation/java/io/BaseIoCallSiteTest.groovy @@ -1,19 +1,14 @@ package datadog.trace.instrumentation.java.io -import datadog.trace.agent.test.AgentTestRunner +import com.datadog.iast.test.IastAgentTestRunner import org.junit.Rule import org.junit.rules.TemporaryFolder -abstract class BaseIoCallSiteTest extends AgentTestRunner { +abstract class BaseIoCallSiteTest extends IastAgentTestRunner { @Rule TemporaryFolder temporaryFolder = new TemporaryFolder(parentFolder()) - @Override - protected void configurePreAgent() { - injectSysConfig("dd.iast.enabled", "true") - } - protected File newFile(final String name) { return temporaryFolder.newFile(name) } diff --git a/dd-java-agent/instrumentation/java-io/src/test/groovy/datadog/trace/instrumentation/java/io/ByteBufferTest.groovy b/dd-java-agent/instrumentation/java-io/src/test/groovy/datadog/trace/instrumentation/java/io/ByteBufferTest.groovy index 2d299f20742..daaeb940f83 100644 --- a/dd-java-agent/instrumentation/java-io/src/test/groovy/datadog/trace/instrumentation/java/io/ByteBufferTest.groovy +++ b/dd-java-agent/instrumentation/java-io/src/test/groovy/datadog/trace/instrumentation/java/io/ByteBufferTest.groovy @@ -1,6 +1,7 @@ package datadog.trace.instrumentation.java.io -import datadog.trace.agent.test.AgentTestRunner +import com.datadog.iast.test.IastAgentTestRunner +import datadog.trace.api.iast.IastContext import datadog.trace.api.iast.InstrumentationBridge import datadog.trace.api.iast.VulnerabilityMarks import datadog.trace.api.iast.propagation.PropagationModule @@ -8,12 +9,7 @@ import foo.bar.TestByteBufferSuite import java.nio.ByteBuffer -class ByteBufferTest extends AgentTestRunner { - - @Override - protected void configurePreAgent() { - injectSysConfig('dd.iast.enabled', 'true') - } +class ByteBufferTest extends IastAgentTestRunner { void 'test wrap method'() { given: @@ -22,10 +18,10 @@ class ByteBufferTest extends AgentTestRunner { InstrumentationBridge.registerIastModule(module) when: - TestByteBufferSuite.wrap(message) + runUnderIastTrace { TestByteBufferSuite.wrap(message) } then: - 1 * module.taintIfTainted(_ as ByteBuffer, message, true, VulnerabilityMarks.NOT_MARKED) + 1 * module.taintIfTainted(_ as IastContext, _ as ByteBuffer, message, true, VulnerabilityMarks.NOT_MARKED) } void 'test array method'() { @@ -35,10 +31,11 @@ class ByteBufferTest extends AgentTestRunner { InstrumentationBridge.registerIastModule(module) when: - final result = TestByteBufferSuite.array(message) + byte[] result = null + runUnderIastTrace { result = TestByteBufferSuite.array(message) } then: result == message.array() - 1 * module.taintIfTainted(_ as byte[], message, true, VulnerabilityMarks.NOT_MARKED) + 1 * module.taintIfTainted(_ as IastContext, _ as byte[], message, true, VulnerabilityMarks.NOT_MARKED) } } diff --git a/dd-java-agent/instrumentation/java-io/src/test/groovy/datadog/trace/instrumentation/java/io/InputStreamInstrumentationTest.groovy b/dd-java-agent/instrumentation/java-io/src/test/groovy/datadog/trace/instrumentation/java/io/InputStreamInstrumentationTest.groovy index a73e5f5e524..4682aff2fe0 100644 --- a/dd-java-agent/instrumentation/java-io/src/test/groovy/datadog/trace/instrumentation/java/io/InputStreamInstrumentationTest.groovy +++ b/dd-java-agent/instrumentation/java-io/src/test/groovy/datadog/trace/instrumentation/java/io/InputStreamInstrumentationTest.groovy @@ -1,17 +1,12 @@ package datadog.trace.instrumentation.java.io -import datadog.trace.agent.test.AgentTestRunner +import com.datadog.iast.test.IastAgentTestRunner +import datadog.trace.api.iast.IastContext import datadog.trace.api.iast.InstrumentationBridge import datadog.trace.api.iast.propagation.PropagationModule import foo.bar.TestInputStreamSuite -class InputStreamInstrumentationTest extends AgentTestRunner { - - - @Override - protected void configurePreAgent() { - injectSysConfig("dd.iast.enabled", "true") - } +class InputStreamInstrumentationTest extends IastAgentTestRunner { def 'test constructor with IS as arg()'() { setup: @@ -20,9 +15,9 @@ class InputStreamInstrumentationTest extends AgentTestRunner { final is = Mock(InputStream) when: - TestInputStreamSuite.pushbackInputStreamFromIS(is) + runUnderIastTrace { TestInputStreamSuite.pushbackInputStreamFromIS(is) } then: - (1.._) * propagationModule.taintIfTainted(_ as InputStream, is) + (1.._) * propagationModule.taintIfTainted(_ as IastContext, _ as InputStream, is) } } diff --git a/dd-java-agent/instrumentation/java-io/src/test/groovy/datadog/trace/instrumentation/java/io/InputStreamReaderCallSiteTest.groovy b/dd-java-agent/instrumentation/java-io/src/test/groovy/datadog/trace/instrumentation/java/io/InputStreamReaderCallSiteTest.groovy index f8b522e3e24..79bef226e30 100644 --- a/dd-java-agent/instrumentation/java-io/src/test/groovy/datadog/trace/instrumentation/java/io/InputStreamReaderCallSiteTest.groovy +++ b/dd-java-agent/instrumentation/java-io/src/test/groovy/datadog/trace/instrumentation/java/io/InputStreamReaderCallSiteTest.groovy @@ -1,5 +1,6 @@ package datadog.trace.instrumentation.java.io +import datadog.trace.api.iast.IastContext import datadog.trace.api.iast.InstrumentationBridge import datadog.trace.api.iast.propagation.PropagationModule import foo.bar.TestInputStreamReaderSuite @@ -14,10 +15,10 @@ class InputStreamReaderCallSiteTest extends BaseIoCallSiteTest{ InstrumentationBridge.registerIastModule(iastModule) when: - TestInputStreamReaderSuite.init(new ByteArrayInputStream("test".getBytes()), Charset.defaultCharset()) + runUnderIastTrace { TestInputStreamReaderSuite.init(new ByteArrayInputStream("test".getBytes()), Charset.defaultCharset()) } then: - 1 * iastModule.taintIfTainted(_ as InputStreamReader, _ as InputStream) + 1 * iastModule.taintIfTainted(_ as IastContext, _ as InputStreamReader, _ as InputStream) 0 * _ } } diff --git a/dd-java-agent/instrumentation/java-io/src/test/groovy/datadog/trace/instrumentation/java/io/StringReaderCallSiteTest.groovy b/dd-java-agent/instrumentation/java-io/src/test/groovy/datadog/trace/instrumentation/java/io/StringReaderCallSiteTest.groovy index f99dca913f4..11faa0fd4a7 100644 --- a/dd-java-agent/instrumentation/java-io/src/test/groovy/datadog/trace/instrumentation/java/io/StringReaderCallSiteTest.groovy +++ b/dd-java-agent/instrumentation/java-io/src/test/groovy/datadog/trace/instrumentation/java/io/StringReaderCallSiteTest.groovy @@ -1,21 +1,22 @@ package datadog.trace.instrumentation.java.io +import datadog.trace.api.iast.IastContext import datadog.trace.api.iast.InstrumentationBridge import datadog.trace.api.iast.propagation.PropagationModule import foo.bar.TestStringReaderSuite -class StringReaderCallSiteTest extends BaseIoCallSiteTest{ +class StringReaderCallSiteTest extends BaseIoCallSiteTest { - void 'test StringReader.'(){ + void 'test StringReader.'() { given: PropagationModule iastModule = Mock(PropagationModule) InstrumentationBridge.registerIastModule(iastModule) final input = 'Test input' when: - TestStringReaderSuite.init(input) + runUnderIastTrace { TestStringReaderSuite.init(input) } then: - 1 * iastModule.taintIfTainted(_ as StringReader, input) + 1 * iastModule.taintIfTainted(_ as IastContext, _ as StringReader, input) } } diff --git a/dd-java-agent/instrumentation/java-lang/src/main/java/datadog/trace/instrumentation/java/lang/StringCallSite.java b/dd-java-agent/instrumentation/java-lang/src/main/java/datadog/trace/instrumentation/java/lang/StringCallSite.java index c48ec0e350c..054f12477b2 100644 --- a/dd-java-agent/instrumentation/java-lang/src/main/java/datadog/trace/instrumentation/java/lang/StringCallSite.java +++ b/dd-java-agent/instrumentation/java-lang/src/main/java/datadog/trace/instrumentation/java/lang/StringCallSite.java @@ -4,6 +4,7 @@ import datadog.trace.agent.tooling.csi.CallSite; import datadog.trace.api.iast.IastCallSites; +import datadog.trace.api.iast.IastContext; import datadog.trace.api.iast.InstrumentationBridge; import datadog.trace.api.iast.Propagation; import datadog.trace.api.iast.propagation.CodecModule; @@ -17,6 +18,7 @@ import javax.annotation.Nonnull; import javax.annotation.Nullable; +@SuppressWarnings("unused") @Propagation @CallSite(spi = IastCallSites.class) public class StringCallSite { @@ -123,7 +125,7 @@ public static String aroundJoin( final StringModule module = InstrumentationBridge.STRING; if (module != null) { try { - module.onStringJoin(result, delimiter, copy.toArray(new CharSequence[copy.size()])); + module.onStringJoin(result, delimiter, copy.toArray(new CharSequence[0])); } catch (final Throwable e) { module.onUnexpectedException("afterSubSequence threw", e); } @@ -384,7 +386,10 @@ public static char[] afterToCharArray( final PropagationModule module = InstrumentationBridge.PROPAGATION; if (module != null) { try { - module.taintIfTainted(result, self, true, NOT_MARKED); + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + module.taintIfTainted(ctx, result, self, true, NOT_MARKED); + } } catch (final Throwable e) { module.onUnexpectedException("afterToCharArray threw", e); } @@ -412,7 +417,7 @@ public static String[] afterSplit( public static String[] afterSplitWithLimit( @CallSite.This @Nonnull final String self, @CallSite.Argument(0) @Nonnull final String regex, - @CallSite.Argument(1) @Nonnull final int pos, + @CallSite.Argument(1) final int pos, @CallSite.Return @Nonnull final String[] result) { final StringModule module = InstrumentationBridge.STRING; if (module != null) { diff --git a/dd-java-agent/instrumentation/java-net/src/main/java/datadog/trace/instrumentation/java/net/URICallSite.java b/dd-java-agent/instrumentation/java-net/src/main/java/datadog/trace/instrumentation/java/net/URICallSite.java index 66bc25d625f..b87697b8238 100644 --- a/dd-java-agent/instrumentation/java-net/src/main/java/datadog/trace/instrumentation/java/net/URICallSite.java +++ b/dd-java-agent/instrumentation/java-net/src/main/java/datadog/trace/instrumentation/java/net/URICallSite.java @@ -2,6 +2,7 @@ import datadog.trace.agent.tooling.csi.CallSite; import datadog.trace.api.iast.IastCallSites; +import datadog.trace.api.iast.IastContext; import datadog.trace.api.iast.InstrumentationBridge; import datadog.trace.api.iast.Propagation; import datadog.trace.api.iast.propagation.PropagationModule; @@ -20,7 +21,10 @@ public static URI afterCreate( final PropagationModule module = InstrumentationBridge.PROPAGATION; if (module != null) { try { - module.taintIfTainted(result, value); + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + module.taintIfTainted(ctx, result, value); + } } catch (final Throwable e) { module.onUnexpectedException("create threw", e); } @@ -43,7 +47,10 @@ public static URI afterCtor( final PropagationModule module = InstrumentationBridge.PROPAGATION; if (module != null) { try { - module.taintIfAnyTainted(result, args); + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + module.taintIfAnyTainted(ctx, result, args); + } } catch (final Throwable e) { module.onUnexpectedException("ctor threw", e); } @@ -57,9 +64,12 @@ public static URI afterCtor( public static String afterToString( @CallSite.This final URI url, @CallSite.Return final String result) { final PropagationModule module = InstrumentationBridge.PROPAGATION; - if (module != null) { + if (result != null && module != null) { try { - module.taintIfTainted(result, url); + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + module.taintIfTainted(ctx, result, url); + } } catch (final Throwable e) { module.onUnexpectedException("After toString threw", e); } @@ -71,9 +81,12 @@ public static String afterToString( public static URI afterNormalize( @CallSite.This final URI url, @CallSite.Return final URI result) { final PropagationModule module = InstrumentationBridge.PROPAGATION; - if (module != null) { + if (result != null && module != null) { try { - module.taintIfTainted(result, url); + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + module.taintIfTainted(ctx, result, url); + } } catch (final Throwable e) { module.onUnexpectedException("After toString threw", e); } diff --git a/dd-java-agent/instrumentation/java-net/src/main/java/datadog/trace/instrumentation/java/net/URLCallSite.java b/dd-java-agent/instrumentation/java-net/src/main/java/datadog/trace/instrumentation/java/net/URLCallSite.java index da954219bf4..c434de2d3ac 100644 --- a/dd-java-agent/instrumentation/java-net/src/main/java/datadog/trace/instrumentation/java/net/URLCallSite.java +++ b/dd-java-agent/instrumentation/java-net/src/main/java/datadog/trace/instrumentation/java/net/URLCallSite.java @@ -2,6 +2,7 @@ import datadog.trace.agent.tooling.csi.CallSite; import datadog.trace.api.iast.IastCallSites; +import datadog.trace.api.iast.IastContext; import datadog.trace.api.iast.InstrumentationBridge; import datadog.trace.api.iast.Propagation; import datadog.trace.api.iast.Sink; @@ -32,7 +33,10 @@ public static URL afterCtor( final PropagationModule module = InstrumentationBridge.PROPAGATION; if (module != null) { try { - module.taintIfAnyTainted(result, args); + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + module.taintIfAnyTainted(ctx, result, args); + } } catch (final Throwable e) { module.onUnexpectedException("ctor threw", e); } @@ -47,9 +51,12 @@ public static URL afterCtor( public static String afterToString( @CallSite.This final URL url, @CallSite.Return final String result) { final PropagationModule module = InstrumentationBridge.PROPAGATION; - if (module != null) { + if (result != null && module != null) { try { - module.taintIfTainted(result, url); + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + module.taintIfTainted(ctx, result, url); + } } catch (final Throwable e) { module.onUnexpectedException("After toString threw", e); } @@ -61,9 +68,12 @@ public static String afterToString( @CallSite.After("java.net.URI java.net.URL.toURI()") public static URI afterToURI(@CallSite.This final URL url, @CallSite.Return final URI result) { final PropagationModule module = InstrumentationBridge.PROPAGATION; - if (module != null) { + if (result != null && module != null) { try { - module.taintIfTainted(result, url); + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + module.taintIfTainted(ctx, result, url); + } } catch (final Throwable e) { module.onUnexpectedException("After toURI threw", e); } diff --git a/dd-java-agent/instrumentation/java-net/src/main/java/datadog/trace/instrumentation/java/net/URLEncoderCallSite.java b/dd-java-agent/instrumentation/java-net/src/main/java/datadog/trace/instrumentation/java/net/URLEncoderCallSite.java index 7c219bcf6ba..227b10363a7 100644 --- a/dd-java-agent/instrumentation/java-net/src/main/java/datadog/trace/instrumentation/java/net/URLEncoderCallSite.java +++ b/dd-java-agent/instrumentation/java-net/src/main/java/datadog/trace/instrumentation/java/net/URLEncoderCallSite.java @@ -2,6 +2,7 @@ import datadog.trace.agent.tooling.csi.CallSite; import datadog.trace.api.iast.IastCallSites; +import datadog.trace.api.iast.IastContext; import datadog.trace.api.iast.InstrumentationBridge; import datadog.trace.api.iast.Propagation; import datadog.trace.api.iast.VulnerabilityMarks; @@ -32,7 +33,10 @@ private static String encode(final String result, final String value) { final PropagationModule module = InstrumentationBridge.PROPAGATION; if (module != null) { try { - module.taintIfTainted(result, value, false, VulnerabilityMarks.XSS_MARK); + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + module.taintIfTainted(ctx, result, value, false, VulnerabilityMarks.XSS_MARK); + } } catch (final Throwable e) { module.onUnexpectedException("afterEncode threw", e); } diff --git a/dd-java-agent/instrumentation/jersey/src/main/java/datadog/trace/instrumentation/jersey/AbstractParamValueExtractorInstrumentation.java b/dd-java-agent/instrumentation/jersey/src/main/java/datadog/trace/instrumentation/jersey/AbstractParamValueExtractorInstrumentation.java index 9fe244af34e..124d641cb74 100644 --- a/dd-java-agent/instrumentation/jersey/src/main/java/datadog/trace/instrumentation/jersey/AbstractParamValueExtractorInstrumentation.java +++ b/dd-java-agent/instrumentation/jersey/src/main/java/datadog/trace/instrumentation/jersey/AbstractParamValueExtractorInstrumentation.java @@ -54,7 +54,7 @@ public static void onExit( final PropagationModule module = InstrumentationBridge.PROPAGATION; if (module != null) { IastContext ctx = reqCtx.getData(RequestContextSlot.IAST); - module.taint(ctx, result, ThreadLocalSourceType.get(), parameterName); + module.taint(ctx, (String) result, ThreadLocalSourceType.get(), parameterName); } } } diff --git a/dd-java-agent/instrumentation/jersey/src/main/java/datadog/trace/instrumentation/jersey/AbstractStringReaderAdvice.java b/dd-java-agent/instrumentation/jersey/src/main/java/datadog/trace/instrumentation/jersey/AbstractStringReaderAdvice.java index 575fbed1100..89274cfee34 100644 --- a/dd-java-agent/instrumentation/jersey/src/main/java/datadog/trace/instrumentation/jersey/AbstractStringReaderAdvice.java +++ b/dd-java-agent/instrumentation/jersey/src/main/java/datadog/trace/instrumentation/jersey/AbstractStringReaderAdvice.java @@ -21,7 +21,7 @@ public static void onExit( final PropagationModule module = InstrumentationBridge.PROPAGATION; if (module != null) { IastContext ctx = reqCtx.getData(RequestContextSlot.IAST); - module.taint(ctx, result, SourceTypes.REQUEST_PARAMETER_VALUE); + module.taint(ctx, (String) result, SourceTypes.REQUEST_PARAMETER_VALUE); } } } diff --git a/dd-java-agent/instrumentation/jose-jwt/src/main/java/datadog/trace/instrumentation/josejwt/JSONObjectUtilsInstrumentation.java b/dd-java-agent/instrumentation/jose-jwt/src/main/java/datadog/trace/instrumentation/josejwt/JSONObjectUtilsInstrumentation.java index 2af1be5219f..baae4938e38 100644 --- a/dd-java-agent/instrumentation/jose-jwt/src/main/java/datadog/trace/instrumentation/josejwt/JSONObjectUtilsInstrumentation.java +++ b/dd-java-agent/instrumentation/jose-jwt/src/main/java/datadog/trace/instrumentation/josejwt/JSONObjectUtilsInstrumentation.java @@ -56,7 +56,7 @@ public static void onEnter( if (value instanceof String) { // TODO: We could represent this source more accurately, perhaps tracking the original // source, or using a special name. - module.taint(ctx, value, SourceTypes.REQUEST_HEADER_VALUE, name); + module.taint(ctx, (String) value, SourceTypes.REQUEST_HEADER_VALUE, name); } } } diff --git a/dd-java-agent/instrumentation/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/KafkaIastHelper.java b/dd-java-agent/instrumentation/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/KafkaIastHelper.java index 2b3c29d3bed..5d08585152f 100644 --- a/dd-java-agent/instrumentation/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/KafkaIastHelper.java +++ b/dd-java-agent/instrumentation/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/KafkaIastHelper.java @@ -3,6 +3,7 @@ import static datadog.trace.api.iast.SourceTypes.KAFKA_MESSAGE_KEY; import static datadog.trace.api.iast.SourceTypes.KAFKA_MESSAGE_VALUE; +import datadog.trace.api.iast.IastContext; import datadog.trace.api.iast.InstrumentationBridge; import datadog.trace.api.iast.propagation.PropagationModule; import datadog.trace.bootstrap.CallDepthThreadLocalMap; @@ -36,8 +37,12 @@ public static void taint( if (module == null) { return; } + final IastContext ctx = IastContext.Provider.get(); + if (ctx == null) { + return; + } final byte source = getSource(store, deserializer); - module.taint(data, source); + module.taint(ctx, data, source); } public static void taint( @@ -54,12 +59,16 @@ public static void taint( if (module == null) { return; } + final IastContext ctx = IastContext.Provider.get(); + if (ctx == null) { + return; + } final byte source = getSource(store, deserializer); int start = data.position(); if (data.hasArray()) { start += data.arrayOffset(); } - module.taint(data, source, start, data.remaining()); + module.taint(ctx, data, source, start, data.remaining()); } public static void afterDeserialize() { diff --git a/dd-java-agent/instrumentation/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/UtilsInstrumentation.java b/dd-java-agent/instrumentation/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/UtilsInstrumentation.java index 38c8a02132b..348e0e0805c 100644 --- a/dd-java-agent/instrumentation/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/UtilsInstrumentation.java +++ b/dd-java-agent/instrumentation/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/UtilsInstrumentation.java @@ -9,6 +9,7 @@ import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.agent.tooling.muzzle.Reference; +import datadog.trace.api.iast.IastContext; import datadog.trace.api.iast.InstrumentationBridge; import datadog.trace.api.iast.Propagation; import datadog.trace.api.iast.propagation.PropagationModule; @@ -58,12 +59,16 @@ public static void toArray( if (propagation == null) { return; } + final IastContext ctx = IastContext.Provider.get(); + if (ctx == null) { + return; + } int start = buffer.position() + offset; if (buffer.hasArray()) { start += buffer.arrayOffset(); } // create a new range shifted to the result byte array coordinates - propagation.taintIfTainted(bytes, buffer, start, length, false, NOT_MARKED); + propagation.taintIfTainted(ctx, bytes, buffer, start, length, false, NOT_MARKED); } } @@ -79,7 +84,11 @@ public static void wrapNullable( if (propagation == null) { return; } - propagation.taintIfTainted(buffer, bytes, true, NOT_MARKED); + final IastContext ctx = IastContext.Provider.get(); + if (ctx == null) { + return; + } + propagation.taintIfTainted(ctx, buffer, bytes, true, NOT_MARKED); } } } diff --git a/dd-java-agent/instrumentation/netty-buffer-4/src/main/java/datadog/trace/instrumentation/netty40/buffer/ByteBufInputStreamInstrumentation.java b/dd-java-agent/instrumentation/netty-buffer-4/src/main/java/datadog/trace/instrumentation/netty40/buffer/ByteBufInputStreamInstrumentation.java index 1de2f01498e..5275f4656c2 100644 --- a/dd-java-agent/instrumentation/netty-buffer-4/src/main/java/datadog/trace/instrumentation/netty40/buffer/ByteBufInputStreamInstrumentation.java +++ b/dd-java-agent/instrumentation/netty-buffer-4/src/main/java/datadog/trace/instrumentation/netty40/buffer/ByteBufInputStreamInstrumentation.java @@ -10,6 +10,7 @@ import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.agent.tooling.bytebuddy.iast.TaintableVisitor; +import datadog.trace.api.iast.IastContext; import datadog.trace.api.iast.InstrumentationBridge; import datadog.trace.api.iast.Propagation; import datadog.trace.api.iast.propagation.PropagationModule; @@ -47,17 +48,16 @@ public void methodAdvice(final MethodTransformer transformer) { public static class ConstructorAdvice { - @Advice.OnMethodExit + @Advice.OnMethodExit(onThrowable = Throwable.class) @Propagation public static void onExit( @Advice.This final Object self, @Advice.Argument(0) final Object buffer) { final PropagationModule module = InstrumentationBridge.PROPAGATION; - try { - if (module != null) { - module.taintIfTainted(self, buffer); + if (buffer != null && module != null) { + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + module.taintIfTainted(ctx, self, buffer); } - } catch (final Throwable e) { - module.onUnexpectedException("ByteBufInputStream ctor threw", e); } } } diff --git a/dd-java-agent/instrumentation/okhttp-2/src/main/java/datadog/trace/instrumentation/okhttp2/IastHttpUrlInstrumentation.java b/dd-java-agent/instrumentation/okhttp-2/src/main/java/datadog/trace/instrumentation/okhttp2/IastHttpUrlInstrumentation.java index 3a894505e36..829c17f8f7a 100644 --- a/dd-java-agent/instrumentation/okhttp-2/src/main/java/datadog/trace/instrumentation/okhttp2/IastHttpUrlInstrumentation.java +++ b/dd-java-agent/instrumentation/okhttp-2/src/main/java/datadog/trace/instrumentation/okhttp2/IastHttpUrlInstrumentation.java @@ -10,6 +10,7 @@ import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.agent.tooling.bytebuddy.iast.TaintableVisitor; +import datadog.trace.api.iast.IastContext; import datadog.trace.api.iast.InstrumentationBridge; import datadog.trace.api.iast.Propagation; import datadog.trace.api.iast.propagation.PropagationModule; @@ -76,8 +77,11 @@ public static class ParseAdvice { public static void onParse( @Advice.Argument(0) final Object argument, @Advice.Return final Object result) { final PropagationModule module = InstrumentationBridge.PROPAGATION; - if (module != null) { - module.taintIfTainted(result, argument); + if (result != null && module != null) { + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + module.taintIfTainted(ctx, result, argument); + } } } } @@ -89,8 +93,11 @@ public static class PropagationAdvice { public static void onPropagation( @Advice.This final Object self, @Advice.Return final Object result) { final PropagationModule module = InstrumentationBridge.PROPAGATION; - if (module != null) { - module.taintIfTainted(result, self); + if (result != null && module != null) { + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + module.taintIfTainted(ctx, result, self); + } } } } diff --git a/dd-java-agent/instrumentation/okhttp-3/src/main/java/datadog/trace/instrumentation/okhttp3/IastHttpUrlInstrumentation.java b/dd-java-agent/instrumentation/okhttp-3/src/main/java/datadog/trace/instrumentation/okhttp3/IastHttpUrlInstrumentation.java index 309f0778039..7803cda9582 100644 --- a/dd-java-agent/instrumentation/okhttp-3/src/main/java/datadog/trace/instrumentation/okhttp3/IastHttpUrlInstrumentation.java +++ b/dd-java-agent/instrumentation/okhttp-3/src/main/java/datadog/trace/instrumentation/okhttp3/IastHttpUrlInstrumentation.java @@ -10,6 +10,7 @@ import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.agent.tooling.bytebuddy.iast.TaintableVisitor; +import datadog.trace.api.iast.IastContext; import datadog.trace.api.iast.InstrumentationBridge; import datadog.trace.api.iast.Propagation; import datadog.trace.api.iast.propagation.PropagationModule; @@ -76,8 +77,11 @@ public static class ParseAdvice { public static void onParse( @Advice.Argument(0) final Object arg, @Advice.Return final Object result) { final PropagationModule module = InstrumentationBridge.PROPAGATION; - if (module != null) { - module.taintIfTainted(result, arg); + if (result != null && module != null) { + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + module.taintIfTainted(ctx, result, arg); + } } } } diff --git a/dd-java-agent/instrumentation/org-json/build.gradle b/dd-java-agent/instrumentation/org-json/build.gradle index ff630aa67dd..4458d70b16e 100644 --- a/dd-java-agent/instrumentation/org-json/build.gradle +++ b/dd-java-agent/instrumentation/org-json/build.gradle @@ -14,6 +14,7 @@ addTestSuiteForDir('latestDepTest', 'test') dependencies { compileOnly group: 'org.json', name: 'json', version: '20230227' testImplementation group: 'org.json', name: 'json', version: '20230227' + testImplementation(testFixtures(project(':dd-java-agent:agent-iast'))) testRuntimeOnly project(':dd-java-agent:instrumentation:iast-instrumenter') diff --git a/dd-java-agent/instrumentation/org-json/src/main/java/datadog/trace/instrumentation/json/JSONArrayInstrumentation.java b/dd-java-agent/instrumentation/org-json/src/main/java/datadog/trace/instrumentation/json/JSONArrayInstrumentation.java index 6ac7e8956e7..85c90f92c08 100644 --- a/dd-java-agent/instrumentation/org-json/src/main/java/datadog/trace/instrumentation/json/JSONArrayInstrumentation.java +++ b/dd-java-agent/instrumentation/org-json/src/main/java/datadog/trace/instrumentation/json/JSONArrayInstrumentation.java @@ -10,10 +10,13 @@ import com.google.auto.service.AutoService; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.api.iast.IastContext; import datadog.trace.api.iast.InstrumentationBridge; import datadog.trace.api.iast.Propagation; import datadog.trace.api.iast.propagation.PropagationModule; import net.bytebuddy.asm.Advice; +import org.json.JSONArray; +import org.json.JSONObject; @AutoService(InstrumenterModule.class) public class JSONArrayInstrumentation extends InstrumenterModule.Iast @@ -51,7 +54,10 @@ public static class ConstructorAdvice { public static void afterInit(@Advice.This Object self, @Advice.Argument(0) final Object input) { final PropagationModule iastModule = InstrumentationBridge.PROPAGATION; if (iastModule != null && input != null) { - iastModule.taintIfTainted(self, input); + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + iastModule.taintIfTainted(ctx, self, input); + } } } } @@ -60,15 +66,21 @@ public static class GetAdvice { @Advice.OnMethodExit(suppress = Throwable.class) @Propagation public static void afterMethod(@Advice.This Object self, @Advice.Return final Object result) { - if (result instanceof Integer - || result instanceof Long - || result instanceof Double - || result instanceof Boolean) { + boolean isString = result instanceof String; + boolean isJson = !isString && (result instanceof JSONObject || result instanceof JSONArray); + if (!isString && !isJson) { return; } final PropagationModule iastModule = InstrumentationBridge.PROPAGATION; - if (iastModule != null && result != null && result instanceof String) { - iastModule.taintIfTainted(result, self); + if (iastModule != null) { + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + if (isString) { + iastModule.taintIfTainted(ctx, (String) result, self); + } else { + iastModule.taintIfTainted(ctx, result, self); + } + } } } } @@ -77,15 +89,21 @@ public static class OptAdvice { @Advice.OnMethodExit(suppress = Throwable.class) @Propagation public static void afterMethod(@Advice.This Object self, @Advice.Return final Object result) { - if (result instanceof Integer - || result instanceof Long - || result instanceof Double - || result instanceof Boolean) { + boolean isString = result instanceof String; + boolean isJson = !isString && (result instanceof JSONObject || result instanceof JSONArray); + if (!isString && !isJson) { return; } final PropagationModule iastModule = InstrumentationBridge.PROPAGATION; - if (iastModule != null && result != null) { - iastModule.taintIfTainted(result, self); + if (iastModule != null) { + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + if (isString) { + iastModule.taintIfTainted(ctx, (String) result, self); + } else { + iastModule.taintIfTainted(ctx, result, self); + } + } } } } diff --git a/dd-java-agent/instrumentation/org-json/src/main/java/datadog/trace/instrumentation/json/JSONCookieInstrumentation.java b/dd-java-agent/instrumentation/org-json/src/main/java/datadog/trace/instrumentation/json/JSONCookieInstrumentation.java index cb3e2ce397d..b390dced85a 100644 --- a/dd-java-agent/instrumentation/org-json/src/main/java/datadog/trace/instrumentation/json/JSONCookieInstrumentation.java +++ b/dd-java-agent/instrumentation/org-json/src/main/java/datadog/trace/instrumentation/json/JSONCookieInstrumentation.java @@ -6,6 +6,7 @@ import com.google.auto.service.AutoService; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.api.iast.IastContext; import datadog.trace.api.iast.InstrumentationBridge; import datadog.trace.api.iast.Propagation; import datadog.trace.api.iast.propagation.PropagationModule; @@ -38,7 +39,10 @@ public static void onExit( @Advice.Return Object retValue, @Advice.Argument(0) final String input) { final PropagationModule iastModule = InstrumentationBridge.PROPAGATION; if (iastModule != null && input != null) { - iastModule.taintIfTainted(retValue, input); + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + iastModule.taintIfTainted(ctx, retValue, input); + } } } } diff --git a/dd-java-agent/instrumentation/org-json/src/main/java/datadog/trace/instrumentation/json/JSONObjectInstrumentation.java b/dd-java-agent/instrumentation/org-json/src/main/java/datadog/trace/instrumentation/json/JSONObjectInstrumentation.java index cae6f15906d..79536d66432 100644 --- a/dd-java-agent/instrumentation/org-json/src/main/java/datadog/trace/instrumentation/json/JSONObjectInstrumentation.java +++ b/dd-java-agent/instrumentation/org-json/src/main/java/datadog/trace/instrumentation/json/JSONObjectInstrumentation.java @@ -10,10 +10,13 @@ import com.google.auto.service.AutoService; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.api.iast.IastContext; import datadog.trace.api.iast.InstrumentationBridge; import datadog.trace.api.iast.Propagation; import datadog.trace.api.iast.propagation.PropagationModule; import net.bytebuddy.asm.Advice; +import org.json.JSONArray; +import org.json.JSONObject; @AutoService(InstrumenterModule.class) public class JSONObjectInstrumentation extends InstrumenterModule.Iast @@ -53,7 +56,10 @@ public static class ConstructorAdvice { public static void afterInit(@Advice.This Object self, @Advice.Argument(0) final Object input) { final PropagationModule iastModule = InstrumentationBridge.PROPAGATION; if (iastModule != null && input != null) { - iastModule.taintIfTainted(self, input); + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + iastModule.taintIfTainted(ctx, self, input); + } } } } @@ -62,15 +68,21 @@ public static class GetAdvice { @Advice.OnMethodExit(suppress = Throwable.class) @Propagation public static void afterMethod(@Advice.This Object self, @Advice.Return final Object result) { - if (result instanceof Integer - || result instanceof Long - || result instanceof Double - || result instanceof Boolean) { + boolean isString = result instanceof String; + boolean isJson = !isString && (result instanceof JSONObject || result instanceof JSONArray); + if (!isString && !isJson) { return; } final PropagationModule iastModule = InstrumentationBridge.PROPAGATION; - if (iastModule != null && result != null) { - iastModule.taintIfTainted(result, self); + if (iastModule != null) { + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + if (isString) { + iastModule.taintIfTainted(ctx, (String) result, self); + } else { + iastModule.taintIfTainted(ctx, result, self); + } + } } } } @@ -79,15 +91,21 @@ public static class OptAdvice { @Advice.OnMethodExit(suppress = Throwable.class) @Propagation public static void afterMethod(@Advice.This Object self, @Advice.Return final Object result) { - if (result instanceof Integer - || result instanceof Long - || result instanceof Double - || result instanceof Boolean) { + boolean isString = result instanceof String; + boolean isJson = !isString && (result instanceof JSONObject || result instanceof JSONArray); + if (!isString && !isJson) { return; } final PropagationModule iastModule = InstrumentationBridge.PROPAGATION; - if (iastModule != null && result != null) { - iastModule.taintIfTainted(result, self); + if (iastModule != null) { + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + if (isString) { + iastModule.taintIfTainted(ctx, (String) result, self); + } else { + iastModule.taintIfTainted(ctx, result, self); + } + } } } } diff --git a/dd-java-agent/instrumentation/org-json/src/main/java/datadog/trace/instrumentation/json/JSONTokenerInstrumentation.java b/dd-java-agent/instrumentation/org-json/src/main/java/datadog/trace/instrumentation/json/JSONTokenerInstrumentation.java index 60b4d25d490..1789929cc19 100644 --- a/dd-java-agent/instrumentation/org-json/src/main/java/datadog/trace/instrumentation/json/JSONTokenerInstrumentation.java +++ b/dd-java-agent/instrumentation/org-json/src/main/java/datadog/trace/instrumentation/json/JSONTokenerInstrumentation.java @@ -6,6 +6,7 @@ import com.google.auto.service.AutoService; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.api.iast.IastContext; import datadog.trace.api.iast.InstrumentationBridge; import datadog.trace.api.iast.Propagation; import datadog.trace.api.iast.propagation.PropagationModule; @@ -37,7 +38,10 @@ public static class ConstructorAdvice { public static void afterInit(@Advice.This Object self, @Advice.Argument(0) final Object input) { final PropagationModule iastModule = InstrumentationBridge.PROPAGATION; if (iastModule != null && input != null) { - iastModule.taintIfTainted(self, input); + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + iastModule.taintIfTainted(ctx, self, input); + } } } } diff --git a/dd-java-agent/instrumentation/org-json/src/test/groovy/JSONArrayInstrumentationTest.groovy b/dd-java-agent/instrumentation/org-json/src/test/groovy/JSONArrayInstrumentationTest.groovy index 9ab462b021b..ad97575180d 100644 --- a/dd-java-agent/instrumentation/org-json/src/test/groovy/JSONArrayInstrumentationTest.groovy +++ b/dd-java-agent/instrumentation/org-json/src/test/groovy/JSONArrayInstrumentationTest.groovy @@ -1,11 +1,12 @@ -import datadog.trace.agent.test.AgentTestRunner +import com.datadog.iast.test.IastAgentTestRunner +import datadog.trace.api.iast.IastContext import datadog.trace.api.iast.InstrumentationBridge import datadog.trace.api.iast.propagation.PropagationModule import org.json.JSONArray import org.json.JSONObject import org.json.JSONTokener -class JSONArrayInstrumentationTest extends AgentTestRunner { +class JSONArrayInstrumentationTest extends IastAgentTestRunner { @Override void configurePreAgent() { @@ -27,17 +28,20 @@ class JSONArrayInstrumentationTest extends AgentTestRunner { }}""" when: - final jsonObject = new JSONObject(json) - final name = jsonObject.getJSONObject("menu").getJSONArray("labels").get(0) + def name = null + runUnderIastTrace { + final jsonObject = new JSONObject(json) + name = jsonObject.getJSONObject("menu").getJSONArray("labels").get(0) + } then: name == "File" - 1 * module.taintIfTainted(_ as JSONObject, json) - 2 * module.taintIfTainted(_ as JSONObject, _ as JSONTokener) - 2 * module.taintIfTainted(_ as JSONObject, _ as JSONObject) - 1 * module.taintIfTainted(_ as JSONTokener, json) - 2 * module.taintIfTainted(_ as JSONArray, _ as JSONObject) - 2 * module.taintIfTainted("File", _ as JSONArray) + 1 * module.taintObjectIfTainted(_ as IastContext, _ as JSONObject, json) + 2 * module.taintObjectIfTainted(_ as IastContext, _ as JSONObject, _ as JSONTokener) + 2 * module.taintObjectIfTainted(_ as IastContext, _ as JSONObject, _ as JSONObject) + 1 * module.taintObjectIfTainted(_ as IastContext, _ as JSONTokener, json) + 2 * module.taintObjectIfTainted(_ as IastContext, _ as JSONArray, _ as JSONObject) + 2 * module.taintStringIfTainted(_ as IastContext, "File", _ as JSONArray) 0 * _ } @@ -56,17 +60,20 @@ class JSONArrayInstrumentationTest extends AgentTestRunner { }}""" when: - final jsonObject = new JSONObject(json) - final name = jsonObject.getJSONObject("menu").getJSONArray("labels").optString(0, "defaultvalue") + def name = null + runUnderIastTrace { + final jsonObject = new JSONObject(json) + name = jsonObject.getJSONObject("menu").getJSONArray("labels").optString(0, "defaultvalue") + } then: name == "File" - 1 * module.taintIfTainted(_ as JSONObject, json) - 2 * module.taintIfTainted(_ as JSONObject, _ as JSONTokener) - 2 * module.taintIfTainted(_ as JSONObject, _ as JSONObject) - 1 * module.taintIfTainted(_ as JSONTokener, json) - 2 * module.taintIfTainted(_ as JSONArray, _ as JSONObject) - 1 * module.taintIfTainted("File", _ as JSONArray) + 1 * module.taintObjectIfTainted(_ as IastContext, _ as JSONObject, json) + 2 * module.taintObjectIfTainted(_ as IastContext, _ as JSONObject, _ as JSONTokener) + 2 * module.taintObjectIfTainted(_ as IastContext, _ as JSONObject, _ as JSONObject) + 1 * module.taintObjectIfTainted(_ as IastContext, _ as JSONTokener, json) + 2 * module.taintObjectIfTainted(_ as IastContext, _ as JSONArray, _ as JSONObject) + 1 * module.taintStringIfTainted(_ as IastContext, "File", _ as JSONArray) 0 * _ } } diff --git a/dd-java-agent/instrumentation/org-json/src/test/groovy/JSONCookieInstrumentationTest.groovy b/dd-java-agent/instrumentation/org-json/src/test/groovy/JSONCookieInstrumentationTest.groovy index 3849b01e4aa..96c9d4f09f3 100644 --- a/dd-java-agent/instrumentation/org-json/src/test/groovy/JSONCookieInstrumentationTest.groovy +++ b/dd-java-agent/instrumentation/org-json/src/test/groovy/JSONCookieInstrumentationTest.groovy @@ -1,11 +1,12 @@ -import datadog.trace.agent.test.AgentTestRunner +import com.datadog.iast.test.IastAgentTestRunner +import datadog.trace.api.iast.IastContext import datadog.trace.api.iast.InstrumentationBridge import datadog.trace.api.iast.propagation.PropagationModule import org.json.Cookie import org.json.JSONObject import org.json.JSONTokener -class JSONCookieInstrumentationTest extends AgentTestRunner { +class JSONCookieInstrumentationTest extends IastAgentTestRunner { @Override void configurePreAgent() { injectSysConfig("dd.iast.enabled", "true") @@ -18,12 +19,14 @@ class JSONCookieInstrumentationTest extends AgentTestRunner { final cookie = "username = Datadog; expires = Thu, 15 Jun 2020 12:00:00 UTC; path = /" when: - Cookie.toJSONObject(cookie) + runUnderIastTrace { + Cookie.toJSONObject(cookie) + } then: - 1 * module.taintIfTainted(_ as JSONObject, cookie) - 1 * module.taintIfTainted(_ as JSONTokener, cookie) + 1 * module.taintObjectIfTainted(_ as IastContext, _ as JSONObject, cookie) + 1 * module.taintObjectIfTainted(_ as IastContext, _ as JSONTokener, cookie) 0 * _ } } diff --git a/dd-java-agent/instrumentation/org-json/src/test/groovy/JSONObjectInstrumentationTest.groovy b/dd-java-agent/instrumentation/org-json/src/test/groovy/JSONObjectInstrumentationTest.groovy index e1befa3f682..806bade5493 100644 --- a/dd-java-agent/instrumentation/org-json/src/test/groovy/JSONObjectInstrumentationTest.groovy +++ b/dd-java-agent/instrumentation/org-json/src/test/groovy/JSONObjectInstrumentationTest.groovy @@ -1,10 +1,11 @@ -import datadog.trace.agent.test.AgentTestRunner +import com.datadog.iast.test.IastAgentTestRunner +import datadog.trace.api.iast.IastContext import datadog.trace.api.iast.InstrumentationBridge import datadog.trace.api.iast.propagation.PropagationModule import org.json.JSONObject import org.json.JSONTokener -class JSONObjectInstrumentationTest extends AgentTestRunner { +class JSONObjectInstrumentationTest extends IastAgentTestRunner { @Override void configurePreAgent() { injectSysConfig("dd.iast.enabled", "true") @@ -25,16 +26,19 @@ class JSONObjectInstrumentationTest extends AgentTestRunner { }}""" when: - final jsonObject = new JSONObject(json) - final name = jsonObject.getJSONObject("menu").get("name") + String name = null + runUnderIastTrace { + final jsonObject = new JSONObject(json) + name = jsonObject.getJSONObject("menu").get("name") + } then: name == "nameTest" - 1 * module.taintIfTainted(_ as JSONObject, json) - 2 * module.taintIfTainted(_ as JSONObject, _ as JSONTokener) - 2 * module.taintIfTainted(_ as JSONObject, _ as JSONObject) - 1 * module.taintIfTainted(_ as JSONTokener, json) - 2 * module.taintIfTainted("nameTest", _ as JSONObject) + 1 * module.taintObjectIfTainted(_ as IastContext, _ as JSONObject, json) + 2 * module.taintObjectIfTainted(_ as IastContext, _ as JSONObject, _ as JSONTokener) + 2 * module.taintObjectIfTainted(_ as IastContext, _ as JSONObject, _ as JSONObject) + 1 * module.taintObjectIfTainted(_ as IastContext, _ as JSONTokener, json) + 2 * module.taintStringIfTainted(_ as IastContext, "nameTest", _ as JSONObject) 0 * _ } @@ -53,16 +57,19 @@ class JSONObjectInstrumentationTest extends AgentTestRunner { }}""" when: - final jsonObject = new JSONObject(json) - final name = jsonObject.getJSONObject("menu").optString("name") + String name = null + runUnderIastTrace { + final jsonObject = new JSONObject(json) + name = jsonObject.getJSONObject("menu").optString("name") + } then: name == "nameTest" - 1 * module.taintIfTainted(_ as JSONObject, json) - 2 * module.taintIfTainted(_ as JSONObject, _ as JSONTokener) - 2 * module.taintIfTainted(_ as JSONObject, _ as JSONObject) - 1 * module.taintIfTainted(_ as JSONTokener, json) - 1 * module.taintIfTainted("nameTest", _ as JSONObject) + 1 * module.taintObjectIfTainted(_ as IastContext, _ as JSONObject, json) + 2 * module.taintObjectIfTainted(_ as IastContext, _ as JSONObject, _ as JSONTokener) + 2 * module.taintObjectIfTainted(_ as IastContext, _ as JSONObject, _ as JSONObject) + 1 * module.taintObjectIfTainted(_ as IastContext, _ as JSONTokener, json) + 1 * module.taintStringIfTainted(_ as IastContext, "nameTest", _ as JSONObject) 0 * _ } @@ -75,14 +82,17 @@ class JSONObjectInstrumentationTest extends AgentTestRunner { final json = '{"name": "nameTest", "value" : "valueTest"}' when: - final jsonObject = new JSONObject(new JSONTokener(json)) - final name = jsonObject.get("name") + def name = null + runUnderIastTrace { + final jsonObject = new JSONObject(new JSONTokener(json)) + name = jsonObject.get("name") + } then: name == "nameTest" - 1 * module.taintIfTainted(_ as JSONObject, _ as JSONTokener) - 1 * module.taintIfTainted(_ as JSONTokener, json) - 2 * module.taintIfTainted("nameTest", _ as JSONObject) + 1 * module.taintObjectIfTainted(_ as IastContext, _ as JSONObject, _ as JSONTokener) + 1 * module.taintObjectIfTainted(_ as IastContext, _ as JSONTokener, json) + 2 * module.taintStringIfTainted(_ as IastContext, "nameTest", _ as JSONObject) 0 * _ } @@ -96,12 +106,14 @@ class JSONObjectInstrumentationTest extends AgentTestRunner { InstrumentationBridge.registerIastModule(module) when: - final jsonObject = new JSONObject(map) - jsonObject.get("name") + runUnderIastTrace { + final jsonObject = new JSONObject(map) + jsonObject.get("name") + } then: - 1 * module.taintIfTainted(_ as JSONObject, map) - 2 * module.taintIfTainted("nameTest", _ as JSONObject) + 1 * module.taintObjectIfTainted(_ as IastContext, _ as JSONObject, map) + 2 * module.taintStringIfTainted(_ as IastContext, "nameTest", _ as JSONObject) 0 * _ } } diff --git a/dd-java-agent/instrumentation/org-json/src/test/groovy/JSONTokenerInstrumentationTest.groovy b/dd-java-agent/instrumentation/org-json/src/test/groovy/JSONTokenerInstrumentationTest.groovy index b995d34072a..c5cd9aa36de 100644 --- a/dd-java-agent/instrumentation/org-json/src/test/groovy/JSONTokenerInstrumentationTest.groovy +++ b/dd-java-agent/instrumentation/org-json/src/test/groovy/JSONTokenerInstrumentationTest.groovy @@ -1,9 +1,10 @@ -import datadog.trace.agent.test.AgentTestRunner +import com.datadog.iast.test.IastAgentTestRunner +import datadog.trace.api.iast.IastContext import datadog.trace.api.iast.InstrumentationBridge import datadog.trace.api.iast.propagation.PropagationModule import org.json.JSONTokener -class JSONTokenerInstrumentationTest extends AgentTestRunner { +class JSONTokenerInstrumentationTest extends IastAgentTestRunner { @Override void configurePreAgent() { injectSysConfig("dd.iast.enabled", "true") @@ -16,10 +17,12 @@ class JSONTokenerInstrumentationTest extends AgentTestRunner { final json = '{"name": "nameTest", "value" : "valueTest"}' when: - new JSONTokener(json) + runUnderIastTrace { + new JSONTokener(json) + } then: - 1 * module.taintIfTainted(_ as JSONTokener, json) + 1 * module.taintObjectIfTainted(_ as IastContext, _ as JSONTokener, json) 0 * _ } } diff --git a/dd-java-agent/instrumentation/owasp-esapi-2/src/main/java/datadog/trace/instrumentation/owasp/esapi/EncoderCallSite.java b/dd-java-agent/instrumentation/owasp-esapi-2/src/main/java/datadog/trace/instrumentation/owasp/esapi/EncoderCallSite.java index 5f711e0a191..37956f11258 100644 --- a/dd-java-agent/instrumentation/owasp-esapi-2/src/main/java/datadog/trace/instrumentation/owasp/esapi/EncoderCallSite.java +++ b/dd-java-agent/instrumentation/owasp-esapi-2/src/main/java/datadog/trace/instrumentation/owasp/esapi/EncoderCallSite.java @@ -2,6 +2,7 @@ import datadog.trace.agent.tooling.csi.CallSite; import datadog.trace.api.iast.IastCallSites; +import datadog.trace.api.iast.IastContext; import datadog.trace.api.iast.InstrumentationBridge; import datadog.trace.api.iast.Propagation; import datadog.trace.api.iast.VulnerabilityMarks; @@ -20,9 +21,12 @@ public static String afterEncodeForHTML( @CallSite.Argument(0) @Nonnull final String input, @CallSite.Return final String result) { final PropagationModule module = InstrumentationBridge.PROPAGATION; - if (module != null) { + if (result != null && module != null) { try { - module.taintIfTainted(result, input, false, VulnerabilityMarks.XSS_MARK); + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + module.taintIfTainted(ctx, result, input, false, VulnerabilityMarks.XSS_MARK); + } } catch (final Throwable e) { module.onUnexpectedException("afterEncodeForHTML threw", e); } @@ -36,9 +40,12 @@ public static String afterCanonicalize1( @CallSite.Argument(0) @Nonnull final String input, @CallSite.Return final String result) { final PropagationModule module = InstrumentationBridge.PROPAGATION; - if (module != null) { + if (result != null && module != null) { try { - module.taintIfTainted(result, input, false, VulnerabilityMarks.XSS_MARK); + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + module.taintIfTainted(ctx, result, input, false, VulnerabilityMarks.XSS_MARK); + } } catch (final Throwable e) { module.onUnexpectedException("afterCanonicalize1 threw", e); } @@ -54,9 +61,12 @@ public static String afterCanonicalize2( @CallSite.Argument(1) final boolean strict, @CallSite.Return final String result) { final PropagationModule module = InstrumentationBridge.PROPAGATION; - if (module != null) { + if (result != null && module != null) { try { - module.taintIfTainted(result, input, false, VulnerabilityMarks.XSS_MARK); + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + module.taintIfTainted(ctx, result, input, false, VulnerabilityMarks.XSS_MARK); + } } catch (final Throwable e) { module.onUnexpectedException("afterCanonicalize2 threw", e); } @@ -73,9 +83,12 @@ public static String afterCanonicalize3( @CallSite.Argument(2) final boolean restrictMixed, @CallSite.Return final String result) { final PropagationModule module = InstrumentationBridge.PROPAGATION; - if (module != null) { + if (result != null && module != null) { try { - module.taintIfTainted(result, input, false, VulnerabilityMarks.XSS_MARK); + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + module.taintIfTainted(ctx, result, input, false, VulnerabilityMarks.XSS_MARK); + } } catch (final Throwable e) { module.onUnexpectedException("afterCanonicalize3 threw", e); } @@ -89,9 +102,12 @@ public static String afterEncodeForLDAP( @CallSite.Argument(0) @Nonnull final String input, @CallSite.Return final String result) { final PropagationModule module = InstrumentationBridge.PROPAGATION; - if (module != null) { + if (result != null && module != null) { try { - module.taintIfTainted(result, input, false, VulnerabilityMarks.LDAP_INJECTION_MARK); + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + module.taintIfTainted(ctx, result, input, false, VulnerabilityMarks.LDAP_INJECTION_MARK); + } } catch (final Throwable e) { module.onUnexpectedException("afterEncodeForLDAP threw", e); } @@ -107,9 +123,13 @@ public static String afterEncodeForOS( @CallSite.Argument(1) @Nonnull final String input, @CallSite.Return final String result) { final PropagationModule module = InstrumentationBridge.PROPAGATION; - if (module != null) { + if (result != null && module != null) { try { - module.taintIfTainted(result, input, false, VulnerabilityMarks.COMMAND_INJECTION_MARK); + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + module.taintIfTainted( + ctx, result, input, false, VulnerabilityMarks.COMMAND_INJECTION_MARK); + } } catch (final Throwable e) { module.onUnexpectedException("afterEncodeForOS threw", e); } @@ -125,9 +145,12 @@ public static String afterEncodeForSQL( @CallSite.Argument(1) @Nonnull final String input, @CallSite.Return final String result) { final PropagationModule module = InstrumentationBridge.PROPAGATION; - if (module != null) { + if (result != null && module != null) { try { - module.taintIfTainted(result, input, false, VulnerabilityMarks.SQL_INJECTION_MARK); + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + module.taintIfTainted(ctx, result, input, false, VulnerabilityMarks.SQL_INJECTION_MARK); + } } catch (final Throwable e) { module.onUnexpectedException("afterEncodeForSQL threw", e); } diff --git a/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/PathMatcherInstrumentation.java b/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/PathMatcherInstrumentation.java index c71c39b7ae3..0f1ad5d1ba0 100644 --- a/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/PathMatcherInstrumentation.java +++ b/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/PathMatcherInstrumentation.java @@ -56,24 +56,23 @@ static void onExit( return; } + scala.Tuple1 tuple = (scala.Tuple1) extractions; + Object value = tuple._1(); + PropagationModule module = InstrumentationBridge.PROPAGATION; - if (module == null) { + if (module == null || !(value instanceof String)) { return; } - - scala.Tuple1 tuple = (scala.Tuple1) extractions; - Object value = tuple._1(); + final String stringValue = (String) value; IastContext ctx = reqCtx.getData(RequestContextSlot.IAST); // in the test, 4 instances of PathMatcher$Match are created, all with the same value - if (module.isTainted(ctx, value)) { + if (module.isTainted(ctx, stringValue)) { return; } - if (value instanceof String) { - module.taint(ctx, value, SourceTypes.REQUEST_PATH_PARAMETER); - } + module.taint(ctx, stringValue, SourceTypes.REQUEST_PATH_PARAMETER); } } } diff --git a/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintParametersFunction.java b/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintParametersFunction.java index ba805419050..e77667a6f72 100644 --- a/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintParametersFunction.java +++ b/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintParametersFunction.java @@ -44,11 +44,11 @@ public Tuple1 apply(Tuple1 v1) { while (iterator.hasNext()) { Object o = iterator.next(); if (o instanceof String) { - mod.taint(ctx, o, SourceTypes.REQUEST_PARAMETER_VALUE, paramName); + mod.taint(ctx, (String) o, SourceTypes.REQUEST_PARAMETER_VALUE, paramName); } } } else if (value instanceof String) { - mod.taint(ctx, value, SourceTypes.REQUEST_PARAMETER_VALUE, paramName); + mod.taint(ctx, (String) value, SourceTypes.REQUEST_PARAMETER_VALUE, paramName); } return v1; diff --git a/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintRequestContextFunction.java b/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintRequestContextFunction.java index 2acbb5a0778..a7e307d5d3e 100644 --- a/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintRequestContextFunction.java +++ b/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintRequestContextFunction.java @@ -18,7 +18,7 @@ public Tuple1 apply(Tuple1 v1) { RequestContext reqCtx = v1._1(); PropagationModule mod = InstrumentationBridge.PROPAGATION; - if (mod == null) { + if (mod == null || reqCtx == null) { return v1; } IastContext ctx = IastContext.Provider.get(AgentTracer.activeSpan()); diff --git a/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintRequestFunction.java b/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintRequestFunction.java index b350d9de377..37b223272d0 100644 --- a/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintRequestFunction.java +++ b/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintRequestFunction.java @@ -19,7 +19,7 @@ public Tuple1 apply(Tuple1 v1) { HttpRequest httpRequest = v1._1(); PropagationModule mod = InstrumentationBridge.PROPAGATION; - if (mod == null) { + if (mod == null || httpRequest == null) { return v1; } IastContext ctx = IastContext.Provider.get(AgentTracer.activeSpan()); diff --git a/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintSingleParameterFunction.java b/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintSingleParameterFunction.java index e22a8729e4d..710cca86977 100644 --- a/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintSingleParameterFunction.java +++ b/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintSingleParameterFunction.java @@ -51,11 +51,11 @@ public Tuple1 apply(Tuple1 v1) { while (iterator.hasNext()) { Object o = iterator.next(); if (o instanceof String) { - mod.taint(ctx, o, SourceTypes.REQUEST_PARAMETER_VALUE, paramName); + mod.taint(ctx, (String) o, SourceTypes.REQUEST_PARAMETER_VALUE, paramName); } } } else if (value instanceof String) { - mod.taint(ctx, value, SourceTypes.REQUEST_PARAMETER_VALUE, paramName); + mod.taint(ctx, (String) value, SourceTypes.REQUEST_PARAMETER_VALUE, paramName); } return v1; diff --git a/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintUnmarshaller.java b/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintUnmarshaller.java index 3c377b3d05d..c9019a5130b 100644 --- a/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintUnmarshaller.java +++ b/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintUnmarshaller.java @@ -30,9 +30,11 @@ public TaintUnmarshaller(PropagationModule propagationModule, Unmarshaller @Override public Future apply(A value, ExecutionContext ec, Materializer materializer) { - IastContext ctx = IastContext.Provider.get(AgentTracer.activeSpan()); - if (ctx != null) { - propagationModule.taint(ctx, value, SourceTypes.REQUEST_BODY); + if (value != null) { + IastContext ctx = IastContext.Provider.get(AgentTracer.activeSpan()); + if (ctx != null) { + propagationModule.taint(ctx, value, SourceTypes.REQUEST_BODY); + } } return delegate.apply(value, ec, materializer); } diff --git a/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintUriFunction.java b/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintUriFunction.java index e51133878b0..3bbdd71a662 100644 --- a/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintUriFunction.java +++ b/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintUriFunction.java @@ -17,7 +17,7 @@ public Tuple1 apply(Tuple1 v1) { Uri uri = v1._1(); PropagationModule mod = InstrumentationBridge.PROPAGATION; - if (mod == null) { + if (mod == null || uri == null) { return v1; } IastContext ctx = IastContext.Provider.get(AgentTracer.activeSpan()); diff --git a/dd-java-agent/instrumentation/resteasy-appsec/src/main/java/datadog/trace/instrumentation/resteasy/CookieParamInjectorAdvice.java b/dd-java-agent/instrumentation/resteasy-appsec/src/main/java/datadog/trace/instrumentation/resteasy/CookieParamInjectorAdvice.java index 96fcbca06ca..81c7f536105 100644 --- a/dd-java-agent/instrumentation/resteasy-appsec/src/main/java/datadog/trace/instrumentation/resteasy/CookieParamInjectorAdvice.java +++ b/dd-java-agent/instrumentation/resteasy-appsec/src/main/java/datadog/trace/instrumentation/resteasy/CookieParamInjectorAdvice.java @@ -28,11 +28,11 @@ public static void onExit( Collection collection = (Collection) result; for (Object o : collection) { if (o instanceof String) { - module.taint(ctx, o, SourceTypes.REQUEST_COOKIE_VALUE, paramName); + module.taint(ctx, (String) o, SourceTypes.REQUEST_COOKIE_VALUE, paramName); } } } else { - module.taint(ctx, result, SourceTypes.REQUEST_COOKIE_VALUE, paramName); + module.taint(ctx, (String) result, SourceTypes.REQUEST_COOKIE_VALUE, paramName); } } } diff --git a/dd-java-agent/instrumentation/resteasy-appsec/src/main/java/datadog/trace/instrumentation/resteasy/FormParamInjectorAdvice.java b/dd-java-agent/instrumentation/resteasy-appsec/src/main/java/datadog/trace/instrumentation/resteasy/FormParamInjectorAdvice.java index 5c481697db6..9f1fff61394 100644 --- a/dd-java-agent/instrumentation/resteasy-appsec/src/main/java/datadog/trace/instrumentation/resteasy/FormParamInjectorAdvice.java +++ b/dd-java-agent/instrumentation/resteasy-appsec/src/main/java/datadog/trace/instrumentation/resteasy/FormParamInjectorAdvice.java @@ -28,11 +28,11 @@ public static void onExit( Collection collection = (Collection) result; for (Object o : collection) { if (o instanceof String) { - module.taint(ctx, o, SourceTypes.REQUEST_PARAMETER_VALUE, paramName); + module.taint(ctx, (String) o, SourceTypes.REQUEST_PARAMETER_VALUE, paramName); } } } else { - module.taint(ctx, result, SourceTypes.REQUEST_PARAMETER_VALUE, paramName); + module.taint(ctx, (String) result, SourceTypes.REQUEST_PARAMETER_VALUE, paramName); } } } diff --git a/dd-java-agent/instrumentation/resteasy-appsec/src/main/java/datadog/trace/instrumentation/resteasy/HeaderParamInjectorAdvice.java b/dd-java-agent/instrumentation/resteasy-appsec/src/main/java/datadog/trace/instrumentation/resteasy/HeaderParamInjectorAdvice.java index 25c46edb3d5..e3a9799739a 100644 --- a/dd-java-agent/instrumentation/resteasy-appsec/src/main/java/datadog/trace/instrumentation/resteasy/HeaderParamInjectorAdvice.java +++ b/dd-java-agent/instrumentation/resteasy-appsec/src/main/java/datadog/trace/instrumentation/resteasy/HeaderParamInjectorAdvice.java @@ -23,20 +23,16 @@ public static void onExit( if (result instanceof String || result instanceof Collection) { final PropagationModule module = InstrumentationBridge.PROPAGATION; if (module != null) { - try { - IastContext ctx = reqCtx.getData(RequestContextSlot.IAST); - if (result instanceof Collection) { - Collection collection = (Collection) result; - for (Object o : collection) { - if (o instanceof String) { - module.taint(ctx, o, SourceTypes.REQUEST_HEADER_VALUE, paramName); - } + IastContext ctx = reqCtx.getData(RequestContextSlot.IAST); + if (result instanceof Collection) { + Collection collection = (Collection) result; + for (Object o : collection) { + if (o instanceof String) { + module.taint(ctx, (String) o, SourceTypes.REQUEST_HEADER_VALUE, paramName); } - } else { - module.taint(ctx, result, SourceTypes.REQUEST_HEADER_VALUE, paramName); } - } catch (final Throwable e) { - module.onUnexpectedException("HeaderParamInjectorAdvice.onExit threw", e); + } else { + module.taint(ctx, (String) result, SourceTypes.REQUEST_HEADER_VALUE, paramName); } } } diff --git a/dd-java-agent/instrumentation/resteasy-appsec/src/main/java/datadog/trace/instrumentation/resteasy/PathParamInjectorAdvice.java b/dd-java-agent/instrumentation/resteasy-appsec/src/main/java/datadog/trace/instrumentation/resteasy/PathParamInjectorAdvice.java index 96fd981adc1..1469d06996c 100644 --- a/dd-java-agent/instrumentation/resteasy-appsec/src/main/java/datadog/trace/instrumentation/resteasy/PathParamInjectorAdvice.java +++ b/dd-java-agent/instrumentation/resteasy-appsec/src/main/java/datadog/trace/instrumentation/resteasy/PathParamInjectorAdvice.java @@ -28,11 +28,11 @@ public static void onExit( Collection collection = (Collection) result; for (Object o : collection) { if (o instanceof String) { - module.taint(ctx, o, SourceTypes.REQUEST_PATH_PARAMETER, paramName); + module.taint(ctx, (String) o, SourceTypes.REQUEST_PATH_PARAMETER, paramName); } } } else { - module.taint(ctx, result, SourceTypes.REQUEST_PATH_PARAMETER, paramName); + module.taint(ctx, (String) result, SourceTypes.REQUEST_PATH_PARAMETER, paramName); } } } diff --git a/dd-java-agent/instrumentation/resteasy-appsec/src/main/java/datadog/trace/instrumentation/resteasy/QueryParamInjectorAdvice.java b/dd-java-agent/instrumentation/resteasy-appsec/src/main/java/datadog/trace/instrumentation/resteasy/QueryParamInjectorAdvice.java index c5c07b27ef2..1271746f377 100644 --- a/dd-java-agent/instrumentation/resteasy-appsec/src/main/java/datadog/trace/instrumentation/resteasy/QueryParamInjectorAdvice.java +++ b/dd-java-agent/instrumentation/resteasy-appsec/src/main/java/datadog/trace/instrumentation/resteasy/QueryParamInjectorAdvice.java @@ -28,11 +28,11 @@ public static void onExit( Collection collection = (Collection) result; for (Object o : collection) { if (o instanceof String) { - module.taint(ctx, o, SourceTypes.REQUEST_PARAMETER_VALUE, paramName); + module.taint(ctx, (String) o, SourceTypes.REQUEST_PARAMETER_VALUE, paramName); } } } else { - module.taint(ctx, result, SourceTypes.REQUEST_PARAMETER_VALUE, paramName); + module.taint(ctx, (String) result, SourceTypes.REQUEST_PARAMETER_VALUE, paramName); } } } diff --git a/dd-java-agent/instrumentation/servlet-common/src/main/java/datadog/trace/instrumentation/servlet/HttpServletResponseInstrumentation.java b/dd-java-agent/instrumentation/servlet-common/src/main/java/datadog/trace/instrumentation/servlet/HttpServletResponseInstrumentation.java index 79213fde074..ba478591a6e 100644 --- a/dd-java-agent/instrumentation/servlet-common/src/main/java/datadog/trace/instrumentation/servlet/HttpServletResponseInstrumentation.java +++ b/dd-java-agent/instrumentation/servlet-common/src/main/java/datadog/trace/instrumentation/servlet/HttpServletResponseInstrumentation.java @@ -12,6 +12,7 @@ import com.google.auto.service.AutoService; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.api.iast.IastContext; import datadog.trace.api.iast.InstrumentationBridge; import datadog.trace.api.iast.Propagation; import datadog.trace.api.iast.Sink; @@ -120,9 +121,10 @@ public static class EncodeURLAdvice { @Propagation public static void onExit(@Advice.Argument(0) final String url, @Advice.Return String encoded) { final PropagationModule module = InstrumentationBridge.PROPAGATION; - if (module != null) { - if (null != url && !url.isEmpty() && null != encoded && !encoded.isEmpty()) { - module.taintIfTainted(encoded, url); + if (encoded != null && module != null) { + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + module.taintIfTainted(ctx, encoded, url); } } } diff --git a/dd-java-agent/instrumentation/servlet-common/src/test/groovy/HttpServletResponseInstrumentationTest.groovy b/dd-java-agent/instrumentation/servlet-common/src/test/groovy/HttpServletResponseInstrumentationTest.groovy index b11e9aee9ab..697e042d349 100644 --- a/dd-java-agent/instrumentation/servlet-common/src/test/groovy/HttpServletResponseInstrumentationTest.groovy +++ b/dd-java-agent/instrumentation/servlet-common/src/test/groovy/HttpServletResponseInstrumentationTest.groovy @@ -1,9 +1,12 @@ import datadog.trace.agent.test.AgentTestRunner +import datadog.trace.api.iast.IastContext import datadog.trace.api.iast.InstrumentationBridge import datadog.trace.api.iast.propagation.PropagationModule import datadog.trace.api.iast.sink.HttpResponseHeaderModule import datadog.trace.api.iast.sink.UnvalidatedRedirectModule import datadog.trace.api.iast.util.Cookie as IastCookie +import datadog.trace.bootstrap.instrumentation.api.AgentTracer +import datadog.trace.bootstrap.instrumentation.api.TagContext import foo.bar.DummyResponse import javax.servlet.http.Cookie @@ -178,11 +181,11 @@ class HttpServletResponseInstrumentationTest extends AgentTestRunner { final response = new DummyResponse() when: - response.encodeRedirectURL("http://dummy.url.com") + runUnderIastTrace { response.encodeRedirectURL("http://dummy.url.com") } then: noExceptionThrown() - 1 * module.taintIfTainted(_, "http://dummy.url.com") + 1 * module.taintIfTainted(_ as IastContext, _, "http://dummy.url.com") 0 * _ } @@ -193,11 +196,11 @@ class HttpServletResponseInstrumentationTest extends AgentTestRunner { final response = new DummyResponse() when: - response.encodeURL("http://dummy.url.com") + runUnderIastTrace { response.encodeURL("http://dummy.url.com") } then: noExceptionThrown() - 1 * module.taintIfTainted(_, "http://dummy.url.com") + 1 * module.taintIfTainted(_ as IastContext, _, "http://dummy.url.com") 0 * _ } @@ -225,4 +228,14 @@ class HttpServletResponseInstrumentationTest extends AgentTestRunner { then: 0 * _ } + + protected E runUnderIastTrace(Closure cl) { + final ddctx = new TagContext().withRequestContextDataIast(Stub(IastContext)) + final span = TEST_TRACER.startSpan("test", "test-iast-span", ddctx) + try { + return AgentTracer.activateSpan(span).withCloseable(cl) + } finally { + span.finish() + } + } } diff --git a/dd-java-agent/instrumentation/servlet/request-5/src/main/java/datadog/trace/instrumentation/servlet5/JakartaHttpServletResponseInstrumentation.java b/dd-java-agent/instrumentation/servlet/request-5/src/main/java/datadog/trace/instrumentation/servlet5/JakartaHttpServletResponseInstrumentation.java index 5dac167fec3..22e926ecff9 100644 --- a/dd-java-agent/instrumentation/servlet/request-5/src/main/java/datadog/trace/instrumentation/servlet5/JakartaHttpServletResponseInstrumentation.java +++ b/dd-java-agent/instrumentation/servlet/request-5/src/main/java/datadog/trace/instrumentation/servlet5/JakartaHttpServletResponseInstrumentation.java @@ -12,6 +12,7 @@ import com.google.auto.service.AutoService; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.api.iast.IastContext; import datadog.trace.api.iast.InstrumentationBridge; import datadog.trace.api.iast.Propagation; import datadog.trace.api.iast.Sink; @@ -118,9 +119,10 @@ public static class EncodeURLAdvice { @Propagation public static void onExit(@Advice.Argument(0) final String url, @Advice.Return String encoded) { final PropagationModule module = InstrumentationBridge.PROPAGATION; - if (module != null) { - if (null != url && !url.isEmpty() && null != encoded && !encoded.isEmpty()) { - module.taintIfTainted(encoded, url); + if (encoded != null && module != null) { + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + module.taintIfTainted(ctx, encoded, url); } } } diff --git a/dd-java-agent/instrumentation/spring-core/src/main/java/datadog/trace/instrumentation/springcore/StreamUtilsInstrumentation.java b/dd-java-agent/instrumentation/spring-core/src/main/java/datadog/trace/instrumentation/springcore/StreamUtilsInstrumentation.java index 232faaa742b..acc820fed9c 100644 --- a/dd-java-agent/instrumentation/spring-core/src/main/java/datadog/trace/instrumentation/springcore/StreamUtilsInstrumentation.java +++ b/dd-java-agent/instrumentation/spring-core/src/main/java/datadog/trace/instrumentation/springcore/StreamUtilsInstrumentation.java @@ -8,6 +8,7 @@ import com.google.auto.service.AutoService; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.api.iast.IastContext; import datadog.trace.api.iast.InstrumentationBridge; import datadog.trace.api.iast.Propagation; import datadog.trace.api.iast.propagation.PropagationModule; @@ -49,8 +50,11 @@ public static class SpringAdvice { public static void checkReturnedObject( @Advice.Return String string, @Advice.Argument(0) final InputStream in) { final PropagationModule module = InstrumentationBridge.PROPAGATION; - if (in != null && string != null && !string.isEmpty()) { - module.taintIfTainted(string, in); + if (string != null && module != null) { + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + module.taintIfTainted(ctx, string, in); + } } } diff --git a/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/DataBufferAsInputStreamAdvice.java b/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/DataBufferAsInputStreamAdvice.java index 3322c5a4461..9a65d7aa5d7 100644 --- a/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/DataBufferAsInputStreamAdvice.java +++ b/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/DataBufferAsInputStreamAdvice.java @@ -1,5 +1,6 @@ package datadog.trace.instrumentation.springwebflux.server.iast; +import datadog.trace.api.iast.IastContext; import datadog.trace.api.iast.InstrumentationBridge; import datadog.trace.api.iast.Propagation; import datadog.trace.api.iast.propagation.PropagationModule; @@ -18,6 +19,9 @@ public static void after(@Advice.This DataBuffer dataBuffer, @Advice.Return Inpu return; } - mod.taintIfTainted(is, dataBuffer); + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + mod.taintIfTainted(ctx, is, dataBuffer); + } } } diff --git a/dd-java-agent/instrumentation/spring-webmvc-5.3/src/main/java/datadog/trace/instrumentation/springweb/EscapeUtilsCallSite.java b/dd-java-agent/instrumentation/spring-webmvc-5.3/src/main/java/datadog/trace/instrumentation/springweb/EscapeUtilsCallSite.java index 5db0c0e4492..98aa9d271c4 100644 --- a/dd-java-agent/instrumentation/spring-webmvc-5.3/src/main/java/datadog/trace/instrumentation/springweb/EscapeUtilsCallSite.java +++ b/dd-java-agent/instrumentation/spring-webmvc-5.3/src/main/java/datadog/trace/instrumentation/springweb/EscapeUtilsCallSite.java @@ -2,6 +2,7 @@ import datadog.trace.agent.tooling.csi.CallSite; import datadog.trace.api.iast.IastCallSites; +import datadog.trace.api.iast.IastContext; import datadog.trace.api.iast.InstrumentationBridge; import datadog.trace.api.iast.Propagation; import datadog.trace.api.iast.VulnerabilityMarks; @@ -19,9 +20,12 @@ public class EscapeUtilsCallSite { public static String afterEscape( @CallSite.Argument(0) @Nullable final String input, @CallSite.Return final String result) { final PropagationModule module = InstrumentationBridge.PROPAGATION; - if (module != null) { + if (result != null && module != null) { try { - module.taintIfTainted(result, input, false, VulnerabilityMarks.XSS_MARK); + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + module.taintIfTainted(ctx, result, input, false, VulnerabilityMarks.XSS_MARK); + } } catch (final Throwable e) { module.onUnexpectedException("afterEscape threw", e); } @@ -36,9 +40,12 @@ public static String afterHtmlEscape2( @CallSite.Argument(1) @Nullable final String encoding, @CallSite.Return final String result) { final PropagationModule module = InstrumentationBridge.PROPAGATION; - if (module != null) { + if (result != null && module != null) { try { - module.taintIfTainted(result, input, false, VulnerabilityMarks.XSS_MARK); + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + module.taintIfTainted(ctx, result, input, false, VulnerabilityMarks.XSS_MARK); + } } catch (final Throwable e) { module.onUnexpectedException("afterHtmlEscape2 threw", e); } diff --git a/dd-java-agent/instrumentation/unbescape/src/main/java/datadog/trace/instrumentation/unbescape/EscapeUtilsCallSite.java b/dd-java-agent/instrumentation/unbescape/src/main/java/datadog/trace/instrumentation/unbescape/EscapeUtilsCallSite.java index 6a70cf3c71b..be1e3e3796e 100644 --- a/dd-java-agent/instrumentation/unbescape/src/main/java/datadog/trace/instrumentation/unbescape/EscapeUtilsCallSite.java +++ b/dd-java-agent/instrumentation/unbescape/src/main/java/datadog/trace/instrumentation/unbescape/EscapeUtilsCallSite.java @@ -2,6 +2,7 @@ import datadog.trace.agent.tooling.csi.CallSite; import datadog.trace.api.iast.IastCallSites; +import datadog.trace.api.iast.IastContext; import datadog.trace.api.iast.InstrumentationBridge; import datadog.trace.api.iast.Propagation; import datadog.trace.api.iast.VulnerabilityMarks; @@ -20,14 +21,15 @@ public class EscapeUtilsCallSite { "java.lang.String org.unbescape.javascript.JavaScriptEscape.escapeJavaScript(java.lang.String)") public static String afterEscape( @CallSite.Argument(0) @Nullable final String input, @CallSite.Return final String result) { - if (input != null && result != null) { - final PropagationModule module = InstrumentationBridge.PROPAGATION; - if (module != null) { - try { - module.taintIfTainted(result, input, false, VulnerabilityMarks.XSS_MARK); - } catch (final Throwable e) { - module.onUnexpectedException("afterEscape threw", e); + final PropagationModule module = InstrumentationBridge.PROPAGATION; + if (result != null && module != null) { + try { + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + module.taintIfTainted(ctx, result, input, false, VulnerabilityMarks.XSS_MARK); } + } catch (final Throwable e) { + module.onUnexpectedException("afterEscape threw", e); } } return result; diff --git a/dd-java-agent/instrumentation/vertx-web-3.4/src/main/java/datadog/trace/instrumentation/vertx_3_4/core/BufferInstrumentation.java b/dd-java-agent/instrumentation/vertx-web-3.4/src/main/java/datadog/trace/instrumentation/vertx_3_4/core/BufferInstrumentation.java index 2ce7e7c9de9..3ca1be1ea9e 100644 --- a/dd-java-agent/instrumentation/vertx-web-3.4/src/main/java/datadog/trace/instrumentation/vertx_3_4/core/BufferInstrumentation.java +++ b/dd-java-agent/instrumentation/vertx-web-3.4/src/main/java/datadog/trace/instrumentation/vertx_3_4/core/BufferInstrumentation.java @@ -13,6 +13,7 @@ import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.agent.tooling.bytebuddy.iast.TaintableVisitor; import datadog.trace.agent.tooling.muzzle.Reference; +import datadog.trace.api.iast.IastContext; import datadog.trace.api.iast.InstrumentationBridge; import datadog.trace.api.iast.Propagation; import datadog.trace.api.iast.propagation.PropagationModule; @@ -60,46 +61,43 @@ public void methodAdvice(final MethodTransformer transformer) { } public static class ToStringAdvice { - @Advice.OnMethodExit + @Advice.OnMethodExit(suppress = Throwable.class) @Propagation public static void get(@Advice.This final Object self, @Advice.Return final String result) { final PropagationModule module = InstrumentationBridge.PROPAGATION; - if (module != null) { - try { - module.taintIfTainted(result, self); - } catch (final Throwable e) { - module.onUnexpectedException("toString threw", e); + if (result != null && module != null) { + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + module.taintIfTainted(ctx, result, self); } } } } public static class GetByteBuffAdvice { - @Advice.OnMethodExit + @Advice.OnMethodExit(suppress = Throwable.class) @Propagation public static void get(@Advice.This final Object self, @Advice.Return final Object result) { final PropagationModule module = InstrumentationBridge.PROPAGATION; - if (module != null) { - try { - module.taintIfTainted(result, self); - } catch (final Throwable e) { - module.onUnexpectedException("getByteBuf threw", e); + if (result != null && module != null) { + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + module.taintIfTainted(ctx, result, self); } } } } public static class AppendBufferAdvice { - @Advice.OnMethodExit + @Advice.OnMethodExit(suppress = Throwable.class) @Propagation public static void get( @Advice.Argument(0) final Object buffer, @Advice.Return final Object result) { final PropagationModule module = InstrumentationBridge.PROPAGATION; - if (module != null) { - try { - module.taintIfTainted(result, buffer); - } catch (final Throwable e) { - module.onUnexpectedException("appendBuffer threw", e); + if (result != null && module != null) { + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + module.taintIfTainted(ctx, result, buffer); } } } diff --git a/dd-java-agent/instrumentation/vertx-web-3.4/src/main/java/datadog/trace/instrumentation/vertx_3_4/core/HttpServerRequestInstrumentation.java b/dd-java-agent/instrumentation/vertx-web-3.4/src/main/java/datadog/trace/instrumentation/vertx_3_4/core/HttpServerRequestInstrumentation.java index 383d1e55056..b75563f3cd4 100644 --- a/dd-java-agent/instrumentation/vertx-web-3.4/src/main/java/datadog/trace/instrumentation/vertx_3_4/core/HttpServerRequestInstrumentation.java +++ b/dd-java-agent/instrumentation/vertx-web-3.4/src/main/java/datadog/trace/instrumentation/vertx_3_4/core/HttpServerRequestInstrumentation.java @@ -59,7 +59,7 @@ public static void onExit( @Advice.Return final Object multiMap, @ActiveRequestContext RequestContext reqCtx) { // only taint the map the first time - if (beforeHeaders != multiMap) { + if (beforeHeaders != multiMap && multiMap != null) { final PropagationModule module = InstrumentationBridge.PROPAGATION; if (module != null) { final IastContext ctx = reqCtx.getData(RequestContextSlot.IAST); diff --git a/dd-java-agent/instrumentation/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/core/BufferInstrumentation.java b/dd-java-agent/instrumentation/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/core/BufferInstrumentation.java index a7774fcf3f9..e3c4ecab3a7 100644 --- a/dd-java-agent/instrumentation/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/core/BufferInstrumentation.java +++ b/dd-java-agent/instrumentation/vertx-web-4.0/src/main/java/datadog/trace/instrumentation/vertx_4_0/core/BufferInstrumentation.java @@ -12,6 +12,7 @@ import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.agent.tooling.bytebuddy.iast.TaintableVisitor; import datadog.trace.agent.tooling.muzzle.Reference; +import datadog.trace.api.iast.IastContext; import datadog.trace.api.iast.InstrumentationBridge; import datadog.trace.api.iast.Propagation; import datadog.trace.api.iast.propagation.PropagationModule; @@ -62,9 +63,16 @@ public static class ToStringAdvice { @Advice.OnMethodExit(suppress = Throwable.class) @Propagation public static void get(@Advice.This final Object self, @Advice.Return final String result) { + if (result == null) { + return; + } final PropagationModule module = InstrumentationBridge.PROPAGATION; - if (module != null) { - module.taintIfTainted(result, self); + if (module == null) { + return; + } + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + module.taintIfTainted(ctx, result, self); } } } @@ -73,9 +81,16 @@ public static class GetByteBuffAdvice { @Advice.OnMethodExit(suppress = Throwable.class) @Propagation public static void get(@Advice.This final Object self, @Advice.Return final Object result) { + if (result == null) { + return; + } final PropagationModule module = InstrumentationBridge.PROPAGATION; - if (module != null) { - module.taintIfTainted(result, self); + if (module == null) { + return; + } + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + module.taintIfTainted(ctx, result, self); } } } @@ -85,9 +100,16 @@ public static class AppendBufferAdvice { @Propagation public static void get( @Advice.Argument(0) final Object buffer, @Advice.Return final Object result) { + if (result == null) { + return; + } final PropagationModule module = InstrumentationBridge.PROPAGATION; - if (module != null) { - module.taintIfTainted(result, buffer); + if (module == null) { + return; + } + final IastContext ctx = IastContext.Provider.get(); + if (ctx != null) { + module.taintIfTainted(ctx, result, buffer); } } } diff --git a/internal-api/src/main/java/datadog/trace/api/iast/propagation/PropagationModule.java b/internal-api/src/main/java/datadog/trace/api/iast/propagation/PropagationModule.java index 250a3962a69..bd276ced00d 100644 --- a/internal-api/src/main/java/datadog/trace/api/iast/propagation/PropagationModule.java +++ b/internal-api/src/main/java/datadog/trace/api/iast/propagation/PropagationModule.java @@ -1,73 +1,109 @@ package datadog.trace.api.iast.propagation; +import static datadog.trace.api.iast.VulnerabilityMarks.NOT_MARKED; + import datadog.trace.api.iast.IastContext; import datadog.trace.api.iast.IastModule; import datadog.trace.api.iast.Taintable.Source; import java.util.function.Predicate; +import javax.annotation.Nonnull; import javax.annotation.Nullable; /** Main API for propagation of tainted values, */ @SuppressWarnings("unused") public interface PropagationModule extends IastModule { - /** @see #taint(IastContext, Object, byte) */ - void taint(@Nullable Object target, byte origin); - /** * Taints the object with a source with the selected origin and no name, if target is a char * sequence it will be used as value */ - void taint(@Nullable IastContext ctx, @Nullable Object target, byte origin); - - /** @see #taint(IastContext, Object, byte, int, int) */ - void taint(@Nullable Object target, byte origin, int start, int length); - - /** - * Taints the object with a source with the selected origin, range and no name. If target is a - * char sequence it will be used as value. - * - *

If the value is already tainted this method will append a new range. - */ - void taint( - @Nullable IastContext ctx, @Nullable Object target, byte origin, int start, int length); + default void taint(@Nonnull IastContext ctx, @Nullable Object target, byte origin) { + taint(ctx, target, origin, null); + } - /** @see #taint(IastContext, Object, byte, CharSequence) */ - void taint(@Nullable Object target, byte origin, @Nullable CharSequence name); + /** @see #taint(IastContext, Object, byte) */ + default void taint(@Nonnull IastContext ctx, @Nullable String target, byte origin) { + taint(ctx, target, origin, null); + } /** * Taints the object with a source with the selected origin and name, if target is a char sequence * it will be used as value */ - void taint( - @Nullable IastContext ctx, @Nullable Object target, byte origin, @Nullable CharSequence name); + default void taint( + @Nonnull IastContext ctx, @Nullable Object target, byte origin, @Nullable CharSequence name) { + taint(ctx, target, origin, name, target); + } - /** @see #taint(IastContext, Object, byte, CharSequence, Object) */ - void taint( - @Nullable Object target, byte origin, @Nullable CharSequence name, @Nullable Object value); + /** @see #taint(IastContext, Object, byte, CharSequence) */ + default void taint( + @Nonnull IastContext ctx, @Nullable String target, byte origin, @Nullable CharSequence name) { + taint(ctx, target, origin, name, target); + } /** Taints the object with a source with the selected origin, name and value */ void taint( - @Nullable IastContext ctx, + @Nonnull IastContext ctx, @Nullable Object target, byte origin, @Nullable CharSequence name, @Nullable Object value); - /** @see #taintIfTainted(IastContext, Object, Object) */ - void taintIfTainted(@Nullable Object target, @Nullable Object input); + /** @see #taint(IastContext, Object, byte, CharSequence, Object) */ + void taint( + @Nonnull IastContext ctx, + @Nullable String target, + byte origin, + @Nullable CharSequence name, + @Nullable CharSequence value); + + /** + * Taints the object with a source with the selected origin, range and no name. If target is a + * char sequence it will be used as value. + * + *

If the value is already tainted this method will append a new range. + */ + void taint(@Nonnull IastContext ctx, @Nullable Object target, byte origin, int start, int length); + + /** @see #taint(IastContext, Object, byte, int, int) */ + void taint(@Nonnull IastContext ctx, @Nullable String target, byte origin, int start, int length); /** * Taints the object only if the input value is tainted. If tainted, it will use the highest * priority source of the input to taint the object. */ - void taintIfTainted(@Nullable IastContext ctx, @Nullable Object target, @Nullable Object input); + default void taintIfTainted( + @Nonnull IastContext ctx, @Nullable Object target, @Nullable Object input) { + taintIfTainted(ctx, target, input, false, NOT_MARKED); + } + + /** @see #taintIfTainted(IastContext, Object, Object) */ + default void taintIfTainted( + @Nonnull IastContext ctx, @Nullable String target, @Nullable Object input) { + taintIfTainted(ctx, target, input, false, NOT_MARKED); + } - /** @see #taintIfTainted(IastContext, Object, Object, int, int, boolean, int) */ + /** + * Taints the object only if the input value is tainted. It will try to reuse sources from the + * input value according to: + * + *

    + *
  • keepRanges=true will reuse the ranges from the input tainted value and mark them + *
  • keepRanges=false will use the highest priority source from the input ranges and mark it + *
+ */ void taintIfTainted( + @Nonnull IastContext ctx, @Nullable Object target, @Nullable Object input, - int start, - int length, + boolean keepRanges, + int mark); + + /** @see #taintIfTainted(IastContext, Object, Object, boolean, int) */ + void taintIfTainted( + @Nonnull IastContext ctx, + @Nullable String target, + @Nullable Object input, boolean keepRanges, int mark); @@ -81,7 +117,7 @@ void taintIfTainted( * */ void taintIfTainted( - @Nullable IastContext ctx, + @Nonnull IastContext ctx, @Nullable Object target, @Nullable Object input, int start, @@ -91,73 +127,74 @@ void taintIfTainted( /** @see #taintIfTainted(IastContext, Object, Object, boolean, int) */ void taintIfTainted( - @Nullable Object target, @Nullable Object input, boolean keepRanges, int mark); - - /** - * Taints the object only if the input value is tainted. It will try to reuse sources from the - * input value according to: - * - *
    - *
  • keepRanges=true will reuse the ranges from the input tainted value and mark them - *
  • keepRanges=false will use the highest priority source from the input ranges and mark it - *
- */ - void taintIfTainted( - @Nullable IastContext ctx, - @Nullable Object target, + @Nonnull IastContext ctx, + @Nullable String target, @Nullable Object input, + int start, + int length, boolean keepRanges, int mark); - /** @see #taintIfTainted(IastContext, Object, Object, byte) */ - void taintIfTainted(@Nullable Object target, @Nullable Object input, byte origin); - /** * Taints the object only if the input value is tainted, the resulting value will be tainted using * a source with the specified origin and no name, if target is a char sequence it will be used as * value */ - void taintIfTainted( - @Nullable IastContext ctx, @Nullable Object target, @Nullable Object input, byte origin); + default void taintIfTainted( + @Nonnull IastContext ctx, @Nullable Object target, @Nullable Object input, byte origin) { + taintIfTainted(ctx, target, input, origin, null, target); + } - /** @see #taintIfTainted(IastContext, Object, Object, byte, CharSequence) */ - void taintIfTainted( - @Nullable Object target, @Nullable Object input, byte origin, @Nullable CharSequence name); + /** @see #taintIfTainted(IastContext, Object, Object, byte) */ + default void taintIfTainted( + @Nonnull IastContext ctx, @Nullable String target, @Nullable Object input, byte origin) { + taintIfTainted(ctx, target, input, origin, null, target); + } /** * Taints the object only if the input value is tainted, the resulting value will be tainted using * a source with the specified origin and name, if target is a char sequence it will be used as * value */ - void taintIfTainted( - @Nullable IastContext ctx, + default void taintIfTainted( + @Nonnull IastContext ctx, @Nullable Object target, @Nullable Object input, byte origin, - @Nullable CharSequence name); + @Nullable CharSequence name) { + taintIfTainted(ctx, target, input, origin, name, target); + } - /** @see #taintIfTainted(IastContext, Object, Object, byte, CharSequence, Object) */ - void taintIfTainted( - @Nullable Object target, + /** @see #taintIfTainted(IastContext, Object, Object, byte, CharSequence) */ + default void taintIfTainted( + @Nonnull IastContext ctx, + @Nullable String target, @Nullable Object input, byte origin, - @Nullable CharSequence name, - @Nullable Object value); + @Nullable CharSequence name) { + taintIfTainted(ctx, target, input, origin, name, target); + } /** * Taints the object only if the input value is tainted, the resulting value will be tainted using * a source with the specified origin, name and value. */ void taintIfTainted( - @Nullable IastContext ctx, + @Nonnull IastContext ctx, @Nullable Object target, @Nullable Object input, byte origin, @Nullable CharSequence name, @Nullable Object value); - /** @see #taintIfAnyTainted(IastContext, Object, Object[]) */ - void taintIfAnyTainted(@Nullable Object target, @Nullable Object[] inputs); + /** @see #taintIfTainted(IastContext, Object, Object, byte, CharSequence, Object) */ + void taintIfTainted( + @Nonnull IastContext ctx, + @Nullable String target, + @Nullable Object input, + byte origin, + @Nullable CharSequence name, + @Nullable Object value); /** * Taints the object if any of the inputs is tainted. When a tainted input is found the logic is @@ -165,12 +202,16 @@ void taintIfTainted( * * @see #taintIfTainted(IastContext, Object, Object) */ - void taintIfAnyTainted( - @Nullable IastContext ctx, @Nullable Object target, @Nullable Object[] inputs); + default void taintIfAnyTainted( + @Nonnull IastContext ctx, @Nullable Object target, @Nullable Object[] inputs) { + taintIfAnyTainted(ctx, target, inputs, false, NOT_MARKED); + } - /** @see #taintIfAnyTainted(IastContext, Object, Object[], boolean, int) */ - void taintIfAnyTainted( - @Nullable Object target, @Nullable Object[] inputs, boolean keepRanges, int mark); + /** @see #taintIfAnyTainted(IastContext, Object, Object[]) */ + default void taintIfAnyTainted( + @Nonnull IastContext ctx, @Nullable String target, @Nullable Object[] inputs) { + taintIfAnyTainted(ctx, target, inputs, false, NOT_MARKED); + } /** * Taints the object if any of the inputs is tainted. When a tainted input is found the logic is @@ -179,14 +220,19 @@ void taintIfAnyTainted( * @see #taintIfTainted(IastContext, Object, Object, boolean, int) */ void taintIfAnyTainted( - @Nullable IastContext ctx, + @Nonnull IastContext ctx, @Nullable Object target, @Nullable Object[] inputs, boolean keepRanges, int mark); - /** @see #taintDeeply(IastContext, Object, byte, Predicate) */ - int taintDeeply(@Nullable Object target, byte origin, Predicate> classFilter); + /** @see #taintIfAnyTainted(IastContext, Object, Object[], boolean, int) */ + void taintIfAnyTainted( + @Nonnull IastContext ctx, + @Nullable String target, + @Nullable Object[] inputs, + boolean keepRanges, + int mark); /** * Visit the graph of the object and taints all the string properties found using a source with @@ -196,24 +242,20 @@ void taintIfAnyTainted( * @return number of tainted elements */ int taintDeeply( - @Nullable IastContext ctx, + @Nonnull IastContext ctx, @Nullable Object target, byte origin, Predicate> classFilter); - /** @see #isTainted(IastContext, Object) */ - boolean isTainted(@Nullable Object target); - /** Checks if an arbitrary object is tainted */ - boolean isTainted(@Nullable IastContext ctx, @Nullable Object target); + boolean isTainted(@Nonnull IastContext ctx, @Nullable Object target); - /** @see #findSource(IastContext, Object) */ - @Nullable - Source findSource(@Nullable Object target); + /** Checks if the string is tainted */ + boolean isTainted(@Nonnull IastContext ctx, @Nullable String target); /** * Returns the source with the highest priority if the object is tainted, {@code null} otherwise */ @Nullable - Source findSource(@Nullable IastContext ctx, @Nullable Object target); + Source findSource(@Nonnull IastContext ctx, @Nullable Object target); }