Rust: Handle async blocks in SSA analysis

paldepind · paldepind · commit a5e13ee71b8e · 2024-11-10T11:59:37.000+01:00
diff --git a/rust/ql/lib/codeql/rust/controlflow/BasicBlocks.qll b/rust/ql/lib/codeql/rust/controlflow/BasicBlocks.qll
@@ -12,7 +12,7 @@ final class BasicBlock = BasicBlockImpl;
  * without branches or joins.
  */
 private class BasicBlockImpl extends TBasicBlockStart {
-  /** Gets the scope of this basic block. */
+  /** Gets the CFG scope of this basic block. */
   CfgScope getScope() { result = this.getAPredecessor().getScope() }
 
   /** Gets an immediate successor of this basic block, if any. */
diff --git a/rust/ql/lib/codeql/rust/controlflow/ControlFlowGraph.qll b/rust/ql/lib/codeql/rust/controlflow/ControlFlowGraph.qll
@@ -9,6 +9,8 @@ private import BasicBlocks
 
 final class CfgScope = Scope::CfgScope;
 
+predicate getEnclosingCfgScope = Scope::getEnclosingCfgScope/1;
+
 final class SuccessorType = SuccessorTypeImpl;
 
 final class NormalSuccessor = NormalSuccessorImpl;
@@ -52,3 +54,6 @@ final class CfgNode extends Node {
   /** Gets the basic block that this control flow node belongs to. */
   BasicBlock getBasicBlock() { result.getANode() = this }
 }
+
+/** Holds if evaluating `e` jumps to the evaluation of a different CFG scope. */
+predicate isControlFlowJump(Expr e) { e instanceof CallExprBase or e instanceof AwaitExpr }
diff --git a/rust/ql/lib/codeql/rust/controlflow/internal/ControlFlowGraphImpl.qll b/rust/ql/lib/codeql/rust/controlflow/internal/ControlFlowGraphImpl.qll
@@ -23,7 +23,7 @@ private module CfgInput implements InputSig<Location> {
   class CfgScope = Scope::CfgScope;
 
   CfgScope getCfgScope(AstNode n) {
-    result = n.getEnclosingCallable() and
+    result = Scope::getEnclosingCfgScope(n) and
     Stages::CfgStage::ref()
   }
 
diff --git a/rust/ql/lib/codeql/rust/controlflow/internal/Scope.qll b/rust/ql/lib/codeql/rust/controlflow/internal/Scope.qll
@@ -14,9 +14,7 @@ abstract private class CfgScopeImpl extends AstNode {
 
 final class CfgScope = CfgScopeImpl;
 
-class AsyncBlockScope extends CfgScopeImpl, BlockExpr {
-  AsyncBlockScope() { this.isAsync() }
-
+final class AsyncBlockScope extends CfgScopeImpl, AsyncBlockExpr {
   override predicate scopeFirst(AstNode first) {
     first(this.(ExprTrees::AsyncBlockExprTree).getFirstChildNode(), first)
   }
@@ -29,7 +27,7 @@ class AsyncBlockScope extends CfgScopeImpl, BlockExpr {
 /**
  * A CFG scope for a callable (a function or a closure) with a body.
  */
-class CallableScope extends CfgScopeImpl, Callable {
+final class CallableScope extends CfgScopeImpl, Callable {
   CallableScope() {
     // A function without a body corresponds to a trait method signature and
     // should not have a CFG scope.
@@ -52,3 +50,13 @@ class CallableScope extends CfgScopeImpl, Callable {
   /** Holds if `scope` is exited when `last` finishes with completion `c`. */
   override predicate scopeLast(AstNode last, Completion c) { last(this.getBody(), last, c) }
 }
+
+/** Gets the CFG scope that encloses `node`, if any. */
+CfgScope getEnclosingCfgScope(AstNode node) {
+  exists(AstNode p | p = node.getParentNode() |
+    result = p
+    or
+    not p instanceof CfgScope and
+    result = getEnclosingCfgScope(p)
+  )
+}
diff --git a/rust/ql/lib/codeql/rust/dataflow/internal/DataFlowImpl.qll b/rust/ql/lib/codeql/rust/dataflow/internal/DataFlowImpl.qll
@@ -134,7 +134,7 @@ module Node {
 
     ExprNode() { this = TExprNode(n) }
 
-    override CfgScope getCfgScope() { result = this.asExpr().getEnclosingCallable() }
+    override CfgScope getCfgScope() { result = getEnclosingCfgScope(this.asExpr()) }
 
     override Location getLocation() { result = n.getExpr().getLocation() }
 
@@ -154,7 +154,7 @@ module Node {
 
     ParameterNode() { this = TParameterNode(parameter) }
 
-    override CfgScope getCfgScope() { result = parameter.getEnclosingCallable() }
+    override CfgScope getCfgScope() { result = getEnclosingCfgScope(parameter) }
 
     override Location getLocation() { result = parameter.getLocation() }
 
diff --git a/rust/ql/lib/codeql/rust/dataflow/internal/SsaImpl.qll b/rust/ql/lib/codeql/rust/dataflow/internal/SsaImpl.qll
@@ -156,19 +156,19 @@ private predicate variableReadActual(BasicBlock bb, int i, Variable v) {
  */
 pragma[noinline]
 private predicate hasCapturedWrite(Variable v, Cfg::CfgScope scope) {
-  any(VariableWriteAccess write | write.getVariable() = v and scope = write.getEnclosingCallable+())
+  any(VariableWriteAccess write | write.getVariable() = v and scope = getEnclosingCfgScope+(write))
       .isCapture()
 }
 
 /**
  * Holds if `v` is read inside basic block `bb` at index `i`, which is in the
- * immediate outer scope of `scope`.
+ * immediate outer CFG scope of `scope`.
  */
 pragma[noinline]
 private predicate variableReadActualInOuterScope(
   BasicBlock bb, int i, Variable v, Cfg::CfgScope scope
 ) {
-  variableReadActual(bb, i, v) and bb.getScope() = scope.getEnclosingCallable()
+  variableReadActual(bb, i, v) and bb.getScope() = getEnclosingCfgScope(scope)
 }
 
 pragma[noinline]
@@ -263,7 +263,7 @@ private predicate readsCapturedVariable(BasicBlock bb, Variable v) {
  */
 pragma[noinline]
 private predicate hasCapturedRead(Variable v, Cfg::CfgScope scope) {
-  any(VariableReadAccess read | read.getVariable() = v and scope = read.getEnclosingCallable+())
+  any(VariableReadAccess read | read.getVariable() = v and scope = getEnclosingCfgScope+(read))
       .isCapture()
 }
 
@@ -273,14 +273,15 @@ private predicate hasCapturedRead(Variable v, Cfg::CfgScope scope) {
  */
 pragma[noinline]
 private predicate variableWriteInOuterScope(BasicBlock bb, int i, Variable v, Cfg::CfgScope scope) {
-  SsaInput::variableWrite(bb, i, v, _) and scope.getEnclosingCallable() = bb.getScope()
+  SsaInput::variableWrite(bb, i, v, _) and getEnclosingCfgScope(scope) = bb.getScope()
 }
 
 /**
  * Holds if the call `call` at index `i` in basic block `bb` may reach
  * a callable that reads captured variable `v`.
  */
-private predicate capturedCallRead(CallExprBase call, BasicBlock bb, int i, Variable v) {
+private predicate capturedCallRead(Expr call, BasicBlock bb, int i, Variable v) {
+  Cfg::isControlFlowJump(call) and
   exists(Cfg::CfgScope scope |
     hasCapturedRead(v, scope) and
     (
@@ -295,7 +296,8 @@ private predicate capturedCallRead(CallExprBase call, BasicBlock bb, int i, Vari
  * Holds if the call `call` at index `i` in basic block `bb` may reach a callable
  * that writes captured variable `v`.
  */
-predicate capturedCallWrite(CallExprBase call, BasicBlock bb, int i, Variable v) {
+predicate capturedCallWrite(Expr call, BasicBlock bb, int i, Variable v) {
+  Cfg::isControlFlowJump(call) and
   call = bb.getNode(i).getAstNode() and
   exists(Cfg::CfgScope scope |
     hasVariableReadWithCapturedWrite(bb, any(int j | j > i), v, scope)
diff --git a/rust/ql/lib/codeql/rust/elements/internal/VariableImpl.qll b/rust/ql/lib/codeql/rust/elements/internal/VariableImpl.qll
@@ -1,4 +1,5 @@
 private import rust
+private import codeql.rust.controlflow.ControlFlowGraph
 private import codeql.rust.elements.internal.generated.ParentChild
 private import codeql.rust.elements.internal.PathExprBaseImpl::Impl as PathExprBaseImpl
 private import codeql.rust.elements.internal.FormatTemplateVariableAccessImpl::Impl as FormatTemplateVariableAccessImpl
@@ -445,7 +446,7 @@ module Impl {
     Variable getVariable() { result = v }
 
     /** Holds if this access is a capture. */
-    predicate isCapture() { this.getEnclosingCallable() != v.getPat().getEnclosingCallable() }
+    predicate isCapture() { getEnclosingCfgScope(this) != getEnclosingCfgScope(v.getPat()) }
 
     override string toString() { result = name }
 
diff --git a/rust/ql/test/library-tests/variables/Ssa.expected b/rust/ql/test/library-tests/variables/Ssa.expected
@@ -124,6 +124,8 @@ definition
 | variables.rs:428:9:428:20 | closure3 | variables.rs:428:13:428:20 | closure3 |
 | variables.rs:436:9:436:13 | i | variables.rs:436:13:436:13 | i |
 | variables.rs:437:9:437:13 | block | variables.rs:437:9:437:13 | block |
+| variables.rs:438:9:438:9 | i | variables.rs:436:13:436:13 | i |
+| variables.rs:441:5:441:15 | AwaitExpr | variables.rs:436:13:436:13 | i |
 | variables.rs:445:8:445:8 | b | variables.rs:445:8:445:8 | b |
 | variables.rs:446:9:446:13 | x | variables.rs:446:13:446:13 | x |
 | variables.rs:449:5:457:5 | phi | variables.rs:446:13:446:13 | x |
@@ -239,8 +241,8 @@ read
 | variables.rs:420:9:420:20 | closure2 | variables.rs:420:13:420:20 | closure2 | variables.rs:423:5:423:12 | closure2 |
 | variables.rs:423:5:423:14 | CallExpr | variables.rs:418:13:418:13 | y | variables.rs:424:15:424:15 | y |
 | variables.rs:428:9:428:20 | closure3 | variables.rs:428:13:428:20 | closure3 | variables.rs:431:5:431:12 | closure3 |
-| variables.rs:436:9:436:13 | i | variables.rs:436:13:436:13 | i | variables.rs:442:15:442:15 | i |
 | variables.rs:437:9:437:13 | block | variables.rs:437:9:437:13 | block | variables.rs:441:5:441:9 | block |
+| variables.rs:441:5:441:15 | AwaitExpr | variables.rs:436:13:436:13 | i | variables.rs:442:15:442:15 | i |
 | variables.rs:445:8:445:8 | b | variables.rs:445:8:445:8 | b | variables.rs:449:8:449:8 | b |
 | variables.rs:446:9:446:13 | x | variables.rs:446:13:446:13 | x | variables.rs:447:15:447:15 | x |
 | variables.rs:446:9:446:13 | x | variables.rs:446:13:446:13 | x | variables.rs:448:15:448:15 | x |
@@ -343,8 +345,8 @@ firstRead
 | variables.rs:420:9:420:20 | closure2 | variables.rs:420:13:420:20 | closure2 | variables.rs:423:5:423:12 | closure2 |
 | variables.rs:423:5:423:14 | CallExpr | variables.rs:418:13:418:13 | y | variables.rs:424:15:424:15 | y |
 | variables.rs:428:9:428:20 | closure3 | variables.rs:428:13:428:20 | closure3 | variables.rs:431:5:431:12 | closure3 |
-| variables.rs:436:9:436:13 | i | variables.rs:436:13:436:13 | i | variables.rs:442:15:442:15 | i |
 | variables.rs:437:9:437:13 | block | variables.rs:437:9:437:13 | block | variables.rs:441:5:441:9 | block |
+| variables.rs:441:5:441:15 | AwaitExpr | variables.rs:436:13:436:13 | i | variables.rs:442:15:442:15 | i |
 | variables.rs:445:8:445:8 | b | variables.rs:445:8:445:8 | b | variables.rs:449:8:449:8 | b |
 | variables.rs:446:9:446:13 | x | variables.rs:446:13:446:13 | x | variables.rs:447:15:447:15 | x |
 | variables.rs:449:5:457:5 | phi | variables.rs:446:13:446:13 | x | variables.rs:458:15:458:15 | x |
@@ -443,8 +445,8 @@ lastRead
 | variables.rs:420:9:420:20 | closure2 | variables.rs:420:13:420:20 | closure2 | variables.rs:423:5:423:12 | closure2 |
 | variables.rs:423:5:423:14 | CallExpr | variables.rs:418:13:418:13 | y | variables.rs:424:15:424:15 | y |
 | variables.rs:428:9:428:20 | closure3 | variables.rs:428:13:428:20 | closure3 | variables.rs:431:5:431:12 | closure3 |
-| variables.rs:436:9:436:13 | i | variables.rs:436:13:436:13 | i | variables.rs:442:15:442:15 | i |
 | variables.rs:437:9:437:13 | block | variables.rs:437:9:437:13 | block | variables.rs:441:5:441:9 | block |
+| variables.rs:441:5:441:15 | AwaitExpr | variables.rs:436:13:436:13 | i | variables.rs:442:15:442:15 | i |
 | variables.rs:445:8:445:8 | b | variables.rs:445:8:445:8 | b | variables.rs:449:8:449:8 | b |
 | variables.rs:446:9:446:13 | x | variables.rs:446:13:446:13 | x | variables.rs:448:15:448:15 | x |
 | variables.rs:449:5:457:5 | phi | variables.rs:446:13:446:13 | x | variables.rs:458:15:458:15 | x |
@@ -554,5 +556,6 @@ assigns
 | variables.rs:23:5:23:6 | x2 | variables.rs:23:10:23:10 | 5 |
 | variables.rs:30:5:30:5 | x | variables.rs:30:9:30:9 | 2 |
 | variables.rs:421:9:421:9 | y | variables.rs:421:13:421:13 | 3 |
+| variables.rs:438:9:438:9 | i | variables.rs:438:13:438:13 | 1 |
 | variables.rs:450:9:450:9 | x | variables.rs:450:13:450:13 | 2 |
 | variables.rs:454:9:454:9 | x | variables.rs:454:13:454:13 | 3 |
diff --git a/rust/ql/test/library-tests/variables/variables.expected b/rust/ql/test/library-tests/variables/variables.expected
@@ -461,8 +461,10 @@ capturedVariable
 | variables.rs:410:13:410:13 | x |
 | variables.rs:418:13:418:13 | y |
 | variables.rs:426:13:426:13 | z |
+| variables.rs:436:13:436:13 | i |
 capturedAccess
 | variables.rs:403:19:403:19 | x |
 | variables.rs:413:19:413:19 | x |
 | variables.rs:421:9:421:9 | y |
 | variables.rs:429:9:429:9 | z |
+| variables.rs:438:9:438:9 | i |

Original file line number	Diff line number	Diff line change
`@@ -23,7 +23,7 @@ private module CfgInput implements InputSig<Location> {`
`23`	`23`	`class CfgScope = Scope::CfgScope;`
`24`	`24`
`25`	`25`	`CfgScope getCfgScope(AstNode n) {`
`26`		`- result = n.getEnclosingCallable() and`
	`26`	`+ result = Scope::getEnclosingCfgScope(n) and`
`27`	`27`	`Stages::CfgStage::ref()`
`28`	`28`	`}`
`29`	`29`