Python: Improve handling of async methods

Author: RasmusWL

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

@@ -51,6 +51,8 @@ private module Cached {
     DataFlowPrivate::iterableUnpackingReadStep(nodeFrom, _, nodeTo)
     or
     DataFlowPrivate::iterableUnpackingStoreStep(nodeFrom, _, nodeTo)
+    or
+    awaitStep(nodeFrom, nodeTo)
   }
 }
 
@@ -201,3 +203,10 @@ predicate copyStep(DataFlow::CfgNode nodeFrom, DataFlow::CfgNode nodeTo) {
     call.getArg(0) = nodeFrom
   )
 }
+
+/**
+ * Holds if taint can flow from `nodeFrom` to `nodeTo` with a step related `await`.
+ */
+predicate awaitStep(DataFlow::Node nodeFrom, DataFlow::Node nodeTo) {
+  nodeTo.asExpr().(Await).getValue() = nodeFrom.asExpr()
+}

python/ql/src/semmle/python/frameworks/internal/InstanceTaintStepsHelper.qll

@@ -33,14 +33,15 @@ private class InstanceAdditionalTaintStep extends TaintTracking::AdditionalTaint
       nodeFrom = helper.getInstance() and
       nodeTo.(DataFlow::MethodCallNode).calls(nodeFrom, helper.getMethodName())
       or
-      // async methods
-      exists(DataFlow::MethodCallNode call, Await await |
-        nodeTo.asExpr() = await and
-        nodeFrom = helper.getInstance()
-      |
-        await.getValue() = any(DataFlow::Node awaitable | call.flowsTo(awaitable)).asExpr() and
-        call.calls(nodeFrom, helper.getAsyncMethodName())
-      )
+      // async methods.
+      //
+      // since we have general taint-step from `foo` in `await foo` to the whole
+      // expression, we simply taint the awaitable that is the result of "calling" the
+      // async method. That also allows such an awaitable to be placed in a list (for
+      // use with `asyncio.gather` for example), and thereby propagate taint to the
+      // list.
+      nodeFrom = helper.getInstance() and
+      nodeTo.(DataFlow::MethodCallNode).calls(nodeFrom, helper.getAsyncMethodName())
       or
       // Attributes
       nodeFrom = helper.getInstance() and

python/ql/test/library-tests/frameworks/aiohttp/taint_test.py

@@ -56,9 +56,9 @@ async def test_taint(request: web.Request): # $ requestHandler
         await request.content.readchunk(), # $ tainted
         (await request.content.readchunk())[0], # $ tainted
         [line async for line in request.content], # $ tainted
-        [data async for data in request.content.iter_chunked(1024)], # $ MISSING: tainted
-        [data async for data in request.content.iter_any()], # $ MISSING: tainted
-        [data async for data, _ in request.content.iter_chunks()], # $ MISSING: tainted
+        [data async for data in request.content.iter_chunked(1024)], # $ tainted
+        [data async for data in request.content.iter_any()], # $ tainted
+        [data async for data, _ in request.content.iter_chunks()], # $ tainted
         request.content.read_nowait(), # $ tainted
 
         # aiohttp.StreamReader