Rust: Integrate SSA into data flow

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

@@ -0,0 +1,20 @@
+/**
+ * Provides subclasses of `CfgNode` that represents different types of nodes in
+ * the control flow graph.
+ */
+
+private import rust
+private import ControlFlowGraph
+
+/** A CFG node that corresponds to an element in the AST. */
+class AstCfgNode extends CfgNode {
+  AstCfgNode() { exists(this.getAstNode()) }
+}
+
+/** A CFG node that corresponds to an expression in the AST. */
+class ExprCfgNode extends AstCfgNode {
+  ExprCfgNode() { this.getAstNode() instanceof Expr }
+
+  /** Gets the underlying expression. */
+  Expr getExpr() { result = this.getAstNode() }
+}

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

@@ -7,7 +7,9 @@ private import codeql.util.Unit
 private import codeql.dataflow.DataFlow
 private import codeql.dataflow.internal.DataFlowImpl
 private import rust
+private import SsaImpl as SsaImpl
 private import codeql.rust.controlflow.ControlFlowGraph
+private import codeql.rust.controlflow.CfgNodes
 private import codeql.rust.dataflow.Ssa
 
 module Node {
@@ -52,18 +54,43 @@ module Node {
     override Location getLocation() { none() }
   }
 
+  /**
+   * A node in the data flow graph that corresponds to an expression in the
+   * AST.
+   *
+   * Note that because of control-flow splitting, one `Expr` may correspond
+   * to multiple `ExprNode`s, just like it may correspond to multiple
+   * `ControlFlow::Node`s.
+   */
+  final class ExprNode extends Node, TExprNode {
+    ExprCfgNode n;
+
+    ExprNode() { this = TExprNode(n) }
+
+    override Location getLocation() { result = n.getExpr().getLocation() }
+
+    override string toString() { result = n.getExpr().toString() }
+
+    override Expr asExpr() { result = n.getExpr() }
+
+    override CfgNode getCfgNode() { result = n }
+  }
+
   /**
    * The value of a parameter at function entry, viewed as a node in a data
    * flow graph.
    */
-  final class ParameterNode extends Node {
-    Param param;
+  final class ParameterNode extends Node, TParameterNode {
+    Param parameter;
+
+    ParameterNode() { this = TParameterNode(parameter) }
 
-    ParameterNode() { this = TSourceParameterNode(param) }
+    override Location getLocation() { result = parameter.getLocation() }
 
-    override Location getLocation() { result = param.getLocation() }
+    override string toString() { result = parameter.toString() }
 
-    override string toString() { result = param.toString() }
+    /** Gets the parameter in the AST that this node corresponds to. */
+    Param getParameter() { result = parameter }
   }
 
   final class ArgumentNode = NaNode;
@@ -93,6 +120,32 @@ module Node {
   final class CastNode = NaNode;
 }
 
+final class Node = Node::Node;
+
+/** Provides logic related to SSA. */
+module SsaFlow {
+  private module Impl = SsaImpl::DataFlowIntegration;
+
+  private Node::ParameterNode toParameterNode(Param p) { result = TParameterNode(p) }
+
+  /** Converts a control flow node into an SSA control flow node. */
+  Impl::Node asNode(Node n) {
+    n = TSsaNode(result)
+    or
+    result.(Impl::ExprNode).getExpr() = n.(Node::ExprNode).getCfgNode()
+    or
+    n = toParameterNode(result.(Impl::ParameterNode).getParameter())
+  }
+
+  predicate localFlowStep(SsaImpl::DefinitionExt def, Node nodeFrom, Node nodeTo, boolean isUseStep) {
+    Impl::localFlowStep(def, asNode(nodeFrom), asNode(nodeTo), isUseStep)
+  }
+
+  predicate localMustFlowStep(SsaImpl::DefinitionExt def, Node nodeFrom, Node nodeTo) {
+    Impl::localMustFlowStep(def, asNode(nodeFrom), asNode(nodeTo))
+  }
+}
+
 module RustDataFlow implements InputSig<Location> {
   /**
    * An element, viewed as a node in a data flow graph. Either an expression
@@ -122,10 +175,10 @@ module RustDataFlow implements InputSig<Location> {
 
   predicate nodeIsHidden(Node node) { none() }
 
-  class DataFlowExpr = Void;
+  class DataFlowExpr = ExprCfgNode;
 
   /** Gets the node corresponding to `e`. */
-  Node exprNode(DataFlowExpr e) { none() }
+  Node exprNode(DataFlowExpr e) { result.getCfgNode() = e }
 
   final class DataFlowCall extends TNormalCall {
     private CallExpr c;
@@ -191,7 +244,7 @@ module RustDataFlow implements InputSig<Location> {
    * Holds if there is a simple local flow step from `node1` to `node2`. These
    * are the value-preserving intra-callable flow steps.
    */
-  predicate simpleLocalFlowStep(Node node1, Node node2, string model) { none() }
+  predicate simpleLocalFlowStep(Node nodeFrom, Node nodeTo, string model) { none() }
 
   /**
    * Holds if data can flow from `node1` to `node2` through a non-local step
@@ -256,7 +309,9 @@ module RustDataFlow implements InputSig<Location> {
    * `node2` must be visited along a flow path, then any type known for `node2`
    * must also apply to `node1`.
    */
-  predicate localMustFlowStep(Node node1, Node node2) { none() }
+  predicate localMustFlowStep(Node node1, Node node2) {
+    SsaFlow::localMustFlowStep(_, node1, node2)
+  }
 
   class LambdaCallKind = Void;
 
@@ -267,7 +322,7 @@ module RustDataFlow implements InputSig<Location> {
   /** Holds if `call` is a lambda call of kind `kind` where `receiver` is the lambda expression. */
   predicate lambdaCall(DataFlowCall call, LambdaCallKind kind, Node receiver) { none() }
 
-  /** Extra data-flow steps needed for lambda flow analysis. */
+  /** Extra data flow steps needed for lambda flow analysis. */
   predicate additionalLambdaFlowStep(Node nodeFrom, Node nodeTo, boolean preservesValue) { none() }
 
   predicate knownSourceModel(Node source, string model) { none() }
@@ -286,8 +341,9 @@ cached
 private module Cached {
   cached
   newtype TNode =
-    TExprNode(CfgNode n, Expr e) { n.getAstNode() = e } or
-    TSourceParameterNode(Param param)
+    TExprNode(ExprCfgNode n) or
+    TParameterNode(Param p) or
+    TSsaNode(SsaImpl::DataFlowIntegration::SsaNode node)
 
   cached
   newtype TDataFlowCall = TNormalCall(CallExpr c)
@@ -302,7 +358,9 @@ private module Cached {
 
   /** This is the local flow predicate that is exposed. */
   cached
-  predicate localFlowStepImpl(Node::Node nodeFrom, Node::Node nodeTo) { none() }
+  predicate localFlowStepImpl(Node::Node nodeFrom, Node::Node nodeTo) {
+    SsaFlow::localFlowStep(_, nodeFrom, nodeTo, _)
+  }
 }
 
 import Cached

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

@@ -2,6 +2,7 @@ private import rust
 private import codeql.rust.controlflow.BasicBlocks as BasicBlocks
 private import BasicBlocks
 private import codeql.rust.controlflow.ControlFlowGraph as Cfg
+private import codeql.rust.controlflow.CfgNodes as CfgNodes
 private import Cfg
 private import codeql.rust.controlflow.internal.ControlFlowGraphImpl as ControlFlowGraphImpl
 private import codeql.ssa.Ssa as SsaImplCommon
@@ -395,6 +396,38 @@ private module Cached {
   Definition uncertainWriteDefinitionInput(UncertainWriteDefinition def) {
     Impl::uncertainWriteDefinitionInput(def, result)
   }
+
+  cached
+  module DataFlowIntegration {
+    import DataFlowIntegrationImpl
+
+    cached
+    predicate localFlowStep(DefinitionExt def, Node nodeFrom, Node nodeTo, boolean isUseStep) {
+      DataFlowIntegrationImpl::localFlowStep(def, nodeFrom, nodeTo, isUseStep)
+    }
+
+    cached
+    predicate localMustFlowStep(DefinitionExt def, Node nodeFrom, Node nodeTo) {
+      DataFlowIntegrationImpl::localMustFlowStep(def, nodeFrom, nodeTo)
+    }
+
+    signature predicate guardChecksSig(CfgNodes::AstCfgNode g, Cfg::CfgNode e, boolean branch);
+
+    cached // nothing is actually cached
+    module BarrierGuard<guardChecksSig/3 guardChecks> {
+      private predicate guardChecksAdjTypes(
+        DataFlowIntegrationInput::Guard g, DataFlowIntegrationInput::Expr e, boolean branch
+      ) {
+        guardChecks(g, e, branch)
+      }
+
+      private Node getABarrierNodeImpl() {
+        result = DataFlowIntegrationImpl::BarrierGuard<guardChecksAdjTypes/3>::getABarrierNode()
+      }
+
+      predicate getABarrierNode = getABarrierNodeImpl/0;
+    }
+  }
 }
 
 import Cached
@@ -426,3 +459,44 @@ class PhiReadNode extends DefinitionExt, Impl::PhiReadNode {
 
   override Location getLocation() { result = Impl::PhiReadNode.super.getLocation() }
 }
+
+private module DataFlowIntegrationInput implements Impl::DataFlowIntegrationInputSig {
+  class Expr extends CfgNodes::ExprCfgNode {
+    predicate hasCfgNode(SsaInput::BasicBlock bb, int i) { this = bb.getNode(i) }
+  }
+
+  Expr getARead(Definition def) { result = Cached::getARead(def) }
+
+  /** Holds if SSA definition `def` assigns `value` to the underlying variable. */
+  predicate ssaDefAssigns(WriteDefinition def, Expr value) { none() }
+
+  class Parameter = Param;
+
+  /** Holds if SSA definition `def` initializes parameter `p` at function entry. */
+  predicate ssaDefInitializesParam(WriteDefinition def, Parameter p) {
+    exists(BasicBlock bb, int i | bb.getNode(i).getAstNode() = p and def.definesAt(_, bb, i))
+  }
+
+  class Guard extends CfgNodes::AstCfgNode {
+    predicate hasCfgNode(SsaInput::BasicBlock bb, int i) { this = bb.getNode(i) }
+  }
+
+  /** Holds if the guard `guard` controls block `bb` upon evaluating to `branch`. */
+  predicate guardControlsBlock(Guard guard, SsaInput::BasicBlock bb, boolean branch) {
+    exists(ConditionBlock conditionBlock, ConditionalSuccessor s |
+      guard = conditionBlock.getLastNode() and
+      s.getValue() = branch and
+      conditionBlock.controls(bb, s)
+    )
+  }
+
+  /** Gets an immediate conditional successor of basic block `bb`, if any. */
+  SsaInput::BasicBlock getAConditionalBasicBlockSuccessor(SsaInput::BasicBlock bb, boolean branch) {
+    exists(Cfg::ConditionalSuccessor s |
+      result = bb.getASuccessor(s) and
+      s.getValue() = branch
+    )
+  }
+}
+
+private module DataFlowIntegrationImpl = Impl::DataFlowIntegration<DataFlowIntegrationInput>;