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..5b11cc7c020 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.taintString(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.taintString(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..a1e29ffa6f0 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.taintString(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.taintString(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..2b585b249e2 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) { + taintStringIfTainted(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) { + taintStringIfRangeTainted(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) { + taintObjectIfTainted(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) { + taintObjectIfTainted(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) { + taintObjectIfTainted(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..bc097e00b23 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( + public void taintObjectRange( + @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, + public void taintStringRange( + @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, + public void taintObject( + @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 taintString( + @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, + public void taintObjectIfTainted( + @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, + public void taintStringIfTainted( + @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( + public void taintObjectIfRangeTainted( + @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, + public void taintStringIfRangeTainted( + @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( + public void taintObjectIfTainted( + @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, + public void taintStringIfTainted( + @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( + public void taintObjectIfAnyTainted( + @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, + public void taintStringIfAnyTainted( + @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,19 @@ 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); - } - - @Override - public boolean isTainted(@Nullable final Object target) { - return target != null && isTainted(LazyContext.build(), target); + final TaintedObjects to = ctx.getTaintedObjects(); + return highestPrioritySource(to, target); } @Override - public boolean isTainted(@Nullable final IastContext ctx, @Nullable final Object target) { + public boolean isTainted(@Nonnull final IastContext ctx, @Nullable final Object target) { return target != null && findSource(ctx, target) != null; } @@ -353,10 +304,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 +349,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 +362,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 +389,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 +409,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 +419,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,23 +461,46 @@ 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 @@ -549,33 +531,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 +547,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..f763607a384 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.taintObject(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.taintObjectRange(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.taintObjectIfTainted(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.taintObjectIfTainted(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.taintObjectIfRangeTainted(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.taintObjectIfRangeTainted(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.taintObjectIfRangeTainted(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.taintObjectIfTainted(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.taintObjectIfTainted(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.taintObjectIfTainted(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.taintObjectIfTainted(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.taintObjectIfTainted(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.taintObjectIfTainted(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.taintObjectIfAnyTainted(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.taintObjectIfAnyTainted(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.taintObjectIfAnyTainted(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.taintObjectIfAnyTainted(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.taintObjectIfAnyTainted(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.taintObjectIfAnyTainted(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.taintObjectIfAnyTainted(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.taintObjectIfAnyTainted(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.taintObject(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.taintObject(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.taintObject(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.taintStringIfTainted(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..a5e38217693 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.taintObjectIfTainted(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/IastAgentTestRunner.groovy b/dd-java-agent/agent-iast/src/testFixtures/groovy/com/datadog/iast/test/IastAgentTestRunner.groovy index 27ed9ee244d..c34279312f0 100644 --- a/dd-java-agent/agent-iast/src/testFixtures/groovy/com/datadog/iast/test/IastAgentTestRunner.groovy +++ b/dd-java-agent/agent-iast/src/testFixtures/groovy/com/datadog/iast/test/IastAgentTestRunner.groovy @@ -45,18 +45,25 @@ class IastAgentTestRunner extends AgentTestRunner implements IastRequestContextP } protected DDSpan runUnderIastTrace(Closure cl) { - CallbackProvider iastCbp = TEST_TRACER.getCallbackProvider(RequestContextSlot.IAST) - Supplier> reqStartCb = iastCbp.getCallback(Events.EVENTS.requestStarted()) + return withIastTrace(cl).v1 + } + + protected E computeUnderIastTrace(Closure cl) { + return withIastTrace(cl).v2 + } - def iastCtx = reqStartCb.get().result - def ddctx = new TagContext().withRequestContextDataIast(iastCtx) - AgentSpan span = TEST_TRACER.startSpan("test", "test-iast-span", ddctx) + private Tuple2 withIastTrace(Closure cl) { + final iastCbp = TEST_TRACER.getCallbackProvider(RequestContextSlot.IAST) + final reqStartCb = iastCbp.getCallback(Events.EVENTS.requestStarted()) + + final iastCtx = reqStartCb.get().result + final ddctx = new TagContext().withRequestContextDataIast(iastCtx) + final span = TEST_TRACER.startSpan("test", "test-iast-span", ddctx) try { - AgentTracer.activateSpan(span).withCloseable cl + final E result = AgentTracer.activateSpan(span).withCloseable cl + return Tuple.tuple(span, result) } finally { span.finish() } - - span } } 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/agent-tooling/src/main/java/datadog/trace/agent/tooling/iast/TaintableEnumeration.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/iast/TaintableEnumeration.java index 81864874966..f1d67e3a8a0 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/iast/TaintableEnumeration.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/iast/TaintableEnumeration.java @@ -58,7 +58,7 @@ public String nextElement() { throw e; } try { - module.taint(context, next, origin, name(next)); + module.taintString(context, next, origin, name(next)); } catch (final Throwable e) { module.onUnexpectedException("Failed to taint enumeration", e); } diff --git a/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/iast/TaintableEnumerationTest.groovy b/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/iast/TaintableEnumerationTest.groovy index b0f5b50f48e..4c3395004d7 100644 --- a/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/iast/TaintableEnumerationTest.groovy +++ b/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/iast/TaintableEnumerationTest.groovy @@ -36,7 +36,7 @@ class TaintableEnumerationTest extends DDSpecification { then: result == values - values.each { 1 * module.taint(iastCtx, it, origin, name) } + values.each { 1 * module.taintString(iastCtx, it, origin, name) } } void 'underlying enumerated values are tainted with the value as a name'() { @@ -50,7 +50,7 @@ class TaintableEnumerationTest extends DDSpecification { then: result == values - values.each { 1 * module.taint(iastCtx, it, origin, it) } + values.each { 1 * module.taintString(iastCtx, it, origin, it) } } void 'taintable enumeration leaves no trace in case of error'() { diff --git a/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/CookieHeaderInstrumentation.java b/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/CookieHeaderInstrumentation.java index 57e0b0e092e..5539fb450c6 100644 --- a/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/CookieHeaderInstrumentation.java +++ b/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/CookieHeaderInstrumentation.java @@ -75,8 +75,8 @@ static void after( while (iterator.hasNext()) { HttpCookiePair pair = iterator.next(); final String name = pair.name(), value = pair.value(); - prop.taint(ctx, name, SourceTypes.REQUEST_COOKIE_NAME, name); - prop.taint(ctx, value, SourceTypes.REQUEST_COOKIE_VALUE, name); + prop.taintString(ctx, name, SourceTypes.REQUEST_COOKIE_NAME, name); + prop.taintString(ctx, value, SourceTypes.REQUEST_COOKIE_VALUE, name); } } } diff --git a/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/HeaderNameCallSite.java b/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/HeaderNameCallSite.java index eaad764a2d2..143f2829989 100644 --- a/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/HeaderNameCallSite.java +++ b/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/HeaderNameCallSite.java @@ -32,7 +32,7 @@ public static String after(@CallSite.This HttpHeader header, @CallSite.Return St if (ctx == null) { return result; } - module.taintIfTainted(ctx, result, header, SourceTypes.REQUEST_HEADER_NAME, result); + module.taintStringIfTainted(ctx, result, header, SourceTypes.REQUEST_HEADER_NAME, result); } catch (final Throwable e) { module.onUnexpectedException("onHeaderNames threw", e); } diff --git a/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/HttpHeaderSubclassesInstrumentation.java b/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/HttpHeaderSubclassesInstrumentation.java index 576496b1c95..e87100071cc 100644 --- a/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/HttpHeaderSubclassesInstrumentation.java +++ b/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/HttpHeaderSubclassesInstrumentation.java @@ -72,7 +72,7 @@ static void onExit( return; } IastContext ctx = reqCtx.getData(RequestContextSlot.IAST); - propagation.taintIfTainted(ctx, retVal, h); + propagation.taintStringIfTainted(ctx, retVal, h); } } } diff --git a/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/HttpRequestInstrumentation.java b/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/HttpRequestInstrumentation.java index cc3c4ddfea1..49553f13bcb 100644 --- a/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/HttpRequestInstrumentation.java +++ b/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/HttpRequestInstrumentation.java @@ -88,7 +88,7 @@ static void onExit( } // unfortunately, the call to h.value() is instrumented, but // because the call to taint() only happens after, the call is a noop - propagation.taint(ctx, h, SourceTypes.REQUEST_HEADER_VALUE, h.name(), h.value()); + propagation.taintObject(ctx, h, SourceTypes.REQUEST_HEADER_VALUE, h.name(), h.value()); } } } @@ -112,7 +112,7 @@ static void onExit( return; } - propagation.taintIfTainted(ctx, entity, thiz); + propagation.taintObjectIfTainted(ctx, entity, thiz); } } } 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..f3b26731b54 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.taintString(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/RequestContextInstrumentation.java b/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/RequestContextInstrumentation.java index e45dd3268d3..90e897fcc7b 100644 --- a/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/RequestContextInstrumentation.java +++ b/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/RequestContextInstrumentation.java @@ -67,7 +67,7 @@ static void onExit( return; } - propagation.taintIfTainted(ctx, request, requestContext); + propagation.taintObjectIfTainted(ctx, request, requestContext); } } } diff --git a/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/UriInstrumentation.java b/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/UriInstrumentation.java index 77c0edf27a7..33ab666f6c3 100644 --- a/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/UriInstrumentation.java +++ b/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/UriInstrumentation.java @@ -80,7 +80,7 @@ static void after( return; } final IastContext ctx = reqCtx.getData(RequestContextSlot.IAST); - mod.taintIfTainted(ctx, ret.get(), uri); + mod.taintStringIfTainted(ctx, ret.get(), uri); } } @@ -109,8 +109,8 @@ static void after( while (iterator.hasNext()) { Tuple2 pair = iterator.next(); final String name = pair._1(), value = pair._2(); - prop.taint(ctx, name, SourceTypes.REQUEST_PARAMETER_NAME, name); - prop.taint(ctx, value, SourceTypes.REQUEST_PARAMETER_VALUE, name); + prop.taintString(ctx, name, SourceTypes.REQUEST_PARAMETER_NAME, name); + prop.taintString(ctx, value, SourceTypes.REQUEST_PARAMETER_VALUE, name); } } } diff --git a/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/helpers/TaintCookieFunction.java b/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/helpers/TaintCookieFunction.java index bc3fc920535..5cbc18d95e1 100644 --- a/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/helpers/TaintCookieFunction.java +++ b/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/helpers/TaintCookieFunction.java @@ -27,8 +27,8 @@ public Tuple1 apply(Tuple1 v1) { } final String name = httpCookiePair.name(); final String value = httpCookiePair.value(); - mod.taint(ctx, name, SourceTypes.REQUEST_COOKIE_NAME, name); - mod.taint(ctx, value, SourceTypes.REQUEST_COOKIE_VALUE, name); + mod.taintString(ctx, name, SourceTypes.REQUEST_COOKIE_NAME, name); + mod.taintString(ctx, value, SourceTypes.REQUEST_COOKIE_VALUE, name); return v1; } } diff --git a/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/helpers/TaintFutureHelper.java b/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/helpers/TaintFutureHelper.java index 7203f9141fe..b7153bef168 100644 --- a/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/helpers/TaintFutureHelper.java +++ b/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/helpers/TaintFutureHelper.java @@ -14,7 +14,7 @@ public static Future wrapFuture( t -> { IastContext ctx = IastContext.Provider.get(AgentTracer.activeSpan()); if (ctx != null) { - mod.taintIfTainted(ctx, t, input); + mod.taintObjectIfTainted(ctx, t, input); } return t; }; diff --git a/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/helpers/TaintMapFunction.java b/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/helpers/TaintMapFunction.java index 68a8056ab79..d8485bfa4b9 100644 --- a/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/helpers/TaintMapFunction.java +++ b/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/helpers/TaintMapFunction.java @@ -32,8 +32,8 @@ public Tuple1> apply(Tuple1> v1) { while (iterator.hasNext()) { Tuple2 e = iterator.next(); final String name = e._1(), value = e._2(); - prop.taint(ctx, name, SourceTypes.REQUEST_PARAMETER_NAME, name); - prop.taint(ctx, value, SourceTypes.REQUEST_PARAMETER_VALUE, name); + prop.taintString(ctx, name, SourceTypes.REQUEST_PARAMETER_NAME, name); + prop.taintString(ctx, value, SourceTypes.REQUEST_PARAMETER_VALUE, name); } return v1; diff --git a/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/helpers/TaintMultiMapFunction.java b/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/helpers/TaintMultiMapFunction.java index a4c9718b8df..e98b45866b7 100644 --- a/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/helpers/TaintMultiMapFunction.java +++ b/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/helpers/TaintMultiMapFunction.java @@ -33,10 +33,10 @@ public Tuple1>> apply(Tuple1>> while (entriesIterator.hasNext()) { Tuple2> e = entriesIterator.next(); final String name = e._1(); - mod.taint(ctx, name, SourceTypes.REQUEST_PARAMETER_NAME, name); + mod.taintString(ctx, name, SourceTypes.REQUEST_PARAMETER_NAME, name); List values = e._2(); for (final String value : ScalaToJava.listAsList(values)) { - mod.taint(ctx, value, SourceTypes.REQUEST_PARAMETER_VALUE, name); + mod.taintString(ctx, value, SourceTypes.REQUEST_PARAMETER_VALUE, name); } } diff --git a/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/helpers/TaintOptionalCookieFunction.java b/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/helpers/TaintOptionalCookieFunction.java index f9db9557664..1aeb51bec1b 100644 --- a/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/helpers/TaintOptionalCookieFunction.java +++ b/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/helpers/TaintOptionalCookieFunction.java @@ -29,8 +29,8 @@ public Tuple1> apply(Tuple1> v1) { final HttpCookiePair cookie = httpCookiePair.get(); final String name = cookie.name(); final String value = cookie.value(); - mod.taint(ctx, name, SourceTypes.REQUEST_COOKIE_NAME, name); - mod.taint(ctx, value, SourceTypes.REQUEST_COOKIE_VALUE, name); + mod.taintString(ctx, name, SourceTypes.REQUEST_COOKIE_NAME, name); + mod.taintString(ctx, value, SourceTypes.REQUEST_COOKIE_VALUE, name); return v1; } } diff --git a/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/helpers/TaintRequestContextFunction.java b/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/helpers/TaintRequestContextFunction.java index 1ca581a7364..4839211d196 100644 --- a/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/helpers/TaintRequestContextFunction.java +++ b/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/helpers/TaintRequestContextFunction.java @@ -25,7 +25,7 @@ public Tuple1 apply(Tuple1 v1) { if (ctx == null) { return v1; } - mod.taint(ctx, reqCtx, SourceTypes.REQUEST_BODY); + mod.taintObject(ctx, reqCtx, SourceTypes.REQUEST_BODY); return v1; } diff --git a/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/helpers/TaintRequestFunction.java b/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/helpers/TaintRequestFunction.java index 14889998d51..c881037cf1c 100644 --- a/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/helpers/TaintRequestFunction.java +++ b/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/helpers/TaintRequestFunction.java @@ -26,7 +26,7 @@ public Tuple1 apply(Tuple1 v1) { if (ctx == null) { return v1; } - mod.taint(ctx, httpRequest, SourceTypes.REQUEST_BODY); + mod.taintObject(ctx, httpRequest, SourceTypes.REQUEST_BODY); return v1; } diff --git a/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/helpers/TaintSeqFunction.java b/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/helpers/TaintSeqFunction.java index 9b58763b34b..1269d8ef448 100644 --- a/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/helpers/TaintSeqFunction.java +++ b/dd-java-agent/instrumentation/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/helpers/TaintSeqFunction.java @@ -39,9 +39,9 @@ public Tuple1>> apply(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.taintString(ctx, (String) o, SourceTypes.REQUEST_PARAMETER_VALUE, paramName); } } } else if (value instanceof String) { - mod.taint(ctx, value, SourceTypes.REQUEST_PARAMETER_VALUE, paramName); + mod.taintString(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..fa6ccd50307 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.taintObject(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..a1cf5168649 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,14 +17,14 @@ 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()); if (ctx == null) { return v1; } - mod.taint(ctx, uri, SourceTypes.REQUEST_QUERY); + mod.taintObject(ctx, uri, SourceTypes.REQUEST_QUERY); return v1; } 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..a5d800fc236 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.taintString(ctx, (String) o, SourceTypes.REQUEST_PARAMETER_VALUE, paramName); } } } else if (value instanceof String) { - mod.taint(ctx, value, SourceTypes.REQUEST_PARAMETER_VALUE, paramName); + mod.taintString(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..1d58ab9000e 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.taintObjectIfTainted(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..f60296d17ac 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.taintObjectIfTainted(_ 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-fileupload/src/main/java/datadog/trace/instrumentation/commons/fileupload/CommonsFileuploadInstrumenter.java b/dd-java-agent/instrumentation/commons-fileupload/src/main/java/datadog/trace/instrumentation/commons/fileupload/CommonsFileuploadInstrumenter.java index 7c251428296..8c5ab41bb30 100644 --- a/dd-java-agent/instrumentation/commons-fileupload/src/main/java/datadog/trace/instrumentation/commons/fileupload/CommonsFileuploadInstrumenter.java +++ b/dd-java-agent/instrumentation/commons-fileupload/src/main/java/datadog/trace/instrumentation/commons/fileupload/CommonsFileuploadInstrumenter.java @@ -60,7 +60,7 @@ public static Map onExit( final IastContext ctx = reqCtx.getData(RequestContextSlot.IAST); for (final Map.Entry entry : map.entrySet()) { if (entry.getValue() != null) { - module.taint( + module.taintString( ctx, entry.getValue(), SourceTypes.REQUEST_MULTIPART_PARAMETER, entry.getKey()); } } diff --git a/dd-java-agent/instrumentation/commons-fileupload/src/test/groovy/MultipartInstrumentationTest.groovy b/dd-java-agent/instrumentation/commons-fileupload/src/test/groovy/MultipartInstrumentationTest.groovy index 6f7da062305..5d1df9dae71 100644 --- a/dd-java-agent/instrumentation/commons-fileupload/src/test/groovy/MultipartInstrumentationTest.groovy +++ b/dd-java-agent/instrumentation/commons-fileupload/src/test/groovy/MultipartInstrumentationTest.groovy @@ -41,8 +41,8 @@ class MultipartInstrumentationTest extends AgentTestRunner { } then: - 1 * module.taint(iastCtx, 'file', SourceTypes.REQUEST_MULTIPART_PARAMETER, 'name') - 1 * module.taint(iastCtx, _, SourceTypes.REQUEST_MULTIPART_PARAMETER, 'filename') + 1 * module.taintObject(iastCtx, 'file', SourceTypes.REQUEST_MULTIPART_PARAMETER, 'name') + 1 * module.taintObject(iastCtx, _, SourceTypes.REQUEST_MULTIPART_PARAMETER, 'filename') 0 * _ where: diff --git a/dd-java-agent/instrumentation/commons-httpclient-2/build.gradle b/dd-java-agent/instrumentation/commons-httpclient-2/build.gradle index b406619e198..58796bf5020 100644 --- a/dd-java-agent/instrumentation/commons-httpclient-2/build.gradle +++ b/dd-java-agent/instrumentation/commons-httpclient-2/build.gradle @@ -26,6 +26,7 @@ dependencies { compileOnly group: 'commons-httpclient', name: 'commons-httpclient', version: '2.0' testImplementation group: 'commons-httpclient', name: 'commons-httpclient', version: '2.0' + testImplementation(testFixtures(project(':dd-java-agent:agent-iast'))) latestDepTestImplementation group: 'commons-httpclient', name: 'commons-httpclient', version: '(2.0,20000000]' } 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..863469688da 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; @@ -46,13 +47,17 @@ public void methodAdvice(MethodTransformer transformer) { } public static class CtorAdvice { - @Advice.OnMethodExit(suppress = Throwable.class) + @Advice.OnMethodExit(suppress = Throwable.class, inline = false) @Propagation 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.taintObjectIfTainted(ctx, self, argument); } } } diff --git a/dd-java-agent/instrumentation/commons-httpclient-2/src/test/groovy/IastCommonsHttpClientInstrumentationTest.groovy b/dd-java-agent/instrumentation/commons-httpclient-2/src/test/groovy/IastCommonsHttpClientInstrumentationTest.groovy index 0a3f5ba5bc7..ffe371ec7ab 100644 --- a/dd-java-agent/instrumentation/commons-httpclient-2/src/test/groovy/IastCommonsHttpClientInstrumentationTest.groovy +++ b/dd-java-agent/instrumentation/commons-httpclient-2/src/test/groovy/IastCommonsHttpClientInstrumentationTest.groovy @@ -1,20 +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.propagation.PropagationModule import datadog.trace.api.iast.sink.SsrfModule import org.apache.commons.httpclient.HttpClient +import org.apache.commons.httpclient.HttpMethod import org.apache.commons.httpclient.methods.GetMethod import spock.lang.AutoCleanup import spock.lang.Shared import static datadog.trace.agent.test.server.http.TestHttpServer.httpServer -class IastCommonsHttpClientInstrumentationTest extends AgentTestRunner { - - @Override - protected void configurePreAgent() { - injectSysConfig('dd.iast.enabled', 'true') - } +class IastCommonsHttpClientInstrumentationTest extends IastAgentTestRunner { @AutoCleanup @Shared @@ -43,7 +40,7 @@ class IastCommonsHttpClientInstrumentationTest extends AgentTestRunner { InstrumentationBridge.registerIastModule(ssrf) when: - new HttpClient().executeMethod(new GetMethod(url)) + runUnderIastTrace { new HttpClient().executeMethod(new GetMethod(url)) } then: 1 * ssrf.onURLConnection({ value -> tainteds.containsKey(value) }) @@ -51,9 +48,11 @@ class IastCommonsHttpClientInstrumentationTest extends AgentTestRunner { private void mockPropagation() { final propagation = Mock(PropagationModule) { - taintIfTainted(_, _) >> { - if (tainteds.containsKey(it[1])) { - tainteds.put(it[0], null) + taintObjectIfTainted(_ as IastContext, _ as HttpMethod, _ as String) >> { + final method = it[1] as HttpMethod + final url = it[2] as String + if (tainteds.containsKey(url)) { + tainteds.put(method, null) } } } 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..da8a7168bbe 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.taintStringIfTainted(ctx, result, input, false, VulnerabilityMarks.XSS_MARK); + } } catch (final Throwable e) { module.onUnexpectedException("afterEscape threw", e); } @@ -38,9 +42,13 @@ 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.taintStringIfTainted( + 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..4988156440c 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,11 @@ class StringEscapeUtilsCallSiteTest extends AgentTestRunner { InstrumentationBridge.registerIastModule(module) when: - final result = TestStringEscapeUtilsSuite.&"$method".call(args) + def result = computeUnderIastTrace { TestStringEscapeUtilsSuite.&"$method".call(args) } then: result == expected - 1 * module.taintIfTainted(_ as String, args[0], false, mark) + 1 * module.taintStringIfTainted(_ 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..a92e6d0a1e5 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.taintStringIfTainted(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.taintStringIfTainted(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..7b0c702a3cf 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,11 @@ class StringEscapeUtilsCallSiteTest extends AgentTestRunner { InstrumentationBridge.registerIastModule(module) when: - final result = TestStringEscapeUtilsSuite.&"$method".call(args) + def result = computeUnderIastTrace { TestStringEscapeUtilsSuite.&"$method".call(args) } then: result == expected - 1 * module.taintIfTainted(_ as String, args[0], false, VulnerabilityMarks.XSS_MARK) + 1 * module.taintStringIfTainted(_ as IastContext, _ as String, args[0], false, VulnerabilityMarks.XSS_MARK) 0 * _ where: @@ -43,11 +37,11 @@ class StringEscapeUtilsCallSiteTest extends AgentTestRunner { InstrumentationBridge.registerIastModule(module) when: - final result = TestStringEscapeUtilsSuite.&"$method".call(args) + def result = computeUnderIastTrace { TestStringEscapeUtilsSuite.&"$method".call(args) } then: result == expected - 1 * module.taintIfTainted(_ as String, args[0]) + 1 * module.taintStringIfTainted(_ 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..bb99b2b7324 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.taintStringIfTainted(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.taintStringIfTainted(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..09dcad73e3b 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,11 @@ class StringEscapeUtilsCallSiteTest extends AgentTestRunner { InstrumentationBridge.registerIastModule(module) when: - final result = TestStringEscapeUtilsSuite.&"$method".call(args) + def result = computeUnderIastTrace { TestStringEscapeUtilsSuite.&"$method".call(args) } then: result == expected - 1 * module.taintIfTainted(_ as String, args[0], false, VulnerabilityMarks.XSS_MARK) + 1 * module.taintStringIfTainted(_ as IastContext, _ as String, args[0], false, VulnerabilityMarks.XSS_MARK) 0 * _ where: @@ -44,11 +38,11 @@ class StringEscapeUtilsCallSiteTest extends AgentTestRunner { InstrumentationBridge.registerIastModule(module) when: - final result = TestStringEscapeUtilsSuite.&"$method".call(args) + def result = computeUnderIastTrace { TestStringEscapeUtilsSuite.&"$method".call(args) } then: result == expected - 1 * module.taintIfTainted(_ as String, args[0]) + 1 * module.taintStringIfTainted(_ 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..b063600ffdf 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.taintStringIfTainted(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..8410318348c 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,11 @@ class StringUtilCallSiteTest extends AgentTestRunner { InstrumentationBridge.registerIastModule(module) when: - final result = TestStringUtilSuite.&"$method".call(args) + def result = computeUnderIastTrace { TestStringUtilSuite.&"$method".call(args) } then: result == expected - 1 * module.taintIfTainted(_ as String, args[0], false, VulnerabilityMarks.XSS_MARK) + 1 * module.taintStringIfTainted(_ as IastContext, _ as String, args[0], false, VulnerabilityMarks.XSS_MARK) 0 * _ where: @@ -42,10 +38,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..d6b7071b743 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.taintObjectIfTainted(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.taintStringIfTainted(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..d33c6f4e416 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,16 @@ class JsonReaderInstrumentationTest extends AgentTestRunner { final gson = new Gson() when: - final reader = new JsonReader(new StringReader(json)) + final reader = computeUnderIastTrace { new JsonReader(new StringReader(json)) } then: - 1 * module.taintIfTainted(_ as JsonReader, _ as StringReader) + 1 * module.taintObjectIfTainted(_ 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.taintStringIfTainted(_ 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..717aed5bf99 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.taintObjectIfTainted(ctx, parser, input); + } } if (input instanceof URL) { final SsrfModule ssrf = InstrumentationBridge.SSRF; @@ -82,10 +86,14 @@ 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.taintObjectIfRangeTainted( + 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..1e868e56dd5 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,11 @@ class Json1FactoryInstrumentationTest extends AgentTestRunner { final content = '{"key":"value"}' when: - final result = new JsonFactory().createJsonParser(content) + final result = computeUnderIastTrace { new JsonFactory().createJsonParser(content) } then: result != null - 1 * propagationModule.taintIfTainted(_ as JsonParser, content) + 1 * propagationModule.taintObjectIfTainted(_ as IastContext, _ as JsonParser, content) 0 * _ } @@ -52,11 +48,11 @@ class Json1FactoryInstrumentationTest extends AgentTestRunner { final is = Mock(InputStream) when: - final result = new JsonFactory().createJsonParser(is) + final result = computeUnderIastTrace { new JsonFactory().createJsonParser(is) } then: result != null - 1 * propagationModule.taintIfTainted(_ as JsonParser, is) + 1 * propagationModule.taintObjectIfTainted(_ as IastContext, _ as JsonParser, is) 2 * is.read(_, _, _) 0 * _ } @@ -68,12 +64,12 @@ class Json1FactoryInstrumentationTest extends AgentTestRunner { final reader = Mock(Reader) when: - final result = new JsonFactory().createJsonParser(reader) + final result = computeUnderIastTrace { new JsonFactory().createJsonParser(reader) } then: result != null - 1 * propagationModule.taintIfTainted(_ as JsonParser, reader) + 1 * propagationModule.taintObjectIfTainted(_ as IastContext, _ as JsonParser, reader) 0 * _ } @@ -85,11 +81,11 @@ class Json1FactoryInstrumentationTest extends AgentTestRunner { final url = new URL("${clientServer.address}/json") when: - final parser = new JsonFactory().createJsonParser(url) + final parser = computeUnderIastTrace { new JsonFactory().createJsonParser(url) } then: parser != null - 1 * propagationModule.taintIfTainted(_ as JsonParser, url) + 1 * propagationModule.taintObjectIfTainted(_ as IastContext, _ as JsonParser, url) 1 * ssrfModule.onURLConnection(url) 0 * _ } @@ -102,12 +98,12 @@ class Json1FactoryInstrumentationTest extends AgentTestRunner { final bytes = '{}'.bytes when: - final result = new JsonFactory().createJsonParser(bytes) + final result = computeUnderIastTrace{ new JsonFactory().createJsonParser(bytes) } then: result != null - 1 * propagationModule.taintIfTainted(_ as JsonParser, bytes) + 1 * propagationModule.taintObjectIfTainted(_ as IastContext, _ as JsonParser, bytes) 0 * _ } @@ -118,11 +114,11 @@ class Json1FactoryInstrumentationTest extends AgentTestRunner { final bytes = '{}'.bytes when: - final parser = new JsonFactory().createJsonParser(bytes, 0, 2) + final parser = computeUnderIastTrace { new JsonFactory().createJsonParser(bytes, 0, 2) } then: parser != null - 1 * propagationModule.taintIfTainted(_ as JsonParser, bytes, 0, 2, false, VulnerabilityMarks.NOT_MARKED) + 1 * propagationModule.taintObjectIfRangeTainted(_ 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..9e8456f9253 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,20 @@ class Json1ParserInstrumentationTest extends AgentTestRunner { final reader = new ObjectMapper().reader(Map) when: - final taintedResult = reader.readValue(target) as Map + final taintedResult = computeUnderIastTrace { + 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.taintObjectIfTainted(_ as IastContext, _ as JsonParser, _) + _ * module.findSource(_ as IastContext, _ as JsonParser) >> source + 1 * module.taintString(_ as IastContext, 'root', source.origin, 'root', JSON_STRING) + 1 * module.taintString(_ as IastContext, 'root_value', source.origin, 'root', JSON_STRING) + 1 * module.taintString(_ as IastContext, 'nested', source.origin, 'nested', JSON_STRING) + 1 * module.taintString(_ as IastContext, 'nested_array', source.origin, 'nested_array', JSON_STRING) + 1 * module.taintString(_ as IastContext, 'array_0', source.origin, 'nested_array', JSON_STRING) + 1 * module.taintString(_ as IastContext, 'array_1', source.origin, 'nested_array', JSON_STRING) 0 * _ where: @@ -54,12 +51,12 @@ class Json1ParserInstrumentationTest extends AgentTestRunner { final reader = new ObjectMapper().reader(Map) when: - final taintedResult = reader.readValue(target) as Map + final taintedResult = computeUnderIastTrace { reader.readValue(target) as Map } then: JsonOutput.toJson(taintedResult) == JSON_STRING - _ * module.taintIfTainted(_, _) - _ * module.findSource(_) >> null + _ * module.taintObjectIfTainted(_ as IastContext, _ as JsonParser, _) + _ * module.findSource(_ as IastContext, _ as JsonParser) >> null 0 * _ where: @@ -67,7 +64,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..062eafa6872 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.taintObjectIfTainted(ctx, parser, input); + } } if (input instanceof URL) { final SsrfModule ssrf = InstrumentationBridge.SSRF; @@ -84,10 +88,14 @@ 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.taintObjectIfRangeTainted( + 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..433bb8499a8 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.taintObjectIfTainted(ctx, parser, tokenBuffer); + } } } } diff --git a/dd-java-agent/instrumentation/jackson-core/src/test/groovy/Json2FactoryInstrumentationTest.groovy b/dd-java-agent/instrumentation/jackson-core/src/test/groovy/Json2FactoryInstrumentationTest.groovy index 6f602c1c8e4..df5a1e629b5 100644 --- a/dd-java-agent/instrumentation/jackson-core/src/test/groovy/Json2FactoryInstrumentationTest.groovy +++ b/dd-java-agent/instrumentation/jackson-core/src/test/groovy/Json2FactoryInstrumentationTest.groovy @@ -1,19 +1,19 @@ +import com.datadog.iast.test.IastAgentTestRunner import com.fasterxml.jackson.core.JsonFactory import com.fasterxml.jackson.core.JsonParser import com.fasterxml.jackson.databind.ObjectMapper -import datadog.trace.agent.test.AgentTestRunner 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 import datadog.trace.api.iast.sink.SsrfModule -import groovy.transform.CompileDynamic import spock.lang.AutoCleanup import spock.lang.Shared import static datadog.trace.agent.test.server.http.TestHttpServer.httpServer -class Json2FactoryInstrumentationTest extends AgentTestRunner { +class Json2FactoryInstrumentationTest extends IastAgentTestRunner { @Shared @AutoCleanup @@ -27,12 +27,6 @@ class Json2FactoryInstrumentationTest extends AgentTestRunner { } } - @CompileDynamic - @Override - protected void configurePreAgent() { - injectSysConfig("dd.iast.enabled", "true") - } - void 'test createParser(String)'() { setup: final propagationModule = Mock(PropagationModule) @@ -40,11 +34,11 @@ class Json2FactoryInstrumentationTest extends AgentTestRunner { final content = '{"key":"value"}' when: - final result = new JsonFactory().createParser(content) + final result = computeUnderIastTrace { new JsonFactory().createParser(content) } then: result != null - 1 * propagationModule.taintIfTainted(_ as JsonParser, content) + 1 * propagationModule.taintObjectIfTainted(_ as IastContext, _ as JsonParser, content) 0 * _ } @@ -55,12 +49,12 @@ class Json2FactoryInstrumentationTest extends AgentTestRunner { final is = Mock(InputStream) when: - final result = new JsonFactory().createParser(is) + final result = computeUnderIastTrace { new JsonFactory().createParser(is) } then: result != null - 1 * propagationModule.taintIfTainted(_ as JsonParser, is) + 1 * propagationModule.taintObjectIfTainted(_ as IastContext, _ as JsonParser, is) 2 * is.read(_, _, _) 0 * _ } @@ -72,12 +66,12 @@ class Json2FactoryInstrumentationTest extends AgentTestRunner { final reader = Mock(Reader) when: - final result = new JsonFactory().createParser(reader) + final result = computeUnderIastTrace { new JsonFactory().createParser(reader) } then: result != null - 1 * propagationModule.taintIfTainted(_ as JsonParser, reader) + 1 * propagationModule.taintObjectIfTainted(_ as IastContext, _ as JsonParser, reader) 0 * _ } @@ -88,12 +82,12 @@ class Json2FactoryInstrumentationTest extends AgentTestRunner { final bytes = '{}'.bytes when: - final result = new JsonFactory().createParser(bytes) + final result = computeUnderIastTrace { new JsonFactory().createParser(bytes) } then: result != null - 1 * propagationModule.taintIfTainted(_ as JsonParser, bytes) + 1 * propagationModule.taintObjectIfTainted(_ as IastContext, _ as JsonParser, bytes) 0 * _ } @@ -104,11 +98,11 @@ class Json2FactoryInstrumentationTest extends AgentTestRunner { final bytes = '{}'.bytes when: - final parser = new JsonFactory().createParser(bytes, 0, 2) + final parser = computeUnderIastTrace { new JsonFactory().createParser(bytes, 0, 2) } then: parser != null - 1 * propagationModule.taintIfTainted(_ as JsonParser, bytes, 0, 2, false, VulnerabilityMarks.NOT_MARKED) + 1 * propagationModule.taintObjectIfRangeTainted(_ as IastContext, _ as JsonParser, bytes, 0, 2, false, VulnerabilityMarks.NOT_MARKED) 0 * _ } @@ -120,15 +114,16 @@ class Json2FactoryInstrumentationTest extends AgentTestRunner { final url = new URL("${clientServer.address}/json") when: - final parser = new JsonFactory().createParser(url) - parser.setCodec(new ObjectMapper()) - final json = parser.readValueAs(Map) + final json = computeUnderIastTrace { + final parser = new JsonFactory().createParser(url) + parser.setCodec(new ObjectMapper()) + return parser.readValueAs(Map) + } then: - parser != null json == [key: 'value'] - 1 * propagationModule.taintIfTainted(_ as JsonParser, url) - 1 * propagationModule.findSource(_ as JsonParser) >> null + 1 * propagationModule.taintObjectIfTainted(_ as IastContext, _ as JsonParser, url) + 1 * propagationModule.findSource(_ as IastContext, _ as JsonParser) >> null 1 * ssrfModule.onURLConnection(url) 0 * _ } diff --git a/dd-java-agent/instrumentation/jackson-core/src/test/groovy/Json2ParserInstrumentationTest.groovy b/dd-java-agent/instrumentation/jackson-core/src/test/groovy/Json2ParserInstrumentationTest.groovy index 0b8f51af5c7..b8521f06852 100644 --- a/dd-java-agent/instrumentation/jackson-core/src/test/groovy/Json2ParserInstrumentationTest.groovy +++ b/dd-java-agent/instrumentation/jackson-core/src/test/groovy/Json2ParserInstrumentationTest.groovy @@ -1,5 +1,7 @@ +import com.datadog.iast.test.IastAgentTestRunner +import com.fasterxml.jackson.core.JsonParser import com.fasterxml.jackson.databind.ObjectMapper -import datadog.trace.agent.test.AgentTestRunner +import datadog.trace.api.iast.IastContext import datadog.trace.api.iast.InstrumentationBridge import datadog.trace.api.iast.SourceTypes import datadog.trace.api.iast.Taintable @@ -8,15 +10,10 @@ import groovy.json.JsonOutput import java.nio.charset.Charset -class Json2ParserInstrumentationTest extends AgentTestRunner { +class Json2ParserInstrumentationTest 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 +24,18 @@ class Json2ParserInstrumentationTest extends AgentTestRunner { final reader = new ObjectMapper().readerFor(Map) when: - final taintedResult = reader.readValue(target) as Map + final taintedResult = computeUnderIastTrace { 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.taintObjectIfTainted(_ as IastContext, _ as JsonParser, _) + _ * module.findSource(_ as IastContext, _ as JsonParser) >> source + 1 * module.taintString(_ as IastContext, 'root', source.origin, 'root', JSON_STRING) + 1 * module.taintString(_ as IastContext, 'root_value', source.origin, 'root', JSON_STRING) + 1 * module.taintString(_ as IastContext, 'nested', source.origin, 'nested', JSON_STRING) + 1 * module.taintString(_ as IastContext, 'nested_array', source.origin, 'nested_array', JSON_STRING) + 1 * module.taintString(_ as IastContext, 'array_0', source.origin, 'nested_array', JSON_STRING) + 1 * module.taintString(_ as IastContext, 'array_1', source.origin, 'nested_array', JSON_STRING) 0 * _ where: @@ -54,12 +51,12 @@ class Json2ParserInstrumentationTest extends AgentTestRunner { final reader = new ObjectMapper().readerFor(Map) when: - final taintedResult = reader.readValue(target) as Map + final taintedResult = computeUnderIastTrace { reader.readValue(target) as Map } then: JsonOutput.toJson(taintedResult) == JSON_STRING - _ * module.taintIfTainted(_, _) - _ * module.findSource(_) >> null + _ * module.taintObjectIfTainted(_ as IastContext, _ as JsonParser, _) + _ * module.findSource(_ as IastContext, _ as JsonParser) >> null 0 * _ where: 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..dd0ab703e78 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.taintObjectIfTainted(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.taintObjectIfTainted(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..aa51c8ff0ba 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.taintObjectIfTainted(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..00be6e3e2d5 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.taintObjectIfTainted(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..a4f3db16857 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.taintObjectIfTainted(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..1f89ab1a541 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.taintObjectIfTainted(_ as IastContext, _ as ByteBuffer, message, true, VulnerabilityMarks.NOT_MARKED) } void 'test array method'() { @@ -35,10 +31,10 @@ class ByteBufferTest extends AgentTestRunner { InstrumentationBridge.registerIastModule(module) when: - final result = TestByteBufferSuite.array(message) + final result = computeUnderIastTrace { TestByteBufferSuite.array(message) } then: result == message.array() - 1 * module.taintIfTainted(_ as byte[], message, true, VulnerabilityMarks.NOT_MARKED) + 1 * module.taintObjectIfTainted(_ 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..6eb8aa168b2 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.taintObjectIfTainted(_ 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..59e4e1cf351 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.taintObjectIfTainted(_ 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..14aba3c94f4 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.taintObjectIfTainted(_ as IastContext, _ as StringReader, input) } } diff --git a/dd-java-agent/instrumentation/java-lang/build.gradle b/dd-java-agent/instrumentation/java-lang/build.gradle index 8258d36ba37..f8541d3cb6c 100644 --- a/dd-java-agent/instrumentation/java-lang/build.gradle +++ b/dd-java-agent/instrumentation/java-lang/build.gradle @@ -11,6 +11,7 @@ addTestSuiteForDir('latestDepTest', 'test') dependencies { testRuntimeOnly project(':dd-java-agent:instrumentation:iast-instrumenter') + testImplementation(testFixtures(project(':dd-java-agent:agent-iast'))) } tasks.compileTestJava.configure { 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..f75a2710050 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.taintObjectIfTainted(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-lang/src/test/groovy/datadog/trace/instrumentation/java/lang/StringCallSiteTest.groovy b/dd-java-agent/instrumentation/java-lang/src/test/groovy/datadog/trace/instrumentation/java/lang/StringCallSiteTest.groovy index f1a47b0ec83..2031553f30d 100644 --- a/dd-java-agent/instrumentation/java-lang/src/test/groovy/datadog/trace/instrumentation/java/lang/StringCallSiteTest.groovy +++ b/dd-java-agent/instrumentation/java-lang/src/test/groovy/datadog/trace/instrumentation/java/lang/StringCallSiteTest.groovy @@ -1,23 +1,17 @@ package datadog.trace.instrumentation.java.lang -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.CodecModule import datadog.trace.api.iast.propagation.PropagationModule import datadog.trace.api.iast.propagation.StringModule import foo.bar.TestStringSuite -import groovy.transform.CompileDynamic import java.nio.charset.Charset -@CompileDynamic -class StringCallSiteTest extends AgentTestRunner { - - @Override - protected void configurePreAgent() { - injectSysConfig("dd.iast.enabled", "true") - } +class StringCallSiteTest extends IastAgentTestRunner { def 'test string concat call site'() { setup: @@ -272,11 +266,11 @@ class StringCallSiteTest extends AgentTestRunner { final string = 'test' when: - final char[] result = TestStringSuite.toCharArray(string) + final char[] result = computeUnderIastTrace { TestStringSuite.toCharArray(string) } then: result != null && result.length > 0 - 1 * module.taintIfTainted(_ as char[], string, true, VulnerabilityMarks.NOT_MARKED) + 1 * module.taintObjectIfTainted(_ as IastContext, _ as char[], string, true, VulnerabilityMarks.NOT_MARKED) 0 * _ } 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..26d61c3d322 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.taintObjectIfTainted(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.taintObjectIfAnyTainted(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.taintStringIfTainted(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.taintObjectIfTainted(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..668ffcd733e 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.taintObjectIfAnyTainted(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.taintStringIfTainted(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.taintObjectIfTainted(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..cf7934d2dd6 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.taintStringIfTainted(ctx, result, value, false, VulnerabilityMarks.XSS_MARK); + } } catch (final Throwable e) { module.onUnexpectedException("afterEncode threw", e); } diff --git a/dd-java-agent/instrumentation/java-net/src/test/groovy/datadog/trace/instrumentation/java/net/URICallSIteTest.groovy b/dd-java-agent/instrumentation/java-net/src/test/groovy/datadog/trace/instrumentation/java/net/URICallSIteTest.groovy index 398d94779d7..502c5c9315f 100644 --- a/dd-java-agent/instrumentation/java-net/src/test/groovy/datadog/trace/instrumentation/java/net/URICallSIteTest.groovy +++ b/dd-java-agent/instrumentation/java-net/src/test/groovy/datadog/trace/instrumentation/java/net/URICallSIteTest.groovy @@ -24,7 +24,7 @@ class URICallSIteTest extends AgentTestRunner { then: uri.toString() == expected - 1 * module.taintIfAnyTainted(_ as URI, args as Object[]) + 1 * module.taintObjectIfAnyTainted(_ as URI, args as Object[]) where: method | args | expected @@ -45,7 +45,7 @@ class URICallSIteTest extends AgentTestRunner { then: uri.toString() == expected - 1 * module.taintIfTainted(_ as URI, args[0]) + 1 * module.taintObjectIfTainted(_ as URI, args[0]) where: method | args | expected @@ -61,7 +61,7 @@ class URICallSIteTest extends AgentTestRunner { TestURICallSiteSuite.&"$method".call(args as Object[]) then: - 1 * module.taintIfTainted(_, _ as URI) + 1 * module.taintObjectIfTainted(_, _ as URI) where: method | args diff --git a/dd-java-agent/instrumentation/java-net/src/test/groovy/datadog/trace/instrumentation/java/net/URLCallSiteTest.groovy b/dd-java-agent/instrumentation/java-net/src/test/groovy/datadog/trace/instrumentation/java/net/URLCallSiteTest.groovy index a51038e769c..b7b03304903 100644 --- a/dd-java-agent/instrumentation/java-net/src/test/groovy/datadog/trace/instrumentation/java/net/URLCallSiteTest.groovy +++ b/dd-java-agent/instrumentation/java-net/src/test/groovy/datadog/trace/instrumentation/java/net/URLCallSiteTest.groovy @@ -23,7 +23,7 @@ class URLCallSiteTest extends AgentTestRunner { then: uri.toString() == expected - 1 * module.taintIfAnyTainted(_ as URL, args as Object[]) + 1 * module.taintObjectIfAnyTainted(_ as URL, args as Object[]) where: method | args | expected @@ -44,7 +44,7 @@ class URLCallSiteTest extends AgentTestRunner { TestURLCallSiteSuite.&"$method".call(args as Object[]) then: - 1 * module.taintIfTainted(_, _ as URL) + 1 * module.taintObjectIfTainted(_, _ as URL) where: method | args diff --git a/dd-java-agent/instrumentation/java-net/src/test/groovy/datadog/trace/instrumentation/java/net/URLEncoderCallSiteTest.groovy b/dd-java-agent/instrumentation/java-net/src/test/groovy/datadog/trace/instrumentation/java/net/URLEncoderCallSiteTest.groovy index aefc594c745..221cc45a0a8 100644 --- a/dd-java-agent/instrumentation/java-net/src/test/groovy/datadog/trace/instrumentation/java/net/URLEncoderCallSiteTest.groovy +++ b/dd-java-agent/instrumentation/java-net/src/test/groovy/datadog/trace/instrumentation/java/net/URLEncoderCallSiteTest.groovy @@ -23,7 +23,7 @@ class URLEncoderCallSiteTest extends AgentTestRunner { then: result == expected - 1 * iastModule.taintIfTainted(_ as String, args[0], false, VulnerabilityMarks.XSS_MARK) + 1 * iastModule.taintObjectIfTainted(_ as String, args[0], false, VulnerabilityMarks.XSS_MARK) 0 * _ where: diff --git a/dd-java-agent/instrumentation/jersey/src/main/java/datadog/trace/instrumentation/jersey/AbstractFormProviderInstrumentation.java b/dd-java-agent/instrumentation/jersey/src/main/java/datadog/trace/instrumentation/jersey/AbstractFormProviderInstrumentation.java index 3299a704b9c..2303e70e123 100644 --- a/dd-java-agent/instrumentation/jersey/src/main/java/datadog/trace/instrumentation/jersey/AbstractFormProviderInstrumentation.java +++ b/dd-java-agent/instrumentation/jersey/src/main/java/datadog/trace/instrumentation/jersey/AbstractFormProviderInstrumentation.java @@ -54,9 +54,9 @@ public static void onExit( final IastContext ctx = reqCtx.getData(RequestContextSlot.IAST); for (Map.Entry> entry : result.entrySet()) { final String name = entry.getKey(); - prop.taint(ctx, name, SourceTypes.REQUEST_PARAMETER_NAME, name); + prop.taintString(ctx, name, SourceTypes.REQUEST_PARAMETER_NAME, name); for (String value : entry.getValue()) { - prop.taint(ctx, value, SourceTypes.REQUEST_PARAMETER_VALUE, name); + prop.taintString(ctx, value, SourceTypes.REQUEST_PARAMETER_VALUE, name); } } } 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..c98a917c554 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.taintString(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..1a07e1103c9 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.taintString(ctx, (String) result, SourceTypes.REQUEST_PARAMETER_VALUE); } } } diff --git a/dd-java-agent/instrumentation/jersey/src/main/java/datadog/trace/instrumentation/jersey/CookieInstrumentation.java b/dd-java-agent/instrumentation/jersey/src/main/java/datadog/trace/instrumentation/jersey/CookieInstrumentation.java index 863ffe5cdd3..5e95ec601a3 100644 --- a/dd-java-agent/instrumentation/jersey/src/main/java/datadog/trace/instrumentation/jersey/CookieInstrumentation.java +++ b/dd-java-agent/instrumentation/jersey/src/main/java/datadog/trace/instrumentation/jersey/CookieInstrumentation.java @@ -53,7 +53,8 @@ public static void onExit( final PropagationModule module = InstrumentationBridge.PROPAGATION; if (module != null) { IastContext ctx = reqCtx.getData(RequestContextSlot.IAST); - module.taintIfTainted(ctx, cookieName, self, SourceTypes.REQUEST_COOKIE_NAME, cookieName); + module.taintStringIfTainted( + ctx, cookieName, self, SourceTypes.REQUEST_COOKIE_NAME, cookieName); } } } @@ -70,7 +71,7 @@ public static void onExit( final PropagationModule module = InstrumentationBridge.PROPAGATION; if (module != null) { IastContext ctx = reqCtx.getData(RequestContextSlot.IAST); - module.taintIfTainted(ctx, cookieValue, self, SourceTypes.REQUEST_COOKIE_VALUE, name); + module.taintStringIfTainted(ctx, cookieValue, self, SourceTypes.REQUEST_COOKIE_VALUE, name); } } } diff --git a/dd-java-agent/instrumentation/jersey/src/main/java/datadog/trace/instrumentation/jersey/InboundMessageContextInstrumentation.java b/dd-java-agent/instrumentation/jersey/src/main/java/datadog/trace/instrumentation/jersey/InboundMessageContextInstrumentation.java index 46e7df5ec25..153ce8618d5 100644 --- a/dd-java-agent/instrumentation/jersey/src/main/java/datadog/trace/instrumentation/jersey/InboundMessageContextInstrumentation.java +++ b/dd-java-agent/instrumentation/jersey/src/main/java/datadog/trace/instrumentation/jersey/InboundMessageContextInstrumentation.java @@ -57,9 +57,9 @@ public static void onExit( final IastContext ctx = reqCtx.getData(RequestContextSlot.IAST); for (Map.Entry> entry : headers.entrySet()) { final String name = entry.getKey(); - prop.taint(ctx, name, SourceTypes.REQUEST_HEADER_NAME, name); + prop.taintString(ctx, name, SourceTypes.REQUEST_HEADER_NAME, name); for (String value : entry.getValue()) { - prop.taint(ctx, value, SourceTypes.REQUEST_HEADER_VALUE, name); + prop.taintString(ctx, value, SourceTypes.REQUEST_HEADER_VALUE, name); } } } @@ -77,8 +77,8 @@ public static void onExit( final IastContext ctx = reqCtx.getData(RequestContextSlot.IAST); for (Map.Entry entry : cookies.entrySet()) { final String name = entry.getKey(); - module.taint(ctx, name, SourceTypes.REQUEST_COOKIE_NAME, name); - module.taint(ctx, entry.getValue(), SourceTypes.REQUEST_COOKIE_VALUE, name); + module.taintString(ctx, name, SourceTypes.REQUEST_COOKIE_NAME, name); + module.taintObject(ctx, entry.getValue(), SourceTypes.REQUEST_COOKIE_VALUE, name); } } } diff --git a/dd-java-agent/instrumentation/jersey/src/main/java/datadog/trace/instrumentation/jersey/ReaderInterceptorExecutorInstrumentation.java b/dd-java-agent/instrumentation/jersey/src/main/java/datadog/trace/instrumentation/jersey/ReaderInterceptorExecutorInstrumentation.java index 8d52a38c4d7..76509b9a85c 100644 --- a/dd-java-agent/instrumentation/jersey/src/main/java/datadog/trace/instrumentation/jersey/ReaderInterceptorExecutorInstrumentation.java +++ b/dd-java-agent/instrumentation/jersey/src/main/java/datadog/trace/instrumentation/jersey/ReaderInterceptorExecutorInstrumentation.java @@ -45,7 +45,7 @@ static void after( final PropagationModule module = InstrumentationBridge.PROPAGATION; if (module != null) { IastContext ctx = reqCtx.getData(RequestContextSlot.IAST); - module.taint(ctx, inputStream, SourceTypes.REQUEST_BODY); + module.taintObject(ctx, inputStream, SourceTypes.REQUEST_BODY); } } } 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..4c2ea64527f 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.taintString(ctx, (String) value, SourceTypes.REQUEST_HEADER_VALUE, name); } } } diff --git a/dd-java-agent/instrumentation/jose-jwt/src/main/java/datadog/trace/instrumentation/josejwt/JWTParserInstrumentation.java b/dd-java-agent/instrumentation/jose-jwt/src/main/java/datadog/trace/instrumentation/josejwt/JWTParserInstrumentation.java index efefcc9b581..5fcd4340c09 100644 --- a/dd-java-agent/instrumentation/jose-jwt/src/main/java/datadog/trace/instrumentation/josejwt/JWTParserInstrumentation.java +++ b/dd-java-agent/instrumentation/jose-jwt/src/main/java/datadog/trace/instrumentation/josejwt/JWTParserInstrumentation.java @@ -51,7 +51,7 @@ public static void onEnter( // TODO: We could represent this source more accurately, perhaps tracking the original // source, or using a special name. IastContext ctx = reqCtx.getData(RequestContextSlot.IAST); - module.taint(ctx, json, SourceTypes.REQUEST_HEADER_VALUE); + module.taintString(ctx, json, SourceTypes.REQUEST_HEADER_VALUE); } } } diff --git a/dd-java-agent/instrumentation/jose-jwt/src/test/groovy/JWTParserInstrumentationTest.groovy b/dd-java-agent/instrumentation/jose-jwt/src/test/groovy/JWTParserInstrumentationTest.groovy index 97e7d6222a0..13900185d41 100644 --- a/dd-java-agent/instrumentation/jose-jwt/src/test/groovy/JWTParserInstrumentationTest.groovy +++ b/dd-java-agent/instrumentation/jose-jwt/src/test/groovy/JWTParserInstrumentationTest.groovy @@ -29,7 +29,7 @@ class JWTParserInstrumentationTest extends AgentTestRunner { runUnderIastTrace { new com.auth0.jwt.impl.JWTParser().parsePayload(payload) } then: - 1 * propagationModule.taint(iastCtx, payload, SourceTypes.REQUEST_HEADER_VALUE) + 1 * propagationModule.taintObject(iastCtx, payload, SourceTypes.REQUEST_HEADER_VALUE) 0 * _ } @@ -43,11 +43,11 @@ class JWTParserInstrumentationTest extends AgentTestRunner { runUnderIastTrace { com.nimbusds.jose.util.JSONObjectUtils.parse(json) } then: - 1 * propagationModule.taint(iastCtx, 'http://foobar.com', SourceTypes.REQUEST_HEADER_VALUE, 'iss') - 1 * propagationModule.taint(iastCtx, 'foo', SourceTypes.REQUEST_HEADER_VALUE, 'sub') - 1 * propagationModule.taint(iastCtx, 'foobar', SourceTypes.REQUEST_HEADER_VALUE, 'aud') - 1 * propagationModule.taint(iastCtx, 'Mr Foo Bar', SourceTypes.REQUEST_HEADER_VALUE, 'name') - 1 * propagationModule.taint(iastCtx, 'read', SourceTypes.REQUEST_HEADER_VALUE, 'scope') + 1 * propagationModule.taintObject(iastCtx, 'http://foobar.com', SourceTypes.REQUEST_HEADER_VALUE, 'iss') + 1 * propagationModule.taintString(iastCtx, 'foo', SourceTypes.REQUEST_HEADER_VALUE, 'sub') + 1 * propagationModule.taintString(iastCtx, 'foobar', SourceTypes.REQUEST_HEADER_VALUE, 'aud') + 1 * propagationModule.taintString(iastCtx, 'Mr Foo Bar', SourceTypes.REQUEST_HEADER_VALUE, 'name') + 1 * propagationModule.taintString(iastCtx, 'read', SourceTypes.REQUEST_HEADER_VALUE, 'scope') 0 * _ } diff --git a/dd-java-agent/instrumentation/kafka-clients-0.11/build.gradle b/dd-java-agent/instrumentation/kafka-clients-0.11/build.gradle index b03714d31ec..871df95c904 100644 --- a/dd-java-agent/instrumentation/kafka-clients-0.11/build.gradle +++ b/dd-java-agent/instrumentation/kafka-clients-0.11/build.gradle @@ -30,6 +30,7 @@ dependencies { testRuntimeOnly project(':dd-java-agent:instrumentation:java-io') testRuntimeOnly project(':dd-java-agent:instrumentation:jackson-core') testImplementation(group: 'com.fasterxml.jackson.core', name: 'jackson-databind', version: '2.9.10') + testImplementation(testFixtures(project(':dd-java-agent:agent-iast'))) // Include latest version of kafka itself along with latest version of client libs. // This seems to help with jar compatibility hell. diff --git a/dd-java-agent/instrumentation/kafka-clients-0.11/src/iastLatestDepTest3/groovy/iast/KafkaIastDeserializerTest.groovy b/dd-java-agent/instrumentation/kafka-clients-0.11/src/iastLatestDepTest3/groovy/iast/KafkaIastDeserializerTest.groovy index 43ccb4f069d..0eb07610484 100644 --- a/dd-java-agent/instrumentation/kafka-clients-0.11/src/iastLatestDepTest3/groovy/iast/KafkaIastDeserializerTest.groovy +++ b/dd-java-agent/instrumentation/kafka-clients-0.11/src/iastLatestDepTest3/groovy/iast/KafkaIastDeserializerTest.groovy @@ -45,21 +45,21 @@ class KafkaIastDeserializerTest extends AgentTestRunner { then: switch (test.method) { case Method.DEFAULT: - 1 * propagationModule.taint(payload, source) // taint byte[] + 1 * propagationModule.taintObject(payload, source) // taint byte[] 1 * codecModule.onStringFromBytes(payload, 0, payload.length, _, _ as String) // taint byte[] => string break case Method.WITH_HEADERS: - 1 * propagationModule.taint(payload, source) // taint byte[] + 1 * propagationModule.taintObject(payload, source) // taint byte[] 1 * codecModule.onStringFromBytes(payload, 0, payload.length, _, _ as String) // taint byte[] => string break case Method.WITH_BYTE_BUFFER: - 1 * propagationModule.taint(_ as ByteBuffer, source, 0, payload.length) // taint ByteBuffer - 1 * propagationModule.taintIfTainted(payload, _ as ByteBuffer, 0, payload.length, false, NOT_MARKED) // taint ByteBuffer => byte[] + 1 * propagationModule.taintObject(_ as ByteBuffer, source, 0, payload.length) // taint ByteBuffer + 1 * propagationModule.taintStringIfTainted(payload, _ as ByteBuffer, 0, payload.length, false, NOT_MARKED) // taint ByteBuffer => byte[] 1 * codecModule.onStringFromBytes(payload, 0, payload.length, _, _ as String) // taint byte[] => string break case Method.WITH_BYTE_BUFFER_OFFSET: - 1 * propagationModule.taint(_ as ByteBuffer, source, BUFF_OFFSET, payload.length) // taint ByteBuffer - 1 * propagationModule.taintIfTainted(_ as byte[], _ as ByteBuffer, true, NOT_MARKED) // taint ByteBuffer => byte[] + 1 * propagationModule.taintObject(_ as ByteBuffer, source, BUFF_OFFSET, payload.length) // taint ByteBuffer + 1 * propagationModule.taintObjectIfTainted(_ as byte[], _ as ByteBuffer, true, NOT_MARKED) // taint ByteBuffer => byte[] 1 * codecModule.onStringFromBytes(_ as byte[], BUFF_OFFSET, payload.length, _, _ as String) // taint byte[] => string break } @@ -86,18 +86,18 @@ class KafkaIastDeserializerTest extends AgentTestRunner { then: switch (test.method) { case Method.DEFAULT: - 1 * propagationModule.taint(payload, source) // taint byte[] + 1 * propagationModule.taintObject(payload, source) // taint byte[] break case Method.WITH_HEADERS: - 1 * propagationModule.taint(payload, source) // taint byte[] + 1 * propagationModule.taintObject(payload, source) // taint byte[] break case Method.WITH_BYTE_BUFFER: - 1 * propagationModule.taint(_ as ByteBuffer, source, 0, payload.length) // taint ByteBuffer - 1 * propagationModule.taintIfTainted(payload, _ as ByteBuffer, 0, payload.length, false, NOT_MARKED) // taint ByteBuffer => byte[] + 1 * propagationModule.taintObject(_ as ByteBuffer, source, 0, payload.length) // taint ByteBuffer + 1 * propagationModule.taintStringIfTainted(payload, _ as ByteBuffer, 0, payload.length, false, NOT_MARKED) // taint ByteBuffer => byte[] break case Method.WITH_BYTE_BUFFER_OFFSET: - 1 * propagationModule.taint(_ as ByteBuffer, source, BUFF_OFFSET, payload.length) // taint ByteBuffer - 1 * propagationModule.taintIfTainted(payload, _ as ByteBuffer, BUFF_OFFSET, payload.length, false, NOT_MARKED) // taint ByteBuffer => byte[] + 1 * propagationModule.taintObject(_ as ByteBuffer, source, BUFF_OFFSET, payload.length) // taint ByteBuffer + 1 * propagationModule.taintStringIfTainted(payload, _ as ByteBuffer, BUFF_OFFSET, payload.length, false, NOT_MARKED) // taint ByteBuffer => byte[] break } 0 * _ @@ -123,18 +123,18 @@ class KafkaIastDeserializerTest extends AgentTestRunner { then: switch (test.method) { case Method.DEFAULT: - 1 * propagationModule.taint(payload, source) // taint byte[] - 1 * propagationModule.taintIfTainted(_ as ByteBuffer, payload, true, NOT_MARKED) // taint byte[] => ByteBuffer + 1 * propagationModule.taintObject(payload, source) // taint byte[] + 1 * propagationModule.taintObjectIfTainted(_ as ByteBuffer, payload, true, NOT_MARKED) // taint byte[] => ByteBuffer break case Method.WITH_HEADERS: - 1 * propagationModule.taint(payload, source) // taint byte[] - 1 * propagationModule.taintIfTainted(_ as ByteBuffer, payload, true, NOT_MARKED) // taint byte[] => ByteBuffer + 1 * propagationModule.taintObject(payload, source) // taint byte[] + 1 * propagationModule.taintObjectIfTainted(_ as ByteBuffer, payload, true, NOT_MARKED) // taint byte[] => ByteBuffer break case Method.WITH_BYTE_BUFFER: - 1 * propagationModule.taint(_ as ByteBuffer, source, 0, payload.length) // taint ByteBuffer + 1 * propagationModule.taintObject(_ as ByteBuffer, source, 0, payload.length) // taint ByteBuffer break case Method.WITH_BYTE_BUFFER_OFFSET: - 1 * propagationModule.taint(_ as ByteBuffer, source, BUFF_OFFSET, payload.length) // taint ByteBuffer + 1 * propagationModule.taintObject(_ as ByteBuffer, source, BUFF_OFFSET, payload.length) // taint ByteBuffer break } 0 * _ @@ -161,28 +161,28 @@ class KafkaIastDeserializerTest extends AgentTestRunner { then: switch (test.method) { case Method.DEFAULT: - 1 * propagationModule.taint(payload, source) // taint byte[] + 1 * propagationModule.taintObject(payload, source) // taint byte[] break case Method.WITH_HEADERS: - 1 * propagationModule.taint(payload, source) // taint byte[] + 1 * propagationModule.taintObject(payload, source) // taint byte[] break case Method.WITH_BYTE_BUFFER: - 1 * propagationModule.taint(_ as ByteBuffer, source, 0, payload.length) // taint ByteBuffer - 1 * propagationModule.taintIfTainted(payload, _ as ByteBuffer, 0, payload.length, false, NOT_MARKED) // taint byte[] => ByteBuffer + 1 * propagationModule.taintObject(_ as ByteBuffer, source, 0, payload.length) // taint ByteBuffer + 1 * propagationModule.taintStringIfTainted(payload, _ as ByteBuffer, 0, payload.length, false, NOT_MARKED) // taint byte[] => ByteBuffer break case Method.WITH_BYTE_BUFFER_OFFSET: - 1 * propagationModule.taint(_ as ByteBuffer, source, BUFF_OFFSET, payload.length) // taint ByteBuffer - 1 * propagationModule.taintIfTainted(payload, _ as ByteBuffer, BUFF_OFFSET, payload.length, false, NOT_MARKED) // taint byte[] => ByteBuffer + 1 * propagationModule.taintObject(_ as ByteBuffer, source, BUFF_OFFSET, payload.length) // taint ByteBuffer + 1 * propagationModule.taintStringIfTainted(payload, _ as ByteBuffer, BUFF_OFFSET, payload.length, false, NOT_MARKED) // taint byte[] => ByteBuffer break } // taint JSON - 1 * propagationModule.taintIfTainted(_ as JsonParser, payload) + 1 * propagationModule.taintObjectIfTainted(_ as JsonParser, payload) 1 * propagationModule.findSource(_) >> Stub(Taintable.Source) { getOrigin() >> source getValue() >> json } - 1 * propagationModule.taint(_, 'name', source, 'name', json) - 1 * propagationModule.taint(_, 'Mr Bean', source, 'name', json) + 1 * propagationModule.taintObject(_, 'name', source, 'name', json) + 1 * propagationModule.taintString(_, 'Mr Bean', source, 'name', json) 0 * _ where: 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..66e417f6807 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.taintObject(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.taintObjectRange(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..cfd643e28af 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.taintObjectIfRangeTainted(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.taintObjectIfTainted(ctx, buffer, bytes, true, NOT_MARKED); } } } diff --git a/dd-java-agent/instrumentation/kafka-clients-0.11/src/test/groovy/iast/KafkaIastDeserializerForkedTest.groovy b/dd-java-agent/instrumentation/kafka-clients-0.11/src/test/groovy/iast/KafkaIastDeserializerForkedTest.groovy index 15729faf655..cdc1a310dd0 100644 --- a/dd-java-agent/instrumentation/kafka-clients-0.11/src/test/groovy/iast/KafkaIastDeserializerForkedTest.groovy +++ b/dd-java-agent/instrumentation/kafka-clients-0.11/src/test/groovy/iast/KafkaIastDeserializerForkedTest.groovy @@ -1,7 +1,8 @@ package iast +import com.datadog.iast.test.IastAgentTestRunner import com.fasterxml.jackson.core.JsonParser -import datadog.trace.agent.test.AgentTestRunner +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 @@ -13,12 +14,7 @@ import org.apache.kafka.common.serialization.ByteBufferDeserializer import org.apache.kafka.common.serialization.StringDeserializer import org.springframework.kafka.support.serializer.JsonDeserializer -class KafkaIastDeserializerForkedTest extends AgentTestRunner { - - @Override - protected void configurePreAgent() { - injectSysConfig('dd.iast.enabled', 'true') - } +class KafkaIastDeserializerForkedTest extends IastAgentTestRunner { void 'test string deserializer'() { given: @@ -31,11 +27,13 @@ class KafkaIastDeserializerForkedTest extends AgentTestRunner { final deserializer = new StringDeserializer() when: - deserializer.configure([:], source == SourceTypes.KAFKA_MESSAGE_KEY) - deserializer.deserialize("test", payload) + runUnderIastTrace { + deserializer.configure([:], source == SourceTypes.KAFKA_MESSAGE_KEY) + deserializer.deserialize("test", payload) + } then: - 1 * propagationModule.taint(payload, source) + 1 * propagationModule.taintObject(_ as IastContext, payload, source) 1 * codecModule.onStringFromBytes(payload, _, _, _, _) 0 * _ @@ -55,11 +53,13 @@ class KafkaIastDeserializerForkedTest extends AgentTestRunner { final deserializer = new ByteArrayDeserializer() when: - deserializer.configure([:], source == SourceTypes.KAFKA_MESSAGE_KEY) - deserializer.deserialize("test", payload) + runUnderIastTrace { + deserializer.configure([:], source == SourceTypes.KAFKA_MESSAGE_KEY) + deserializer.deserialize("test", payload) + } then: - 1 * propagationModule.taint(payload, source) + 1 * propagationModule.taintObject(_ as IastContext, payload, source) 0 * _ where: @@ -78,12 +78,14 @@ class KafkaIastDeserializerForkedTest extends AgentTestRunner { final deserializer = new ByteBufferDeserializer() when: - deserializer.configure([:], source == SourceTypes.KAFKA_MESSAGE_KEY) - deserializer.deserialize("test", payload) + runUnderIastTrace { + deserializer.configure([:], source == SourceTypes.KAFKA_MESSAGE_KEY) + deserializer.deserialize("test", payload) + } then: - 1 * propagationModule.taint(payload, source) - 1 * propagationModule.taintIfTainted(_, payload, true, VulnerabilityMarks.NOT_MARKED) + 1 * propagationModule.taintObject(_ as IastContext, payload, source) + 1 * propagationModule.taintObjectIfTainted(_ as IastContext, payload, true, VulnerabilityMarks.NOT_MARKED) 0 * _ where: @@ -103,18 +105,20 @@ class KafkaIastDeserializerForkedTest extends AgentTestRunner { final deserializer = new JsonDeserializer(TestBean) when: - deserializer.configure([:], source == SourceTypes.KAFKA_MESSAGE_KEY) - deserializer.deserialize('test', payload) + runUnderIastTrace { + deserializer.configure([:], source == SourceTypes.KAFKA_MESSAGE_KEY) + deserializer.deserialize('test', payload) + } then: - 1 * propagationModule.taint(payload, source) - 1 * propagationModule.taintIfTainted(_ as JsonParser, payload) - 1 * propagationModule.findSource(_) >> Stub(Source) { + 1 * propagationModule.taintObject(_ as IastContext, payload, source) + 1 * propagationModule.taintObjectIfTainted(_ as IastContext, _ as JsonParser, payload) + 1 * propagationModule.findSource(_ as IastContext, _) >> Stub(Source) { getOrigin() >> source getValue() >> json } - 1 * propagationModule.taint(_, 'name', source, 'name', json) - 1 * propagationModule.taint(_, 'Mr Bean', source, 'name', json) + 1 * propagationModule.taintString(_ as IastContext, 'name', source, 'name', json) + 1 * propagationModule.taintString(_ as IastContext, 'Mr Bean', source, 'name', json) 0 * _ where: diff --git a/dd-java-agent/instrumentation/kafka-clients-0.11/src/test/groovy/iast/UtilsInstrumentationForkedTest.groovy b/dd-java-agent/instrumentation/kafka-clients-0.11/src/test/groovy/iast/UtilsInstrumentationForkedTest.groovy index d3332fc6105..b669dc2763e 100644 --- a/dd-java-agent/instrumentation/kafka-clients-0.11/src/test/groovy/iast/UtilsInstrumentationForkedTest.groovy +++ b/dd-java-agent/instrumentation/kafka-clients-0.11/src/test/groovy/iast/UtilsInstrumentationForkedTest.groovy @@ -1,6 +1,7 @@ package iast -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 org.apache.kafka.common.utils.Utils import java.nio.ByteBuffer -class UtilsInstrumentationForkedTest extends AgentTestRunner { - - @Override - protected void configurePreAgent() { - injectSysConfig('dd.iast.enabled', 'true') - } +class UtilsInstrumentationForkedTest extends IastAgentTestRunner { void 'test toArray'() { setup: @@ -22,11 +18,11 @@ class UtilsInstrumentationForkedTest extends AgentTestRunner { final buffer = args[0] when: - final bytes = Utils.&"$method".call(args as Object[]) + final bytes = computeUnderIastTrace { Utils.&"$method".call(args as Object[]) } then: bytes != null - 1 * propagationModule.taintIfTainted(_ as byte[], buffer, offset, length, false, VulnerabilityMarks.NOT_MARKED) + 1 * propagationModule.taintObjectIfRangeTainted(_ as IastContext, _ as byte[], buffer, offset, length, false, VulnerabilityMarks.NOT_MARKED) 0 * _ where: @@ -47,11 +43,11 @@ class UtilsInstrumentationForkedTest extends AgentTestRunner { final bytes = 'Hello World!'.bytes when: - final buffer = Utils.wrapNullable(bytes) + final buffer = computeUnderIastTrace { Utils.wrapNullable(bytes) } then: buffer != null - 1 * propagationModule.taintIfTainted(_ as ByteBuffer, bytes, true, VulnerabilityMarks.NOT_MARKED) + 1 * propagationModule.taintObjectIfTainted(_ as IastContext, _ as ByteBuffer, bytes, true, VulnerabilityMarks.NOT_MARKED) 0 * _ } 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..feb0689bcf8 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.taintObjectIfTainted(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..3783866ef50 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.taintObjectIfTainted(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.taintObjectIfTainted(ctx, result, self); + } } } } diff --git a/dd-java-agent/instrumentation/okhttp-2/src/test/groovy/IastOkHttp2InstrumentationTest.groovy b/dd-java-agent/instrumentation/okhttp-2/src/test/groovy/IastOkHttp2InstrumentationTest.groovy index c77fef63dcd..26ce58f2648 100644 --- a/dd-java-agent/instrumentation/okhttp-2/src/test/groovy/IastOkHttp2InstrumentationTest.groovy +++ b/dd-java-agent/instrumentation/okhttp-2/src/test/groovy/IastOkHttp2InstrumentationTest.groovy @@ -62,12 +62,12 @@ class IastOkHttp2InstrumentationTest extends AgentTestRunner { private void mockPropagation() { final propagation = Mock(PropagationModule) { - taintIfAnyTainted(_, _) >> { + taintObjectIfAnyTainted(_, _) >> { if ((it[1] as List).any { input -> tainteds.containsKey(input) }) { tainteds.put(it[0], null) } } - taintIfTainted(_, _) >> { + taintObjectIfTainted(_, _) >> { if (tainteds.containsKey(it[1])) { tainteds.put(it[0], null) } 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..6db61b7c583 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.taintObjectIfTainted(ctx, result, arg); + } } } } diff --git a/dd-java-agent/instrumentation/okhttp-3/src/test/groovy/IastOkHttp3InstrumentationTest.groovy b/dd-java-agent/instrumentation/okhttp-3/src/test/groovy/IastOkHttp3InstrumentationTest.groovy index 988b112604c..8be29667071 100644 --- a/dd-java-agent/instrumentation/okhttp-3/src/test/groovy/IastOkHttp3InstrumentationTest.groovy +++ b/dd-java-agent/instrumentation/okhttp-3/src/test/groovy/IastOkHttp3InstrumentationTest.groovy @@ -67,7 +67,7 @@ class IastOkHttp3InstrumentationTest extends AgentTestRunner { private void mockPropagation() { final propagation = Mock(PropagationModule) { - taintIfTainted(_, _) >> { + taintObjectIfTainted(_, _) >> { if (tainteds.containsKey(it[1])) { tainteds.put(it[0], null) } 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..256ac219fa5 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.taintObjectIfTainted(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.taintStringIfTainted(ctx, (String) result, self); + } else { + iastModule.taintObjectIfTainted(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.taintStringIfTainted(ctx, (String) result, self); + } else { + iastModule.taintObjectIfTainted(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..6ea15840c66 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.taintObjectIfTainted(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..64f74eb8878 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.taintObjectIfTainted(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.taintStringIfTainted(ctx, (String) result, self); + } else { + iastModule.taintObjectIfTainted(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.taintStringIfTainted(ctx, (String) result, self); + } else { + iastModule.taintObjectIfTainted(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..18c6a775390 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.taintObjectIfTainted(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..0f92d728453 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,16 +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 { - - @Override - void configurePreAgent() { - injectSysConfig("dd.iast.enabled", "true") - } +class JSONArrayInstrumentationTest extends IastAgentTestRunner { void 'test JSONObject returning an array'() { given: @@ -27,17 +23,19 @@ class JSONArrayInstrumentationTest extends AgentTestRunner { }}""" when: - final jsonObject = new JSONObject(json) - final name = jsonObject.getJSONObject("menu").getJSONArray("labels").get(0) + final name = computeUnderIastTrace { + final jsonObject = new JSONObject(json) + return 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 +54,19 @@ class JSONArrayInstrumentationTest extends AgentTestRunner { }}""" when: - final jsonObject = new JSONObject(json) - final name = jsonObject.getJSONObject("menu").getJSONArray("labels").optString(0, "defaultvalue") + final name = computeUnderIastTrace { + final jsonObject = new JSONObject(json) + return 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..8dd1dbcb6be 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,15 +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 { - - @Override void configurePreAgent() { - injectSysConfig("dd.iast.enabled", "true") - } +class JSONCookieInstrumentationTest extends IastAgentTestRunner { void 'test JSon Cookie toJSONObject'() { given: @@ -18,12 +15,11 @@ 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..398b17bdcbd 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,14 +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 { - - @Override void configurePreAgent() { - injectSysConfig("dd.iast.enabled", "true") - } +class JSONObjectInstrumentationTest extends IastAgentTestRunner { void 'test JSONObject string constructor'() { given: @@ -25,16 +22,18 @@ class JSONObjectInstrumentationTest extends AgentTestRunner { }}""" when: - final jsonObject = new JSONObject(json) - final name = jsonObject.getJSONObject("menu").get("name") + final name = computeUnderIastTrace { + final jsonObject = new JSONObject(json) + return 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 +52,18 @@ class JSONObjectInstrumentationTest extends AgentTestRunner { }}""" when: - final jsonObject = new JSONObject(json) - final name = jsonObject.getJSONObject("menu").optString("name") + final name = computeUnderIastTrace { + final jsonObject = new JSONObject(json) + return 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 +76,16 @@ class JSONObjectInstrumentationTest extends AgentTestRunner { final json = '{"name": "nameTest", "value" : "valueTest"}' when: - final jsonObject = new JSONObject(new JSONTokener(json)) - final name = jsonObject.get("name") + final name = computeUnderIastTrace { + final jsonObject = new JSONObject(new JSONTokener(json)) + return 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 +99,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..c15efe7eee6 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,13 +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 { - - @Override void configurePreAgent() { - injectSysConfig("dd.iast.enabled", "true") - } +class JSONTokenerInstrumentationTest extends IastAgentTestRunner { void 'test JSONTokener string constructor'() { given: @@ -16,10 +13,10 @@ 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/build.gradle b/dd-java-agent/instrumentation/owasp-esapi-2/build.gradle index 01d0a112cee..e175753f535 100644 --- a/dd-java-agent/instrumentation/owasp-esapi-2/build.gradle +++ b/dd-java-agent/instrumentation/owasp-esapi-2/build.gradle @@ -16,6 +16,7 @@ dependencies { compileOnly group: 'org.owasp.esapi', name: 'esapi', version: '2.1.0' testImplementation group: 'org.owasp.esapi', name: 'esapi', version: '2.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/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..0bb1736bb10 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.taintStringIfTainted(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.taintStringIfTainted(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.taintStringIfTainted(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.taintStringIfTainted(ctx, result, input, false, VulnerabilityMarks.XSS_MARK); + } } catch (final Throwable e) { module.onUnexpectedException("afterCanonicalize3 threw", e); } @@ -89,9 +102,13 @@ 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.taintStringIfTainted( + ctx, result, input, false, VulnerabilityMarks.LDAP_INJECTION_MARK); + } } catch (final Throwable e) { module.onUnexpectedException("afterEncodeForLDAP threw", e); } @@ -107,9 +124,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.taintStringIfTainted( + ctx, result, input, false, VulnerabilityMarks.COMMAND_INJECTION_MARK); + } } catch (final Throwable e) { module.onUnexpectedException("afterEncodeForOS threw", e); } @@ -125,9 +146,13 @@ 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.taintStringIfTainted( + ctx, result, input, false, VulnerabilityMarks.SQL_INJECTION_MARK); + } } catch (final Throwable e) { module.onUnexpectedException("afterEncodeForSQL threw", e); } diff --git a/dd-java-agent/instrumentation/owasp-esapi-2/src/test/groovy/datadog/trace/instrumentation/owasp/esapi/EncoderCallSiteTest.groovy b/dd-java-agent/instrumentation/owasp-esapi-2/src/test/groovy/datadog/trace/instrumentation/owasp/esapi/EncoderCallSiteTest.groovy index e7bd272fd6d..e37c023f69c 100644 --- a/dd-java-agent/instrumentation/owasp-esapi-2/src/test/groovy/datadog/trace/instrumentation/owasp/esapi/EncoderCallSiteTest.groovy +++ b/dd-java-agent/instrumentation/owasp-esapi-2/src/test/groovy/datadog/trace/instrumentation/owasp/esapi/EncoderCallSiteTest.groovy @@ -1,6 +1,7 @@ package datadog.trace.instrumentation.owasp.esapi -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,24 +9,24 @@ import foo.bar.TestEncoderSuite import org.owasp.esapi.Encoder import org.owasp.esapi.codecs.Codec -class EncoderCallSiteTest extends AgentTestRunner { - - @Override - protected void configurePreAgent() { - injectSysConfig("dd.iast.enabled", "true") - } +class EncoderCallSiteTest extends IastAgentTestRunner { void 'test #method propagation with mark #mark'() { given: final module = Mock(PropagationModule) - final testSuite = new TestEncoderSuite(Mock(Encoder)) + final encoder = Stub(Encoder) { + "$method"(_) >> { + return it[0] + "Encoded" + } + } + final testSuite = new TestEncoderSuite(encoder) InstrumentationBridge.registerIastModule(module) when: - testSuite.&"$method".call(args) + runUnderIastTrace { testSuite.&"$method".call(args) } then: - 1 * module.taintIfTainted(_, _, false, mark) + 1 * module.taintObjectIfTainted(_ as IastContext, _, _, false, mark) 0 * module._ where: @@ -35,7 +36,7 @@ class EncoderCallSiteTest extends AgentTestRunner { 'canonicalize' | ['Ø-This is a quote', true] | VulnerabilityMarks.XSS_MARK 'canonicalize' | ['Ø-This is a quote', true, true] | VulnerabilityMarks.XSS_MARK 'encodeForLDAP' | ['Ø-This is a quote'] | VulnerabilityMarks.LDAP_INJECTION_MARK - 'encodeForOS' | [Mock(Codec), 'Ø-This is a quote'] | VulnerabilityMarks.COMMAND_INJECTION_MARK - 'encodeForSQL' | [Mock(Codec), 'Ø-This is a quote'] | VulnerabilityMarks.SQL_INJECTION_MARK + 'encodeForOS' | [Stub(Codec), 'Ø-This is a quote'] | VulnerabilityMarks.COMMAND_INJECTION_MARK + 'encodeForSQL' | [Stub(Codec), 'Ø-This is a quote'] | VulnerabilityMarks.SQL_INJECTION_MARK } } diff --git a/dd-java-agent/instrumentation/owasp-esapi-2/src/test/java/foo/bar/TestEncoderSuite.java b/dd-java-agent/instrumentation/owasp-esapi-2/src/test/java/foo/bar/TestEncoderSuite.java index 21f91858fc0..3dec41606a0 100644 --- a/dd-java-agent/instrumentation/owasp-esapi-2/src/test/java/foo/bar/TestEncoderSuite.java +++ b/dd-java-agent/instrumentation/owasp-esapi-2/src/test/java/foo/bar/TestEncoderSuite.java @@ -7,7 +7,7 @@ public class TestEncoderSuite { private Encoder encoder; - TestEncoderSuite(Encoder encoder) { + public TestEncoderSuite(Encoder encoder) { this.encoder = encoder; } diff --git a/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/CookieHeaderInstrumentation.java b/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/CookieHeaderInstrumentation.java index ddf71a8f5a4..a3103aa58f4 100644 --- a/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/CookieHeaderInstrumentation.java +++ b/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/CookieHeaderInstrumentation.java @@ -75,8 +75,8 @@ static void after( while (iterator.hasNext()) { HttpCookiePair pair = iterator.next(); final String name = pair.name(), value = pair.value(); - prop.taint(ctx, name, SourceTypes.REQUEST_COOKIE_NAME, name); - prop.taint(ctx, value, SourceTypes.REQUEST_COOKIE_VALUE, name); + prop.taintString(ctx, name, SourceTypes.REQUEST_COOKIE_NAME, name); + prop.taintString(ctx, value, SourceTypes.REQUEST_COOKIE_VALUE, name); } } } diff --git a/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/HeaderNameCallSite.java b/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/HeaderNameCallSite.java index 39ec2647462..f0de41cd3da 100644 --- a/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/HeaderNameCallSite.java +++ b/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/HeaderNameCallSite.java @@ -33,7 +33,7 @@ public static String after(@CallSite.This HttpHeader header, @CallSite.Return St if (ctx == null) { return result; } - module.taintIfTainted(ctx, result, header, SourceTypes.REQUEST_HEADER_NAME, result); + module.taintStringIfTainted(ctx, result, header, SourceTypes.REQUEST_HEADER_NAME, result); } catch (final Throwable e) { module.onUnexpectedException("onHeaderNames threw", e); } diff --git a/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/HttpHeaderSubclassesInstrumentation.java b/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/HttpHeaderSubclassesInstrumentation.java index abea74fca74..1e7c9d2a72e 100644 --- a/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/HttpHeaderSubclassesInstrumentation.java +++ b/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/HttpHeaderSubclassesInstrumentation.java @@ -72,7 +72,7 @@ static void onExit( return; } IastContext ctx = reqCtx.getData(RequestContextSlot.IAST); - propagation.taintIfTainted(ctx, retVal, h); + propagation.taintStringIfTainted(ctx, retVal, h); } } } diff --git a/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/HttpRequestInstrumentation.java b/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/HttpRequestInstrumentation.java index 01e4c352961..db02dc8a6c3 100644 --- a/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/HttpRequestInstrumentation.java +++ b/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/HttpRequestInstrumentation.java @@ -89,7 +89,7 @@ static void onExit( } // unfortunately, the call to h.value() is instrumented, but // because the call to taint() only happens after, the call is a noop - propagation.taint(ctx, h, SourceTypes.REQUEST_HEADER_VALUE, h.name(), h.value()); + propagation.taintObject(ctx, h, SourceTypes.REQUEST_HEADER_VALUE, h.name(), h.value()); } } } @@ -113,7 +113,7 @@ static void onExit( return; } - propagation.taintIfTainted(ctx, entity, thiz); + propagation.taintObjectIfTainted(ctx, entity, thiz); } } } 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..92d4298066e 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.taintString(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/RequestContextInstrumentation.java b/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/RequestContextInstrumentation.java index 6d252fe3564..25a3c5b44c5 100644 --- a/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/RequestContextInstrumentation.java +++ b/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/RequestContextInstrumentation.java @@ -67,7 +67,7 @@ static void onExit( return; } - propagation.taintIfTainted(ctx, request, requestContext); + propagation.taintObjectIfTainted(ctx, request, requestContext); } } } diff --git a/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/UriInstrumentation.java b/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/UriInstrumentation.java index ef7354d358d..663daebf2f3 100644 --- a/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/UriInstrumentation.java +++ b/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/UriInstrumentation.java @@ -80,7 +80,7 @@ static void after( return; } IastContext ctx = reqCtx.getData(RequestContextSlot.IAST); - mod.taintIfTainted(ctx, ret.get(), uri); + mod.taintStringIfTainted(ctx, ret.get(), uri); } } @@ -108,8 +108,8 @@ static void after( while (iterator.hasNext()) { Tuple2 pair = iterator.next(); final String name = pair._1(), value = pair._2(); - prop.taint(ctx, name, SourceTypes.REQUEST_PARAMETER_NAME, name); - prop.taint(ctx, value, SourceTypes.REQUEST_PARAMETER_VALUE, name); + prop.taintString(ctx, name, SourceTypes.REQUEST_PARAMETER_NAME, name); + prop.taintString(ctx, value, SourceTypes.REQUEST_PARAMETER_VALUE, name); } } } diff --git a/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintCookieFunction.java b/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintCookieFunction.java index 57246093509..bc3e83b493b 100644 --- a/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintCookieFunction.java +++ b/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintCookieFunction.java @@ -27,8 +27,8 @@ public Tuple1 apply(Tuple1 v1) { } final String name = httpCookiePair.name(); final String value = httpCookiePair.value(); - mod.taint(ctx, name, SourceTypes.REQUEST_COOKIE_NAME, name); - mod.taint(ctx, value, SourceTypes.REQUEST_COOKIE_VALUE, name); + mod.taintString(ctx, name, SourceTypes.REQUEST_COOKIE_NAME, name); + mod.taintString(ctx, value, SourceTypes.REQUEST_COOKIE_VALUE, name); return v1; } } diff --git a/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintFutureHelper.java b/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintFutureHelper.java index 5ad320a442c..fe1459d632f 100644 --- a/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintFutureHelper.java +++ b/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintFutureHelper.java @@ -14,7 +14,7 @@ public static Future wrapFuture( t -> { IastContext ctx = IastContext.Provider.get(AgentTracer.activeSpan()); if (ctx != null) { - mod.taintIfTainted(ctx, t, input); + mod.taintObjectIfTainted(ctx, t, input); } return t; }; diff --git a/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintMapFunction.java b/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintMapFunction.java index 44343912c37..aaf2b8ee79a 100644 --- a/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintMapFunction.java +++ b/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintMapFunction.java @@ -32,8 +32,8 @@ public Tuple1> apply(Tuple1> v1) { while (iterator.hasNext()) { Tuple2 e = iterator.next(); final String name = e._1(), value = e._2(); - prop.taint(ctx, name, SourceTypes.REQUEST_PARAMETER_NAME, name); - prop.taint(ctx, value, SourceTypes.REQUEST_PARAMETER_VALUE, name); + prop.taintString(ctx, name, SourceTypes.REQUEST_PARAMETER_NAME, name); + prop.taintString(ctx, value, SourceTypes.REQUEST_PARAMETER_VALUE, name); } return v1; diff --git a/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintMultiMapFunction.java b/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintMultiMapFunction.java index f9c6d2a0948..58dcebd688b 100644 --- a/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintMultiMapFunction.java +++ b/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintMultiMapFunction.java @@ -33,10 +33,10 @@ public Tuple1>> apply(Tuple1>> while (entriesIterator.hasNext()) { Tuple2> e = entriesIterator.next(); final String name = e._1(); - mod.taint(ctx, name, SourceTypes.REQUEST_PARAMETER_NAME, name); + mod.taintString(ctx, name, SourceTypes.REQUEST_PARAMETER_NAME, name); List values = e._2(); for (final String value : ScalaToJava.listAsList(values)) { - mod.taint(ctx, value, SourceTypes.REQUEST_PARAMETER_VALUE, name); + mod.taintString(ctx, value, SourceTypes.REQUEST_PARAMETER_VALUE, name); } } diff --git a/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintOptionalCookieFunction.java b/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintOptionalCookieFunction.java index b50c22bccf8..a8b1569945c 100644 --- a/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintOptionalCookieFunction.java +++ b/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintOptionalCookieFunction.java @@ -29,8 +29,8 @@ public Tuple1> apply(Tuple1> v1) { final HttpCookiePair cookie = httpCookiePair.get(); final String name = cookie.name(); final String value = cookie.value(); - mod.taint(ctx, name, SourceTypes.REQUEST_COOKIE_NAME, name); - mod.taint(ctx, value, SourceTypes.REQUEST_COOKIE_VALUE, name); + mod.taintString(ctx, name, SourceTypes.REQUEST_COOKIE_NAME, name); + mod.taintString(ctx, value, SourceTypes.REQUEST_COOKIE_VALUE, name); return v1; } } 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..1ae24ca9b3c 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.taintString(ctx, (String) o, SourceTypes.REQUEST_PARAMETER_VALUE, paramName); } } } else if (value instanceof String) { - mod.taint(ctx, value, SourceTypes.REQUEST_PARAMETER_VALUE, paramName); + mod.taintString(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..14eab48cad4 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,14 +18,14 @@ 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()); if (ctx == null) { return v1; } - mod.taint(ctx, reqCtx, SourceTypes.REQUEST_BODY); + mod.taintObject(ctx, reqCtx, SourceTypes.REQUEST_BODY); return v1; } 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..2c31ced0d07 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,14 +19,14 @@ 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()); if (ctx == null) { return v1; } - mod.taint(ctx, httpRequest, SourceTypes.REQUEST_BODY); + mod.taintObject(ctx, httpRequest, SourceTypes.REQUEST_BODY); return v1; } diff --git a/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintSeqFunction.java b/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintSeqFunction.java index 04a843008f7..dc49d4f5efc 100644 --- a/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintSeqFunction.java +++ b/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintSeqFunction.java @@ -39,9 +39,9 @@ public Tuple1>> apply(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.taintString(ctx, (String) o, SourceTypes.REQUEST_PARAMETER_VALUE, paramName); } } } else if (value instanceof String) { - mod.taint(ctx, value, SourceTypes.REQUEST_PARAMETER_VALUE, paramName); + mod.taintString(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..60938488b27 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.taintObject(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..b74addc2a99 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,14 +17,14 @@ 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()); if (ctx == null) { return v1; } - mod.taint(ctx, uri, SourceTypes.REQUEST_QUERY); + mod.taintObject(ctx, uri, SourceTypes.REQUEST_QUERY); return v1; } 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..282c2ec34af 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.taintString(ctx, (String) o, SourceTypes.REQUEST_COOKIE_VALUE, paramName); } } } else { - module.taint(ctx, result, SourceTypes.REQUEST_COOKIE_VALUE, paramName); + module.taintString(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..5068bd1b3f9 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.taintString(ctx, (String) o, SourceTypes.REQUEST_PARAMETER_VALUE, paramName); } } } else { - module.taint(ctx, result, SourceTypes.REQUEST_PARAMETER_VALUE, paramName); + module.taintString(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..1dc2e4afc07 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.taintString(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.taintString(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..0edfa333cf9 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.taintString(ctx, (String) o, SourceTypes.REQUEST_PATH_PARAMETER, paramName); } } } else { - module.taint(ctx, result, SourceTypes.REQUEST_PATH_PARAMETER, paramName); + module.taintString(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..3a3c7e323b0 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.taintString(ctx, (String) o, SourceTypes.REQUEST_PARAMETER_VALUE, paramName); } } } else { - module.taint(ctx, result, SourceTypes.REQUEST_PARAMETER_VALUE, paramName); + module.taintString(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..a93e24ade0a 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.taintStringIfTainted(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..7f155984c4e 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.taintObjectIfTainted(_ 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.taintObjectIfTainted(_ 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-3/src/main/java/datadog/trace/instrumentation/servlet3/MultipartInstrumentation.java b/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/MultipartInstrumentation.java index 850c70de5b4..814626f119f 100644 --- a/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/MultipartInstrumentation.java +++ b/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/MultipartInstrumentation.java @@ -69,7 +69,8 @@ public static String onExit( final PropagationModule module = InstrumentationBridge.PROPAGATION; if (module != null) { IastContext ctx = reqCtx.getData(RequestContextSlot.IAST); - module.taint(ctx, name, SourceTypes.REQUEST_MULTIPART_PARAMETER, "Content-Disposition"); + module.taintString( + ctx, name, SourceTypes.REQUEST_MULTIPART_PARAMETER, "Content-Disposition"); } return name; } @@ -86,7 +87,7 @@ public static String onExit( final PropagationModule module = InstrumentationBridge.PROPAGATION; if (module != null) { IastContext ctx = reqCtx.getData(RequestContextSlot.IAST); - module.taint(ctx, value, SourceTypes.REQUEST_MULTIPART_PARAMETER, name); + module.taintString(ctx, value, SourceTypes.REQUEST_MULTIPART_PARAMETER, name); } return value; } @@ -107,7 +108,7 @@ public static void onExit( if (module != null) { final IastContext ctx = reqCtx.getData(RequestContextSlot.IAST); for (final String value : headerValues) { - module.taint(ctx, value, SourceTypes.REQUEST_MULTIPART_PARAMETER, headerName); + module.taintString(ctx, value, SourceTypes.REQUEST_MULTIPART_PARAMETER, headerName); } } } @@ -127,7 +128,7 @@ public static void onExit( if (module != null) { final IastContext ctx = reqCtx.getData(RequestContextSlot.IAST); for (final String name : headerNames) { - module.taint(ctx, name, SourceTypes.REQUEST_MULTIPART_PARAMETER); + module.taintString(ctx, name, SourceTypes.REQUEST_MULTIPART_PARAMETER); } } } diff --git a/dd-java-agent/instrumentation/servlet/request-3/src/test/groovy/MultipartInstrumentationForkedTest.groovy b/dd-java-agent/instrumentation/servlet/request-3/src/test/groovy/MultipartInstrumentationForkedTest.groovy index 4800e1f3197..83e7221e6ff 100644 --- a/dd-java-agent/instrumentation/servlet/request-3/src/test/groovy/MultipartInstrumentationForkedTest.groovy +++ b/dd-java-agent/instrumentation/servlet/request-3/src/test/groovy/MultipartInstrumentationForkedTest.groovy @@ -36,7 +36,7 @@ class MultipartInstrumentationForkedTest extends AgentTestRunner { runUnderIastTrace { part.getName() } then: - 1 * module.taint(iastCtx, 'partName', SourceTypes.REQUEST_MULTIPART_PARAMETER, 'Content-Disposition') + 1 * module.taintObject(iastCtx, 'partName', SourceTypes.REQUEST_MULTIPART_PARAMETER, 'Content-Disposition') 0 * _ } @@ -50,7 +50,7 @@ class MultipartInstrumentationForkedTest extends AgentTestRunner { runUnderIastTrace { part.getHeader('headerName') } then: - 1 * module.taint(iastCtx, 'headerValue', SourceTypes.REQUEST_MULTIPART_PARAMETER, 'headerName') + 1 * module.taintObject(iastCtx, 'headerValue', SourceTypes.REQUEST_MULTIPART_PARAMETER, 'headerName') 0 * _ } @@ -64,7 +64,7 @@ class MultipartInstrumentationForkedTest extends AgentTestRunner { runUnderIastTrace { part.getHeaders('headerName') } then: - 1 * module.taint(iastCtx, 'headerValue', SourceTypes.REQUEST_MULTIPART_PARAMETER, 'headerName') + 1 * module.taintObject(iastCtx, 'headerValue', SourceTypes.REQUEST_MULTIPART_PARAMETER, 'headerName') 0 * _ } @@ -78,7 +78,7 @@ class MultipartInstrumentationForkedTest extends AgentTestRunner { runUnderIastTrace { part.getHeaderNames() } then: - 1 * module.taint(iastCtx, 'headerName', SourceTypes.REQUEST_MULTIPART_PARAMETER) + 1 * module.taintObject(iastCtx, 'headerName', SourceTypes.REQUEST_MULTIPART_PARAMETER) 0 * _ } diff --git a/dd-java-agent/instrumentation/servlet/request-5/src/main/java/datadog/trace/instrumentation/servlet5/JakartaHttpServletRequestCallSite.java b/dd-java-agent/instrumentation/servlet/request-5/src/main/java/datadog/trace/instrumentation/servlet5/JakartaHttpServletRequestCallSite.java index c4f9a7d4bac..b1dc62cb58a 100644 --- a/dd-java-agent/instrumentation/servlet/request-5/src/main/java/datadog/trace/instrumentation/servlet5/JakartaHttpServletRequestCallSite.java +++ b/dd-java-agent/instrumentation/servlet/request-5/src/main/java/datadog/trace/instrumentation/servlet5/JakartaHttpServletRequestCallSite.java @@ -33,7 +33,7 @@ public static String afterPath( try { final IastContext ctx = IastContext.Provider.get(AgentTracer.activeSpan()); if (ctx != null) { - module.taint(ctx, retValue, SourceTypes.REQUEST_PATH); + module.taintString(ctx, retValue, SourceTypes.REQUEST_PATH); } } catch (final Throwable e) { module.onUnexpectedException("afterPath threw", e); @@ -55,7 +55,7 @@ public static StringBuffer afterGetRequestURL( try { final IastContext ctx = IastContext.Provider.get(AgentTracer.activeSpan()); if (ctx != null) { - module.taint(ctx, retValue, SourceTypes.REQUEST_URI); + module.taintObject(ctx, retValue, SourceTypes.REQUEST_URI); } } catch (final Throwable e) { module.onUnexpectedException("afterGetRequestURL threw", e); diff --git a/dd-java-agent/instrumentation/servlet/request-5/src/main/java/datadog/trace/instrumentation/servlet5/JakartaHttpServletRequestInstrumentation.java b/dd-java-agent/instrumentation/servlet/request-5/src/main/java/datadog/trace/instrumentation/servlet5/JakartaHttpServletRequestInstrumentation.java index 58a65edbc62..f487a5ed35e 100644 --- a/dd-java-agent/instrumentation/servlet/request-5/src/main/java/datadog/trace/instrumentation/servlet5/JakartaHttpServletRequestInstrumentation.java +++ b/dd-java-agent/instrumentation/servlet/request-5/src/main/java/datadog/trace/instrumentation/servlet5/JakartaHttpServletRequestInstrumentation.java @@ -114,7 +114,7 @@ public static void onExit( return; } IastContext ctx = reqCtx.getData(RequestContextSlot.IAST); - module.taint(ctx, value, SourceTypes.REQUEST_HEADER_VALUE, name); + module.taintString(ctx, value, SourceTypes.REQUEST_HEADER_VALUE, name); } } @@ -177,7 +177,7 @@ public static void onExit( return; } IastContext ctx = reqCtx.getData(RequestContextSlot.IAST); - module.taint(ctx, value, SourceTypes.REQUEST_PARAMETER_VALUE, name); + module.taintString(ctx, value, SourceTypes.REQUEST_PARAMETER_VALUE, name); } } @@ -198,7 +198,7 @@ public static void onExit( } final IastContext ctx = reqCtx.getData(RequestContextSlot.IAST); for (final String value : values) { - module.taint(ctx, value, SourceTypes.REQUEST_PARAMETER_VALUE, name); + module.taintString(ctx, value, SourceTypes.REQUEST_PARAMETER_VALUE, name); } } } @@ -220,11 +220,11 @@ public static void onExit( final IastContext ctx = reqCtx.getData(RequestContextSlot.IAST); for (final Map.Entry entry : parameters.entrySet()) { final String name = entry.getKey(); - module.taint(ctx, name, SourceTypes.REQUEST_PARAMETER_NAME, name); + module.taintString(ctx, name, SourceTypes.REQUEST_PARAMETER_NAME, name); final String[] values = entry.getValue(); if (values != null) { for (final String value : entry.getValue()) { - module.taint(ctx, value, SourceTypes.REQUEST_PARAMETER_VALUE, name); + module.taintString(ctx, value, SourceTypes.REQUEST_PARAMETER_VALUE, name); } } } @@ -268,7 +268,7 @@ public static void onExit( } final IastContext ctx = reqCtx.getData(RequestContextSlot.IAST); for (final Cookie cookie : cookies) { - module.taint(ctx, cookie, SourceTypes.REQUEST_COOKIE_VALUE); + module.taintObject(ctx, cookie, SourceTypes.REQUEST_COOKIE_VALUE); } } } @@ -287,7 +287,7 @@ public static void onExit( return; } IastContext ctx = reqCtx.getData(RequestContextSlot.IAST); - module.taint(ctx, queryString, SourceTypes.REQUEST_QUERY); + module.taintString(ctx, queryString, SourceTypes.REQUEST_QUERY); } } @@ -305,7 +305,7 @@ public static void onExit( return; } IastContext ctx = reqCtx.getData(RequestContextSlot.IAST); - module.taint(ctx, body, SourceTypes.REQUEST_BODY); + module.taintObject(ctx, body, SourceTypes.REQUEST_BODY); } } 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..324a3d4361b 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.taintStringIfTainted(ctx, encoded, url); } } } diff --git a/dd-java-agent/instrumentation/servlet/request-5/src/main/java/datadog/trace/instrumentation/servlet5/JakartaMultipartInstrumentation.java b/dd-java-agent/instrumentation/servlet/request-5/src/main/java/datadog/trace/instrumentation/servlet5/JakartaMultipartInstrumentation.java index 05bf05ab568..32eb2732631 100644 --- a/dd-java-agent/instrumentation/servlet/request-5/src/main/java/datadog/trace/instrumentation/servlet5/JakartaMultipartInstrumentation.java +++ b/dd-java-agent/instrumentation/servlet/request-5/src/main/java/datadog/trace/instrumentation/servlet5/JakartaMultipartInstrumentation.java @@ -65,7 +65,8 @@ public static String onExit( final PropagationModule module = InstrumentationBridge.PROPAGATION; if (module != null) { final IastContext ctx = reqCtx.getData(RequestContextSlot.IAST); - module.taint(ctx, name, SourceTypes.REQUEST_MULTIPART_PARAMETER, "Content-Disposition"); + module.taintString( + ctx, name, SourceTypes.REQUEST_MULTIPART_PARAMETER, "Content-Disposition"); } return name; } @@ -82,7 +83,7 @@ public static String onExit( final PropagationModule module = InstrumentationBridge.PROPAGATION; if (module != null) { final IastContext ctx = reqCtx.getData(RequestContextSlot.IAST); - module.taint(ctx, value, SourceTypes.REQUEST_MULTIPART_PARAMETER, name); + module.taintString(ctx, value, SourceTypes.REQUEST_MULTIPART_PARAMETER, name); } return value; } @@ -103,7 +104,7 @@ public static void onExit( if (module != null) { final IastContext ctx = reqCtx.getData(RequestContextSlot.IAST); for (final String value : headerValues) { - module.taint(ctx, value, SourceTypes.REQUEST_MULTIPART_PARAMETER, headerName); + module.taintString(ctx, value, SourceTypes.REQUEST_MULTIPART_PARAMETER, headerName); } } } @@ -123,7 +124,7 @@ public static void onExit( if (module != null) { final IastContext ctx = reqCtx.getData(RequestContextSlot.IAST); for (final String name : headerNames) { - module.taint(ctx, name, SourceTypes.REQUEST_MULTIPART_PARAMETER); + module.taintString(ctx, name, SourceTypes.REQUEST_MULTIPART_PARAMETER); } } } diff --git a/dd-java-agent/instrumentation/servlet/request-5/src/test/groovy/JakartaHttpServletRequestInstrumentationTest.groovy b/dd-java-agent/instrumentation/servlet/request-5/src/test/groovy/JakartaHttpServletRequestInstrumentationTest.groovy index 932ccaa8dab..83287c46c53 100644 --- a/dd-java-agent/instrumentation/servlet/request-5/src/test/groovy/JakartaHttpServletRequestInstrumentationTest.groovy +++ b/dd-java-agent/instrumentation/servlet/request-5/src/test/groovy/JakartaHttpServletRequestInstrumentationTest.groovy @@ -47,7 +47,7 @@ class JakartaHttpServletRequestInstrumentationTest extends AgentTestRunner { then: result == 'value' 1 * mock.getHeader('header') >> 'value' - 1 * iastModule.taint(iastCtx, 'value', SourceTypes.REQUEST_HEADER_VALUE, 'header') + 1 * iastModule.taintObject(iastCtx, 'value', SourceTypes.REQUEST_HEADER_VALUE, 'header') 0 * _ where: @@ -68,7 +68,7 @@ class JakartaHttpServletRequestInstrumentationTest extends AgentTestRunner { then: result == headers 1 * mock.getHeaders('headers') >> Collections.enumeration(headers) - headers.each { 1 * iastModule.taint(iastCtx, it, SourceTypes.REQUEST_HEADER_VALUE, 'headers') } + headers.each { 1 * iastModule.taintObject(iastCtx, it, SourceTypes.REQUEST_HEADER_VALUE, 'headers') } 0 * _ where: @@ -89,7 +89,7 @@ class JakartaHttpServletRequestInstrumentationTest extends AgentTestRunner { then: result == headers 1 * mock.getHeaderNames() >> Collections.enumeration(headers) - headers.each { 1 * iastModule.taint(iastCtx, it, SourceTypes.REQUEST_HEADER_NAME, it) } + headers.each { 1 * iastModule.taintObject(iastCtx, it, SourceTypes.REQUEST_HEADER_NAME, it) } 0 * _ where: @@ -109,7 +109,7 @@ class JakartaHttpServletRequestInstrumentationTest extends AgentTestRunner { then: result == 'value' 1 * mock.getParameter('parameter') >> 'value' - 1 * iastModule.taint(iastCtx, 'value', SourceTypes.REQUEST_PARAMETER_VALUE, 'parameter') + 1 * iastModule.taintObject(iastCtx, 'value', SourceTypes.REQUEST_PARAMETER_VALUE, 'parameter') 0 * _ where: @@ -130,7 +130,7 @@ class JakartaHttpServletRequestInstrumentationTest extends AgentTestRunner { then: result == values 1 * mock.getParameterValues('parameter') >> { values as String[] } - values.each { 1 * iastModule.taint(iastCtx, it, SourceTypes.REQUEST_PARAMETER_VALUE, 'parameter') } + values.each { 1 * iastModule.taintObject(iastCtx, it, SourceTypes.REQUEST_PARAMETER_VALUE, 'parameter') } 0 * _ where: @@ -152,9 +152,9 @@ class JakartaHttpServletRequestInstrumentationTest extends AgentTestRunner { result == parameters 1 * mock.getParameterMap() >> parameters parameters.each { key, values -> - 1 * iastModule.taint(iastCtx, key, SourceTypes.REQUEST_PARAMETER_NAME, key) + 1 * iastModule.taintObject(iastCtx, key, SourceTypes.REQUEST_PARAMETER_NAME, key) values.each { value -> - 1 * iastModule.taint(iastCtx, value, SourceTypes.REQUEST_PARAMETER_VALUE, key) + 1 * iastModule.taintString(iastCtx, value, SourceTypes.REQUEST_PARAMETER_VALUE, key) } } 0 * _ @@ -178,7 +178,7 @@ class JakartaHttpServletRequestInstrumentationTest extends AgentTestRunner { then: result == parameters 1 * mock.getParameterNames() >> Collections.enumeration(parameters) - parameters.each { 1 * iastModule.taint(iastCtx, it, SourceTypes.REQUEST_PARAMETER_NAME, it) } + parameters.each { 1 * iastModule.taintObject(iastCtx, it, SourceTypes.REQUEST_PARAMETER_NAME, it) } 0 * _ where: @@ -199,7 +199,7 @@ class JakartaHttpServletRequestInstrumentationTest extends AgentTestRunner { then: result == cookies 1 * mock.getCookies() >> cookies - cookies.each { 1 * iastModule.taint(iastCtx, it, SourceTypes.REQUEST_COOKIE_VALUE) } + cookies.each { 1 * iastModule.taintObject(iastCtx, it, SourceTypes.REQUEST_COOKIE_VALUE) } 0 * _ where: @@ -276,7 +276,7 @@ class JakartaHttpServletRequestInstrumentationTest extends AgentTestRunner { then: result == queryString 1 * mock.getQueryString() >> queryString - 1 * iastModule.taint(iastCtx, queryString, SourceTypes.REQUEST_QUERY) + 1 * iastModule.taintObject(iastCtx, queryString, SourceTypes.REQUEST_QUERY) 0 * _ where: @@ -297,7 +297,7 @@ class JakartaHttpServletRequestInstrumentationTest extends AgentTestRunner { then: result == is 1 * mock.getInputStream() >> is - 1 * iastModule.taint(iastCtx, is, SourceTypes.REQUEST_BODY) + 1 * iastModule.taintObject(iastCtx, is, SourceTypes.REQUEST_BODY) 0 * _ where: @@ -318,7 +318,7 @@ class JakartaHttpServletRequestInstrumentationTest extends AgentTestRunner { then: result == reader 1 * mock.getReader() >> reader - 1 * iastModule.taint(iastCtx, reader, SourceTypes.REQUEST_BODY) + 1 * iastModule.taintObject(iastCtx, reader, SourceTypes.REQUEST_BODY) 0 * _ where: @@ -361,7 +361,7 @@ class JakartaHttpServletRequestInstrumentationTest extends AgentTestRunner { then: result == uri 1 * mock.getRequestURI() >> uri - 1 * iastModule.taint(iastCtx, uri, SourceTypes.REQUEST_PATH) + 1 * iastModule.taintObject(iastCtx, uri, SourceTypes.REQUEST_PATH) 0 * _ where: @@ -382,7 +382,7 @@ class JakartaHttpServletRequestInstrumentationTest extends AgentTestRunner { then: result == pathInfo 1 * mock.getPathInfo() >> pathInfo - 1 * iastModule.taint(iastCtx, pathInfo, SourceTypes.REQUEST_PATH) + 1 * iastModule.taintObject(iastCtx, pathInfo, SourceTypes.REQUEST_PATH) 0 * _ where: @@ -403,7 +403,7 @@ class JakartaHttpServletRequestInstrumentationTest extends AgentTestRunner { then: result == pathTranslated 1 * mock.getPathTranslated() >> pathTranslated - 1 * iastModule.taint(iastCtx, pathTranslated, SourceTypes.REQUEST_PATH) + 1 * iastModule.taintObject(iastCtx, pathTranslated, SourceTypes.REQUEST_PATH) 0 * _ where: @@ -424,7 +424,7 @@ class JakartaHttpServletRequestInstrumentationTest extends AgentTestRunner { then: result == url 1 * mock.getRequestURL() >> url - 1 * iastModule.taint(iastCtx, url, SourceTypes.REQUEST_URI) + 1 * iastModule.taintObject(iastCtx, url, SourceTypes.REQUEST_URI) 0 * _ where: diff --git a/dd-java-agent/instrumentation/servlet/request-5/src/test/groovy/JakartaHttpServletResponseInstrumentationTest.groovy b/dd-java-agent/instrumentation/servlet/request-5/src/test/groovy/JakartaHttpServletResponseInstrumentationTest.groovy index e4b7cbdcb04..d15d975fa89 100644 --- a/dd-java-agent/instrumentation/servlet/request-5/src/test/groovy/JakartaHttpServletResponseInstrumentationTest.groovy +++ b/dd-java-agent/instrumentation/servlet/request-5/src/test/groovy/JakartaHttpServletResponseInstrumentationTest.groovy @@ -219,7 +219,7 @@ class JakartaHttpServletResponseInstrumentationTest extends AgentTestRunner { result = response.encodeRedirectURL(url) then: - 1 * module.taintIfTainted(_, url) >> { args -> expected = args[0] } + 1 * module.taintObjectIfTainted(_, url) >> { args -> expected = args[0] } 0 * _ result == expected } @@ -236,7 +236,7 @@ class JakartaHttpServletResponseInstrumentationTest extends AgentTestRunner { result = response.encodeURL(url) then: - 1 * module.taintIfTainted(_, url) >> { args -> expected = args[0] } + 1 * module.taintObjectIfTainted(_, url) >> { args -> expected = args[0] } 0 * _ expected == result } diff --git a/dd-java-agent/instrumentation/servlet/request-5/src/test/groovy/JakartaMultipartInstrumentationTest.groovy b/dd-java-agent/instrumentation/servlet/request-5/src/test/groovy/JakartaMultipartInstrumentationTest.groovy index 287bdfe506c..af82b7c4019 100644 --- a/dd-java-agent/instrumentation/servlet/request-5/src/test/groovy/JakartaMultipartInstrumentationTest.groovy +++ b/dd-java-agent/instrumentation/servlet/request-5/src/test/groovy/JakartaMultipartInstrumentationTest.groovy @@ -35,7 +35,7 @@ class JakartaMultipartInstrumentationTest extends AgentTestRunner { runUnderIastTrace { part.getName() } then: - 1 * module.taint(iastCtx, 'partName', SourceTypes.REQUEST_MULTIPART_PARAMETER, 'Content-Disposition') + 1 * module.taintObject(iastCtx, 'partName', SourceTypes.REQUEST_MULTIPART_PARAMETER, 'Content-Disposition') 0 * _ } @@ -49,7 +49,7 @@ class JakartaMultipartInstrumentationTest extends AgentTestRunner { runUnderIastTrace { part.getHeader('headerName') } then: - 1 * module.taint(iastCtx, 'headerValue', SourceTypes.REQUEST_MULTIPART_PARAMETER, 'headerName') + 1 * module.taintObject(iastCtx, 'headerValue', SourceTypes.REQUEST_MULTIPART_PARAMETER, 'headerName') 0 * _ } @@ -63,7 +63,7 @@ class JakartaMultipartInstrumentationTest extends AgentTestRunner { runUnderIastTrace { part.getHeaders('headerName') } then: - 1 * module.taint(iastCtx, 'headerValue', SourceTypes.REQUEST_MULTIPART_PARAMETER, 'headerName') + 1 * module.taintObject(iastCtx, 'headerValue', SourceTypes.REQUEST_MULTIPART_PARAMETER, 'headerName') 0 * _ } @@ -77,7 +77,7 @@ class JakartaMultipartInstrumentationTest extends AgentTestRunner { runUnderIastTrace { part.getHeaderNames() } then: - 1 * module.taint(iastCtx, 'headerName', SourceTypes.REQUEST_MULTIPART_PARAMETER) + 1 * module.taintObject(iastCtx, 'headerName', SourceTypes.REQUEST_MULTIPART_PARAMETER) 0 * _ } diff --git a/dd-java-agent/instrumentation/servlet/src/main/java/datadog/trace/instrumentation/servlet/http/CookieInstrumentation.java b/dd-java-agent/instrumentation/servlet/src/main/java/datadog/trace/instrumentation/servlet/http/CookieInstrumentation.java index 004b1aacdff..5f2d997ccff 100644 --- a/dd-java-agent/instrumentation/servlet/src/main/java/datadog/trace/instrumentation/servlet/http/CookieInstrumentation.java +++ b/dd-java-agent/instrumentation/servlet/src/main/java/datadog/trace/instrumentation/servlet/http/CookieInstrumentation.java @@ -59,7 +59,7 @@ public static void afterGetName( if (module != null) { try { IastContext ctx = reqCtx.getData(RequestContextSlot.IAST); - module.taintIfTainted(ctx, result, self, SourceTypes.REQUEST_COOKIE_NAME, result); + module.taintStringIfTainted(ctx, result, self, SourceTypes.REQUEST_COOKIE_NAME, result); } catch (final Throwable e) { module.onUnexpectedException("afterGetName threw", e); } @@ -81,7 +81,7 @@ public static void afterGetValue( if (module != null) { try { IastContext ctx = reqCtx.getData(RequestContextSlot.IAST); - module.taintIfTainted(ctx, result, self, SourceTypes.REQUEST_COOKIE_VALUE, name); + module.taintStringIfTainted(ctx, result, self, SourceTypes.REQUEST_COOKIE_VALUE, name); } catch (final Throwable e) { module.onUnexpectedException("getValue threw", e); } diff --git a/dd-java-agent/instrumentation/servlet/src/main/java/datadog/trace/instrumentation/servlet/http/HttpServletRequestCallSite.java b/dd-java-agent/instrumentation/servlet/src/main/java/datadog/trace/instrumentation/servlet/http/HttpServletRequestCallSite.java index c0f45f7cb5c..6eb24635b92 100644 --- a/dd-java-agent/instrumentation/servlet/src/main/java/datadog/trace/instrumentation/servlet/http/HttpServletRequestCallSite.java +++ b/dd-java-agent/instrumentation/servlet/src/main/java/datadog/trace/instrumentation/servlet/http/HttpServletRequestCallSite.java @@ -33,7 +33,7 @@ public static String afterPath( try { final IastContext ctx = IastContext.Provider.get(AgentTracer.activeSpan()); if (ctx != null) { - module.taint(ctx, retValue, SourceTypes.REQUEST_PATH); + module.taintString(ctx, retValue, SourceTypes.REQUEST_PATH); } } catch (final Throwable e) { module.onUnexpectedException("afterPath threw", e); @@ -55,7 +55,7 @@ public static StringBuffer afterGetRequestURL( try { final IastContext ctx = IastContext.Provider.get(AgentTracer.activeSpan()); if (ctx != null) { - module.taint(ctx, retValue, SourceTypes.REQUEST_URI); + module.taintObject(ctx, retValue, SourceTypes.REQUEST_URI); } } catch (final Throwable e) { module.onUnexpectedException("afterGetRequestURL threw", e); diff --git a/dd-java-agent/instrumentation/servlet/src/main/java/datadog/trace/instrumentation/servlet/http/HttpServletRequestInstrumentation.java b/dd-java-agent/instrumentation/servlet/src/main/java/datadog/trace/instrumentation/servlet/http/HttpServletRequestInstrumentation.java index cb2d1eebb80..5ae518435d5 100644 --- a/dd-java-agent/instrumentation/servlet/src/main/java/datadog/trace/instrumentation/servlet/http/HttpServletRequestInstrumentation.java +++ b/dd-java-agent/instrumentation/servlet/src/main/java/datadog/trace/instrumentation/servlet/http/HttpServletRequestInstrumentation.java @@ -114,7 +114,7 @@ public static void onExit( return; } IastContext ctx = reqCtx.getData(RequestContextSlot.IAST); - module.taint(ctx, value, SourceTypes.REQUEST_HEADER_VALUE, name); + module.taintString(ctx, value, SourceTypes.REQUEST_HEADER_VALUE, name); } } @@ -177,7 +177,7 @@ public static void onExit( return; } IastContext ctx = reqCtx.getData(RequestContextSlot.IAST); - module.taint(ctx, value, SourceTypes.REQUEST_PARAMETER_VALUE, name); + module.taintString(ctx, value, SourceTypes.REQUEST_PARAMETER_VALUE, name); } } @@ -198,7 +198,7 @@ public static void onExit( } final IastContext ctx = reqCtx.getData(RequestContextSlot.IAST); for (final String value : values) { - module.taint(ctx, value, SourceTypes.REQUEST_PARAMETER_VALUE, name); + module.taintString(ctx, value, SourceTypes.REQUEST_PARAMETER_VALUE, name); } } } @@ -220,11 +220,11 @@ public static void onExit( final IastContext ctx = reqCtx.getData(RequestContextSlot.IAST); for (final Map.Entry entry : parameters.entrySet()) { final String name = entry.getKey(); - module.taint(ctx, name, SourceTypes.REQUEST_PARAMETER_NAME, name); + module.taintString(ctx, name, SourceTypes.REQUEST_PARAMETER_NAME, name); final String[] values = entry.getValue(); if (values != null) { for (final String value : entry.getValue()) { - module.taint(ctx, value, SourceTypes.REQUEST_PARAMETER_VALUE, name); + module.taintString(ctx, value, SourceTypes.REQUEST_PARAMETER_VALUE, name); } } } @@ -268,7 +268,7 @@ public static void onExit( } final IastContext ctx = reqCtx.getData(RequestContextSlot.IAST); for (final Cookie cookie : cookies) { - module.taint(ctx, cookie, SourceTypes.REQUEST_COOKIE_VALUE); + module.taintObject(ctx, cookie, SourceTypes.REQUEST_COOKIE_VALUE); } } } @@ -287,7 +287,7 @@ public static void onExit( return; } IastContext ctx = reqCtx.getData(RequestContextSlot.IAST); - module.taint(ctx, queryString, SourceTypes.REQUEST_QUERY); + module.taintString(ctx, queryString, SourceTypes.REQUEST_QUERY); } } @@ -305,7 +305,7 @@ public static void onExit( return; } IastContext ctx = reqCtx.getData(RequestContextSlot.IAST); - module.taint(ctx, body, SourceTypes.REQUEST_BODY); + module.taintObject(ctx, body, SourceTypes.REQUEST_BODY); } } diff --git a/dd-java-agent/instrumentation/servlet/src/test/groovy/CookieInstrumentationTest.groovy b/dd-java-agent/instrumentation/servlet/src/test/groovy/CookieInstrumentationTest.groovy index 2bef1727ae5..8d78ef7404c 100644 --- a/dd-java-agent/instrumentation/servlet/src/test/groovy/CookieInstrumentationTest.groovy +++ b/dd-java-agent/instrumentation/servlet/src/test/groovy/CookieInstrumentationTest.groovy @@ -38,7 +38,7 @@ class CookieInstrumentationTest extends AgentTestRunner { then: result == NAME - 1 * iastModule.taintIfTainted(iastCtx, NAME, cookie, SourceTypes.REQUEST_COOKIE_NAME, NAME) + 1 * iastModule.taintObjectIfTainted(iastCtx, NAME, cookie, SourceTypes.REQUEST_COOKIE_NAME, NAME) } void 'test getValue'() { @@ -52,7 +52,7 @@ class CookieInstrumentationTest extends AgentTestRunner { then: result == VALUE - 1 * iastModule.taintIfTainted(iastCtx, VALUE, cookie, SourceTypes.REQUEST_COOKIE_VALUE, NAME) + 1 * iastModule.taintObjectIfTainted(iastCtx, VALUE, cookie, SourceTypes.REQUEST_COOKIE_VALUE, NAME) } protected E runUnderIastTrace(Closure cl) { diff --git a/dd-java-agent/instrumentation/servlet/src/test/groovy/HttpServletRequestTest.groovy b/dd-java-agent/instrumentation/servlet/src/test/groovy/HttpServletRequestTest.groovy index 6b3f535fb81..05c715d1410 100644 --- a/dd-java-agent/instrumentation/servlet/src/test/groovy/HttpServletRequestTest.groovy +++ b/dd-java-agent/instrumentation/servlet/src/test/groovy/HttpServletRequestTest.groovy @@ -49,7 +49,7 @@ class HttpServletRequestTest extends AgentTestRunner { then: result == 'value' 1 * mock.getHeader('header') >> 'value' - 1 * iastModule.taint(iastCtx, 'value', SourceTypes.REQUEST_HEADER_VALUE, 'header') + 1 * iastModule.taintObject(iastCtx, 'value', SourceTypes.REQUEST_HEADER_VALUE, 'header') 0 * _ where: @@ -70,7 +70,7 @@ class HttpServletRequestTest extends AgentTestRunner { then: result == headers 1 * mock.getHeaders('headers') >> Collections.enumeration(headers) - headers.each { 1 * iastModule.taint(iastCtx, it, SourceTypes.REQUEST_HEADER_VALUE, 'headers') } + headers.each { 1 * iastModule.taintObject(iastCtx, it, SourceTypes.REQUEST_HEADER_VALUE, 'headers') } 0 * _ where: @@ -91,7 +91,7 @@ class HttpServletRequestTest extends AgentTestRunner { then: result == headers 1 * mock.getHeaderNames() >> Collections.enumeration(headers) - headers.each { 1 * iastModule.taint(iastCtx, it, SourceTypes.REQUEST_HEADER_NAME, it) } + headers.each { 1 * iastModule.taintObject(iastCtx, it, SourceTypes.REQUEST_HEADER_NAME, it) } 0 * _ where: @@ -111,7 +111,7 @@ class HttpServletRequestTest extends AgentTestRunner { then: result == 'value' 1 * mock.getParameter('parameter') >> 'value' - 1 * iastModule.taint(iastCtx, 'value', SourceTypes.REQUEST_PARAMETER_VALUE, 'parameter') + 1 * iastModule.taintObject(iastCtx, 'value', SourceTypes.REQUEST_PARAMETER_VALUE, 'parameter') 0 * _ where: @@ -132,7 +132,7 @@ class HttpServletRequestTest extends AgentTestRunner { then: result == values 1 * mock.getParameterValues('parameter') >> { values as String[] } - values.each { 1 * iastModule.taint(iastCtx, it, SourceTypes.REQUEST_PARAMETER_VALUE, 'parameter') } + values.each { 1 * iastModule.taintObject(iastCtx, it, SourceTypes.REQUEST_PARAMETER_VALUE, 'parameter') } 0 * _ where: @@ -154,9 +154,9 @@ class HttpServletRequestTest extends AgentTestRunner { result == parameters 1 * mock.getParameterMap() >> parameters parameters.each { key, values -> - 1 * iastModule.taint(iastCtx, key, SourceTypes.REQUEST_PARAMETER_NAME, key) + 1 * iastModule.taintObject(iastCtx, key, SourceTypes.REQUEST_PARAMETER_NAME, key) values.each { value -> - 1 * iastModule.taint(iastCtx, value, SourceTypes.REQUEST_PARAMETER_VALUE, key) + 1 * iastModule.taintString(iastCtx, value, SourceTypes.REQUEST_PARAMETER_VALUE, key) } } 0 * _ @@ -179,7 +179,7 @@ class HttpServletRequestTest extends AgentTestRunner { then: result == parameters 1 * mock.getParameterNames() >> Collections.enumeration(parameters) - parameters.each { 1 * iastModule.taint(iastCtx, it, SourceTypes.REQUEST_PARAMETER_NAME, it) } + parameters.each { 1 * iastModule.taintObject(iastCtx, it, SourceTypes.REQUEST_PARAMETER_NAME, it) } 0 * _ where: @@ -200,7 +200,7 @@ class HttpServletRequestTest extends AgentTestRunner { then: result == cookies 1 * mock.getCookies() >> cookies - cookies.each { 1 * iastModule.taint(iastCtx, it, SourceTypes.REQUEST_COOKIE_VALUE) } + cookies.each { 1 * iastModule.taintObject(iastCtx, it, SourceTypes.REQUEST_COOKIE_VALUE) } 0 * _ where: @@ -277,7 +277,7 @@ class HttpServletRequestTest extends AgentTestRunner { then: result == queryString 1 * mock.getQueryString() >> queryString - 1 * iastModule.taint(iastCtx, queryString, SourceTypes.REQUEST_QUERY) + 1 * iastModule.taintObject(iastCtx, queryString, SourceTypes.REQUEST_QUERY) 0 * _ where: @@ -298,7 +298,7 @@ class HttpServletRequestTest extends AgentTestRunner { then: result == is 1 * mock.getInputStream() >> is - 1 * iastModule.taint(iastCtx, is, SourceTypes.REQUEST_BODY) + 1 * iastModule.taintObject(iastCtx, is, SourceTypes.REQUEST_BODY) 0 * _ where: @@ -319,7 +319,7 @@ class HttpServletRequestTest extends AgentTestRunner { then: result == reader 1 * mock.getReader() >> reader - 1 * iastModule.taint(iastCtx, reader, SourceTypes.REQUEST_BODY) + 1 * iastModule.taintObject(iastCtx, reader, SourceTypes.REQUEST_BODY) 0 * _ where: @@ -362,7 +362,7 @@ class HttpServletRequestTest extends AgentTestRunner { then: result == uri 1 * mock.getRequestURI() >> uri - 1 * iastModule.taint(iastCtx, uri, SourceTypes.REQUEST_PATH) + 1 * iastModule.taintObject(iastCtx, uri, SourceTypes.REQUEST_PATH) 0 * _ where: @@ -383,7 +383,7 @@ class HttpServletRequestTest extends AgentTestRunner { then: result == pathInfo 1 * mock.getPathInfo() >> pathInfo - 1 * iastModule.taint(iastCtx, pathInfo, SourceTypes.REQUEST_PATH) + 1 * iastModule.taintObject(iastCtx, pathInfo, SourceTypes.REQUEST_PATH) 0 * _ where: @@ -404,7 +404,7 @@ class HttpServletRequestTest extends AgentTestRunner { then: result == pathTranslated 1 * mock.getPathTranslated() >> pathTranslated - 1 * iastModule.taint(iastCtx, pathTranslated, SourceTypes.REQUEST_PATH) + 1 * iastModule.taintObject(iastCtx, pathTranslated, SourceTypes.REQUEST_PATH) 0 * _ where: @@ -425,7 +425,7 @@ class HttpServletRequestTest extends AgentTestRunner { then: result == url 1 * mock.getRequestURL() >> url - 1 * iastModule.taint(iastCtx, url, SourceTypes.REQUEST_URI) + 1 * iastModule.taintObject(iastCtx, url, SourceTypes.REQUEST_URI) 0 * _ where: diff --git a/dd-java-agent/instrumentation/spring-core/build.gradle b/dd-java-agent/instrumentation/spring-core/build.gradle index 70258ccb85b..5a3486ee736 100644 --- a/dd-java-agent/instrumentation/spring-core/build.gradle +++ b/dd-java-agent/instrumentation/spring-core/build.gradle @@ -14,6 +14,7 @@ addTestSuiteForDir('latestDepTest', 'test') dependencies { compileOnly group: 'org.springframework', name: 'spring-beans', version: '3.2.2.RELEASE' testImplementation group: 'org.springframework', name: 'spring-beans', version: '3.2.2.RELEASE' + testImplementation(testFixtures(project(':dd-java-agent:agent-iast'))) testRuntimeOnly project(':dd-java-agent:instrumentation:iast-instrumenter') latestDepTestImplementation group: 'org.unbescape', name: 'unbescape', version: '+' } 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..4ffff4d953b 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.taintStringIfTainted(ctx, string, in); + } } } diff --git a/dd-java-agent/instrumentation/spring-core/src/test/groovy/datadog/trace/instrumentation/springcore/StreamUtilsInstrumentationTest.groovy b/dd-java-agent/instrumentation/spring-core/src/test/groovy/datadog/trace/instrumentation/springcore/StreamUtilsInstrumentationTest.groovy index fe58485ce0b..aec55aaa66b 100644 --- a/dd-java-agent/instrumentation/spring-core/src/test/groovy/datadog/trace/instrumentation/springcore/StreamUtilsInstrumentationTest.groovy +++ b/dd-java-agent/instrumentation/spring-core/src/test/groovy/datadog/trace/instrumentation/springcore/StreamUtilsInstrumentationTest.groovy @@ -1,18 +1,13 @@ package datadog.trace.instrumentation.springcore -import datadog.trace.agent.test.AgentTestRunner +import com.datadog.iast.test.IastAgentTestRunner import datadog.trace.api.iast.InstrumentationBridge import datadog.trace.api.iast.propagation.PropagationModule import org.springframework.util.StreamUtils import java.nio.charset.StandardCharsets -class StreamUtilsInstrumentationTest extends AgentTestRunner { - - @Override - protected void configurePreAgent() { - injectSysConfig("dd.iast.enabled", "true") - } +class StreamUtilsInstrumentationTest extends IastAgentTestRunner { void 'test'(){ setup: @@ -21,10 +16,10 @@ class StreamUtilsInstrumentationTest extends AgentTestRunner { InstrumentationBridge.registerIastModule(module) when: - StreamUtils.copyToString(new ByteArrayInputStream("test".getBytes()), StandardCharsets.ISO_8859_1) + runUnderIastTrace { StreamUtils.copyToString(new ByteArrayInputStream("test".getBytes()), StandardCharsets.ISO_8859_1) } then: - 1 * module.taintIfTainted(_ as String, _ as InputStream) + 1 * module.taintObjectIfTainted(_ as String, _ as InputStream) 0 * _ } } diff --git a/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/AbstractServerHttpRequestInstrumentation.java b/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/AbstractServerHttpRequestInstrumentation.java index 37811b137f1..80e63208d02 100644 --- a/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/AbstractServerHttpRequestInstrumentation.java +++ b/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/AbstractServerHttpRequestInstrumentation.java @@ -50,7 +50,7 @@ public static void after( return; } IastContext ctx = reqCtx.getData(RequestContextSlot.IAST); - propagation.taint(ctx, object, SourceTypes.REQUEST_HEADER_VALUE); + propagation.taintObject(ctx, object, SourceTypes.REQUEST_HEADER_VALUE); } } } 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..d14c64ff48a 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.taintObjectIfTainted(ctx, is, dataBuffer); + } } } diff --git a/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/HandleMatchAdvice.java b/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/HandleMatchAdvice.java index 5ac2054dec7..c3f1f992e91 100644 --- a/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/HandleMatchAdvice.java +++ b/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/HandleMatchAdvice.java @@ -40,7 +40,7 @@ public static void after( if (parameterName == null || value == null) { continue; // should not happen } - module.taint( + module.taintString( iastRequestContext, value, SourceTypes.REQUEST_PATH_PARAMETER, parameterName); } } @@ -57,7 +57,7 @@ public static void after( for (Map.Entry> ie : value.entrySet()) { String innerKey = ie.getKey(); if (innerKey != null) { - module.taint( + module.taintString( iastRequestContext, innerKey, SourceTypes.REQUEST_MATRIX_PARAMETER, @@ -66,7 +66,7 @@ public static void after( Iterable innerValues = ie.getValue(); if (innerValues != null) { for (String iv : innerValues) { - module.taint( + module.taintString( iastRequestContext, iv, SourceTypes.REQUEST_MATRIX_PARAMETER, parameterName); } } diff --git a/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/HttpMessageInstrumentation.java b/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/HttpMessageInstrumentation.java index 363f4c7b53b..6d800b29445 100644 --- a/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/HttpMessageInstrumentation.java +++ b/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/HttpMessageInstrumentation.java @@ -58,7 +58,7 @@ public static void after( return; } IastContext ctx = reqCtx.getData(RequestContextSlot.IAST); - propagation.taint(ctx, object, SourceTypes.REQUEST_HEADER_VALUE); + propagation.taintObject(ctx, object, SourceTypes.REQUEST_HEADER_VALUE); } } } diff --git a/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/ReactorServerHttpRequestInstrumentation.java b/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/ReactorServerHttpRequestInstrumentation.java index 6b2de237618..4d4a73a8a72 100644 --- a/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/ReactorServerHttpRequestInstrumentation.java +++ b/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/ReactorServerHttpRequestInstrumentation.java @@ -48,7 +48,7 @@ public static void after( return; } IastContext ctx = reqCtx.getData(RequestContextSlot.IAST); - propagation.taint(ctx, object, SourceTypes.REQUEST_HEADER_VALUE); + propagation.taintObject(ctx, object, SourceTypes.REQUEST_HEADER_VALUE); } } } diff --git a/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/RequestHeaderMapResolveAdvice.java b/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/RequestHeaderMapResolveAdvice.java index f72c97ef3b7..75bec42145e 100644 --- a/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/RequestHeaderMapResolveAdvice.java +++ b/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/RequestHeaderMapResolveAdvice.java @@ -32,17 +32,17 @@ public static void after( for (Map.Entry> e : ((MultiValueMap) values).entrySet()) { final String name = e.getKey(); - prop.taint(ctx, name, SourceTypes.REQUEST_HEADER_NAME, name); + prop.taintString(ctx, name, SourceTypes.REQUEST_HEADER_NAME, name); for (String v : e.getValue()) { - prop.taint(ctx, v, SourceTypes.REQUEST_HEADER_VALUE, name); + prop.taintString(ctx, v, SourceTypes.REQUEST_HEADER_VALUE, name); } } } else { for (Map.Entry e : ((Map) values).entrySet()) { final String name = e.getKey(); final String value = e.getValue(); - prop.taint(ctx, name, SourceTypes.REQUEST_HEADER_NAME, name); - prop.taint(ctx, value, SourceTypes.REQUEST_HEADER_VALUE, name); + prop.taintString(ctx, name, SourceTypes.REQUEST_HEADER_NAME, name); + prop.taintString(ctx, value, SourceTypes.REQUEST_HEADER_VALUE, name); } } } diff --git a/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/ServerServletHttpRequestInstrumentation.java b/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/ServerServletHttpRequestInstrumentation.java index bbafdeb8342..0755022ab42 100644 --- a/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/ServerServletHttpRequestInstrumentation.java +++ b/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/ServerServletHttpRequestInstrumentation.java @@ -48,7 +48,7 @@ public static void after( return; } IastContext ctx = reqCtx.getData(RequestContextSlot.IAST); - propagation.taint(ctx, object, SourceTypes.REQUEST_HEADER_VALUE); + propagation.taintObject(ctx, object, SourceTypes.REQUEST_HEADER_VALUE); } } } diff --git a/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/TaintCookiesAdvice.java b/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/TaintCookiesAdvice.java index 3f887639d1a..291abe554d1 100644 --- a/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/TaintCookiesAdvice.java +++ b/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/TaintCookiesAdvice.java @@ -30,8 +30,8 @@ public static void after( for (HttpCookie cookie : cookieList) { final String name = cookie.getName(); final String value = cookie.getValue(); - module.taint(ctx, name, SourceTypes.REQUEST_COOKIE_NAME, name); - module.taint(ctx, value, SourceTypes.REQUEST_COOKIE_VALUE, name); + module.taintString(ctx, name, SourceTypes.REQUEST_COOKIE_NAME, name); + module.taintString(ctx, value, SourceTypes.REQUEST_COOKIE_VALUE, name); } } } diff --git a/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/TaintFluxElementsFunction.java b/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/TaintFluxElementsFunction.java index e532f9392aa..6086d084f34 100644 --- a/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/TaintFluxElementsFunction.java +++ b/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/TaintFluxElementsFunction.java @@ -17,7 +17,7 @@ public TaintFluxElementsFunction(IastContext ctx, PropagationModule propagationM @Override public T apply(T t) { - propagation.taint(ctx, t, SourceTypes.REQUEST_BODY); + propagation.taintObject(ctx, t, SourceTypes.REQUEST_BODY); return t; } } diff --git a/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/TaintGetBodyAdvice.java b/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/TaintGetBodyAdvice.java index 7adf1ca893f..ec2eaeaa155 100644 --- a/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/TaintGetBodyAdvice.java +++ b/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/TaintGetBodyAdvice.java @@ -27,7 +27,7 @@ public static void after( // taint both the flux and the individual DataBuffers IastContext ctx = reqCtx.getData(RequestContextSlot.IAST); - propagation.taint(ctx, flux, SourceTypes.REQUEST_BODY); + propagation.taintObject(ctx, flux, SourceTypes.REQUEST_BODY); flux = flux.map(new TaintFluxElementsFunction<>(ctx, propagation)); } } diff --git a/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/TaintHttpHeadersGetAdvice.java b/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/TaintHttpHeadersGetAdvice.java index 47b22d7e320..58305275194 100644 --- a/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/TaintHttpHeadersGetAdvice.java +++ b/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/TaintHttpHeadersGetAdvice.java @@ -34,7 +34,7 @@ public static void after( final IastContext ctx = reqCtx.getData(RequestContextSlot.IAST); String lc = ((String) arg).toLowerCase(Locale.ROOT); for (String value : values) { - module.taintIfTainted(ctx, value, self, SourceTypes.REQUEST_HEADER_VALUE, lc); + module.taintStringIfTainted(ctx, value, self, SourceTypes.REQUEST_HEADER_VALUE, lc); } } } diff --git a/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/TaintHttpHeadersGetFirstAdvice.java b/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/TaintHttpHeadersGetFirstAdvice.java index 6fa89bd14c3..9a882277008 100644 --- a/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/TaintHttpHeadersGetFirstAdvice.java +++ b/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/TaintHttpHeadersGetFirstAdvice.java @@ -27,6 +27,6 @@ public static void after( return; } IastContext ctx = reqCtx.getData(RequestContextSlot.IAST); - module.taintIfTainted(ctx, value, self, SourceTypes.REQUEST_HEADER_VALUE, arg); + module.taintStringIfTainted(ctx, value, self, SourceTypes.REQUEST_HEADER_VALUE, arg); } } diff --git a/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/TaintHttpHeadersToSingleValueMapAdvice.java b/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/TaintHttpHeadersToSingleValueMapAdvice.java index 4401007c1d8..5a64eb49ff9 100644 --- a/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/TaintHttpHeadersToSingleValueMapAdvice.java +++ b/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/TaintHttpHeadersToSingleValueMapAdvice.java @@ -31,8 +31,8 @@ public static void after( for (Map.Entry e : values.entrySet()) { final String name = e.getKey(); final String value = e.getValue(); - module.taintIfTainted(ctx, name, self, SourceTypes.REQUEST_HEADER_NAME, name); - module.taintIfTainted(ctx, value, self, SourceTypes.REQUEST_HEADER_VALUE, name); + module.taintStringIfTainted(ctx, name, self, SourceTypes.REQUEST_HEADER_NAME, name); + module.taintStringIfTainted(ctx, value, self, SourceTypes.REQUEST_HEADER_VALUE, name); } } } diff --git a/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/TaintQueryParamsAdvice.java b/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/TaintQueryParamsAdvice.java index 692e2e2f890..31e0f23397a 100644 --- a/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/TaintQueryParamsAdvice.java +++ b/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/TaintQueryParamsAdvice.java @@ -31,9 +31,9 @@ public static void after( final IastContext ctx = reqCtx.getData(RequestContextSlot.IAST); for (Map.Entry> e : queryParams.entrySet()) { String name = e.getKey(); - prop.taint(ctx, name, SourceTypes.REQUEST_PARAMETER_NAME, name); + prop.taintString(ctx, name, SourceTypes.REQUEST_PARAMETER_NAME, name); for (String value : e.getValue()) { - prop.taint(ctx, value, SourceTypes.REQUEST_PARAMETER_VALUE, name); + prop.taintString(ctx, value, SourceTypes.REQUEST_PARAMETER_VALUE, name); } } } diff --git a/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/TaintReadOnlyHttpHeadersAdvice.java b/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/TaintReadOnlyHttpHeadersAdvice.java index 397cf87e5ef..2efe52c6c79 100644 --- a/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/TaintReadOnlyHttpHeadersAdvice.java +++ b/dd-java-agent/instrumentation/spring-webflux-5/src/main/java/datadog/trace/instrumentation/springwebflux/server/iast/TaintReadOnlyHttpHeadersAdvice.java @@ -26,6 +26,6 @@ public static void after( return; } final IastContext ctx = reqCtx.getData(RequestContextSlot.IAST); - module.taintIfTainted(ctx, retValue, headers); + module.taintObjectIfTainted(ctx, retValue, headers); } } diff --git a/dd-java-agent/instrumentation/spring-webflux-5/src/test/groovy/dd/trace/instrumentation/springwebflux/client/HeadersAdviceForkedTest.groovy b/dd-java-agent/instrumentation/spring-webflux-5/src/test/groovy/dd/trace/instrumentation/springwebflux/client/HeadersAdviceForkedTest.groovy index ee1e47d5f14..7231985efb6 100644 --- a/dd-java-agent/instrumentation/spring-webflux-5/src/test/groovy/dd/trace/instrumentation/springwebflux/client/HeadersAdviceForkedTest.groovy +++ b/dd-java-agent/instrumentation/spring-webflux-5/src/test/groovy/dd/trace/instrumentation/springwebflux/client/HeadersAdviceForkedTest.groovy @@ -35,7 +35,7 @@ class HeadersAdviceForkedTest extends AgentTestRunner { runUnderIastTrace { request.getHeaders() } then: - 2 * module.taint(iastCtx , _ as Object, SourceTypes.REQUEST_HEADER_VALUE) + 2 * module.taintObject(iastCtx , _ as Object, SourceTypes.REQUEST_HEADER_VALUE) 0 * _ } diff --git a/dd-java-agent/instrumentation/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/TemplateAndMatrixVariablesInstrumentation.java b/dd-java-agent/instrumentation/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/TemplateAndMatrixVariablesInstrumentation.java index 475367e16d9..10d99a65224 100644 --- a/dd-java-agent/instrumentation/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/TemplateAndMatrixVariablesInstrumentation.java +++ b/dd-java-agent/instrumentation/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/TemplateAndMatrixVariablesInstrumentation.java @@ -196,7 +196,7 @@ public static void after( if (parameterName == null || value == null) { continue; // should not happen } - module.taint( + module.taintString( iastRequestContext, value, SourceTypes.REQUEST_PATH_PARAMETER, parameterName); } } @@ -213,7 +213,7 @@ public static void after( for (Map.Entry> ie : value.entrySet()) { String innerKey = ie.getKey(); if (innerKey != null) { - module.taint( + module.taintString( iastRequestContext, innerKey, SourceTypes.REQUEST_MATRIX_PARAMETER, @@ -222,7 +222,7 @@ public static void after( Iterable innerValues = ie.getValue(); if (innerValues != null) { for (String iv : innerValues) { - module.taint( + module.taintString( iastRequestContext, iv, SourceTypes.REQUEST_MATRIX_PARAMETER, diff --git a/dd-java-agent/instrumentation/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/TemplateVariablesUrlHandlerInstrumentation.java b/dd-java-agent/instrumentation/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/TemplateVariablesUrlHandlerInstrumentation.java index ee786e4ce2e..37f1b2d4158 100644 --- a/dd-java-agent/instrumentation/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/TemplateVariablesUrlHandlerInstrumentation.java +++ b/dd-java-agent/instrumentation/spring-webmvc-3.1/src/main/java/datadog/trace/instrumentation/springweb/TemplateVariablesUrlHandlerInstrumentation.java @@ -154,7 +154,7 @@ public static void after( if (parameterName == null || value == null) { continue; // should not happen } - module.taint( + module.taintString( iastRequestContext, value, SourceTypes.REQUEST_PATH_PARAMETER, parameterName); } } diff --git a/dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/test/boot/SpringBootBasedTest.groovy b/dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/test/boot/SpringBootBasedTest.groovy index bf35356ba14..694f5a660a9 100644 --- a/dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/test/boot/SpringBootBasedTest.groovy +++ b/dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/test/boot/SpringBootBasedTest.groovy @@ -259,7 +259,7 @@ class SpringBootBasedTest extends HttpServerTest TEST_WRITER.waitForTraces(1) then: - 1 * mod.taint(_ as IastContext, '123', SourceTypes.REQUEST_PATH_PARAMETER, 'id') + 1 * mod.taintString(_ as IastContext, '123', SourceTypes.REQUEST_PATH_PARAMETER, 'id') 0 * mod._ cleanup: @@ -295,11 +295,11 @@ class SpringBootBasedTest extends HttpServerTest TEST_WRITER.waitForTraces(1) then: - 1 * mod.taint(_ as IastContext, 'a=x,y;a=z', SourceTypes.REQUEST_PATH_PARAMETER, 'var') - 1 * mod.taint(_ as IastContext, 'a', SourceTypes.REQUEST_MATRIX_PARAMETER, 'var') - 1 * mod.taint(_ as IastContext, 'x', SourceTypes.REQUEST_MATRIX_PARAMETER, 'var') - 1 * mod.taint(_ as IastContext, 'y', SourceTypes.REQUEST_MATRIX_PARAMETER, 'var') - 1 * mod.taint(_ as IastContext, 'z', SourceTypes.REQUEST_MATRIX_PARAMETER, 'var') + 1 * mod.taintString(_ as IastContext, 'a=x,y;a=z', SourceTypes.REQUEST_PATH_PARAMETER, 'var') + 1 * mod.taintString(_ as IastContext, 'a', SourceTypes.REQUEST_MATRIX_PARAMETER, 'var') + 1 * mod.taintString(_ as IastContext, 'x', SourceTypes.REQUEST_MATRIX_PARAMETER, 'var') + 1 * mod.taintString(_ as IastContext, 'y', SourceTypes.REQUEST_MATRIX_PARAMETER, 'var') + 1 * mod.taintString(_ as IastContext, 'z', SourceTypes.REQUEST_MATRIX_PARAMETER, 'var') 0 * mod._ cleanup: diff --git a/dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/test/urlhandlermapping/UrlHandlerMappingTest.groovy b/dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/test/urlhandlermapping/UrlHandlerMappingTest.groovy index cbfb014bed4..f1d606f737f 100644 --- a/dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/test/urlhandlermapping/UrlHandlerMappingTest.groovy +++ b/dd-java-agent/instrumentation/spring-webmvc-3.1/src/test/groovy/test/urlhandlermapping/UrlHandlerMappingTest.groovy @@ -203,7 +203,7 @@ class UrlHandlerMappingTest extends HttpServerTest> ie : value.entrySet()) { String innerKey = ie.getKey(); if (innerKey != null) { - module.taint( + module.taintString( iastRequestContext, innerKey, SourceTypes.REQUEST_MATRIX_PARAMETER, @@ -150,7 +150,7 @@ public static void after( Iterable innerValues = ie.getValue(); if (innerValues != null) { for (String iv : innerValues) { - module.taint( + module.taintString( iastRequestContext, iv, SourceTypes.REQUEST_MATRIX_PARAMETER, diff --git a/dd-java-agent/instrumentation/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/InterceptorPreHandleAdvice.java b/dd-java-agent/instrumentation/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/InterceptorPreHandleAdvice.java index 3e293762f66..cb5e563b748 100644 --- a/dd-java-agent/instrumentation/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/InterceptorPreHandleAdvice.java +++ b/dd-java-agent/instrumentation/spring-webmvc-6.0/src/main/java17/datadog/trace/instrumentation/springweb6/InterceptorPreHandleAdvice.java @@ -91,7 +91,7 @@ public static void after( if (parameterName == null || value == null) { continue; // should not happen } - module.taint( + module.taintString( iastRequestContext, value, SourceTypes.REQUEST_PATH_PARAMETER, parameterName); } } diff --git a/dd-java-agent/instrumentation/spring-webmvc-6.0/src/test/groovy/datadog/trace/instrumentation/springweb6/boot/SpringBootBasedTest.groovy b/dd-java-agent/instrumentation/spring-webmvc-6.0/src/test/groovy/datadog/trace/instrumentation/springweb6/boot/SpringBootBasedTest.groovy index 5a360c6f89c..ecdce88fbeb 100644 --- a/dd-java-agent/instrumentation/spring-webmvc-6.0/src/test/groovy/datadog/trace/instrumentation/springweb6/boot/SpringBootBasedTest.groovy +++ b/dd-java-agent/instrumentation/spring-webmvc-6.0/src/test/groovy/datadog/trace/instrumentation/springweb6/boot/SpringBootBasedTest.groovy @@ -305,7 +305,7 @@ class SpringBootBasedTest extends HttpServerTest then: // spring-security filter causes uri matching to happen twice - (1.._) * mod.taint(_, '123', SourceTypes.REQUEST_PATH_PARAMETER, 'id') + (1.._) * mod.taintObject(_, '123', SourceTypes.REQUEST_PATH_PARAMETER, 'id') 0 * mod._ cleanup: @@ -344,11 +344,11 @@ class SpringBootBasedTest extends HttpServerTest then: // spring-security filter (AuthorizationFilter.java:95) causes uri matching to happen twice (or three times in recent spring (6.1+) versions) - (1.._) * mod.taint(_ as IastContext, 'a=x,y', SourceTypes.REQUEST_PATH_PARAMETER, 'var') // this version of spring removes ;a=z - (1.._) * mod.taint(_ as IastContext, 'a', SourceTypes.REQUEST_MATRIX_PARAMETER, 'var') - (1.._) * mod.taint(_ as IastContext, 'x', SourceTypes.REQUEST_MATRIX_PARAMETER, 'var') - (1.._) * mod.taint(_ as IastContext, 'y', SourceTypes.REQUEST_MATRIX_PARAMETER, 'var') - (1.._) * mod.taint(_ as IastContext, 'z', SourceTypes.REQUEST_MATRIX_PARAMETER, 'var') + (1.._) * mod.taintString(_ as IastContext, 'a=x,y', SourceTypes.REQUEST_PATH_PARAMETER, 'var') // this version of spring removes ;a=z + (1.._) * mod.taintString(_ as IastContext, 'a', SourceTypes.REQUEST_MATRIX_PARAMETER, 'var') + (1.._) * mod.taintString(_ as IastContext, 'x', SourceTypes.REQUEST_MATRIX_PARAMETER, 'var') + (1.._) * mod.taintString(_ as IastContext, 'y', SourceTypes.REQUEST_MATRIX_PARAMETER, 'var') + (1.._) * mod.taintString(_ as IastContext, 'z', SourceTypes.REQUEST_MATRIX_PARAMETER, 'var') 0 * mod._ cleanup: diff --git a/dd-java-agent/instrumentation/spring-webmvc-6.0/src/test/groovy/datadog/trace/instrumentation/springweb6/urlhandlermapping/UrlHandlerMappingTest.groovy b/dd-java-agent/instrumentation/spring-webmvc-6.0/src/test/groovy/datadog/trace/instrumentation/springweb6/urlhandlermapping/UrlHandlerMappingTest.groovy index 9c211b34aad..526b1f68a38 100644 --- a/dd-java-agent/instrumentation/spring-webmvc-6.0/src/test/groovy/datadog/trace/instrumentation/springweb6/urlhandlermapping/UrlHandlerMappingTest.groovy +++ b/dd-java-agent/instrumentation/spring-webmvc-6.0/src/test/groovy/datadog/trace/instrumentation/springweb6/urlhandlermapping/UrlHandlerMappingTest.groovy @@ -207,7 +207,7 @@ class UrlHandlerMappingTest extends HttpServerTest"escape this < '] | '<htmlTag>"escape this < </htmlTag>' - 'escapeHtml4' | ['"escape this < '] | '<htmlTag>"escape this < </htmlTag>' - 'escapeHtml5' | ['"escape this < '] | '<htmlTag>"escape this < </htmlTag>' - 'escapeHtml5Xml' | ['"escape this < '] | '<htmlTag>"escape this < </htmlTag>' - 'escapeJavaScript' | ['