Python: Add additional taint steps for for-iteration

RasmusWL · RasmusWL · commit afb160fbbb87 · 2020-08-26T20:18:31.000+02:00
diff --git a/python/ql/src/experimental/dataflow/internal/TaintTrackingPrivate.qll b/python/ql/src/experimental/dataflow/internal/TaintTrackingPrivate.qll
@@ -36,6 +36,8 @@ predicate localAdditionalTaintStep(DataFlow::Node nodeFrom, DataFlow::Node nodeT
   containerStep(nodeFrom, nodeTo)
   or
   copyStep(nodeFrom, nodeTo)
+  or
+  forStep(nodeFrom, nodeTo)
 }
 
 /**
@@ -190,3 +192,16 @@ predicate copyStep(DataFlow::CfgNode nodeFrom, DataFlow::CfgNode nodeTo) {
     call.getArg(0) = nodeFrom.getNode()
   )
 }
+
+
+/**
+ * Holds if taint can flow from `nodeFrom` to `nodeTo` with a step related to `for`-iteration,
+ * for example `for x in xs`, or `for x,y in points`.
+ */
+predicate forStep(DataFlow::CfgNode nodeFrom, DataFlow::EssaNode nodeTo) {
+  exists(EssaNodeDefinition defn, For for |
+    for.getTarget().getAChildNode*() = defn.getDefiningNode().getNode() and
+    nodeTo.getVar() = defn and
+    nodeFrom.getNode().getNode() = for.getIter()
+  )
+}
diff --git a/python/ql/test/experimental/dataflow/tainttracking/defaultAdditionalTaintStep-py3/TestTaint.expected b/python/ql/test/experimental/dataflow/tainttracking/defaultAdditionalTaintStep-py3/TestTaint.expected
@@ -13,8 +13,8 @@
 | unpacking.py:18 | fail | extended_unpacking | last |
 | unpacking.py:18 | fail | extended_unpacking | rest |
 | unpacking.py:23 | fail | also_allowed | a |
-| unpacking.py:31 | fail | also_allowed | b |
-| unpacking.py:31 | fail | also_allowed | c |
+| unpacking.py:31 | ok   | also_allowed | b |
+| unpacking.py:31 | ok   | also_allowed | c |
 | unpacking.py:39 | fail | nested | x |
 | unpacking.py:39 | fail | nested | xs |
 | unpacking.py:39 | fail | nested | ys |
diff --git a/python/ql/test/experimental/dataflow/tainttracking/defaultAdditionalTaintStep/TestTaint.expected b/python/ql/test/experimental/dataflow/tainttracking/defaultAdditionalTaintStep/TestTaint.expected
@@ -23,14 +23,14 @@
 | collections.py:61 | fail | test_access | a |
 | collections.py:61 | fail | test_access | b |
 | collections.py:61 | fail | test_access | c |
-| collections.py:64 | fail | test_access | h |
-| collections.py:66 | fail | test_access | i |
+| collections.py:64 | ok   | test_access | h |
+| collections.py:66 | ok   | test_access | i |
 | collections.py:73 | ok   | test_dict_access | tainted_dict["name"] |
 | collections.py:74 | ok   | test_dict_access | tainted_dict.get(..) |
 | collections.py:75 | ok   | test_dict_access | tainted_dict[x] |
 | collections.py:76 | ok   | test_dict_access | tainted_dict.copy() |
-| collections.py:80 | fail | test_dict_access | v |
-| collections.py:82 | fail | test_dict_access | v |
+| collections.py:80 | ok   | test_dict_access | v |
+| collections.py:82 | ok   | test_dict_access | v |
 | collections.py:90 | fail | test_named_tuple | point[0] |
 | collections.py:91 | fail | test_named_tuple | point.x |
 | collections.py:95 | ok   | test_named_tuple | point[1] |
