Rust: Model more data flow constructs as calls using MaD

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

@@ -90,8 +90,6 @@
  * Holds if `arg` is an argument of `call` at the position `pos`.
  */
 predicate isArgumentForCall(Expr arg, Call call, RustDataFlow::ArgumentPosition pos) {
-  // TODO: Handle index expressions as calls in data flow.
-  not call instanceof IndexExpr and
   arg = pos.getArgument(call)
 }
 
@@ -206,8 +204,11 @@
     )
     or
     // An edge from a pattern/expression to its corresponding SSA definition.
-    nodeFrom.(AstNodeNode).getAstNode() =
-      nodeTo.(SsaNode).asDefinition().(Ssa::WriteDefinition).getWriteAccess()
+    exists(AstNode n |
+      n = nodeTo.(SsaNode).asDefinition().(Ssa::WriteDefinition).getWriteAccess() and
+      n = nodeFrom.(AstNodeNode).getAstNode() and
+      not n = any(CompoundAssignmentExpr cae).getLhs()
+    )
     or
     nodeFrom.(SourceParameterNode).getParameter().(Param).getPat() = nodeTo.asPat()
     or
@@ -261,6 +262,30 @@
   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
@@ -292,19 +317,19 @@
    * arguments positions, so we use the same underlying type for both.
    */
   final class ParameterPosition extends TParameterPosition {
     /** Gets the underlying integer position, if any. */
     int getPosition() { this = TPositionalParameterPosition(result) }
 
     predicate hasPosition() { exists(this.getPosition()) }
 
     /** Holds if this position represents the `self` position. */
     predicate isSelf() { this = TSelfParameterPosition() }
 
     /**
      * Holds if this position represents a reference to a closure itself. Only
      * used for tracking flow through captured variables.
      */
     predicate isClosureSelf() { this = TClosureSelfParameterPosition() }
 
     /** Gets a textual representation of this position. */
     string toString() {
@@ -360,6 +385,7 @@
     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())
   }
@@ -399,13 +425,23 @@
 
   final class ReturnKind = ReturnKindAlias;
 
+  private Function getStaticTargetExt(Call c) {
+    result = c.getStaticTarget()
+    or
+    not exists(c.getStaticTarget()) and
+    exists(TraitItemNode t, string methodName |
+      c.(Operation).isOverloaded(t, methodName, _) and
+      result = t.getAssocItem(methodName)
+    )
+  }
+
   /** Gets a viable implementation of the target of the given `Call`. */
   DataFlowCallable viableCallable(DataFlowCall call) {
     exists(Call c | c = call.asCall() |
       result.asCfgScope() = c.getARuntimeTarget()
       or
       exists(SummarizedCallable sc, Function staticTarget |
-        staticTarget = c.getStaticTarget() and
+        staticTarget = getStaticTargetExt(c) and
         sc = result.asSummarizedCallable()
       |
         sc = staticTarget
@@ -552,12 +588,6 @@
       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
@@ -583,6 +613,12 @@
       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
@@ -602,8 +638,18 @@
     implicitDeref(node1, node2, c)
     or
     // A read step dual to the store step for implicit borrows.
-    implicitBorrow(node2.(PostUpdateNode).getPreUpdateNode(),
-      node1.(PostUpdateNode).getPreUpdateNode(), c)
+    exists(Node n | implicitBorrow(n, node1.(PostUpdateNode).getPreUpdateNode(), c) |
+      node2.(PostUpdateNode).getPreUpdateNode() = n
+      or
+      // For compound assignments into variables like `x += y`, we do not want flow into
+      // `[post] x`, as that would create spurious flow when `x` is a parameter. Instead,
+      // we add the step directly into the SSA definition for `x` after the update.
+      exists(CompoundAssignmentExpr cae, Expr lhs |
+        lhs = cae.getLhs() and
+        lhs = node2.(SsaNode).asDefinition().(Ssa::WriteDefinition).getWriteAccess() and
+        n = TExprNode(lhs)
+      )
+    )
     or
     VariableCapture::readStep(node1, c, node2)
   }
@@ -644,13 +690,27 @@
   }
 
   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
+        )
     )
   }
 
@@ -694,14 +754,14 @@
     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
@@ -738,7 +798,7 @@
   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
@@ -982,9 +1042,7 @@
   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

@@ -20,18 +20,6 @@ module RustTaintTracking implements InputSig<Location, RustDataFlow> {
     Stages::DataFlowStage::ref() and
     model = "" and
     (
-      exists(BinaryExpr binary |
-        binary.getOperatorName() = ["+", "-", "*", "/", "%", "&", "|", "^", "<<", ">>"] and
-        pred.asExpr() = [binary.getLhs(), binary.getRhs()] and
-        succ.asExpr() = binary
-      )
-      or
-      exists(PrefixExpr prefix |
-        prefix.getOperatorName() = ["-", "!"] and
-        pred.asExpr() = prefix.getExpr() and
-        succ.asExpr() = prefix
-      )
-      or
       pred.asExpr() = succ.asExpr().(CastExpr).getExpr()
       or
       exists(IndexExpr index |
@@ -65,6 +53,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/OperationImpl.qll

@@ -131,6 +131,7 @@ module Impl {
      * Holds if this operation is overloaded to the method `methodName` of the
      * trait `trait`.
      */
+    pragma[nomagic]
     predicate isOverloaded(Trait trait, string methodName, int borrows) {
       isOverloaded(this.getOperatorName(), this.getNumberOfOperands(), trait.getCanonicalPath(),
         methodName, borrows)

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. */