Merge pull request #22637 from github/rust/async-block-lambda

Rust: Make await call async blocks as lambdas in dataflow

Author: paldepind

rust/ql/lib/change-notes/2026-09-22-async-block-flow.md

@@ -0,0 +1,4 @@
+---
+category: minorAnalysis
+---
+* Improve data flow for async blocks when used with `await`.

rust/ql/lib/codeql/rust/dataflow/internal/DataFlowImpl.qll

@@ -62,6 +62,9 @@ final class DataFlowCall extends TDataFlowCall {
   /** Gets the underlying call, if any. */
   Call asCall() { this = TCall(result) }
 
+  /** Gets the underlying await expression, if any. */
+  AwaitExpr asAwaitExpr() { this = TAwaitExpr(result) }
+
   predicate isImplicitDerefCall(Expr e, DerefChain derefChain, int i, Function target) {
     this = TImplicitDerefCall(e, derefChain, i, target)
   }
@@ -73,15 +76,16 @@ final class DataFlowCall extends TDataFlowCall {
   }
 
   DataFlowCallable getEnclosingCallable() {
-    result.asCfgScope() = this.asCall().getEnclosingCfgScope()
+    result.asCfgScope() =
+      [this.asCall().getEnclosingCfgScope(), this.asAwaitExpr().getEnclosingCfgScope()]
     or
     result.asCfgScope() = any(Expr e | this.isImplicitDerefCall(e, _, _, _)).getEnclosingCfgScope()
     or
     this.isSummaryCall(result.asSummarizedCallable(), _)
   }
 
   string toString() {
-    result = this.asCall().toString()
+    result = [this.asCall().toString(), this.asAwaitExpr().toString()]
     or
     exists(Expr e, DerefChain derefChain, int i |
       this.isImplicitDerefCall(e, derefChain, i, _) and
@@ -97,7 +101,7 @@ final class DataFlowCall extends TDataFlowCall {
   }
 
   Location getLocation() {
-    result = this.asCall().getLocation()
+    result = [this.asCall().getLocation(), this.asAwaitExpr().getLocation()]
     or
     result = any(Expr e | this.isImplicitDerefCall(e, _, _, _)).getLocation()
   }
@@ -1004,7 +1008,11 @@ module RustDataFlowGen<RustDataFlowInputSig Input> implements InputSig<Location>
    */
   predicate lambdaCall(DataFlowCall call, LambdaCallKind kind, Node receiver) {
     (
-      receiver.asExpr() = call.asCall().(CallExprImpl::DynamicCallExpr).getFunction()
+      receiver.asExpr() =
+        [
+          call.asCall().(CallExprImpl::DynamicCallExpr).getFunction(),
+          call.asAwaitExpr().getExpr()
+        ]
       or
       call.isSummaryCall(_, receiver.(FlowSummaryNode).getSummaryNode())
     ) and
@@ -1176,6 +1184,7 @@ private module Cached {
       Stages::DataFlowStage::ref() and
       call.hasEnclosingCfgScope()
     } or
+    TAwaitExpr(AwaitExpr await) { await.hasEnclosingCfgScope() } or
     TImplicitDerefCall(Expr e, DerefChain derefChain, int i, Function target) {
       TypeInference::implicitDerefChainBorrow(e, derefChain, _) and
       target = derefChain.getElement(i).getDerefFunction() and

rust/ql/lib/codeql/rust/dataflow/internal/Node.qll

@@ -435,6 +435,20 @@ final class ClosureArgumentNode extends ArgumentNode, ExprNode {
   }
 }
 
+/**
+ * A data flow node that represents the run-time representation of an async
+ * block passed into its body when awaited.
+ */
+final class AsyncBlockArgumentNode extends ArgumentNode, ExprNode {
+  private AwaitExpr await;
+
+  AsyncBlockArgumentNode() { this.asExpr() = await.getExpr() }
+
+  override predicate isArgumentOf(DataFlowCall call, RustDataFlow::ArgumentPosition pos) {
+    call.asAwaitExpr() = await and pos.isClosureSelf()
+  }
+}
+
 /** An SSA node. */
 class SsaNode extends Node, TSsaNode {
   SsaImpl::DataFlowIntegration::SsaNode node;
@@ -485,12 +499,17 @@ final private class ExprOutNode extends ExprNode, OutNode {
       not call instanceof DerefExpr and // Handled by `DerefOutNode`
       not call instanceof IndexExpr // Handled by `IndexOutNode`
     )
+    or
+    this.asExpr() instanceof AwaitExpr
   }
 
   /** Gets the underlying call node that includes this out node. */
   override DataFlowCall getCall(ReturnKind kind) {
     result.asCall() = n and
     kind = TNormalReturnKind()
+    or
+    result.asAwaitExpr() = n and
+    kind = TNormalReturnKind()
   }
 }
 

rust/ql/test/library-tests/dataflow/global/inline-flow.expected

@@ -196,7 +196,9 @@ edges
 | main.rs:342:9:342:9 | a | main.rs:343:10:343:10 | a | provenance |  |
 | main.rs:342:13:342:26 | async_source(...) | main.rs:342:9:342:9 | a | provenance |  |
 | main.rs:346:13:346:13 | c | main.rs:347:14:347:14 | c | provenance |  |
+| main.rs:346:13:346:13 | c | main.rs:348:9:348:9 | c | provenance |  |
 | main.rs:346:17:346:25 | source(...) | main.rs:346:13:346:13 | c | provenance |  |
+| main.rs:348:9:348:9 | c | main.rs:350:10:350:16 | await b | provenance |  |
 | main.rs:354:9:354:9 | a | main.rs:355:10:355:10 | a | provenance |  |
 | main.rs:354:13:354:55 | ...::block_on(...) | main.rs:354:9:354:9 | a | provenance |  |
 | main.rs:354:41:354:54 | async_source(...) | main.rs:354:13:354:55 | ...::block_on(...) | provenance | MaD:3 |
@@ -443,6 +445,8 @@ nodes
 | main.rs:346:13:346:13 | c | semmle.label | c |
 | main.rs:346:17:346:25 | source(...) | semmle.label | source(...) |
 | main.rs:347:14:347:14 | c | semmle.label | c |
+| main.rs:348:9:348:9 | c | semmle.label | c |
+| main.rs:350:10:350:16 | await b | semmle.label | await b |
 | main.rs:354:9:354:9 | a | semmle.label | a |
 | main.rs:354:13:354:55 | ...::block_on(...) | semmle.label | ...::block_on(...) |
 | main.rs:354:41:354:54 | async_source(...) | semmle.label | async_source(...) |
@@ -539,6 +543,7 @@ testFailures
 | main.rs:337:10:337:10 | a | main.rs:336:13:336:21 | source(...) | main.rs:337:10:337:10 | a | $@ | main.rs:336:13:336:21 | source(...) | source(...) |
 | main.rs:343:10:343:10 | a | main.rs:336:13:336:21 | source(...) | main.rs:343:10:343:10 | a | $@ | main.rs:336:13:336:21 | source(...) | source(...) |
 | main.rs:347:14:347:14 | c | main.rs:346:17:346:25 | source(...) | main.rs:347:14:347:14 | c | $@ | main.rs:346:17:346:25 | source(...) | source(...) |
+| main.rs:350:10:350:16 | await b | main.rs:346:17:346:25 | source(...) | main.rs:350:10:350:16 | await b | $@ | main.rs:346:17:346:25 | source(...) | source(...) |
 | main.rs:355:10:355:10 | a | main.rs:336:13:336:21 | source(...) | main.rs:355:10:355:10 | a | $@ | main.rs:336:13:336:21 | source(...) | source(...) |
 | main.rs:404:14:404:15 | n1 | main.rs:379:13:379:21 | source(...) | main.rs:404:14:404:15 | n1 | $@ | main.rs:379:13:379:21 | source(...) | source(...) |
 | main.rs:408:14:408:15 | n2 | main.rs:379:13:379:21 | source(...) | main.rs:408:14:408:15 | n2 | $@ | main.rs:379:13:379:21 | source(...) | source(...) |

rust/ql/test/library-tests/dataflow/global/main.rs

@@ -347,7 +347,7 @@ async fn test_async_await_async_part() {
         sink(c); // $ hasValueFlow=2
         c
     };
-    sink(b.await); // $ MISSING: hasValueFlow=2
+    sink(b.await); // $ hasValueFlow=2
 }
 
 fn test_async_await() {

rust/ql/test/library-tests/dataflow/lambdas/Cargo.lock

@@ -63,6 +63,14 @@ edges
 | main.rs:111:26:111:36 | \|...\| ... : ... => .. [captured a] | main.rs:111:35:111:35 | a | provenance | heuristic-callback |
 | main.rs:112:26:112:36 | \|...\| ... : ... => .. [captured a] | main.rs:112:35:112:35 | a | provenance | heuristic-callback |
 | main.rs:113:26:113:36 | \|...\| ... : ... => .. [captured a] | main.rs:113:35:113:35 | a | provenance | heuristic-callback |
+| main.rs:120:30:120:40 | source(...) | main.rs:121:14:121:25 | await future | provenance |  |
+| main.rs:125:20:125:30 | source(...) | main.rs:126:9:128:9 | { ... } : ... => .. [captured data] | provenance |  |
+| main.rs:126:9:128:9 | { ... } : ... => .. [captured data] | main.rs:127:18:127:21 | data | provenance |  |
+| main.rs:133:20:133:30 | source(...) | main.rs:135:14:135:19 | future [captured data] | provenance |  |
+| main.rs:135:14:135:19 | future [captured data] | main.rs:134:30:134:33 | data | provenance |  |
+| main.rs:135:14:135:19 | future [captured data] | main.rs:135:14:135:25 | await future | provenance |  |
+| main.rs:141:9:143:9 | [post] { ... } : ... => .. [captured captured] | main.rs:145:14:145:21 | captured | provenance |  |
+| main.rs:142:24:142:34 | source(...) | main.rs:141:9:143:9 | [post] { ... } : ... => .. [captured captured] | provenance |  |
 nodes
 | main.rs:10:20:10:52 | if cond {...} else {...} | semmle.label | if cond {...} else {...} |
 | main.rs:10:30:10:39 | source(...) | semmle.label | source(...) |
@@ -138,9 +146,22 @@ nodes
 | main.rs:112:35:112:35 | a | semmle.label | a |
 | main.rs:113:26:113:36 | \|...\| ... : ... => .. [captured a] | semmle.label | \|...\| ... : ... => .. [captured a] |
 | main.rs:113:35:113:35 | a | semmle.label | a |
+| main.rs:120:30:120:40 | source(...) | semmle.label | source(...) |
+| main.rs:121:14:121:25 | await future | semmle.label | await future |
+| main.rs:125:20:125:30 | source(...) | semmle.label | source(...) |
+| main.rs:126:9:128:9 | { ... } : ... => .. [captured data] | semmle.label | { ... } : ... => .. [captured data] |
+| main.rs:127:18:127:21 | data | semmle.label | data |
+| main.rs:133:20:133:30 | source(...) | semmle.label | source(...) |
+| main.rs:134:30:134:33 | data | semmle.label | data |
+| main.rs:135:14:135:19 | future [captured data] | semmle.label | future [captured data] |
+| main.rs:135:14:135:25 | await future | semmle.label | await future |
+| main.rs:141:9:143:9 | [post] { ... } : ... => .. [captured captured] | semmle.label | [post] { ... } : ... => .. [captured captured] |
+| main.rs:142:24:142:34 | source(...) | semmle.label | source(...) |
+| main.rs:145:14:145:21 | captured | semmle.label | captured |
 subpaths
 | main.rs:29:21:29:21 | a | main.rs:27:20:27:23 | ... | main.rs:27:26:27:52 | if cond {...} else {...} | main.rs:29:13:29:22 | f(...) |
 | main.rs:77:21:77:21 | a | main.rs:66:24:66:32 | ...: i64 | main.rs:66:42:72:1 | { ... } | main.rs:77:13:77:22 | f(...) |
+| main.rs:135:14:135:19 | future [captured data] | main.rs:134:30:134:33 | data | main.rs:134:30:134:33 | data | main.rs:135:14:135:25 | await future |
 testFailures
 #select
 | main.rs:11:10:11:16 | f(...) | main.rs:10:30:10:39 | source(...) | main.rs:11:10:11:16 | f(...) | $@ | main.rs:10:30:10:39 | source(...) | source(...) |
@@ -161,3 +182,7 @@ testFailures
 | main.rs:111:35:111:35 | a | main.rs:109:13:109:22 | source(...) | main.rs:111:35:111:35 | a | $@ | main.rs:109:13:109:22 | source(...) | source(...) |
 | main.rs:112:35:112:35 | a | main.rs:109:13:109:22 | source(...) | main.rs:112:35:112:35 | a | $@ | main.rs:109:13:109:22 | source(...) | source(...) |
 | main.rs:113:35:113:35 | a | main.rs:109:13:109:22 | source(...) | main.rs:113:35:113:35 | a | $@ | main.rs:109:13:109:22 | source(...) | source(...) |
+| main.rs:121:14:121:25 | await future | main.rs:120:30:120:40 | source(...) | main.rs:121:14:121:25 | await future | $@ | main.rs:120:30:120:40 | source(...) | source(...) |
+| main.rs:127:18:127:21 | data | main.rs:125:20:125:30 | source(...) | main.rs:127:18:127:21 | data | $@ | main.rs:125:20:125:30 | source(...) | source(...) |
+| main.rs:135:14:135:25 | await future | main.rs:133:20:133:30 | source(...) | main.rs:135:14:135:25 | await future | $@ | main.rs:133:20:133:30 | source(...) | source(...) |
+| main.rs:145:14:145:21 | captured | main.rs:142:24:142:34 | source(...) | main.rs:145:14:145:21 | captured | $@ | main.rs:142:24:142:34 | source(...) | source(...) |