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

Author: hvitved

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,30 @@ private module Aliases {
   class LambdaCallKindAlias = LambdaCallKind;
 }
 
+/**
+ * Index assignments like `a[i] = rhs` are treated as `*a.index_mut(i) = rhs`,
+ * so they should in principle be handled by `referenceAssignment`.
+ *
+ * However, this would require support for [generalized reverse flow][1], which
+ * is not yet implemented, so instead we simulate reverse flow where it would
+ * have applied via the model for `<_ as core::ops::index::IndexMut>::index_mut`.
+ *
+ * The same is the case for compound assignments like `a[i] += rhs`, which are
+ * treated as `(*a.index_mut(i)).add_assign(rhs)`.
+ *
+ * [1]: https://github.com/github/codeql/pull/18109
+ */
+predicate indexAssignment(
+  AssignmentOperation assignment, IndexExpr index, Node rhs, PostUpdateNode base, Content c
+) {
+  assignment.getLhs() = index and
+  rhs.asExpr() = assignment.getRhs() and
+  base.getPreUpdateNode().asExpr() = index.getBase() and
+  c instanceof ElementContent and
+  // simulate that the flow summary applies
+  not index.getResolvedTarget().fromSource()
+}
+
 module RustDataFlow implements InputSig<Location> {
   private import Aliases
   private import codeql.rust.dataflow.DataFlow
@@ -363,6 +385,7 @@ 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
     nodeIsHidden(node.(PostUpdateNode).getPreUpdateNode())
   }
@@ -559,12 +582,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 +607,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
@@ -651,13 +674,27 @@ module RustDataFlow implements InputSig<Location> {
   }
 
   pragma[nomagic]
-  private predicate referenceAssignment(Node node1, Node node2, ReferenceContent c) {
-    exists(AssignmentExpr assignment, PrefixExpr deref |
-      assignment.getLhs() = deref and
-      deref.getOperatorName() = "*" and
+  private predicate referenceAssignment(
+    Node node1, Node node2, Expr e, boolean clears, ReferenceContent c
+  ) {
+    exists(AssignmentExpr assignment, Expr lhs |
+      assignment.getLhs() = lhs and
       node1.asExpr() = assignment.getRhs() and
-      node2.asExpr() = deref.getExpr() and
       exists(c)
+    |
+      lhs =
+        any(DerefExpr de |
+          de = node2.(DerefOutNode).getDerefExpr() and
+          e = de.getExpr()
+        ) and
+      clears = true
+      or
+      lhs =
+        any(IndexExpr ie |
+          ie = node2.(IndexOutNode).getIndexExpr() and
+          e = ie.getBase() and
+          clears = false
+        )
     )
   }
 
@@ -701,14 +738,14 @@ module RustDataFlow implements InputSig<Location> {
     or
     fieldAssignment(node1, node2.(PostUpdateNode).getPreUpdateNode(), c)
     or
-    referenceAssignment(node1, node2.(PostUpdateNode).getPreUpdateNode(), c)
+    referenceAssignment(node1, node2.(PostUpdateNode).getPreUpdateNode(), _, _, c)
     or
-    exists(AssignmentExpr assignment, IndexExpr index |
-      c instanceof ElementContent and
-      assignment.getLhs() = index and
-      node1.asExpr() = assignment.getRhs() and
-      node2.(PostUpdateNode).getPreUpdateNode().asExpr() = index.getBase()
-    )
+    indexAssignment(any(AssignmentExpr ae), _, node1, node2, c)
+    or
+    // Compund assignment like `a[i] += rhs` are modeled as a store step from `rhs`
+    // to `[post] a[i]`, followed by a taint step into `[post] a`.
+    indexAssignment(any(CompoundAssignmentExpr cae),
+      node2.(PostUpdateNode).getPreUpdateNode().asExpr(), node1, _, c)
     or
     referenceExprToExpr(node1, node2, c)
     or
@@ -745,7 +782,7 @@ module RustDataFlow implements InputSig<Location> {
   predicate clearsContent(Node n, ContentSet cs) {
     fieldAssignment(_, n, cs.(SingletonContentSet).getContent())
     or
-    referenceAssignment(_, n, cs.(SingletonContentSet).getContent())
+    referenceAssignment(_, _, n.asExpr(), true, cs.(SingletonContentSet).getContent())
     or
     FlowSummaryImpl::Private::Steps::summaryClearsContent(n.(FlowSummaryNode).getSummaryNode(), cs)
     or
@@ -989,9 +1026,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;
 
@@ -514,7 +553,8 @@ newtype TNode =
   TExprPostUpdateNode(Expr e) {
     e.hasEnclosingCfgScope() and
     (
-      isArgumentForCall(e, _, _)
+      isArgumentForCall(e, _, _) and
+      not (e = any(CompoundAssignmentExpr cae).getLhs() and e instanceof VariableAccess)
       or
       lambdaCallExpr(_, _, e)
       or
@@ -526,7 +566,6 @@ newtype TNode =
       or
       e =
         [
-          any(IndexExpr i).getBase(), //
           any(FieldExpr access).getContainer(), //
           any(TryExpr try).getExpr(), //
           any(AwaitExpr a).getExpr(), //
@@ -542,6 +581,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/dataflow/internal/TaintTrackingImpl.qll

@@ -65,6 +65,9 @@ module RustTaintTracking implements InputSig<Location, RustDataFlow> {
       or
       succ.(Node::PostUpdateNode).getPreUpdateNode().asExpr() =
         getPostUpdateReverseStep(pred.(Node::PostUpdateNode).getPreUpdateNode().asExpr(), false)
+      or
+      indexAssignment(any(CompoundAssignmentExpr cae),
+        pred.(Node::PostUpdateNode).getPreUpdateNode().asExpr(), _, succ, _)
     )
     or
     FlowSummaryImpl::Private::Steps::summaryLocalStep(pred.(Node::FlowSummaryNode).getSummaryNode(),

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
@@ -665,7 +666,7 @@ private predicate typeEquality(AstNode n1, TypePath prefix1, AstNode n2, TypePat
 
 /**
  * Holds if `child` is a child of `parent`, and the Rust compiler applies [least
- * upper bound (LUB) coercion](1) to infer the type of `parent` from the type of
+ * upper bound (LUB) coercion][1] to infer the type of `parent` from the type of
  * `child`.
  *
  * In this case, we want type information to only flow from `child` to `parent`,
@@ -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 {