Rust: Support Argument[x] MaD source definitions

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

@@ -90,10 +90,10 @@ final class DataFlowCall extends TDataFlowCall {
 }
 
 /**
- * The position of a parameter or an argument in a function or call.
+ * The position of a parameter in a function.
  *
- * As there is a 1-to-1 correspondence between parameter positions and
- * arguments positions in Rust we use the same type for both.
+ * 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. */
@@ -126,6 +126,22 @@ final class ParameterPosition extends TParameterPosition {
   }
 }
 
+/**
+ * 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 {
+  /** Gets the argument of `call` at this position, if any. */
+  Expr getArgument(CallExprBase call) {
+    result = call.getArgList().getArg(this.getPosition())
+    or
+    this.isSelf() and
+    result = call.(MethodCallExpr).getReceiver()
+  }
+}
+
 /** Holds if `call` invokes a qualified path that resolves to a method. */
 private predicate callToMethod(CallExpr call) {
   exists(Path path |
@@ -432,6 +448,8 @@ private module Aliases {
 
   class ParameterPositionAlias = ParameterPosition;
 
+  class ArgumentPositionAlias = ArgumentPosition;
+
   class ContentAlias = Content;
 
   class ContentSetAlias = ContentSet;
@@ -550,7 +568,7 @@ module RustDataFlow implements InputSig<Location> {
 
   class ParameterPosition = ParameterPositionAlias;
 
-  class ArgumentPosition = ParameterPosition;
+  class ArgumentPosition = ArgumentPositionAlias;
 
   /**
    * Holds if the parameter position `ppos` matches the argument position

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

@@ -58,7 +58,9 @@ module Input implements InputSig<Location, RustDataFlow> {
 
   string encodeParameterPosition(ParameterPosition pos) { result = pos.toString() }
 
-  predicate encodeArgumentPosition = encodeParameterPosition/1;
+  string encodeArgumentPosition(RustDataFlow::ArgumentPosition pos) {
+    result = encodeParameterPosition(pos)
+  }
 
   string encodeContent(ContentSet cs, string arg) {
     exists(Content c | cs = TSingletonContentSet(c) |
@@ -146,17 +148,21 @@ private module StepsInput implements Impl::Private::StepsInputSig {
   RustDataFlow::Node getSourceNode(Input::SourceBase source, Impl::Private::SummaryComponent sc) {
     sc = Impl::Private::SummaryComponent::return(_) and
     result.asExpr().getExpr() = source.getCall()
+    or
+    exists(CallExprBase call, Expr arg, ArgumentPosition pos |
+      result.(RustDataFlow::PostUpdateNode).getPreUpdateNode().asExpr().getExpr() = arg and
+      sc = Impl::Private::SummaryComponent::argument(pos) and
+      call = source.getCall() and
+      arg = pos.getArgument(call)
+    )
   }
 
   RustDataFlow::Node getSinkNode(Input::SinkBase sink, Impl::Private::SummaryComponent sc) {
-    exists(CallExprBase call, Expr arg, ParameterPosition pos |
+    exists(CallExprBase call, Expr arg, ArgumentPosition pos |
       result.asExpr().getExpr() = arg and
       sc = Impl::Private::SummaryComponent::argument(pos) and
-      call = sink.getCall()
-    |
-      arg = call.getArgList().getArg(pos.getPosition())
-      or
-      arg = call.(MethodCallExpr).getReceiver() and pos.isSelf()
+      call = sink.getCall() and
+      arg = pos.getArgument(call)
     )
   }
 }

rust/ql/test/library-tests/dataflow/models/main.rs

@@ -279,6 +279,15 @@ fn test_simple_sink() {
     simple_sink(s); // $ hasValueFlow=17
 }
 
+// has a source model
+fn arg_source(i: i64) {}
+
+fn test_arg_source() {
+    let i = 19;
+    arg_source(i);
+    sink(i) // $ hasValueFlow=i
+}
+
 #[tokio::main]
 async fn main() {
     test_identify();
@@ -299,5 +308,6 @@ async fn main() {
     test_simple_source();
     test_simple_sink();
     test_get_async_number().await;
+    test_arg_source();
     let dummy = Some(0); // ensure that the the `lang:core` crate is extracted
 }