Rust: Data flow through overloaded operators

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

@@ -4,6 +4,7 @@
  */
 
 private import rust
+private import codeql.rust.elements.Call
 private import ControlFlowGraph
 private import internal.ControlFlowGraphImpl as CfgImpl
 private import internal.CfgNodes
@@ -162,6 +163,30 @@ final class CallExprBaseCfgNode extends Nodes::CallExprBaseCfgNode {
  */
 final class MethodCallExprCfgNode extends CallExprBaseCfgNode, Nodes::MethodCallExprCfgNode { }
 
+/**
+ * A CFG node that calls a function.
+ *
+ * This class abstract over the different ways in which a function can be called in Rust.
+ */
+final class CallCfgNode extends ExprCfgNode {
+  private Call node;
+
+  CallCfgNode() { node = this.getAstNode() }
+
+  /** Gets the underlying `Call`. */
+  Call getCall() { result = node }
+
+  /** Gets the receiver of this call if it is a method call. */
+  ExprCfgNode getReceiver() {
+    any(ChildMapping mapping).hasCfgChild(node, node.getReceiver(), this, result)
+  }
+
+  /** Gets the `i`th argument of this call, if any. */
+  ExprCfgNode getArgument(int i) {
+    any(ChildMapping mapping).hasCfgChild(node, node.getArgument(i), this, result)
+  }
+}
+
 /**
  * A function call expression. For example:
  * ```rust

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

@@ -1,5 +1,8 @@
 import codeql.rust.dataflow.DataFlow::DataFlow as DataFlow
 private import rust
+private import codeql.rust.controlflow.ControlFlowGraph
+private import codeql.rust.controlflow.internal.Splitting
+private import codeql.rust.controlflow.CfgNodes as CfgNodes
 private import codeql.rust.dataflow.internal.DataFlowImpl
 private import codeql.rust.dataflow.internal.FlowSummaryImpl as FlowSummaryImpl
 private import codeql.rust.dataflow.internal.Node as Node
@@ -26,6 +29,17 @@ private module Input implements InputSig<Location, RustDataFlow> {
   }
 
   predicate missingLocationExclude(RustDataFlow::Node n) { not exists(n.asExpr().getLocation()) }
+
+  predicate multipleArgumentCallExclude(Node::ArgumentNode arg, DataFlowCall call) {
+    // An argument such as `x` in `if !x { ... }` has two successors (and hence
+    // two calls); one for each Boolean outcome of `x`.
+    exists(CfgNodes::ExprCfgNode n |
+      arg.isArgumentOf(call, _) and
+      n = call.asCallCfgNode() and
+      arg.asExpr().getASuccessor(any(ConditionalSuccessor c)).getASuccessor*() = n and
+      n.getASplit() instanceof ConditionalCompletionSplitting::ConditionalCompletionSplit
+    )
+  }
 }
 
 import MakeConsistency<Location, RustDataFlow, RustTaintTracking, Input>

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

@@ -8,6 +8,7 @@ private import codeql.util.Boolean
 private import codeql.dataflow.DataFlow
 private import codeql.dataflow.internal.DataFlowImpl
 private import rust
+private import codeql.rust.elements.Call
 private import SsaImpl as SsaImpl
 private import codeql.rust.controlflow.internal.Scope as Scope
 private import codeql.rust.internal.PathResolution
@@ -55,11 +56,7 @@ final class DataFlowCallable extends TDataFlowCallable {
 
 final class DataFlowCall extends TDataFlowCall {
   /** Gets the underlying call in the CFG, if any. */
-  CallExprCfgNode asCallExprCfgNode() { result = this.asCallBaseExprCfgNode() }
-
-  MethodCallExprCfgNode asMethodCallExprCfgNode() { result = this.asCallBaseExprCfgNode() }
-
-  CallExprBaseCfgNode asCallBaseExprCfgNode() { this = TCall(result) }
+  CallCfgNode asCallCfgNode() { this = TCall(result) }
 
   predicate isSummaryCall(
     FlowSummaryImpl::Public::SummarizedCallable c, FlowSummaryImpl::Private::SummaryNode receiver
@@ -68,7 +65,7 @@ final class DataFlowCall extends TDataFlowCall {
   }
 
   DataFlowCallable getEnclosingCallable() {
-    result = TCfgScope(this.asCallBaseExprCfgNode().getExpr().getEnclosingCfgScope())
+    result = TCfgScope(this.asCallCfgNode().getExpr().getEnclosingCfgScope())
     or
     exists(FlowSummaryImpl::Public::SummarizedCallable c |
       this.isSummaryCall(c, _) and
@@ -77,7 +74,7 @@ final class DataFlowCall extends TDataFlowCall {
   }
 
   string toString() {
-    result = this.asCallBaseExprCfgNode().toString()
+    result = this.asCallCfgNode().toString()
     or
     exists(
       FlowSummaryImpl::Public::SummarizedCallable c, FlowSummaryImpl::Private::SummaryNode receiver
@@ -87,7 +84,7 @@ final class DataFlowCall extends TDataFlowCall {
     )
   }
 
-  Location getLocation() { result = this.asCallBaseExprCfgNode().getLocation() }
+  Location getLocation() { result = this.asCallCfgNode().getLocation() }
 }
 
 /**
@@ -135,38 +132,23 @@ final class ParameterPosition extends TParameterPosition {
  */
 final class ArgumentPosition extends ParameterPosition {
   /** Gets the argument of `call` at this position, if any. */
-  Expr getArgument(CallExprBase call) {
-    result = call.getArgList().getArg(this.getPosition())
+  Expr getArgument(Call call) {
+    result = call.getArgument(this.getPosition())
     or
-    this.isSelf() and
-    result = call.(MethodCallExpr).getReceiver()
+    result = call.getReceiver() and this.isSelf()
   }
 }
 
-/** Holds if `call` invokes a qualified path that resolves to a method. */
-private predicate callToMethod(CallExpr call) {
-  exists(Path path |
-    path = call.getFunction().(PathExpr).getPath() and
-    path.hasQualifier() and
-    resolvePath(path).(Function).getParamList().hasSelfParam()
-  )
-}
-
 /**
  * Holds if `arg` is an argument of `call` at the position `pos`.
  *
  * Note that this does not hold for the receiever expression of a method call
  * as the synthetic `ReceiverNode` is the argument for the `self` parameter.
  */
-predicate isArgumentForCall(ExprCfgNode arg, CallExprBaseCfgNode call, ParameterPosition pos) {
-  if callToMethod(call.(CallExprCfgNode).getCallExpr())
-  then
-    // The first argument is for the `self` parameter
-    arg = call.getArgument(0) and pos.isSelf()
-    or
-    // Succeeding arguments are shifted left
-    arg = call.getArgument(pos.getPosition() + 1)
-  else arg = call.getArgument(pos.getPosition())
+predicate isArgumentForCall(ExprCfgNode arg, CallCfgNode call, ParameterPosition pos) {
+  call.getArgument(pos.getPosition()) = arg
+  or
+  call.getReceiver() = arg and pos.isSelf() and not call.getCall().receiverImplicitlyBorrowed()
 }
 
 /** Provides logic related to SSA. */
@@ -419,9 +401,9 @@ module RustDataFlow implements InputSig<Location> {
 
   /** Gets a viable implementation of the target of the given `Call`. */
   DataFlowCallable viableCallable(DataFlowCall call) {
-    result.asCfgScope() = call.asCallBaseExprCfgNode().getCallExprBase().getStaticTarget()
+    result.asCfgScope() = call.asCallCfgNode().getCall().getStaticTarget()
     or
-    result.asLibraryCallable().getACall() = call.asCallBaseExprCfgNode().getCallExprBase()
+    result.asLibraryCallable().getACall() = call.asCallCfgNode().getCall()
   }
 
   /**
@@ -812,7 +794,7 @@ module RustDataFlow implements InputSig<Location> {
    */
   predicate lambdaCall(DataFlowCall call, LambdaCallKind kind, Node receiver) {
     (
-      receiver.asExpr() = call.asCallExprCfgNode().getFunction() and
+      receiver.asExpr() = call.asCallCfgNode().(CallExprCfgNode).getFunction() and
       // All calls to complex expressions and local variable accesses are lambda call.
       exists(Expr f | f = receiver.asExpr().getExpr() |
         f instanceof PathExpr implies f = any(Variable v).getAnAccess()
@@ -976,7 +958,7 @@ private module Cached {
 
   cached
   newtype TDataFlowCall =
-    TCall(CallExprBaseCfgNode c) { Stages::DataFlowStage::ref() } or
+    TCall(CallCfgNode c) { Stages::DataFlowStage::ref() } or
     TSummaryCall(
       FlowSummaryImpl::Public::SummarizedCallable c, FlowSummaryImpl::Private::SummaryNode receiver
     ) {

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

@@ -153,7 +153,7 @@ private import Make<Location, RustDataFlow, Input> as Impl
 
 private module StepsInput implements Impl::Private::StepsInputSig {
   DataFlowCall getACall(Public::SummarizedCallable sc) {
-    result.asCallBaseExprCfgNode().getCallExprBase() = sc.(LibraryCallable).getACall()
+    result.asCallCfgNode().getCall() = sc.(LibraryCallable).getACall()
   }
 
   RustDataFlow::Node getSourceNode(Input::SourceBase source, Impl::Private::SummaryComponent sc) {

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

@@ -224,13 +224,13 @@ abstract class ArgumentNode extends Node {
 }
 
 final class ExprArgumentNode extends ArgumentNode, ExprNode {
-  private CallExprBaseCfgNode call_;
+  private CallCfgNode call_;
   private RustDataFlow::ArgumentPosition pos_;
 
   ExprArgumentNode() { isArgumentForCall(n, call_, pos_) }
 
   override predicate isArgumentOf(DataFlowCall call, RustDataFlow::ArgumentPosition pos) {
-    call.asCallBaseExprCfgNode() = call_ and pos = pos_
+    call.asCallCfgNode() = call_ and pos = pos_
   }
 }
 
@@ -239,7 +239,7 @@ final class ExprArgumentNode extends ArgumentNode, ExprNode {
  * has taken place.
  */
 final class ReceiverNode extends ArgumentNode, TReceiverNode {
-  private MethodCallExprCfgNode n;
+  private CallCfgNode n;
 
   ReceiverNode() { this = TReceiverNode(n, false) }
 
@@ -248,7 +248,7 @@ final class ReceiverNode extends ArgumentNode, TReceiverNode {
   MethodCallExprCfgNode getMethodCall() { result = n }
 
   override predicate isArgumentOf(DataFlowCall call, RustDataFlow::ArgumentPosition pos) {
-    call.asMethodCallExprCfgNode() = n and pos = TSelfParameterPosition()
+    call.asCallCfgNode() = n and pos = TSelfParameterPosition()
   }
 
   override CfgScope getCfgScope() { result = n.getAstNode().getEnclosingCfgScope() }
@@ -281,7 +281,7 @@ final class ClosureArgumentNode extends ArgumentNode, ExprNode {
   ClosureArgumentNode() { lambdaCallExpr(call_, _, this.asExpr()) }
 
   override predicate isArgumentOf(DataFlowCall call, RustDataFlow::ArgumentPosition pos) {
-    call.asCallExprCfgNode() = call_ and
+    call.asCallCfgNode() = call_ and
     pos.isClosureSelf()
   }
 }
@@ -330,11 +330,11 @@ abstract class OutNode extends Node {
 }
 
 final private class ExprOutNode extends ExprNode, OutNode {
-  ExprOutNode() { this.asExpr() instanceof CallExprBaseCfgNode }
+  ExprOutNode() { this.asExpr() instanceof CallCfgNode }
 
   /** Gets the underlying call CFG node that includes this out node. */
   override DataFlowCall getCall(ReturnKind kind) {
-    result.asCallBaseExprCfgNode() = this.getCfgNode() and
+    result.asCallCfgNode() = this.getCfgNode() and
     kind = TNormalReturnKind()
   }
 }
@@ -404,7 +404,7 @@ final class ExprPostUpdateNode extends PostUpdateNode, TExprPostUpdateNode {
 }
 
 final class ReceiverPostUpdateNode extends PostUpdateNode, TReceiverNode {
-  private MethodCallExprCfgNode n;
+  private CallCfgNode n;
 
   ReceiverPostUpdateNode() { this = TReceiverNode(n, true) }
 
@@ -467,11 +467,12 @@ newtype TNode =
         any(FieldExprCfgNode access).getContainer(), //
         any(TryExprCfgNode try).getExpr(), //
         any(PrefixExprCfgNode pe | pe.getOperatorName() = "*").getExpr(), //
-        any(AwaitExprCfgNode a).getExpr(), any(MethodCallExprCfgNode mc).getReceiver(), //
+        any(AwaitExprCfgNode a).getExpr(), //
+        any(MethodCallExprCfgNode mc).getReceiver(), //
         getPostUpdateReverseStep(any(PostUpdateNode n).getPreUpdateNode().asExpr(), _)
       ]
   } or
-  TReceiverNode(MethodCallExprCfgNode mc, Boolean isPost) or
+  TReceiverNode(CallCfgNode mc, Boolean isPost) { mc.getCall().receiverImplicitlyBorrowed() } or
   TSsaNode(SsaImpl::DataFlowIntegration::SsaNode node) or
   TFlowSummaryNode(FlowSummaryImpl::Private::SummaryNode sn) or
   TClosureSelfReferenceNode(CfgScope c) { lambdaCreationExpr(c, _) } or

rust/ql/lib/codeql/rust/elements/Call.qll

@@ -0,0 +1,7 @@
+/**
+ * This module provides the public class `Call`.
+ */
+
+private import internal.CallImpl
+
+final class Call = Impl::Call;