Python: Recognize taint for extended iterable unpacking

Author: RasmusWL

python/ql/src/semmle/python/Flow.qll

@@ -1068,6 +1068,17 @@ class NameConstantNode extends NameNode {
     */
 }
 
+/** A control flow node correspoinding to a starred expression, `*a`. */
+class StarredNode extends ControlFlowNode {
+    StarredNode() {
+        toAst(this) instanceof Starred
+    }
+
+    ControlFlowNode getValue() {
+        toAst(result) = toAst(this).(Starred).getValue()
+    }
+}
+
 private module Scopes {
 
     private predicate fast_local(NameNode n) {

python/ql/src/semmle/python/dataflow/Implementation.qll

@@ -727,7 +727,10 @@ private class EssaTaintTracking extends string {
         SequenceNode left_parent, ControlFlowNode left_defn, CollectionKind parent_kind
     ) {
         left_parent.getAnElement() = left_defn and
-        result = parent_kind.getMember()
+        // Handle `a, *b = some_iterable`
+        if left_defn instanceof StarredNode
+        then result = parent_kind
+        else result = parent_kind.getMember()
         or
         result = iterable_unpacking_decent(left_parent.getAnElement(), left_defn,
                 parent_kind.getMember())

python/ql/src/semmle/python/essa/SsaDefinitions.qll

@@ -39,7 +39,11 @@ cached module SsaSource {
 
     /** Holds if `v` is defined by multiple assignment at `defn`. */
     cached predicate multi_assignment_definition(Variable v, ControlFlowNode defn, int n, SequenceNode lhs) {
-        defn.(NameNode).defines(v) and
+        (
+            defn.(NameNode).defines(v)
+            or
+            defn.(StarredNode).getValue().(NameNode).defines(v)
+        ) and
         not exists(defn.(DefinitionNode).getValue()) and
         lhs.getElement(n) = defn and
         lhs.getBasicBlock().dominates(defn.getBasicBlock())

python/ql/test/3/library-tests/taint/unpacking/TestTaint.expected

@@ -1,6 +1,9 @@
 | test.py:11 | extended_unpacking | first | externally controlled string |
 | test.py:11 | extended_unpacking | last | externally controlled string |
-| test.py:11 | extended_unpacking | rest | NO TAINT |
-| test.py:16 | also_allowed | a | NO TAINT |
+| test.py:11 | extended_unpacking | rest | [externally controlled string] |
+| test.py:16 | also_allowed | a | [externally controlled string] |
 | test.py:24 | also_allowed | b | NO TAINT |
 | test.py:24 | also_allowed | c | NO TAINT |
+| test.py:31 | nested | x | externally controlled string |
+| test.py:31 | nested | xs | [externally controlled string] |
+| test.py:31 | nested | ys | [externally controlled string] |

python/ql/test/3/library-tests/taint/unpacking/test.py

@@ -8,12 +8,12 @@ def test(*args):
 
 def extended_unpacking():
     first, *rest, last = TAINTED_LIST
-    test(first, rest, last) # TODO: mark `rest` as [taint]
+    test(first, rest, last)
 
 
 def also_allowed():
     *a, = TAINTED_LIST
-    test(a) # TODO: mark `a` as [taint]
+    test(a)
 
     # for b, *c in [(1, 2, 3), (4, 5, 6, 7)]:
         # print(c)
@@ -22,3 +22,10 @@ def also_allowed():
 
     for b, *c in [TAINTED_LIST, TAINTED_LIST]:
         test(b, c) # TODO: mark `c` as [taint]
+
+def nested():
+    l = TAINTED_LIST
+    ll = [l,l]
+
+    [[x, *xs], ys] = ll
+    test(x, xs, ys)