Attach values to sources coming from non char sequences

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

@@ -16,6 +16,9 @@ public final class Source implements Taintable.Source {
 
   private static final Logger LOGGER = LoggerFactory.getLogger(Source.class);
 
+  /** Placeholder for non char sequence objects */
+  public static final Object PROPAGATION_PLACEHOLDER = new Object();
+
   // value to send in the rare case that the name/value have been garbage collected
   private static final String GARBAGE_COLLECTED_REF =
       "[unknown: original value was garbage collected]";
@@ -51,7 +54,9 @@ public String getValue() {
   @Nullable
   private String asString(@Nullable final Object target) {
     Object value = target;
-    if (value instanceof Reference) {
+    if (value == PROPAGATION_PLACEHOLDER) {
+      value = null;
+    } else if (value instanceof Reference) {
       value = ((Reference<?>) value).get();
       if (value == null) {
         value = GARBAGE_COLLECTED_REF;
@@ -90,4 +95,11 @@ public boolean equals(Object o) {
   public int hashCode() {
     return Objects.hash(origin, getName(), getValue());
   }
+
+  public Source attachValue(final CharSequence result) {
+    if (value != PROPAGATION_PLACEHOLDER) {
+      return this;
+    }
+    return new Source(origin, name, result);
+  }
 }

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

@@ -1,55 +1,36 @@
 package com.datadog.iast.propagation;
 
-import static com.datadog.iast.taint.Ranges.highestPriorityRange;
-
-import com.datadog.iast.model.Range;
-import com.datadog.iast.taint.Ranges;
+import datadog.trace.api.iast.propagation.CodecModule;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 
-public class FastCodecModule extends BaseCodecModule {
+public class FastCodecModule extends PropagationModuleImpl implements CodecModule {
 
   @Override
-  protected Range[] urlDecodeRanges(
-      @Nonnull final String value,
-      @Nullable final String encoding,
-      @Nonnull final String result,
-      @Nonnull final Range[] ranges) {
-    final Range range = highestPriorityRange(ranges);
-    return new Range[] {Ranges.copyWithPosition(range, 0, result.length())};
+  public void onUrlDecode(
+      @Nonnull final String value, @Nullable final String encoding, @Nonnull final String result) {
+    taintIfTainted(result, value);
   }
 
   @Override
-  protected Range[] fromBytesRanges(
-      @Nonnull final byte[] value,
-      @Nullable final String charset,
-      @Nonnull final String result,
-      @Nonnull final Range[] ranges) {
-    final Range range = highestPriorityRange(ranges);
-    return new Range[] {Ranges.copyWithPosition(range, 0, result.length())};
+  public void onStringFromBytes(
+      @Nonnull final byte[] value, @Nullable final String charset, @Nonnull final String result) {
+    taintIfTainted(result, value);
   }
 
   @Override
-  protected Range[] getBytesRanges(
-      @Nonnull final String value,
-      @Nullable final String charset,
-      @Nonnull final byte[] result,
-      @Nonnull final Range[] ranges) {
-    final Range range = highestPriorityRange(ranges);
-    return new Range[] {Ranges.copyWithPosition(range, 0, result.length)};
+  public void onStringGetBytes(
+      @Nonnull final String value, @Nullable final String charset, @Nonnull final byte[] result) {
+    taintIfTainted(result, value);
   }
 
   @Override
-  protected Range[] decodeBase64Ranges(
-      @Nonnull final byte[] value, @Nonnull final byte[] result, @Nonnull final Range[] ranges) {
-    final Range range = highestPriorityRange(ranges);
-    return new Range[] {Ranges.copyWithPosition(range, 0, result.length)};
+  public void onBase64Encode(@Nullable byte[] value, @Nullable byte[] result) {
+    taintIfTainted(result, value);
   }
 
   @Override
-  protected Range[] encodeBase64Ranges(
-      @Nonnull final byte[] value, @Nonnull final byte[] result, @Nonnull final Range[] ranges) {
-    final Range range = highestPriorityRange(ranges);
-    return new Range[] {Ranges.copyWithPosition(range, 0, result.length)};
+  public void onBase64Decode(@Nullable byte[] value, @Nullable byte[] result) {
+    taintIfTainted(result, value);
   }
 }

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

@@ -1,5 +1,6 @@
 package com.datadog.iast.propagation;
 
+import static com.datadog.iast.model.Source.PROPAGATION_PLACEHOLDER;
 import static com.datadog.iast.taint.Ranges.highestPriorityRange;
 import static com.datadog.iast.util.ObjectVisitor.State.CONTINUE;
 import static datadog.trace.api.iast.VulnerabilityMarks.NOT_MARKED;
@@ -16,6 +17,7 @@
 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;
@@ -313,7 +315,8 @@ private static Object sourceString(@Nullable final Object target, final boolean
       }
       return charSequence.toString();
     }
-    return null; // ignore non char-sequence instances (e.g. byte buffers)
+    // ignore non char-sequence instances (e.g. byte buffers)
+    return value ? PROPAGATION_PLACEHOLDER : null;
   }
 
   @Contract("null -> false")
@@ -324,6 +327,9 @@ private static boolean canBeTainted(@Nullable final Object target) {
     if (target instanceof CharSequence) {
       return Tainteds.canBeTainted((CharSequence) target);
     }
+    if (target.getClass().isArray()) {
+      return Array.getLength(target) > 0;
+    }
     return true;
   }
 
@@ -396,7 +402,7 @@ private static Source highestPrioritySource(
   private static void internalTaint(
       @Nullable final IastContext ctx,
       @Nonnull final Object value,
-      @Nullable final Source source,
+      @Nullable Source source,
       int mark) {
     if (source == null) {
       return;
@@ -409,6 +415,7 @@ private static void internalTaint(
         return;
       }
       if (value instanceof CharSequence) {
+        source = source.attachValue((CharSequence) value);
         to.taint(value, Ranges.forCharSequence((CharSequence) value, source, mark));
       } else {
         to.taint(value, Ranges.forObject(source, mark));
@@ -419,7 +426,7 @@ private static void internalTaint(
   private static void internalTaint(
       @Nullable final IastContext ctx,
       @Nonnull final Object value,
-      @Nullable final Range[] ranges,
+      @Nullable Range[] ranges,
       final int mark) {
     if (ranges == null || ranges.length == 0) {
       return;
@@ -429,8 +436,11 @@ private static void internalTaint(
     } else {
       final TaintedObjects to = getTaintedObjects(ctx);
       if (to != null) {
-        final Range[] markedRanges = markRanges(ranges, mark);
-        to.taint(value, markedRanges);
+        if (value instanceof CharSequence) {
+          ranges = attachSourceValue(ranges, (CharSequence) value);
+        }
+        ranges = markRanges(ranges, mark);
+        to.taint(value, ranges);
       }
     }
   }
@@ -449,6 +459,20 @@ private static Range[] markRanges(@Nonnull final Range[] ranges, final int mark)
     return result;
   }
 
+  public static Range[] attachSourceValue(
+      @Nonnull final Range[] ranges, @Nonnull final CharSequence value) {
+    // unbound sources can only occur when there's a single range in the array
+    if (ranges.length != 1) {
+      return ranges;
+    }
+    final Range range = ranges[0];
+    final Source source = range.getSource();
+    final Source newSource = range.getSource().attachValue(value);
+    return newSource == source
+        ? ranges
+        : Ranges.forCharSequence(value, newSource, range.getMarks());
+  }
+
   private static class LazyContext implements IastContext {
 
     private boolean fetched;

dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/propagation/FastCodecModuleTest.groovy

@@ -3,10 +3,7 @@ package com.datadog.iast.propagation
 import com.datadog.iast.taint.Ranges
 import com.datadog.iast.taint.TaintedObject
 import datadog.trace.api.iast.propagation.CodecModule
-import groovy.transform.CompileDynamic
 
-
-@CompileDynamic
 class FastCodecModuleTest extends BaseCodecModuleTest {
 
   @Override
@@ -40,37 +37,34 @@ class FastCodecModuleTest extends BaseCodecModuleTest {
 
   @Override
   protected void assertOnStringGetBytes(final String value, final String charset, final TaintedObject source, final TaintedObject target) {
-    final result = target.get() as byte[]
     assert target.ranges.size() == 1
 
     final sourceRange = Ranges.highestPriorityRange(source.ranges)
     final range = target.ranges.first()
     assert range.start == 0
-    assert range.length == result.length
+    assert range.length == Integer.MAX_VALUE // unbound for non char sequences
     assert range.source == sourceRange.source
   }
 
   @Override
   protected void assertBase64Decode(byte[] value, TaintedObject source, TaintedObject target) {
-    final result = target.get() as byte[]
     assert target.ranges.size() == 1
 
     final sourceRange = Ranges.highestPriorityRange(source.ranges)
     final range = target.ranges.first()
     assert range.start == 0
-    assert range.length == result.length
+    assert range.length == Integer.MAX_VALUE // unbound for non char sequences
     assert range.source == sourceRange.source
   }
 
   @Override
   protected void assertBase64Encode(byte[] value, TaintedObject source, TaintedObject target) {
-    final result = target.get() as byte[]
     assert target.ranges.size() == 1
 
     final sourceRange = Ranges.highestPriorityRange(source.ranges)
     final range = target.ranges.first()
     assert range.start == 0
-    assert range.length == result.length
+    assert range.length == Integer.MAX_VALUE // unbound for non char sequences
     assert range.source == sourceRange.source
   }
 }