Put exprNodeReturnedFrom predicate in DataFlowDispatch.qll

Author: alexrford

ql/lib/codeql/ruby/dataflow/internal/DataFlowDispatch.qll

@@ -314,3 +314,16 @@ predicate mayBenefitFromCallContext(DataFlowCall call, Callable c) { none() }
  * restricted to those `call`s for which a context might make a difference.
  */
 DataFlowCallable viableImplInCallContext(DataFlowCall call, DataFlowCall ctx) { none() }
+
+/**
+ * Holds if `e` is an `ExprNode` that may be returned by a call to `c`.
+ */
+predicate exprNodeReturnedFrom(DataFlow::ExprNode e, DataFlowCallable c) {
+  exists(ReturnNode r |
+    r.getEnclosingCallable() = c and
+    (
+      r.(ExplicitReturnNode).getReturningNode().getReturnedValueNode() = e.asExpr() or
+      r.(ExprReturnNode) = e
+    )
+  )
+}

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

@@ -4552,13 +4552,3 @@ private predicate revPartialFlow(
   sink.isRevSink() and
   node.getASuccessor+() = sink
 }
-
-/**
- * Holds if `n` is a return node from callable `c`.
- */
-predicate nodeReturnedFrom(ReturnNodeExt n, DataFlowCallable c) {
-  exists(RetNodeEx ret |
-    n = ret.asNode() and
-    c = ret.getReturnPosition().getCallable()
-  )
-}

ql/lib/codeql/ruby/security/ReflectedXSSCustomizations.qll

@@ -7,6 +7,7 @@ private import codeql.ruby.frameworks.ActionController
 private import codeql.ruby.frameworks.ActionView
 private import codeql.ruby.dataflow.RemoteFlowSources
 private import codeql.ruby.dataflow.BarrierGuards
+import codeql.ruby.dataflow.internal.DataFlowDispatch
 private import codeql.ruby.typetracking.TypeTracker
 
 /**
@@ -190,8 +191,9 @@ module ReflectedXSS {
       template = node2.getLocation().getFile() and
       helperMethod.getName() = helperMethodCall.getExpr().getMethodName() and
       helperMethod.getControllerClass() = getAssociatedControllerClass(template) and
-      // `node1` is a returned value
-      DataFlow::nodeReturnedFrom(node1, helperMethod) and
+      // `node1` is an expr node that may be returned by the helper method
+      exprNodeReturnedFrom(node1, helperMethod) and
+      // `node2` is a call to the helper method
       node2.asExpr() = helperMethodCall
     )
   }

ql/test/query-tests/security/cwe-079/ReflectedXSS.expected

@@ -1,6 +1,6 @@
 edges
-| app/controllers/foo/bars_controller.rb:9:5:9:29 | return :  | app/views/foo/bars/show.html.erb:47:5:47:13 | call to user_name |
-| app/controllers/foo/bars_controller.rb:9:12:9:17 | call to params :  | app/controllers/foo/bars_controller.rb:9:5:9:29 | return :  |
+| app/controllers/foo/bars_controller.rb:9:12:9:17 | call to params :  | app/controllers/foo/bars_controller.rb:9:12:9:29 | ...[...] :  |
+| app/controllers/foo/bars_controller.rb:9:12:9:29 | ...[...] :  | app/views/foo/bars/show.html.erb:47:5:47:13 | call to user_name |
 | app/controllers/foo/bars_controller.rb:13:5:13:37 | ... = ... :  | app/views/foo/bars/show.html.erb:51:5:51:18 | call to user_name_memo |
 | app/controllers/foo/bars_controller.rb:13:20:13:25 | call to params :  | app/controllers/foo/bars_controller.rb:13:5:13:37 | ... = ... :  |
 | app/controllers/foo/bars_controller.rb:17:21:17:26 | call to params :  | app/controllers/foo/bars_controller.rb:17:21:17:36 | ...[...] :  |
@@ -19,8 +19,8 @@ edges
 | app/views/foo/bars/show.html.erb:54:29:54:34 | call to params :  | app/views/foo/bars/show.html.erb:54:29:54:44 | ...[...] |
 | app/views/foo/bars/show.html.erb:57:13:57:18 | call to params :  | app/views/foo/bars/show.html.erb:57:13:57:28 | ...[...] |
 nodes
-| app/controllers/foo/bars_controller.rb:9:5:9:29 | return :  | semmle.label | return :  |
 | app/controllers/foo/bars_controller.rb:9:12:9:17 | call to params :  | semmle.label | call to params :  |
+| app/controllers/foo/bars_controller.rb:9:12:9:29 | ...[...] :  | semmle.label | ...[...] :  |
 | app/controllers/foo/bars_controller.rb:13:5:13:37 | ... = ... :  | semmle.label | ... = ... :  |
 | app/controllers/foo/bars_controller.rb:13:20:13:25 | call to params :  | semmle.label | call to params :  |
 | app/controllers/foo/bars_controller.rb:17:21:17:26 | call to params :  | semmle.label | call to params :  |