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/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 @@
+/**
+ * Call.
+ */
+
+private import internal.CallImpl
+
+final class Call = Impl::Call;

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

@@ -0,0 +1,128 @@
+private import rust
+private import codeql.rust.internal.PathResolution
+private import codeql.rust.internal.TypeInference as TypeInference
+private import codeql.rust.elements.internal.ExprImpl::Impl as ExprImpl
+private import codeql.rust.elements.Operation
+
+module Impl {
+  /**
+   * An expression that calls a function.
+   *
+   * This class abstract over the different ways in which a function can be called in Rust.
+   */
+  abstract class Call extends ExprImpl::Expr {
+    Call() { this.fromSource() }
+
+    /** Gets the number of arguments _excluding_ any `self` argument. */
+    abstract int getNumberOfArguments();
+
+    /** Gets the receiver of this call if it is a method call. */
+    abstract Expr getReceiver();
+
+    /** Holds if the call has a receiver that might be implicitly borrowed. */
+    abstract predicate receiverImplicitlyBorrowed();
+
+    /** Gets the trait targeted by this call, if any. */
+    abstract Trait getTrait();
+
+    /** Gets the name of the method called if this call is a method call. */
+    abstract string getMethodName();
+
+    /** Gets the `i`th argument of this call, if any. */
+    abstract Expr getArgument(int i);
+
+    /** Gets the static target of this call, if any. */
+    Function getStaticTarget() {
+      result = TypeInference::resolveMethodCallTarget(this)
+      or
+      not exists(TypeInference::resolveMethodCallTarget(this)) and
+      result = this.(CallExpr).getStaticTarget()
+    }
+  }
+
+  /** Holds if the call expression dispatches to a trait method. */
+  private predicate callIsMethodCall(CallExpr call, Path qualifier, string methodName) {
+    exists(Path path, Function f |
+      path = call.getFunction().(PathExpr).getPath() and
+      f = resolvePath(path) and
+      f.getParamList().hasSelfParam() and
+      qualifier = path.getQualifier() and
+      path.getSegment().getIdentifier().getText() = methodName
+    )
+  }
+
+  private class CallExprCall extends Call instanceof CallExpr {
+    CallExprCall() { not callIsMethodCall(this, _, _) }
+
+    override string getMethodName() { none() }
+
+    override Expr getReceiver() { none() }
+
+    override Trait getTrait() { none() }
+
+    override predicate receiverImplicitlyBorrowed() { none() }
+
+    override int getNumberOfArguments() { result = super.getArgList().getNumberOfArgs() }
+
+    override Expr getArgument(int i) { result = super.getArgList().getArg(i) }
+  }
+
+  private class CallExprMethodCall extends Call instanceof CallExpr {
+    Path qualifier;
+    string methodName;
+
+    CallExprMethodCall() { callIsMethodCall(this, qualifier, methodName) }
+
+    override string getMethodName() { result = methodName }
+
+    override Expr getReceiver() { result = super.getArgList().getArg(0) }
+
+    override Trait getTrait() {
+      result = resolvePath(qualifier) and
+      // When the qualifier is `Self` and resolves to a trait, it's inside a
+      // trait method's default implementation. This is not a dispatch whose
+      // target is inferred from the type of the receiver, but should always
+      // resolve to the function in the trait block as path resolution does.
+      qualifier.toString() != "Self"
+    }
+
+    override predicate receiverImplicitlyBorrowed() { none() }
+
+    override int getNumberOfArguments() { result = super.getArgList().getNumberOfArgs() - 1 }
+
+    override Expr getArgument(int i) { result = super.getArgList().getArg(i + 1) }
+  }
+
+  private class MethodCallExprCall extends Call instanceof MethodCallExpr {
+    override string getMethodName() { result = super.getIdentifier().getText() }
+
+    override Expr getReceiver() { result = this.(MethodCallExpr).getReceiver() }
+
+    override Trait getTrait() { none() }
+
+    override predicate receiverImplicitlyBorrowed() { any() }
+
+    override int getNumberOfArguments() { result = super.getArgList().getNumberOfArgs() }
+
+    override Expr getArgument(int i) { result = super.getArgList().getArg(i) }
+  }
+
+  private class OperatorCall extends Call instanceof Operation {
+    Trait trait;
+    string methodName;
+
+    OperatorCall() { super.isOverloaded(trait, methodName) }
+
+    override string getMethodName() { result = methodName }
+
+    override Expr getReceiver() { result = super.getOperand(0) }
+
+    override Trait getTrait() { result = trait }
+
+    override predicate receiverImplicitlyBorrowed() { none() }
+
+    override int getNumberOfArguments() { result = super.getNumberOfOperands() - 1 }
+
+    override Expr getArgument(int i) { result = super.getOperand(1) and i = 0 }
+  }
+}