Rust: Include patterns as data flow nodes

Author: paldepind

rust/ql/lib/codeql/rust/controlflow/CfgNodes.qll

@@ -12,6 +12,9 @@ class AstCfgNode extends CfgNode {
   AstNode node;
 
   AstCfgNode() { node = this.getAstNode() }
+
+  /** Gets the underlying ast node. */
+  AstNode getAstNode() { result = node }
 }
 
 /** A CFG node that corresponds to a parameter in the AST. */
@@ -22,6 +25,14 @@ class ParamCfgNode extends AstCfgNode {
   Param getParam() { result = node }
 }
 
+/** A CFG node that corresponds to a parameter in the AST. */
+class PatCfgNode extends AstCfgNode {
+  override Pat node;
+
+  /** Gets the underlying pattern. */
+  Pat getPat() { result = node }
+}
+
 /** A CFG node that corresponds to an expression in the AST. */
 class ExprCfgNode extends AstCfgNode {
   override Expr node;

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

@@ -123,6 +123,19 @@ module Node {
     override Location getLocation() { none() }
   }
 
+  /** A data flow node that corresponds to a CFG node for an AST node. */
+  abstract private class AstCfgFlowNode extends Node {
+    AstCfgNode n;
+
+    override CfgNode getCfgNode() { result = n }
+
+    override CfgScope getCfgScope() { result = n.getAstNode().getEnclosingCallable() }
+
+    override Location getLocation() { result = n.getAstNode().getLocation() }
+
+    override string toString() { result = n.getAstNode().toString() }
+  }
+
   /**
    * A node in the data flow graph that corresponds to an expression in the
    * AST.
@@ -131,39 +144,34 @@ module Node {
    * to multiple `ExprNode`s, just like it may correspond to multiple
    * `ControlFlow::Node`s.
    */
-  class ExprNode extends Node, TExprNode {
-    ExprCfgNode n;
+  class ExprNode extends AstCfgFlowNode, TExprNode {
+    override ExprCfgNode n;
 
     ExprNode() { this = TExprNode(n) }
 
-    override CfgScope getCfgScope() { result = this.asExpr().getEnclosingCallable() }
-
-    override Location getLocation() { result = n.getExpr().getLocation() }
+    override Expr asExpr() { result = n.getExpr() }
+  }
 
-    override string toString() { result = n.getExpr().toString() }
+  final class PatNode extends AstCfgFlowNode, TPatNode {
+    override PatCfgNode n;
 
-    override Expr asExpr() { result = n.getExpr() }
+    PatNode() { this = TPatNode(n) }
 
-    override CfgNode getCfgNode() { result = n }
+    /** Gets the Pat in the AST that this node corresponds to. */
+    Pat getPat() { result = n.getPat() }
   }
 
   /**
    * The value of a parameter at function entry, viewed as a node in a data
    * flow graph.
    */
-  final class ParameterNode extends Node, TParameterNode {
-    ParamCfgNode parameter;
+  final class ParameterNode extends AstCfgFlowNode, TParameterNode {
+    override ParamCfgNode n;
 
-    ParameterNode() { this = TParameterNode(parameter) }
-
-    override CfgScope getCfgScope() { result = parameter.getParam().getEnclosingCallable() }
-
-    override Location getLocation() { result = parameter.getLocation() }
-
-    override string toString() { result = parameter.toString() }
+    ParameterNode() { this = TParameterNode(n) }
 
     /** Gets the parameter in the AST that this node corresponds to. */
-    Param getParameter() { result = parameter.getParam() }
+    Param getParameter() { result = n.getParam() }
   }
 
   final class ArgumentNode = NaNode;
@@ -281,6 +289,11 @@ module LocalFlow {
   pragma[nomagic]
   predicate localFlowStepCommon(Node nodeFrom, Node nodeTo) {
     nodeFrom.getCfgNode() = getALastEvalNode(nodeTo.getCfgNode())
+    or
+    exists(LetStmt s |
+      nodeFrom.getCfgNode().getAstNode() = s.getInitializer() and
+      nodeTo.getCfgNode().getAstNode() = s.getPat()
+    )
   }
 }
 
@@ -485,6 +498,7 @@ private module Cached {
   newtype TNode =
     TExprNode(ExprCfgNode n) or
     TParameterNode(ParamCfgNode p) or
+    TPatNode(PatCfgNode p) or
     TSsaNode(SsaImpl::DataFlowIntegration::SsaNode node)
 
   cached

rust/ql/test/library-tests/dataflow/local/DataFlowStep.expected

@@ -1,18 +1,25 @@
 | main.rs:2:9:2:9 | [SSA] s | main.rs:3:33:3:33 | s |
+| main.rs:2:13:2:19 | "Hello" | main.rs:2:9:2:9 | s |
 | main.rs:6:18:6:21 | [SSA] cond | main.rs:9:16:9:19 | cond |
 | main.rs:7:9:7:9 | [SSA] a | main.rs:10:9:10:9 | a |
+| main.rs:7:13:7:13 | 1 | main.rs:7:9:7:9 | a |
 | main.rs:8:9:8:9 | [SSA] b | main.rs:12:9:12:9 | b |
+| main.rs:8:13:8:13 | 2 | main.rs:8:9:8:9 | b |
 | main.rs:9:9:9:9 | [SSA] c | main.rs:14:5:14:5 | c |
+| main.rs:9:13:13:5 | IfExpr | main.rs:9:9:9:9 | c |
 | main.rs:9:21:11:5 | BlockExpr | main.rs:9:13:13:5 | IfExpr |
 | main.rs:10:9:10:9 | a | main.rs:9:21:11:5 | BlockExpr |
 | main.rs:11:12:13:5 | BlockExpr | main.rs:9:13:13:5 | IfExpr |
 | main.rs:12:9:12:9 | b | main.rs:11:12:13:5 | BlockExpr |
 | main.rs:14:5:14:5 | c | main.rs:6:37:15:1 | BlockExpr |
 | main.rs:18:9:18:9 | [SSA] a | main.rs:20:15:20:15 | a |
+| main.rs:18:13:18:13 | 1 | main.rs:18:9:18:9 | a |
 | main.rs:19:9:19:9 | [SSA] b | main.rs:22:5:22:5 | b |
+| main.rs:19:13:21:5 | LoopExpr | main.rs:19:9:19:9 | b |
 | main.rs:20:9:20:15 | BreakExpr | main.rs:19:13:21:5 | LoopExpr |
 | main.rs:20:15:20:15 | a | main.rs:20:9:20:15 | BreakExpr |
 | main.rs:22:5:22:5 | b | main.rs:17:29:23:1 | BlockExpr |
+| main.rs:26:17:26:17 | 1 | main.rs:26:9:26:13 | i |
 | main.rs:27:5:27:5 | [SSA] i | main.rs:28:5:28:5 | i |
 | main.rs:27:5:27:5 | i | main.rs:27:5:27:5 | [SSA] i |
 | main.rs:28:5:28:5 | i | main.rs:25:24:29:1 | BlockExpr |