Rust: Create CFG scope for async blocks

Author: paldepind

rust/ql/lib/codeql/rust/controlflow/internal/CfgConsistency.qll

@@ -32,8 +32,12 @@ query predicate nonPostOrderExpr(Expr e, string cls) {
  */
 query predicate scopeNoFirst(CfgScope scope) {
   Consistency::scopeNoFirst(scope) and
-  not scope = any(Function f | not exists(f.getBody())) and
-  not scope = any(ClosureExpr c | not exists(c.getBody()))
+  not scope =
+    [
+      any(AstNode f | not f.(Function).hasBody()),
+      any(ClosureExpr c | not c.hasBody()),
+      any(AsyncBlockExpr b | not b.hasStmtList())
+    ]
 }
 
 /** Holds if  `be` is the `else` branch of a `let` statement that results in a panic. */

rust/ql/lib/codeql/rust/controlflow/internal/ControlFlowGraphImpl.qll

@@ -44,12 +44,10 @@ private module CfgInput implements InputSig<Location> {
   predicate successorTypeIsCondition(SuccessorType t) { t instanceof Cfg::BooleanSuccessor }
 
   /** Holds if `first` is first executed when entering `scope`. */
-  predicate scopeFirst(CfgScope scope, AstNode first) {
-    first(scope.(CfgScopeTree).getFirstChildNode(), first)
-  }
+  predicate scopeFirst(CfgScope scope, AstNode first) { scope.scopeFirst(first) }
 
   /** Holds if `scope` is exited when `last` finishes with completion `c`. */
-  predicate scopeLast(CfgScope scope, AstNode last, Completion c) { last(scope.getBody(), last, c) }
+  predicate scopeLast(CfgScope scope, AstNode last, Completion c) { scope.scopeLast(last, c) }
 }
 
 private module CfgSplittingInput implements SplittingInputSig<Location, CfgInput> {
@@ -71,14 +69,7 @@ private module CfgImpl =
 
 import CfgImpl
 
-class CfgScopeTree extends StandardTree, Scope::CfgScope {
-  override predicate first(AstNode first) { first = this }
-
-  override predicate last(AstNode last, Completion c) {
-    last = this and
-    completionIsValidFor(c, this)
-  }
-
+class CallableScopeTree extends StandardTree, PreOrderTree, PostOrderTree, Scope::CallableScope {
   override predicate propagatesAbnormal(AstNode child) { none() }
 
   override AstNode getChildNode(int i) {
@@ -280,13 +271,23 @@ module ExprTrees {
     }
   }
 
+  private AstNode getBlockChildNode(BlockExpr b, int i) {
+    result = b.getStmtList().getStatement(i)
+    or
+    i = b.getStmtList().getNumberOfStatements() and
+    result = b.getStmtList().getTailExpr()
+  }
+
+  class AsyncBlockExprTree extends StandardTree, PreOrderTree, PostOrderTree, AsyncBlockExpr {
+    override AstNode getChildNode(int i) { result = getBlockChildNode(this, i) }
+
+    override predicate propagatesAbnormal(AstNode child) { none() }
+  }
+
   class BlockExprTree extends StandardPostOrderTree, BlockExpr {
-    override AstNode getChildNode(int i) {
-      result = this.getStmtList().getStatement(i)
-      or
-      i = this.getStmtList().getNumberOfStatements() and
-      result = this.getStmtList().getTailExpr()
-    }
+    BlockExprTree() { not this.isAsync() }
+
+    override AstNode getChildNode(int i) { result = getBlockChildNode(this, i) }
 
     override predicate propagatesAbnormal(AstNode child) { child = this.getChildNode(_) }
   }

rust/ql/lib/codeql/rust/controlflow/internal/Scope.qll

@@ -5,11 +5,32 @@ private import codeql.rust.elements.internal.generated.ParentChild
 
 /**
  * A control-flow graph (CFG) scope.
- *
- * A CFG scope is a callable with a body.
  */
-class CfgScope extends Callable {
-  CfgScope() {
+abstract private class CfgScopeImpl extends AstNode {
+  abstract predicate scopeFirst(AstNode first);
+
+  abstract predicate scopeLast(AstNode last, Completion c);
+}
+
+final class CfgScope = CfgScopeImpl;
+
+class AsyncBlockScope extends CfgScopeImpl, BlockExpr {
+  AsyncBlockScope() { this.isAsync() }
+
+  override predicate scopeFirst(AstNode first) {
+    first(this.(ExprTrees::AsyncBlockExprTree).getFirstChildNode(), first)
+  }
+
+  override predicate scopeLast(AstNode last, Completion c) {
+    last(this.(ExprTrees::AsyncBlockExprTree).getLastChildElement(), last, c)
+  }
+}
+
+/**
+ * A CFG scope for a callable (a function or a closure) with a body.
+ */
+class CallableScope extends CfgScopeImpl, Callable {
+  CallableScope() {
     // A function without a body corresponds to a trait method signature and
     // should not have a CFG scope.
     this.(Function).hasBody()
@@ -23,4 +44,11 @@ class CfgScope extends Callable {
     or
     result = this.(ClosureExpr).getBody()
   }
+
+  override predicate scopeFirst(AstNode first) {
+    first(this.(CallableScopeTree).getFirstChildNode(), first)
+  }
+
+  /** Holds if `scope` is exited when `last` finishes with completion `c`. */
+  override predicate scopeLast(AstNode last, Completion c) { last(this.getBody(), last, c) }
 }

rust/ql/lib/codeql/rust/elements/AsyncBlockExpr.qll

@@ -0,0 +1,13 @@
+private import codeql.rust.elements.BlockExpr
+
+/**
+ * A async block expression. For example:
+ * ```rust
+ * async {
+ *     let x = 42;
+ * }
+ * ```
+ */
+final class AsyncBlockExpr extends BlockExpr {
+  AsyncBlockExpr() { this.isAsync() }
+}

rust/ql/lib/rust.qll

@@ -5,6 +5,7 @@ import codeql.Locations
 import codeql.files.FileSystem
 import codeql.rust.elements.AssignmentOperation
 import codeql.rust.elements.LogicalOperation
+import codeql.rust.elements.AsyncBlockExpr
 import codeql.rust.elements.Variable
 import codeql.rust.elements.NamedFormatArgument
 import codeql.rust.elements.PositionalFormatArgument

rust/ql/test/library-tests/controlflow/Cfg.expected

@@ -781,10 +781,12 @@ edges
 | test.rs:372:5:384:5 | enter async_block | test.rs:373:9:375:10 | LetStmt |  |
 | test.rs:372:5:384:5 | exit async_block (normal) | test.rs:372:5:384:5 | exit async_block |  |
 | test.rs:372:28:384:5 | BlockExpr | test.rs:372:5:384:5 | exit async_block (normal) |  |
-| test.rs:373:9:375:10 | LetStmt | test.rs:374:13:374:42 | ExprStmt |  |
+| test.rs:373:9:375:10 | LetStmt | test.rs:373:26:375:9 | BlockExpr |  |
 | test.rs:373:13:373:22 | say_godbye | test.rs:376:9:378:10 | LetStmt | match |
 | test.rs:373:26:375:9 | BlockExpr | test.rs:373:13:373:22 | say_godbye |  |
-| test.rs:374:13:374:41 | MacroExpr | test.rs:373:26:375:9 | BlockExpr |  |
+| test.rs:373:26:375:9 | enter BlockExpr | test.rs:374:13:374:42 | ExprStmt |  |
+| test.rs:373:26:375:9 | exit BlockExpr (normal) | test.rs:373:26:375:9 | exit BlockExpr |  |
+| test.rs:374:13:374:41 | MacroExpr | test.rs:373:26:375:9 | exit BlockExpr (normal) |  |
 | test.rs:374:13:374:41 | PathExpr | test.rs:374:22:374:40 | "godbye, everyone!\\n" |  |
 | test.rs:374:13:374:42 | ExprStmt | test.rs:374:22:374:40 | MacroStmts |  |
 | test.rs:374:22:374:40 | "godbye, everyone!\\n" | test.rs:374:22:374:40 | FormatArgsExpr |  |
@@ -794,10 +796,12 @@ edges
 | test.rs:374:22:374:40 | FormatArgsExpr | test.rs:374:22:374:40 | MacroExpr |  |
 | test.rs:374:22:374:40 | MacroExpr | test.rs:374:22:374:40 | CallExpr |  |
 | test.rs:374:22:374:40 | MacroStmts | test.rs:374:22:374:40 | ExprStmt |  |
-| test.rs:376:9:378:10 | LetStmt | test.rs:377:13:377:37 | ExprStmt |  |
+| test.rs:376:9:378:10 | LetStmt | test.rs:376:31:378:9 | BlockExpr |  |
 | test.rs:376:13:376:27 | say_how_are_you | test.rs:379:9:379:28 | LetStmt | match |
 | test.rs:376:31:378:9 | BlockExpr | test.rs:376:13:376:27 | say_how_are_you |  |
-| test.rs:377:13:377:36 | MacroExpr | test.rs:376:31:378:9 | BlockExpr |  |
+| test.rs:376:31:378:9 | enter BlockExpr | test.rs:377:13:377:37 | ExprStmt |  |
+| test.rs:376:31:378:9 | exit BlockExpr (normal) | test.rs:376:31:378:9 | exit BlockExpr |  |
+| test.rs:377:13:377:36 | MacroExpr | test.rs:376:31:378:9 | exit BlockExpr (normal) |  |
 | test.rs:377:13:377:36 | PathExpr | test.rs:377:22:377:35 | "how are you?\\n" |  |
 | test.rs:377:13:377:37 | ExprStmt | test.rs:377:22:377:35 | MacroStmts |  |
 | test.rs:377:22:377:35 | "how are you?\\n" | test.rs:377:22:377:35 | FormatArgsExpr |  |