Add debug info to tainted objects to troubleshoot null ranges (#7455)

Add debug info to tainted objects to troubleshoot null ranges

Author: manuel-alvarez-alvarez

dd-java-agent/agent-iast/src/main/java/com/datadog/iast/taint/TaintedObject.java

@@ -5,11 +5,16 @@
 import com.datadog.iast.model.Range;
 import datadog.trace.api.Config;
 import java.lang.ref.WeakReference;
+import java.util.Arrays;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 public class TaintedObject extends WeakReference<Object> {
 
+  private static final Logger LOGGER = LoggerFactory.getLogger(TaintedObject.class);
+
   public static final int MAX_RANGE_COUNT = Config.get().getIastMaxRangeCount();
 
   final int positiveHashCode;
@@ -21,6 +26,7 @@ public class TaintedObject extends WeakReference<Object> {
 
   public TaintedObject(final @Nonnull Object obj, final @Nonnull Range[] ranges) {
     super(obj);
+    validateRanges(ranges);
     this.positiveHashCode = System.identityHashCode(obj) & POSITIVE_MASK;
     // ensure ranges never go over the limit
     if (ranges.length > MAX_RANGE_COUNT) {
@@ -41,7 +47,12 @@ public Range[] getRanges() {
   }
 
   public void setRanges(@Nonnull final Range[] ranges) {
-    this.ranges = ranges;
+    try {
+      validateRanges(ranges);
+      this.ranges = ranges;
+    } catch (Throwable e) {
+      LOGGER.debug("Error tainting object with custom ranges, ranges won't be updated", e);
+    }
   }
 
   @Override
@@ -57,4 +68,15 @@ public String toString() {
         + (ranges == null ? 0 : ranges.length)
         + " ranges)";
   }
+
+  private void validateRanges(final Range[] ranges) {
+    if (ranges == null) {
+      throw new IllegalArgumentException("ranges cannot be null");
+    }
+    for (Range range : ranges) {
+      if (range == null) {
+        throw new IllegalArgumentException("found null range in " + Arrays.toString(ranges));
+      }
+    }
+  }
 }

dd-java-agent/agent-iast/src/main/java/com/datadog/iast/taint/TaintedObjects.java

@@ -18,6 +18,8 @@
 @SuppressWarnings("UnusedReturnValue")
 public interface TaintedObjects extends Iterable<TaintedObject> {
 
+  Logger LOGGER = LoggerFactory.getLogger(TaintedObjects.class);
+
   static TaintedObjects build(@Nonnull final TaintedMap map) {
     final TaintedObjectsImpl taintedObjects = new TaintedObjectsImpl(map);
     return IastSystem.DEBUG ? new TaintedObjectsDebugAdapter(taintedObjects) : taintedObjects;
@@ -41,12 +43,17 @@ private TaintedObjectsImpl(final @Nonnull TaintedMap map) {
       this.map = map;
     }
 
-    @Nonnull
+    @Nullable
     @Override
     public TaintedObject taint(final @Nonnull Object obj, final @Nonnull Range[] ranges) {
-      final TaintedObject tainted = new TaintedObject(obj, ranges);
-      map.put(tainted);
-      return tainted;
+      try {
+        final TaintedObject tainted = new TaintedObject(obj, ranges);
+        map.put(tainted);
+        return tainted;
+      } catch (Throwable e) {
+        LOGGER.debug("Error tainting object, it won't be tainted", e);
+        return null;
+      }
     }
 
     @Nullable
@@ -73,7 +80,6 @@ public Iterator<TaintedObject> iterator() {
   }
 
   final class TaintedObjectsDebugAdapter implements TaintedObjects, Wrapper<TaintedObjectsImpl> {
-    static final Logger LOGGER = LoggerFactory.getLogger(TaintedObjects.class);
 
     private final TaintedObjectsImpl delegated;
     private final UUID id;
@@ -125,10 +131,14 @@ public Iterator<TaintedObject> iterator() {
       return delegated.iterator();
     }
 
-    private void logTainted(final TaintedObject tainted) {
+    private void logTainted(@Nullable final TaintedObject tainted) {
       if (LOGGER.isDebugEnabled()) {
         try {
-          LOGGER.debug("taint {}: tainted={}", id, TaintedObjectEncoding.toJson(tainted));
+          if (tainted == null) {
+            LOGGER.debug("taint {}: ignored", id);
+          } else {
+            LOGGER.debug("taint {}: tainted={}", id, TaintedObjectEncoding.toJson(tainted));
+          }
         } catch (final Throwable e) {
           LOGGER.error("Failed to debug new tainted object", e);
         }

dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/taint/TaintedObjectTest.groovy

@@ -1,5 +1,7 @@
 package com.datadog.iast.taint
 
+import ch.qos.logback.classic.Logger
+import ch.qos.logback.core.Appender
 import com.datadog.iast.model.Range
 import com.datadog.iast.model.Source
 import datadog.trace.api.Config
@@ -14,7 +16,7 @@ class TaintedObjectTest extends Specification {
     given:
     final max = Config.get().iastMaxRangeCount
     final ranges = (0..max + 1)
-      .collect { index -> new Range(index, 1, new Source(REQUEST_HEADER_NAME, 'a', 'b'), NOT_MARKED) }
+    .collect { index -> new Range(index, 1, new Source(REQUEST_HEADER_NAME, 'a', 'b'), NOT_MARKED) }
 
     when:
     final tainted = new TaintedObject('test', ranges.toArray(new Range[0]))
@@ -24,4 +26,36 @@ class TaintedObjectTest extends Specification {
     tainted.ranges.size() == max
     tainted.toString().contains("${tainted.ranges.size()} ranges")
   }
+
+  void 'test that objects are not tainted if null ranges are provided'() {
+    setup:
+    def toTaint = UUID.randomUUID().toString()
+    def to = new TaintedObject(toTaint, Ranges.forCharSequence(toTaint, new Source(1 as byte, 'a', 'b'), NOT_MARKED))
+    def logger = TaintedObject.LOGGER as Logger
+    def appender = Mock(Appender<?>)
+    logger.addAppender(appender)
+
+    when:
+    to.setRanges(ranges as Range[])
+
+    then:
+    noExceptionThrown()
+    if (taint) {
+      to.ranges.length == ranges.size()
+      0 * appender.doAppend(_)
+    } else {
+      to.ranges.length == 1
+      1 * appender.doAppend(_ as Object)
+    }
+
+    cleanup:
+    logger.detachAppender(appender)
+
+    where:
+    ranges                                                                                                       | taint
+    null                                                                                                         | false
+    []                                                                                                           | true
+    [new Range(0, 10, new Source(1 as byte, 'a', 'b'), 1), null]                                                 | false
+    [new Range(0, 10, new Source(1 as byte, 'a', 'b'), 1), new Range(0, 10, new Source(2 as byte, 'a', 'b'), 1)] | true
+  }
 }

dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/taint/TaintedObjectsLogTest.groovy

@@ -79,5 +79,19 @@ class TaintedObjectsLogTest extends DDSpecification {
     taintedObjects.size() == 1
     taintedObjects.iterator().size() == 1
   }
+
+  void 'should not taint null ranges'() {
+    given:
+    IastSystem.DEBUG = true
+    logger.level = Level.ALL
+    TaintedObjects taintedObjects = taintedObjects()
+    final obj = 'A'
+
+    when:
+    taintedObjects.taint(obj, null)
+
+    then:
+    taintedObjects.empty
+  }
 }
 

dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/taint/TaintedObjectsTest.groovy

@@ -0,0 +1,42 @@
+package com.datadog.iast.taint
+
+import ch.qos.logback.classic.Logger
+import ch.qos.logback.core.Appender
+import com.datadog.iast.model.Range
+import com.datadog.iast.model.Source
+import spock.lang.Specification
+
+class TaintedObjectsTest extends Specification {
+
+  void 'test that objects are not tainted if null ranges are provided'() {
+    setup:
+    def toTaint = UUID.randomUUID().toString()
+    def to = TaintedObjects.build(TaintedMap.build(8))
+    def logger = TaintedObjects.LOGGER as Logger
+    def appender = Mock(Appender<?>)
+    logger.addAppender(appender)
+
+    when:
+    to.taint(toTaint, ranges as Range[])
+
+    then:
+    noExceptionThrown()
+    if (taint) {
+      to.get(toTaint) != null
+      0 * appender.doAppend(_)
+    } else {
+      to.get(toTaint) == null
+      1 * appender.doAppend(_ as Object)
+    }
+
+    cleanup:
+    logger.detachAppender(appender)
+
+    where:
+    ranges                                                                                                       | taint
+    null                                                                                                         | false
+    []                                                                                                           | true
+    [new Range(0, 10, new Source(1 as byte, 'a', 'b'), 1), null]                                                 | false
+    [new Range(0, 10, new Source(1 as byte, 'a', 'b'), 1), new Range(0, 10, new Source(2 as byte, 'a', 'b'), 1)] | true
+  }
+}