Python: Add tainttracking step that was removed

when the correpsonding datadlow step was removed.

Author: yoff

python/ql/lib/semmle/python/dataflow/new/internal/TaintTrackingPrivate.qll

@@ -53,6 +53,8 @@ private module Cached {
     DataFlowPrivate::iterableUnpackingStoreStep(nodeFrom, _, nodeTo)
     or
     awaitStep(nodeFrom, nodeTo)
+    or
+    asyncWithStep(nodeFrom, nodeTo)
   }
 }
 
@@ -211,3 +213,24 @@ predicate copyStep(DataFlow::CfgNode nodeFrom, DataFlow::CfgNode nodeTo) {
 predicate awaitStep(DataFlow::Node nodeFrom, DataFlow::Node nodeTo) {
   nodeTo.asExpr().(Await).getValue() = nodeFrom.asExpr()
 }
+
+/**
+ * Holds if taint can flow from `nodeFrom` to `nodeTo` inside an `async with` statement.
+ *
+ * For example in
+ * ```python
+ * async with open("foo") as f:
+ * ```
+ * the variable `f` is tainted if the result of `open("foo")` is tainted.
+ */
+predicate asyncWithStep(DataFlow::Node nodeFrom, DataFlow::Node nodeTo) {
+  exists(With with, ControlFlowNode contextManager, ControlFlowNode var |
+    nodeFrom.(DataFlow::CfgNode).getNode() = contextManager and
+    nodeTo.(DataFlow::EssaNode).getVar().getDefinition().(WithDefinition).getDefiningNode() = var and
+    // see `with_flow` in `python/ql/src/semmle/python/dataflow/Implementation.qll`
+    with.getContextExpr() = contextManager.getNode() and
+    with.getOptionalVars() = var.getNode() and
+    with.isAsync() and
+    contextManager.strictlyDominates(var)
+  )
+}

python/ql/test/experimental/dataflow/tainttracking/defaultAdditionalTaintStep/test_async.py

@@ -31,7 +31,7 @@ async def test_async_with():
     ctx = AsyncContext()
     taint(ctx)
     async with ctx as tainted:
-        ensure_tainted(tainted) # $ MISSING: tainted
+        ensure_tainted(tainted) # $ tainted
 
 
 class AsyncIter: