Rust: Share the returned-value check with callback return sinks

Author: LucaCappelletti94

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

@@ -144,8 +144,6 @@ private Callable getCallable(Expr expr) {
 }
 
 module Input2 implements Impl::Private::InputSig2 {
-  private import codeql.rust.controlflow.CfgNodes
-
   class SourceSinkReportingElement extends AstNode {
     DataFlowCallable getEnclosingCallable() { result.asCfgScope() = this.getEnclosingCfgScope() }
 
@@ -158,7 +156,7 @@ module Input2 implements Impl::Private::InputSig2 {
       exists(Callable c |
         sc = Impl::Private::SummaryComponent::return(_) and
         c = getCallable(this) and
-        result.getACfgNode().getASuccessor() instanceof AnnotatedExitCfgNode and
+        isReturnedValue(result) and
         result.getEnclosingCfgScope() = c
       )
     }

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

@@ -472,14 +472,17 @@ abstract class ReturnNode extends Node {
   abstract ReturnKind getKind();
 }
 
+/** Holds if the value of `n` is returned from its enclosing callable. */
+predicate isReturnedValue(AstNode n) {
+  exists(SuccessorType t | n.getACfgNode().getASuccessor(t) instanceof AnnotatedExitCfgNode |
+    // The early return of `x?` returns the `Err` or `None` held by `x`, which is
+    // a different value from the one `x?` evaluates to
+    not (n instanceof TryExpr and t instanceof ReturnSuccessor)
+  )
+}
+
 final class ExprReturnNode extends ExprNode, ReturnNode {
-  ExprReturnNode() {
-    exists(SuccessorType t | n.getACfgNode().getASuccessor(t) instanceof AnnotatedExitCfgNode |
-      // The early return of `x?` returns the `Err` or `None` held by `x`, which is
-      // a different value from the one `x?` evaluates to
-      not (n instanceof TryExpr and t instanceof ReturnSuccessor)
-    )
-  }
+  ExprReturnNode() { isReturnedValue(n) }
 
   override ReturnKind getKind() { result = TNormalReturnKind() }
 }

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

@@ -581,3 +581,18 @@ async fn main() {
     test_neutrals();
     let dummy = Some(0); // ensure that the the `lang:core` crate is extracted
 }
+
+mod try_in_sink_callback {
+    use super::source;
+
+    // has a sink model
+    fn pass_sink_option(f: impl FnOnce(()) -> Option<i64>) {}
+
+    fn test_try_in_sink_callback() {
+        let a = |_a| {
+            let _s = Some(source(100))?;
+            Some(0)
+        };
+        pass_sink_option(a);
+    }
+}

rust/ql/test/library-tests/dataflow/models/models.ext.yml

@@ -26,6 +26,7 @@ extensions:
       - ["<main::MyFieldEnum>::sink", "Argument[self].Field[main::MyFieldEnum::D::field_d]", "test-sink", "manual"]
       - ["main::sink_out_of_function::pass_sink", "Argument[0].ReturnValue", "test-sink", "manual"]
       - ["main::sink_out_of_function::pass_sink_nested", "Argument[0].ReturnValue.Field[main::MyFieldEnum::E::field_e].Field[core::option::Option::Some(0)]", "test-sink", "manual"]
+      - ["main::try_in_sink_callback::pass_sink_option", "Argument[0].ReturnValue", "test-sink", "manual"]
       - ["main::external_file::generated_sink", "Argument[0]", "test-sink", "dfc-generated"]
       - ["main::external_file::neutral_generated_sink", "Argument[0]", "test-sink", "dfc-generated"]
       - ["main::external_file::neutral_manual_sink", "Argument[0]", "test-sink", "manual"]