Add support for taintable objects

DataDog · Mariovido · Nov 27, 2024 · Nov 25, 2024 · Nov 25, 2024 · Nov 26, 2024
commit 18f19e057487eb653475efccc5cdd1fcff34485d
@@ -8,13 +8,15 @@
 import static datadog.trace.api.telemetry.LogCollector.SEND_TELEMETRY;
 
 import com.datadog.iast.model.Range;
+import com.datadog.iast.model.Source;
 import com.datadog.iast.taint.Ranges;
 import com.datadog.iast.taint.TaintedObject;
 import com.datadog.iast.taint.TaintedObjects;
 import com.datadog.iast.util.RangeBuilder;
 import com.datadog.iast.util.Ranged;
 import com.datadog.iast.util.StringUtils;
 import datadog.trace.api.iast.IastContext;
+import datadog.trace.api.iast.Taintable;
 import datadog.trace.api.iast.propagation.StringModule;
 import de.thetaphi.forbiddenapis.SuppressForbidden;
 import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
@@ -753,17 +755,31 @@ public void onStringValueOf(Object param, @Nonnull String result) {
       return;
     }
     final TaintedObjects taintedObjects = ctx.getTaintedObjects();
-    final TaintedObject taintedParam = taintedObjects.get(param);
-    if (taintedParam == null) {
-      return;
-    }
 
-    final Range[] rangesParam = taintedParam.getRanges();
-    if (rangesParam.length == 0) {
-      return;
-    }
+    if (param instanceof Taintable) {
+      final Taintable taintable = (Taintable) param;
+      if (!taintable.$DD$isTainted()) {
+        return;
+      }
+      final Taintable.Source source = taintable.$$DD$getSource();
+      final Range[] ranges =
+          Ranges.forCharSequence(
+              result, new Source(source.getOrigin(), source.getName(), source.getValue()));
 
-    taintedObjects.taint(result, rangesParam);
+      taintedObjects.taint(result, ranges);
+    } else {
+      final TaintedObject taintedParam = taintedObjects.get(param);
+      if (taintedParam == null) {
+        return;
+      }
+
+      final Range[] rangesParam = taintedParam.getRanges();
+      if (rangesParam.length == 0) {
+        return;
+      }
+
+      taintedObjects.taint(result, rangesParam);
+    }
   }
 
   /**

@@ -1,8 +1,12 @@
 package com.datadog.iast.propagation
 
 import com.datadog.iast.IastModuleImplTestBase
+import com.datadog.iast.model.Source
+import com.datadog.iast.taint.TaintedObjects
 import datadog.trace.api.gateway.RequestContext
 import datadog.trace.api.gateway.RequestContextSlot
+import datadog.trace.api.iast.SourceTypes
+import datadog.trace.api.iast.Taintable
 import datadog.trace.api.iast.propagation.StringModule
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer
@@ -16,6 +20,7 @@ import static com.datadog.iast.taint.TaintUtils.fromTaintFormat
 import static com.datadog.iast.taint.TaintUtils.getStringFromTaintFormat
 import static com.datadog.iast.taint.TaintUtils.taint
 import static com.datadog.iast.taint.TaintUtils.taintFormat
+import static com.datadog.iast.taint.TaintUtils.taintObject
 
 @CompileDynamic
 class StringModuleTest extends IastModuleImplTestBase {
@@ -1319,6 +1324,22 @@ class StringModuleTest extends IastModuleImplTestBase {
     sbf("==>my_input<==") | "==>my_input<=="
   }
 
+  void 'test valueOf with taintable object and make sure IastRequestContext is called'() {
+    given:
+    final taintedObjects = ctx.getTaintedObjects()
+    final source = taintedSource(null)
+    final param = taintable(taintedObjects, source)
+    def result = String.valueOf(param)
+
+    when:
+    module.onStringValueOf(param, result)
+    def taintedObject = taintedObjects.get(result)
+
+    then:
+    1 * tracer.activeSpan() >> span
+    taintFormat(result, taintedObject.getRanges()) == "==>" + param.toString() + "<=="
+  }
+
   private static Date date(final String pattern, final String value) {
     return new SimpleDateFormat(pattern).parse(value)
   }
@@ -1331,11 +1352,44 @@ class StringModuleTest extends IastModuleImplTestBase {
     return new StringBuilder(string)
   }
 
-  private static StringBuilder sbf() {
+  private static StringBuffer sbf() {
     return sbf('')
   }
 
   private static StringBuffer sbf(final String string) {
     return new StringBuffer(string)
   }
+
+  private static Source taintedSource(String value = 'value') {
+    return new Source(SourceTypes.REQUEST_PARAMETER_VALUE, 'name', value)
+  }
+
+  private static Taintable taintable(TaintedObjects tos, Source source = null) {
+    final result = new MockTaintable()
+    if (source != null) {
+      taintObject(tos, result, source)
+    }
+    return result
+  }
+
+  private static class MockTaintable implements Taintable {
+    private Source source
+
+    @SuppressWarnings('CodeNarc')
+    @Override
+    Source $$DD$getSource() {
+      return source
+    }
+
+    @SuppressWarnings('CodeNarc')
+    @Override
+    void $$DD$setSource(Source source) {
+      this.source = source
+    }
+
+    @Override
+    String toString() {
+      return Taintable.name
+    }
+  }
 }
@@ -3,6 +3,7 @@ package com.datadog.iast.taint
 import com.datadog.iast.model.Range
 import com.datadog.iast.model.Source
 import datadog.trace.api.iast.SourceTypes
+import datadog.trace.api.iast.Taintable
 
 import static datadog.trace.api.iast.VulnerabilityMarks.NOT_MARKED
 
@@ -81,6 +82,14 @@ class TaintUtils {
     getStringFromTaintFormat(appendable.toString())
   }
 
+  static TaintedObject getTaintedObject(final TaintedObjects tos, final Object target) {
+    if (target instanceof Taintable) {
+      final source = (target as Taintable).$$DD$getSource() as Source
+      return source == null ? null : new TaintedObject(target, Ranges.forObject(source))
+    }
+    return tos.get(target)
+  }
+
   static <E> E taint(final TaintedObjects tos, final E value) {
     if (value instanceof String) {
       return addFromTaintFormat(tos, value as String)
@@ -89,6 +98,17 @@ class TaintUtils {
     return value
   }
 
+  static TaintedObject taintObject(final TaintedObjects tos, final Object target, Source source, int mark = NOT_MARKED) {
+    if (target instanceof Taintable) {
+      target.$$DD$setSource(source)
+    } else if (target instanceof CharSequence) {
+      tos.taint(target, Ranges.forCharSequence(target, source, mark))
+    } else {
+      tos.taint(target, Ranges.forObject(source, mark))
+    }
+    return getTaintedObject(tos, target)
+  }
+
   static String addFromTaintFormat(final TaintedObjects tos, final String s) {
     return addFromTaintFormat(tos, s, NOT_MARKED)
   }