Rust: Untangle argument/parameter data flow matching from call graph

hvitved · hvitved · commit 56895065315b · 2025-11-20T11:36:57.000+01:00
diff --git a/rust/ql/lib/codeql/rust/dataflow/internal/DataFlowImpl.qll b/rust/ql/lib/codeql/rust/dataflow/internal/DataFlowImpl.qll
@@ -89,15 +89,12 @@ final class DataFlowCall extends TDataFlowCall {
 
 /**
  * The position of a parameter in a function.
- *
- * In Rust there is a 1-to-1 correspondence between parameter positions and
- * 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) }
+  int getPosition(boolean inMethod) { this = TPositionalParameterPosition(result, inMethod) }
 
-  predicate hasPosition() { exists(this.getPosition()) }
+  predicate hasPosition(boolean inMethod) { exists(this.getPosition(inMethod)) }
 
   /** Holds if this position represents the `self` position. */
   predicate isSelf() { this = TSelfParameterPosition() }
@@ -110,32 +107,88 @@ final class ParameterPosition extends TParameterPosition {
 
   /** Gets a textual representation of this position. */
   string toString() {
-    result = this.getPosition().toString()
+    exists(boolean inMethod, string suffix |
+      inMethod = true and suffix = " (method)"
+      or
+      inMethod = false and suffix = ""
+    |
+      result = this.getPosition(inMethod).toString() + suffix
+    )
     or
     result = "self" and this.isSelf()
     or
     result = "closure self" and this.isClosureSelf()
   }
 
   ParamBase getParameterIn(ParamList ps) {
-    result = ps.getParam(this.getPosition())
+    exists(boolean inMethod |
+      result = ps.getParam(this.getPosition(inMethod)) and
+      if ps.hasSelfParam() then inMethod = true else inMethod = false
+    )
     or
     result = ps.getSelfParam() and this.isSelf()
   }
 }
 
 /**
  * The position of an argument in a call.
- *
- * In Rust there is a 1-to-1 correspondence between parameter positions and
- * arguments positions, so we use the same underlying type for both.
  */
-final class ArgumentPosition extends ParameterPosition {
+final class ArgumentPosition extends TArgumentPosition {
   /** Gets the argument of `call` at this position, if any. */
   Expr getArgument(Call call) {
-    result = call.getPositionalArgument(this.getPosition())
+    this.isReceiver() and
+    (
+      result = call.(MethodCallExpr).getReceiver()
+      or
+      result = call.(Operation).getOperand(0)
+      or
+      result = call.(IndexExpr).getBase()
+    )
+    or
+    exists(boolean inMethodCall, int pos | this.getPosition(inMethodCall) = pos |
+      result = call.(CallExpr).getArg(pos) and
+      inMethodCall = false
+      or
+      result = call.(MethodCallExpr).getArg(pos) and
+      inMethodCall = true
+      or
+      result = call.(Operation).getOperand(pos + 1) and
+      pos >= 0 and
+      inMethodCall = true
+      or
+      result = call.(IndexExpr).getIndex() and
+      pos = 1 and
+      inMethodCall = true
+    )
+  }
+
+  /** Gets the underlying integer position, if any. */
+  int getPosition(boolean inMethodCall) { this = TPositionalArgumentPosition(result, inMethodCall) }
+
+  predicate hasPosition(boolean inMethodCall) { exists(this.getPosition(inMethodCall)) }
+
+  /** Holds if this position represents a receiver in a method call. */
+  predicate isReceiver() { this = TReceiverArgumentPosition() }
+
+  /**
+   * Holds if this position represents a reference to a closure itself. Only
+   * used for tracking flow through captured variables.
+   */
+  predicate isClosureReceiver() { this = TClosureReceiverArgumentPosition() }
+
+  /** Gets a textual representation of this position. */
+  string toString() {
+    exists(boolean inMethodCall, string suffix |
+      inMethodCall = true and suffix = " (method call)"
+      or
+      inMethodCall = false and suffix = ""
+    |
+      result = this.getPosition(inMethodCall).toString() + suffix
+    )
     or
-    result = call.getReceiver() and this.isSelf()
+    result = "receiver" and this.isReceiver()
+    or
+    result = "closure receiver" and this.isClosureReceiver()
   }
 }
 
@@ -145,14 +198,11 @@ final class ArgumentPosition extends ParameterPosition {
  * 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(Expr arg, Call call, ParameterPosition pos) {
+predicate isArgumentForCall(Expr arg, Call call, ArgumentPosition pos) {
   // TODO: Handle index expressions as calls in data flow.
   not call instanceof IndexExpr and
-  (
-    call.getPositionalArgument(pos.getPosition()) = arg
-    or
-    call.getReceiver() = arg and pos.isSelf() and not call.receiverImplicitlyBorrowed()
-  )
+  arg = pos.getArgument(call) and
+  not (pos.isReceiver() and call.receiverImplicitlyBorrowed())
 }
 
 /** Provides logic related to SSA. */
@@ -475,7 +525,21 @@ module RustDataFlow implements InputSig<Location> {
    * Holds if the parameter position `ppos` matches the argument position
    * `apos`.
    */
-  predicate parameterMatch(ParameterPosition ppos, ArgumentPosition apos) { ppos = apos }
+  pragma[nomagic]
+  predicate parameterMatch(ParameterPosition ppos, ArgumentPosition apos) {
+    exists(boolean inMethod | ppos.getPosition(inMethod) = apos.getPosition(inMethod))
+    or
+    // `Vec::push(vec, value))`: match `value` against parameter after `self`
+    ppos.getPosition(true) = apos.getPosition(false) - 1
+    or
+    // `vec.push(value)`: match `vec` against `self` parameter
+    ppos.isSelf() and apos.isReceiver()
+    or
+    // `Vec::push(vec, value))`: match `vec` against `self` parameter
+    ppos.isSelf() and apos.getPosition(false) = 0
+    or
+    ppos.isClosureSelf() and apos.isClosureReceiver()
+  }
 
   /**
    * Holds if there is a simple local flow step from `node1` to `node2`. These
@@ -722,7 +786,7 @@ module RustDataFlow implements InputSig<Location> {
     or
     // Store in function argument
     exists(DataFlowCall call, int i |
-      isArgumentNode(node1, call, TPositionalParameterPosition(i)) and
+      isArgumentNode(node1, call, TPositionalArgumentPosition(i, false)) and
       lambdaCall(call, _, node2.(PostUpdateNode).getPreUpdateNode()) and
       c.(FunctionCallArgumentContent).getPosition() = i
     )
@@ -1032,16 +1096,26 @@ private module Cached {
 
   cached
   newtype TParameterPosition =
-    TPositionalParameterPosition(int i) {
-      i in [0 .. max([any(ParamList l).getNumberOfParams(), any(ArgList l).getNumberOfArgs()]) - 1]
+    TPositionalParameterPosition(int i, Boolean inMethod) {
+      i in [0 .. any(ParamList l).getNumberOfParams() - 1]
       or
-      FlowSummaryImpl::ParsePositions::isParsedArgumentPosition(_, i)
-      or
-      FlowSummaryImpl::ParsePositions::isParsedParameterPosition(_, i)
+      FlowSummaryImpl::ParsePositions::isParsedArgumentPosition(_, i) and
+      inMethod = false
     } or
     TClosureSelfParameterPosition() or
     TSelfParameterPosition()
 
+  cached
+  newtype TArgumentPosition =
+    TPositionalArgumentPosition(int i, Boolean inMethodCall) {
+      i in [0 .. any(ArgList l).getNumberOfArgs() - 1]
+      or
+      FlowSummaryImpl::ParsePositions::isParsedParameterPosition(_, i) and
+      inMethodCall = false
+    } or
+    TClosureReceiverArgumentPosition() or
+    TReceiverArgumentPosition()
+
   cached
   newtype TReturnKind = TNormalReturnKind()
 
diff --git a/rust/ql/lib/codeql/rust/dataflow/internal/FlowSummaryImpl.qll b/rust/ql/lib/codeql/rust/dataflow/internal/FlowSummaryImpl.qll
@@ -25,6 +25,24 @@ module Input implements InputSig<Location, RustDataFlow> {
     final predicate callResolvesTo(string path) {
       path = this.getCall().getStaticTarget().(Addressable).getCanonicalPath()
     }
+
+    Expr getArgument(ParameterPosition pos) {
+      exists(ParameterPosition pos0, ArgumentPosition apos |
+        result = apos.getArgument(this.getCall()) and
+        RustDataFlow::parameterMatch(pos0, apos)
+      |
+        not pos.hasPosition(_) and
+        pos0 = pos
+        or
+        // We parse all positions as if they do not belong to a method; here we adjust if needed
+        exists(boolean inMethod |
+          pos0.getPosition(inMethod) = pos.getPosition(_) and
+          if this.getCall().getStaticTarget().(Function).hasSelfParam()
+          then inMethod = true
+          else inMethod = false
+        )
+      )
+    }
   }
 
   abstract class SourceBase extends SourceSinkBase { }
@@ -47,14 +65,18 @@ module Input implements InputSig<Location, RustDataFlow> {
     override MethodCallExpr getCall() { result = call }
   }
 
-  RustDataFlow::ArgumentPosition callbackSelfParameterPosition() { result.isClosureSelf() }
+  RustDataFlow::ArgumentPosition callbackSelfParameterPosition() { result.isClosureReceiver() }
 
   ReturnKind getStandardReturnValueKind() { result = TNormalReturnKind() }
 
-  string encodeParameterPosition(ParameterPosition pos) { result = pos.toString() }
+  string encodeParameterPosition(ParameterPosition pos) {
+    result = pos.toString() and
+    not pos.hasPosition(true)
+  }
 
   string encodeArgumentPosition(RustDataFlow::ArgumentPosition pos) {
-    result = encodeParameterPosition(pos)
+    result = pos.toString() and
+    not pos.hasPosition(true)
   }
 
   string encodeContent(ContentSet cs, string arg) {
@@ -111,14 +133,14 @@ module Input implements InputSig<Location, RustDataFlow> {
   ParameterPosition decodeUnknownParameterPosition(AccessPath::AccessPathTokenBase token) {
     // needed to support `Argument[x..y]` ranges
     token.getName() = "Argument" and
-    result.getPosition() = AccessPath::parseInt(token.getAnArgument())
+    result.getPosition(false) = AccessPath::parseInt(token.getAnArgument())
   }
 
   bindingset[token]
   RustDataFlow::ArgumentPosition decodeUnknownArgumentPosition(AccessPath::AccessPathTokenBase token) {
     // needed to support `Parameter[x..y]` ranges
     token.getName() = "Parameter" and
-    result.getPosition() = AccessPath::parseInt(token.getAnArgument())
+    result.getPosition(false) = AccessPath::parseInt(token.getAnArgument())
   }
 
   bindingset[token]
@@ -135,9 +157,9 @@ private module StepsInput implements Impl::Private::StepsInputSig {
 
   /** Gets the argument of `source` described by `sc`, if any. */
   private Expr getSourceNodeArgument(Input::SourceBase source, Impl::Private::SummaryComponent sc) {
-    exists(ArgumentPosition pos |
+    exists(ParameterPosition pos |
       sc = Impl::Private::SummaryComponent::argument(pos) and
-      result = pos.getArgument(source.getCall())
+      result = source.getArgument(pos)
     )
   }
 
@@ -165,19 +187,18 @@ private module StepsInput implements Impl::Private::StepsInputSig {
     exists(ArgumentPosition pos, Expr arg |
       s.head() = Impl::Private::SummaryComponent::parameter(pos) and
       arg = getSourceNodeArgument(source, s.tail().headOfSingleton()) and
-      result.asParameter() = getCallable(arg).getParam(pos.getPosition())
+      result.asParameter() = getCallable(arg).getParam(pos.getPosition(_))
     )
     or
     result.(RustDataFlow::PostUpdateNode).getPreUpdateNode().asExpr() =
       getSourceNodeArgument(source, s.headOfSingleton())
   }
 
   RustDataFlow::Node getSinkNode(Input::SinkBase sink, Impl::Private::SummaryComponent sc) {
-    exists(CallExprBase call, Expr arg, ArgumentPosition pos |
+    exists(Expr arg, ParameterPosition ppos |
       result.asExpr() = arg and
-      sc = Impl::Private::SummaryComponent::argument(pos) and
-      call = sink.getCall() and
-      arg = pos.getArgument(call)
+      sc = Impl::Private::SummaryComponent::argument(ppos) and
+      arg = sink.getArgument(ppos)
     )
   }
 }
diff --git a/rust/ql/lib/codeql/rust/dataflow/internal/Node.qll b/rust/ql/lib/codeql/rust/dataflow/internal/Node.qll
@@ -190,7 +190,19 @@ final class SummaryParameterNode extends ParameterNode, FlowSummaryNode {
   private ParameterPosition pos_;
 
   SummaryParameterNode() {
-    FlowSummaryImpl::Private::summaryParameterNode(this.getSummaryNode(), pos_)
+    exists(ParameterPosition pos0 |
+      FlowSummaryImpl::Private::summaryParameterNode(this.getSummaryNode(), pos0)
+    |
+      not pos0.hasPosition(_) and
+      pos_ = pos0
+      or
+      // We parse all positions as if they do not belong to a method; here we adjust if needed
+      exists(SummarizedCallable sc, boolean inMethod |
+        sc = this.getSummarizedCallable() and
+        pos_.getPosition(inMethod) = pos0.getPosition(_) and
+        if sc.hasSelfParam() then inMethod = true else inMethod = false
+      )
+    )
   }
 
   override predicate isParameterOf(DataFlowCallable c, ParameterPosition pos) {
@@ -247,7 +259,7 @@ final class ReceiverNode extends ArgumentNode, TReceiverNode {
   MethodCallExpr getMethodCall() { result = n }
 
   override predicate isArgumentOf(DataFlowCall call, RustDataFlow::ArgumentPosition pos) {
-    call.asCall() = n and pos = TSelfParameterPosition()
+    call.asCall() = n and pos = TReceiverArgumentPosition()
   }
 
   override CfgScope getCfgScope() { result = n.getEnclosingCfgScope() }
@@ -281,7 +293,7 @@ final class ClosureArgumentNode extends ArgumentNode, ExprNode {
 
   override predicate isArgumentOf(DataFlowCall call, RustDataFlow::ArgumentPosition pos) {
     call.asCall() = call_ and
-    pos.isClosureSelf()
+    pos.isClosureReceiver()
   }
 }
 
diff --git a/rust/ql/src/utils/modelgenerator/internal/CaptureModels.qll b/rust/ql/src/utils/modelgenerator/internal/CaptureModels.qll
@@ -66,9 +66,15 @@ module ModelGeneratorCommonInput implements
   string qualifierString() { result = "Argument[self]" }
 
   string parameterExactAccess(R::ParamBase p) {
-    result =
-      "Argument[" + any(DataFlowImpl::ParameterPosition pos | p = pos.getParameterIn(_)).toString() +
-        "]"
+    exists(DataFlowImpl::ParameterPosition pos0, DataFlowImpl::ParameterPosition pos |
+      p = pos0.getParameterIn(_) and
+      result = "Argument[" + pos + "]"
+    |
+      not pos0.hasPosition(true) and
+      pos = pos0
+      or
+      pos.getPosition(false) = pos0.getPosition(true)
+    )
   }
 
   string parameterApproximateAccess(R::ParamBase p) { result = parameterExactAccess(p) }
@@ -84,7 +90,7 @@ module ModelGeneratorCommonInput implements
 
   bindingset[c]
   string paramReturnNodeAsExactOutput(QualifiedCallable c, DataFlowImpl::ParameterPosition pos) {
-    result = parameterExactAccess(c.getFunction().getParam(pos.getPosition()))
+    result = parameterExactAccess(c.getFunction().getParam(pos.getPosition(_)))
     or
     pos.isSelf() and result = qualifierString()
   }
