Rust: Handle x[y] expressions as *.index(y) calls in data flow

Author: hvitved

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

@@ -22,6 +22,8 @@ private module Input implements InputSig<Location, RustDataFlow> {
     n.(Node::PostUpdateNode).getPreUpdateNode().asExpr() = getPostUpdateReverseStep(_, _)
     or
     FlowSummaryImpl::Private::Steps::sourceLocalStep(_, n, _)
+    or
+    assignmentPostUpdateNode(_, n)
   }
 
   predicate missingLocationExclude(RustDataFlow::Node n) { not exists(n.asExpr().getLocation()) }

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

@@ -141,8 +141,6 @@ final class ArgumentPosition extends ParameterPosition {
  * Holds if `arg` is an argument of `call` at the position `pos`.
  */
 predicate isArgumentForCall(Expr arg, Call call, ArgumentPosition pos) {
-  // TODO: Handle index expressions as calls in data flow.
-  not call instanceof IndexExpr and
   arg = pos.getArgument(call)
 }
 
@@ -316,6 +314,20 @@ private module Aliases {
   class LambdaCallKindAlias = LambdaCallKind;
 }
 
+/**
+ * Holds if `n` is the post-update node for the left-hand side of `assignment`.
+ *
+ * For such assignments, we add a local flow step from the right-hand side to
+ * `n`, which means that for assignments like `*x = y`, we get an induced
+ * reverse `*`-store step into `[post] x`.
+ *
+ * This also applies to field assignments, `x.field = y`, and index assignments,
+ * `x[index] = y`.
+ */
+predicate assignmentPostUpdateNode(AssignmentExpr assignment, PostUpdateNode n) {
+  n.getPreUpdateNode().asExpr() = assignment.getLhs()
+}
+
 module RustDataFlow implements InputSig<Location> {
   private import Aliases
   private import codeql.rust.dataflow.DataFlow
@@ -363,7 +375,9 @@ module RustDataFlow implements InputSig<Location> {
     node instanceof ClosureParameterNode or
     node instanceof DerefBorrowNode or
     node instanceof DerefOutNode or
+    node instanceof IndexOutNode or
     node.asExpr() instanceof ParenExpr or
+    assignmentPostUpdateNode(_, node) or
     nodeIsHidden(node.(PostUpdateNode).getPreUpdateNode())
   }
 
@@ -386,6 +400,8 @@ module RustDataFlow implements InputSig<Location> {
       node.asPat() = match.getAnArm().getPat()
     )
     or
+    node.asExpr() = any(AssignmentExpr ae).getRhs()
+    or
     FlowSummaryImpl::Private::Steps::sourceLocalStep(_, node, _)
     or
     FlowSummaryImpl::Private::Steps::sinkLocalStep(node, _, _)
@@ -488,6 +504,12 @@ module RustDataFlow implements InputSig<Location> {
           .(Node::FlowSummaryNode)
           .getSummaryNode(), TOptionalBarrier(_), nodeTo.(Node::FlowSummaryNode).getSummaryNode()) and
     model = ""
+    or
+    exists(AssignmentExpr assignment |
+      nodeFrom.asExpr() = assignment.getRhs() and
+      assignmentPostUpdateNode(assignment, nodeTo) and
+      model = ""
+    )
   }
 
   /**
@@ -559,12 +581,6 @@ module RustDataFlow implements InputSig<Location> {
       access = c.(FieldContent).getAnAccess()
     )
     or
-    exists(IndexExpr arr |
-      c instanceof ElementContent and
-      node1.asExpr() = arr.getBase() and
-      node2.asExpr() = arr
-    )
-    or
     exists(ForExpr for |
       c instanceof ElementContent and
       node1.asExpr() = for.getIterable() and
@@ -590,6 +606,12 @@ module RustDataFlow implements InputSig<Location> {
       node2.asExpr() = deref
     )
     or
+    exists(IndexExpr index |
+      c instanceof ReferenceContent and
+      node1.(IndexOutNode).getIndexExpr() = index and
+      node2.asExpr() = index
+    )
+    or
     // Read from function return
     exists(DataFlowCall call |
       lambdaCall(call, _, node1) and
@@ -652,9 +674,8 @@ module RustDataFlow implements InputSig<Location> {
 
   pragma[nomagic]
   private predicate referenceAssignment(Node node1, Node node2, ReferenceContent c) {
-    exists(AssignmentExpr assignment, PrefixExpr deref |
+    exists(AssignmentExpr assignment, DerefExpr deref |
       assignment.getLhs() = deref and
-      deref.getOperatorName() = "*" and
       node1.asExpr() = assignment.getRhs() and
       node2.asExpr() = deref.getExpr() and
       exists(c)
@@ -699,15 +720,17 @@ module RustDataFlow implements InputSig<Location> {
       node2.(NameNode).getName() = p.getName()
     )
     or
-    fieldAssignment(node1, node2.(PostUpdateNode).getPreUpdateNode(), c)
-    or
-    referenceAssignment(node1, node2.(PostUpdateNode).getPreUpdateNode(), c)
-    or
-    exists(AssignmentExpr assignment, IndexExpr index |
+    // Since generalized reverse flow is not yet supported, we add a special case
+    // for assignments into arrays and slices
+    exists(IndexExpr index |
+      node1.(PostUpdateNode).getPreUpdateNode().asExpr() = index and
+      node2.(PostUpdateNode).getPreUpdateNode().asExpr() = index.getBase() and
       c instanceof ElementContent and
-      assignment.getLhs() = index and
-      node1.asExpr() = assignment.getRhs() and
-      node2.(PostUpdateNode).getPreUpdateNode().asExpr() = index.getBase()
+      index.getResolvedTarget().getCanonicalPath() =
+        [
+          "<[;] as core::ops::index::IndexMut>::index_mut",
+          "<[] as core::ops::index::IndexMut>::index_mut"
+        ]
     )
     or
     referenceExprToExpr(node1, node2, c)
@@ -989,9 +1012,7 @@ private module Cached {
   newtype TDataFlowCall =
     TCall(Call call) {
       Stages::DataFlowStage::ref() and
-      call.hasEnclosingCfgScope() and
-      // TODO: Handle index expressions as calls in data flow.
-      not call instanceof IndexExpr
+      call.hasEnclosingCfgScope()
     } or
     TSummaryCall(
       FlowSummaryImpl::Public::SummarizedCallable c, FlowSummaryImpl::Private::SummaryNode receiver

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

@@ -350,7 +350,8 @@ final private class ExprOutNode extends ExprNode, OutNode {
   ExprOutNode() {
     exists(Call call |
       call = this.asExpr() and
-      not call instanceof DerefExpr // Handled by `DerefOutNode`
+      not call instanceof DerefExpr and // Handled by `DerefOutNode`
+      not call instanceof IndexExpr // Handled by `IndexOutNode`
     )
   }
 
@@ -387,6 +388,32 @@ class DerefOutNode extends OutNode, TDerefOutNode {
   override string toString() { result = de.toString() + " [pre-dereferenced]" }
 }
 
+/**
+ * A node that represents the value of a `x[y]` expression _before_ implicit
+ * dereferencing:
+ *
+ * `x[y]` equivalent to `*x.index(y)`, and this node represents the
+ * `x.index(y)` part.
+ */
+class IndexOutNode extends OutNode, TIndexOutNode {
+  IndexExpr ie;
+
+  IndexOutNode() { this = TIndexOutNode(ie, false) }
+
+  IndexExpr getIndexExpr() { result = ie }
+
+  override CfgScope getCfgScope() { result = ie.getEnclosingCfgScope() }
+
+  override DataFlowCall getCall(ReturnKind kind) {
+    result.asCall() = ie and
+    kind = TNormalReturnKind()
+  }
+
+  override Location getLocation() { result = ie.getLocation() }
+
+  override string toString() { result = ie.toString() + " [pre-dereferenced]" }
+}
+
 final class SummaryOutNode extends FlowSummaryNode, OutNode {
   private DataFlowCall call;
   private ReturnKind kind_;
@@ -476,6 +503,18 @@ class DerefOutPostUpdateNode extends PostUpdateNode, TDerefOutNode {
   override Location getLocation() { result = de.getLocation() }
 }
 
+class IndexOutPostUpdateNode extends PostUpdateNode, TIndexOutNode {
+  IndexExpr ie;
+
+  IndexOutPostUpdateNode() { this = TIndexOutNode(ie, true) }
+
+  override IndexOutNode getPreUpdateNode() { result = TIndexOutNode(ie, false) }
+
+  override CfgScope getCfgScope() { result = ie.getEnclosingCfgScope() }
+
+  override Location getLocation() { result = ie.getLocation() }
+}
+
 final class SummaryPostUpdateNode extends FlowSummaryNode, PostUpdateNode {
   private FlowSummaryNode pre;
 
@@ -526,7 +565,9 @@ newtype TNode =
       or
       e =
         [
-          any(IndexExpr i).getBase(), //
+          any(DerefExpr de), //
+          any(IndexExpr ie), //
+          any(FieldExpr fe), //
           any(FieldExpr access).getContainer(), //
           any(TryExpr try).getExpr(), //
           any(AwaitExpr a).getExpr(), //
@@ -542,6 +583,7 @@ newtype TNode =
     borrow = true
   } or
   TDerefOutNode(DerefExpr de, Boolean isPost) or
+  TIndexOutNode(IndexExpr ie, Boolean isPost) or
   TSsaNode(SsaImpl::DataFlowIntegration::SsaNode node) or
   TFlowSummaryNode(FlowSummaryImpl::Private::SummaryNode sn) {
     forall(AstNode n | n = sn.getSinkElement() or n = sn.getSourceElement() |

rust/ql/lib/codeql/rust/elements/internal/VariableImpl.qll

@@ -676,12 +676,12 @@ module Impl {
     predicate isCapture() { this.getEnclosingCfgScope() != v.getEnclosingCfgScope() }
   }
 
-  /** Holds if `e` occurs in the LHS of an assignment or compound assignment. */
-  private predicate assignmentExprDescendant(AssignmentExpr ae, Expr e) {
-    e = ae.getLhs()
+  /** Holds if `e` occurs in the LHS of an assignment operation. */
+  predicate assignmentOperationDescendant(AssignmentOperation ao, Expr e) {
+    e = ao.getLhs()
     or
     exists(Expr mid |
-      assignmentExprDescendant(ae, mid) and
+      assignmentOperationDescendant(ao, mid) and
       getImmediateParentAdj(e) = mid and
       not mid instanceof DerefExpr and
       not mid instanceof FieldExpr and
@@ -696,7 +696,7 @@ module Impl {
     cached
     VariableWriteAccess() {
       Cached::ref() and
-      assignmentExprDescendant(ae, this)
+      assignmentOperationDescendant(ae, this)
     }
 
     /** Gets the assignment expression that has this write access in the left-hand side. */

rust/ql/lib/codeql/rust/frameworks/stdlib/core.model.yml

@@ -6,6 +6,9 @@ extensions:
       # Builtin deref
       - ["<& as core::ops::deref::Deref>::deref", "Argument[self].Reference", "ReturnValue", "value", "manual"]
       - ["<&mut as core::ops::deref::Deref>::deref", "Argument[self].Reference", "ReturnValue", "value", "manual"]
+      # Index
+      - ["<_ as core::ops::index::Index>::index", "Argument[self].Reference.Element", "ReturnValue.Reference", "value", "manual"]
+      - ["<_ as core::ops::index::IndexMut>::index_mut", "Argument[self].Reference.Element", "ReturnValue.Reference", "value", "manual"]
       # Arithmetic
       - ["<_ as core::ops::arith::Add>::add", "Argument[self]", "ReturnValue", "taint", "manual"]
       - ["<_ as core::ops::arith::Add>::add", "Argument[0]", "ReturnValue", "taint", "manual"]

rust/ql/lib/codeql/rust/internal/TypeInference.qll

@@ -9,6 +9,7 @@ private import TypeMention
 private import typeinference.FunctionType
 private import typeinference.FunctionOverloading as FunctionOverloading
 private import typeinference.BlanketImplementation as BlanketImplementation
+private import codeql.rust.elements.internal.VariableImpl::Impl as VariableImpl
 private import codeql.rust.internal.CachedStages
 private import codeql.typeinference.internal.TypeInference
 private import codeql.rust.frameworks.stdlib.Stdlib
@@ -1636,9 +1637,14 @@ private module MethodResolution {
   }
 
   private class MethodCallIndexExpr extends MethodCall instanceof IndexExpr {
+    private predicate isInMutableContext() {
+      // todo: does not handle all cases yet
+      VariableImpl::assignmentOperationDescendant(_, this)
+    }
+
     pragma[nomagic]
     override predicate hasNameAndArity(string name, int arity) {
-      name = "index" and
+      (if this.isInMutableContext() then name = "index_mut" else name = "index") and
       arity = 1
     }
 
@@ -1652,7 +1658,11 @@ private module MethodResolution {
 
     override predicate supportsAutoDerefAndBorrow() { any() }
 
-    override Trait getTrait() { result.getCanonicalPath() = "core::ops::index::Index" }
+    override Trait getTrait() {
+      if this.isInMutableContext()
+      then result.getCanonicalPath() = "core::ops::index::IndexMut"
+      else result.getCanonicalPath() = "core::ops::index::Index"
+    }
   }
 
   private class MethodCallCallExpr extends MethodCall instanceof CallExpr {