Fix rb/reflected-xss flow from helper method return values

Author: alexrford

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

@@ -4552,3 +4552,13 @@ 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

@@ -185,15 +185,13 @@ module ReflectedXSS {
     // flow out of controller helper method into template
     exists(
       ErbFile template, ActionControllerHelperMethod helperMethod,
-      CfgNodes::ExprNodes::MethodCallCfgNode helperMethodCall, ReturnStmt ret
+      CfgNodes::ExprNodes::MethodCallCfgNode helperMethodCall
     |
       template = node2.getLocation().getFile() and
       helperMethod.getName() = helperMethodCall.getExpr().getMethodName() and
       helperMethod.getControllerClass() = getAssociatedControllerClass(template) and
       // `node1` is a returned value
-      // TODO: we don't pick up implicit returns with this approach
-      node1.asExpr().getExpr().getParent() = ret and
-      ret.getParent+() = helperMethod and
+      DataFlow::nodeReturnedFrom(node1, helperMethod) and
       node2.asExpr() = helperMethodCall
     )
   }

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

@@ -1,6 +1,8 @@
 edges
-| 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: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: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 | ...[...] :  |
 | app/controllers/foo/bars_controller.rb:17:21:17:36 | ...[...] :  | app/views/foo/bars/show.html.erb:2:18:2:30 | @user_website |
 | app/controllers/foo/bars_controller.rb:18:10:18:15 | call to params :  | app/controllers/foo/bars_controller.rb:19:22:19:23 | dt :  |
@@ -17,8 +19,10 @@ 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 :  |
 | app/controllers/foo/bars_controller.rb:17:21:17:36 | ...[...] :  | semmle.label | ...[...] :  |
 | app/controllers/foo/bars_controller.rb:18:10:18:15 | call to params :  | semmle.label | call to params :  |
@@ -35,6 +39,7 @@ nodes
 | app/views/foo/bars/show.html.erb:44:64:44:87 | ... + ... :  | semmle.label | ... + ... :  |
 | app/views/foo/bars/show.html.erb:44:76:44:87 | call to display_text :  | semmle.label | call to display_text :  |
 | app/views/foo/bars/show.html.erb:47:5:47:13 | call to user_name | semmle.label | call to user_name |
+| app/views/foo/bars/show.html.erb:51:5:51:18 | call to user_name_memo | semmle.label | call to user_name_memo |
 | app/views/foo/bars/show.html.erb:54:29:54:34 | call to params :  | semmle.label | call to params :  |
 | app/views/foo/bars/show.html.erb:54:29:54:44 | ...[...] | semmle.label | ...[...] |
 | app/views/foo/bars/show.html.erb:57:13:57:18 | call to params :  | semmle.label | call to params :  |
@@ -49,5 +54,6 @@ nodes
 | app/views/foo/bars/show.html.erb:36:3:36:14 | call to display_text | app/controllers/foo/bars_controller.rb:18:10:18:15 | call to params :  | app/views/foo/bars/show.html.erb:36:3:36:14 | call to display_text | Cross-site scripting vulnerability due to $@. | app/controllers/foo/bars_controller.rb:18:10:18:15 | call to params | a user-provided value |
 | app/views/foo/bars/show.html.erb:41:3:41:16 | @instance_text | app/controllers/foo/bars_controller.rb:18:10:18:15 | call to params :  | app/views/foo/bars/show.html.erb:41:3:41:16 | @instance_text | Cross-site scripting vulnerability due to $@. | app/controllers/foo/bars_controller.rb:18:10:18:15 | call to params | a user-provided value |
 | 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/views/foo/bars/show.html.erb:47:5:47:13 | call to user_name | Cross-site scripting vulnerability due to $@. | app/controllers/foo/bars_controller.rb:9:12:9:17 | call to params | a user-provided value |
+| 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/views/foo/bars/show.html.erb:51:5:51:18 | call to user_name_memo | Cross-site scripting vulnerability due to $@. | app/controllers/foo/bars_controller.rb:13:20:13:25 | call to params | a user-provided value |
 | app/views/foo/bars/show.html.erb:54:29:54:44 | ...[...] | 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 | ...[...] | Cross-site scripting vulnerability due to $@. | app/views/foo/bars/show.html.erb:54:29:54:34 | call to params | a user-provided value |
 | app/views/foo/bars/show.html.erb:57:13:57:28 | ...[...] | 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 | ...[...] | Cross-site scripting vulnerability due to $@. | app/views/foo/bars/show.html.erb:57:13:57:18 | call to params | a user-provided value |