Python: Attempt at taint step for list.append/set.add

RasmusWL · RasmusWL · commit d0081dfbfa95 · 2020-08-27T10:57:07.000+02:00
diff --git a/python/ql/src/experimental/dataflow/internal/TaintTrackingPrivate.qll b/python/ql/src/experimental/dataflow/internal/TaintTrackingPrivate.qll
@@ -177,6 +177,15 @@ predicate containerStep(DataFlow::CfgNode nodeFrom, DataFlow::CfgNode nodeTo) {
           "values", "items", "get", "popitem"] and
     call.getFunction().(AttrNode).getObject(name) = nodeFrom.getNode()
   )
+  or
+  // list.append, set.add
+  // NOTE: this currently doesn't work, since there are no PostUpdateNodes
+  exists(CallNode call, string name |
+    name in ["append", "add"] and
+    call.getFunction().(AttrNode).getObject(name) =
+      nodeTo.(PostUpdateNode).getPreUpdateNode().(DataFlow::CfgNode).getNode() and
+    call.getArg(0) = nodeFrom.getNode()
+  )
 }
 
 /**
diff --git a/python/ql/test/experimental/dataflow/tainttracking/defaultAdditionalTaintStep/TestTaint.expected b/python/ql/test/experimental/dataflow/tainttracking/defaultAdditionalTaintStep/TestTaint.expected
@@ -49,6 +49,8 @@
 | collections_.py:137 | fail | list_index_aug_assign | my_list |
 | collections_.py:144 | ok   | list_append | my_list |
 | collections_.py:147 | fail | list_append | my_list |
+| collections_.py:154 | ok   | set_add | my_set |
+| collections_.py:157 | fail | set_add | my_set |
 | json_.py:26 | ok   | test | json.dumps(..) |
 | json_.py:27 | ok   | test | json.loads(..) |
 | json_.py:34 | fail | test | tainted_filelike |
diff --git a/python/ql/test/experimental/dataflow/tainttracking/defaultAdditionalTaintStep/collections_.py b/python/ql/test/experimental/dataflow/tainttracking/defaultAdditionalTaintStep/collections_.py
@@ -147,6 +147,16 @@ def list_append():
     ensure_tainted(my_list)
 
 
+def set_add():
+    tainted_string = TAINTED_STRING
+    my_set = {"safe"}
+
+    ensure_not_tainted(my_set)
+
+    my_set.add(tainted_string)
+    ensure_tainted(my_set)
+
+
 # Make tests runable
 
 test_construction()
@@ -158,3 +168,4 @@ def list_append():
 list_index_assign()
 list_index_aug_assign()
 list_append()
+set_add()
