Python: Align implementations of awaited.

yoff · yoff · commit a5912ff76de3 · 2021-09-28T16:42:19.000+02:00
diff --git a/python/ql/lib/semmle/python/ApiGraphs.qll b/python/ql/lib/semmle/python/ApiGraphs.qll
@@ -500,21 +500,21 @@ module API {
       or
       // `async for x in l`
       // - `awaitedValue` is `l`
-      // - `result` is `l` (should perhaps be `x`)
+      // - `result` is `l` (should perhaps be `x`, but that should really be a read)
       exists(AsyncFor asyncFor |
         result.asExpr() = asyncFor.getTarget() and
-        // Morally, we should perhaps use asyncFor.getIter() = awaitedValue.asExpr()
-        // but that does not work.
+        // Morally, we should perhaps use asyncFor.getIter() = awaitedValue.asExpr(),
+        // but that is actually behind a read step rather than a flow step.
         asyncFor.getTarget() = awaitedValue.asExpr()
       )
       or
       // `async with x as y`
       // - `awaitedValue` is `x`
-      // - `result` is `x` (should probably be `y`)
+      // - `result` is `x` (should probably be `y` but it might not exist)
       exists(AsyncWith asyncWith |
         result.asExpr() = asyncWith.getContextExpr() and
-        // Morally, we should perhaps use asyncWith.getOptionalVars() = awaitedValue.asExpr()
-        // but that does not work.
+        // Morally, we should perhaps use asyncWith.getOptionalVars() = awaitedValue.asExpr(),
+        // but that might not exist.
         asyncWith.getContextExpr() = awaitedValue.asExpr()
       )
     }
diff --git a/python/ql/lib/semmle/python/frameworks/Asyncpg.qll b/python/ql/lib/semmle/python/frameworks/Asyncpg.qll
@@ -42,29 +42,31 @@ private module Asyncpg {
    * Holds if `result` is the result of awaiting `n`.
    */
   pragma[inline]
-  DataFlow::Node awaited(DataFlow::Node n) {
+  DataFlow::Node awaited(DataFlow::Node awaitedValue) {
     // `await` x
     // - `awaitedValue` is `x`
     // - `result` is `await x`
     exists(Await await |
       result.asExpr() = await and
-      await.getValue() = n.asExpr()
+      await.getValue() = awaitedValue.asExpr()
     )
     or
     // `async for x in l`
-    // - `awaitedValue` is `l`
-    // - `result` is `x`
-    exists(AsyncFor asyncFor |
+    // - `awaitedValue` is local source of `l`
+    // - `result` is `l`
+    exists(AsyncFor asyncFor, DataFlow::Node awaited |
       result.asExpr() = asyncFor.getTarget() and
-      asyncFor.getIter() = n.asExpr()
+      asyncFor.getIter() = awaited.asExpr() and
+      awaited.getALocalSource() = awaitedValue
     )
     or
     // `async with x as y`
-    // - `awaitedValue` is `x`
-    // - `result` is `y`
-    exists(AsyncWith asyncWith |
+    // - `awaitedValue` is local source of `x`
+    // - `result` is `x`
+    exists(AsyncWith asyncWith, DataFlow::Node awaited |
       result.asExpr() = asyncWith.getContextExpr() and
-      asyncWith.getOptionalVars() = n.asExpr()
+      asyncWith.getOptionalVars() = awaited.asExpr() and
+      awaited.getALocalSource() = awaitedValue
     )
   }
 
