Rust: Assume callbacks will be invoked in library functions

rust/ql/lib/codeql/rust/controlflow/internal/Scope.qll

@@ -35,7 +35,8 @@ final class CallableScope extends CfgScopeImpl, Callable {
   CallableScope() {
     // A function without a body corresponds to a trait method signature and
     // should not have a CFG scope.
-    this.hasBody()
+    this.hasBody() and
+    this.fromSource() // exclude stubs in tests defined using `additionalExternalFile`
   }
 
   override predicate scopeFirst(AstNode first) {

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

@@ -440,7 +440,7 @@ module RustDataFlowGen<RustDataFlowInputSig Input> implements InputSig<Location>
       or
       result = "self" and this.isSelf()
       or
-      result = "closure self" and this.isClosureSelf()
+      result = "closure-self" and this.isClosureSelf()
     }
 
     ParamBase getParameterIn(ParamList ps) {

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

@@ -58,6 +58,10 @@ private import codeql.rust.dataflow.FlowBarrier
 private import codeql.rust.dataflow.FlowSummary
 private import codeql.rust.dataflow.FlowSource
 private import codeql.rust.dataflow.FlowSink
+private import codeql.rust.internal.CachedStages
+private import codeql.rust.internal.typeinference.FunctionType
+private import codeql.rust.internal.typeinference.TypeMention
+private import codeql.rust.frameworks.stdlib.Stdlib
 
 /**
  * Holds if in a call to the function with canonical path `path`, the value referred
@@ -207,6 +211,39 @@ private class SummarizedCallableFromModel extends SummarizedCallable::Range {
   }
 }
 
+/**
+ * Holds if library function `f` has a callback at position `n`. In this case we
+ * add a flow model that achieves the effect of simulating that the callback is
+ * invoked, which is needed for flow through captured variables to work.
+ */
+cached
+predicate mayInvokeCallback(Function f, int n) {
+  Stages::TypeInferenceStage::ref() and
+  exists(TypeMention tm, Trait trait |
+    tm = f.getParam(n).getTypeRepr() and
+    trait = getALookupTrait(f, tm.getType()) and
+    trait.getSupertrait*() instanceof FnOnceTrait and
+    not f.fromSource()
+  )
+}
+
+private class SummarizedCallableWithCallback extends SummarizedCallable::Range {
+  private int pos;
+
+  SummarizedCallableWithCallback() { mayInvokeCallback(this, pos) }
+
+  override predicate propagatesFlow(
+    string input, string output, boolean preservesValue, Provenance p, boolean isExact, string model
+  ) {
+    input = "Argument[" + pos + "]" and
+    output = "Argument[" + pos + "].Parameter[closure-self]" and
+    preservesValue = true and
+    p = "hq-generated" and
+    isExact = true and
+    model = "heuristic-callback"
+  }
+}
+
 private class FlowSourceFromModel extends FlowSource::Range {
   private string path;
 

rust/ql/lib/codeql/rust/internal/CachedStages.qll

@@ -130,6 +130,7 @@ module Stages {
   module TypeInferenceStage {
     private import codeql.rust.internal.typeinference.Type
     private import codeql.rust.internal.typeinference.TypeInference
+    private import codeql.rust.dataflow.internal.ModelsAsData
 
     /**
      * Always holds.
@@ -150,6 +151,8 @@ module Stages {
       exists(Type t)
       or
       exists(inferType(_))
+      or
+      mayInvokeCallback(_, _)
     }
   }
 

rust/ql/test/library-tests/dataflow/lambdas/external_file.rs

@@ -0,0 +1,11 @@
+pub fn may_invoke_callback1<F: Fn(i64)>(f: F) {}
+
+pub fn may_invoke_callback2<F: FnOnce(i64)>(f: F) {}
+
+pub fn may_invoke_callback3(f: impl Fn(i64)) {}
+
+pub fn may_invoke_callback4<T>(f: T)
+where
+    T: for<'a> FnOnce(&'a mut i64),
+{
+}

rust/ql/test/library-tests/dataflow/lambdas/inline-flow.expected

@@ -55,6 +55,14 @@ edges
 | main.rs:100:13:100:22 | source(...) | main.rs:100:9:100:9 | b | provenance |  |
 | main.rs:101:17:101:17 | ... | main.rs:101:25:101:25 | x | provenance |  |
 | main.rs:101:29:101:29 | b | main.rs:93:33:93:38 | ...: i64 | provenance |  |
+| main.rs:109:13:109:22 | source(...) | main.rs:110:26:110:36 | \|...\| ... : ... => .. [captured a] | provenance |  |
+| main.rs:109:13:109:22 | source(...) | main.rs:111:26:111:36 | \|...\| ... : ... => .. [captured a] | provenance |  |
+| main.rs:109:13:109:22 | source(...) | main.rs:112:26:112:36 | \|...\| ... : ... => .. [captured a] | provenance |  |
+| main.rs:109:13:109:22 | source(...) | main.rs:113:26:113:36 | \|...\| ... : ... => .. [captured a] | provenance |  |
+| main.rs:110:26:110:36 | \|...\| ... : ... => .. [captured a] | main.rs:110:35:110:35 | a | provenance | heuristic-callback |
+| main.rs:111:26:111:36 | \|...\| ... : ... => .. [captured a] | main.rs:111:35:111:35 | a | provenance | heuristic-callback |
+| main.rs:112:26:112:36 | \|...\| ... : ... => .. [captured a] | main.rs:112:35:112:35 | a | provenance | heuristic-callback |
+| main.rs:113:26:113:36 | \|...\| ... : ... => .. [captured a] | main.rs:113:35:113:35 | a | provenance | heuristic-callback |
 nodes
 | main.rs:10:20:10:52 | if cond {...} else {...} | semmle.label | if cond {...} else {...} |
 | main.rs:10:30:10:39 | source(...) | semmle.label | source(...) |
@@ -121,6 +129,15 @@ nodes
 | main.rs:101:17:101:17 | ... | semmle.label | ... |
 | main.rs:101:25:101:25 | x | semmle.label | x |
 | main.rs:101:29:101:29 | b | semmle.label | b |
+| main.rs:109:13:109:22 | source(...) | semmle.label | source(...) |
+| main.rs:110:26:110:36 | \|...\| ... : ... => .. [captured a] | semmle.label | \|...\| ... : ... => .. [captured a] |
+| main.rs:110:35:110:35 | a | semmle.label | a |
+| main.rs:111:26:111:36 | \|...\| ... : ... => .. [captured a] | semmle.label | \|...\| ... : ... => .. [captured a] |
+| main.rs:111:35:111:35 | a | semmle.label | a |
+| main.rs:112:26:112:36 | \|...\| ... : ... => .. [captured a] | semmle.label | \|...\| ... : ... => .. [captured a] |
+| main.rs:112:35:112:35 | a | semmle.label | a |
+| main.rs:113:26:113:36 | \|...\| ... : ... => .. [captured a] | semmle.label | \|...\| ... : ... => .. [captured a] |
+| main.rs:113:35:113:35 | a | semmle.label | a |
 subpaths
 | main.rs:29:21:29:21 | a | main.rs:27:20:27:23 | ... | main.rs:27:26:27:52 | if cond {...} else {...} | main.rs:29:13:29:22 | f(...) |
 | main.rs:77:21:77:21 | a | main.rs:66:24:66:32 | ...: i64 | main.rs:66:42:72:1 | { ... } | main.rs:77:13:77:22 | f(...) |
@@ -140,3 +157,7 @@ testFailures
 | main.rs:99:25:99:25 | x | main.rs:100:13:100:22 | source(...) | main.rs:99:25:99:25 | x | $@ | main.rs:100:13:100:22 | source(...) | source(...) |
 | main.rs:101:25:101:25 | x | main.rs:98:13:98:22 | source(...) | main.rs:101:25:101:25 | x | $@ | main.rs:98:13:98:22 | source(...) | source(...) |
 | main.rs:101:25:101:25 | x | main.rs:100:13:100:22 | source(...) | main.rs:101:25:101:25 | x | $@ | main.rs:100:13:100:22 | source(...) | source(...) |
+| main.rs:110:35:110:35 | a | main.rs:109:13:109:22 | source(...) | main.rs:110:35:110:35 | a | $@ | main.rs:109:13:109:22 | source(...) | source(...) |
+| main.rs:111:35:111:35 | a | main.rs:109:13:109:22 | source(...) | main.rs:111:35:111:35 | a | $@ | main.rs:109:13:109:22 | source(...) | source(...) |
+| main.rs:112:35:112:35 | a | main.rs:109:13:109:22 | source(...) | main.rs:112:35:112:35 | a | $@ | main.rs:109:13:109:22 | source(...) | source(...) |
+| main.rs:113:35:113:35 | a | main.rs:109:13:109:22 | source(...) | main.rs:113:35:113:35 | a | $@ | main.rs:109:13:109:22 | source(...) | source(...) |

rust/ql/test/library-tests/dataflow/lambdas/inline-flow.ext.yml

@@ -0,0 +1,6 @@
+extensions:
+  - addsTo:
+      pack: codeql/rust-all
+      extensible: additionalExternalFile
+    data:
+      - ["external_file.rs"]

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

@@ -102,6 +102,17 @@ fn test_apply_wrap() {
     apply_wrap(|x| sink(x), 0);
 }
 
+mod external_file;
+use external_file::*;
+
+fn test_external_call() {
+    let a = source(81);
+    may_invoke_callback1(|x| sink(a)); // $ hasValueFlow=81
+    may_invoke_callback2(|x| sink(a)); // $ hasValueFlow=81
+    may_invoke_callback3(|x| sink(a)); // $ hasValueFlow=81
+    may_invoke_callback4(|x| sink(a)); // $ hasValueFlow=81
+}
+
 fn main() {
     closure_flow_out();
     closure_flow_in();