Shared: Pass SummaryComponentStack to isSource and getSourceType

Author: paldepind

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

@@ -6,7 +6,10 @@ private import rust
 private import codeql.dataflow.internal.FlowSummaryImpl
 private import codeql.dataflow.internal.AccessPathSyntax as AccessPath
 private import codeql.rust.dataflow.internal.DataFlowImpl
+private import codeql.rust.internal.PathResolution
 private import codeql.rust.dataflow.FlowSummary
+private import codeql.rust.dataflow.Ssa
+private import codeql.rust.controlflow.CfgNodes
 private import Content
 
 module Input implements InputSig<Location, RustDataFlow> {
@@ -133,16 +136,40 @@ private module StepsInput implements Impl::Private::StepsInputSig {
     result.asCallCfgNode().getCall().getStaticTarget() = sc
   }
 
-  RustDataFlow::Node getSourceNode(Input::SourceBase source, Impl::Private::SummaryComponent sc) {
-    sc = Impl::Private::SummaryComponent::return(_) and
+  /** Gets the argument of `source` described by `sc`, if any. */
+  private Expr getSourceNodeArgument(Input::SourceBase source, Impl::Private::SummaryComponent sc) {
+    exists(ArgumentPosition pos |
+      sc = Impl::Private::SummaryComponent::argument(pos) and
+      result = pos.getArgument(source.getCall())
+    )
+  }
+
+  /** Get the callable that `expr` refers to. */
+  private Callable getCallable(Expr expr) {
+    result = resolvePath(expr.(PathExpr).getPath()).(Function)
+    or
+    result = expr.(ClosureExpr)
+    or
+    // The expression is an SSA read of an assignment of a closure
+    exists(Ssa::Definition def, ExprCfgNode value |
+      def.getARead().getAstNode() = expr and
+      def.getAnUltimateDefinition().(Ssa::WriteDefinition).assigns(value) and
+      result = value.getExpr().(ClosureExpr)
+    )
+  }
+
+  RustDataFlow::Node getSourceNode(Input::SourceBase source, Impl::Private::SummaryComponentStack s) {
+    s.head() = 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)
+    exists(ArgumentPosition pos, Expr arg |
+      s.head() = Impl::Private::SummaryComponent::parameter(pos) and
+      arg = getSourceNodeArgument(source, s.tail().head()) and
+      result.asParameter() = getCallable(arg).getParam(pos.getPosition())
     )
+    or
+    result.(RustDataFlow::PostUpdateNode).getPreUpdateNode().asExpr().getExpr() =
+      getSourceNodeArgument(source, s.head())
   }
 
   RustDataFlow::Node getSinkNode(Input::SinkBase sink, Impl::Private::SummaryComponent sc) {

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

@@ -258,20 +258,20 @@ mod source_into_function {
     }
 
     fn test_source_into_function() {
-        let a = |a| sink(a); // $ MISSING: hasValueFlow=1
+        let a = |a| sink(a); // $ hasValueFlow=1
         pass_source(1, a);
 
         pass_source(2, |a| {
-            sink(a); // $ MISSING: hasValueFlow=2
+            sink(a); // $ hasValueFlow=2
         });
 
         fn f(a: i64) {
-            sink(a) // $ MISSING: hasValueFlow=3
+            sink(a) // $ hasValueFlow=3
         }
         pass_source(3, f);
 
         pass_source(4, async move |a| {
-            sink(a); // $ MISSING: hasValueFlow=4
+            sink(a); // $ hasValueFlow=4
         });
     }
 }