Ruby: Fix ExprReturnNode to handle Rescue, Ensure, and Else.

Author: aschackmull

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

@@ -1643,16 +1643,24 @@ private module ReturnNodes {
     }
   }
 
-  pragma[noinline]
-  private AstNode implicitReturn(Callable c, ExprNode n) {
-    exists(CfgNodes::ExprCfgNode en |
-      en = n.getExprNode() and
-      en.getASuccessor().(CfgNodes::AnnotatedExitNode).isNormal() and
-      n.(NodeImpl).getCfgScope() = c and
-      result = en.getExpr()
-    )
+  private AstNode desugar(AstNode n) {
+    result = n.getDesugared()
+    or
+    not exists(n.getDesugared()) and
+    result = n
+  }
+
+  private Expr getLast(StmtSequence s) {
+    result = getLast(s.(BodyStmt).getElse())
     or
-    result = implicitReturn(c, n).getParent()
+    result = getLast(s.(BodyStmt).getARescue().getBody())
+    or
+    not exists(s.(BodyStmt).getElse()) and
+    exists(Stmt last | last = s.getLastStmt() |
+      result = getLast(last)
+      or
+      result = last and not last instanceof StmtSequence
+    )
   }
 
   /**
@@ -1661,7 +1669,7 @@ private module ReturnNodes {
    * last thing that is evaluated in the body of the callable.
    */
   class ExprReturnNode extends SourceReturnNode, ExprNode {
-    ExprReturnNode() { exists(Callable c | implicitReturn(c, this) = c.getBody().getAStmt()) }
+    ExprReturnNode() { this.getExprNode().getExpr() = desugar(getLast(any(Callable c).getBody())) }
 
     override ReturnKind getKindSource() {
       exists(CfgScope scope | scope = this.(NodeImpl).getCfgScope() |

ruby/ql/test/library-tests/dataflow/implicit-return/implicit-return.expected

@@ -1,10 +1,30 @@
 models
 edges
 | implicit_return.rb:12:3:12:11 | call to source | implicit_return.rb:15:6:15:11 | call to m_body | provenance |  |
+| implicit_return.rb:19:3:19:11 | call to source | implicit_return.rb:24:6:24:18 | call to m_body_ensure | provenance |  |
+| implicit_return.rb:30:3:30:11 | call to source | implicit_return.rb:33:6:33:13 | call to m_rescue | provenance |  |
+| implicit_return.rb:39:3:39:11 | call to source | implicit_return.rb:44:6:44:20 | call to m_rescue_ensure | provenance |  |
+| implicit_return.rb:52:3:52:11 | call to source | implicit_return.rb:55:6:55:11 | call to m_else | provenance |  |
+| implicit_return.rb:63:3:63:11 | call to source | implicit_return.rb:68:6:68:18 | call to m_else_ensure | provenance |  |
 nodes
 | implicit_return.rb:12:3:12:11 | call to source | semmle.label | call to source |
 | implicit_return.rb:15:6:15:11 | call to m_body | semmle.label | call to m_body |
+| implicit_return.rb:19:3:19:11 | call to source | semmle.label | call to source |
+| implicit_return.rb:24:6:24:18 | call to m_body_ensure | semmle.label | call to m_body_ensure |
+| implicit_return.rb:30:3:30:11 | call to source | semmle.label | call to source |
+| implicit_return.rb:33:6:33:13 | call to m_rescue | semmle.label | call to m_rescue |
+| implicit_return.rb:39:3:39:11 | call to source | semmle.label | call to source |
+| implicit_return.rb:44:6:44:20 | call to m_rescue_ensure | semmle.label | call to m_rescue_ensure |
+| implicit_return.rb:52:3:52:11 | call to source | semmle.label | call to source |
+| implicit_return.rb:55:6:55:11 | call to m_else | semmle.label | call to m_else |
+| implicit_return.rb:63:3:63:11 | call to source | semmle.label | call to source |
+| implicit_return.rb:68:6:68:18 | call to m_else_ensure | semmle.label | call to m_else_ensure |
 subpaths
 testFailures
 #select
 | implicit_return.rb:15:6:15:11 | call to m_body | implicit_return.rb:12:3:12:11 | call to source | implicit_return.rb:15:6:15:11 | call to m_body | $@ | implicit_return.rb:12:3:12:11 | call to source | call to source |
+| implicit_return.rb:24:6:24:18 | call to m_body_ensure | implicit_return.rb:19:3:19:11 | call to source | implicit_return.rb:24:6:24:18 | call to m_body_ensure | $@ | implicit_return.rb:19:3:19:11 | call to source | call to source |
+| implicit_return.rb:33:6:33:13 | call to m_rescue | implicit_return.rb:30:3:30:11 | call to source | implicit_return.rb:33:6:33:13 | call to m_rescue | $@ | implicit_return.rb:30:3:30:11 | call to source | call to source |
+| implicit_return.rb:44:6:44:20 | call to m_rescue_ensure | implicit_return.rb:39:3:39:11 | call to source | implicit_return.rb:44:6:44:20 | call to m_rescue_ensure | $@ | implicit_return.rb:39:3:39:11 | call to source | call to source |
+| implicit_return.rb:55:6:55:11 | call to m_else | implicit_return.rb:52:3:52:11 | call to source | implicit_return.rb:55:6:55:11 | call to m_else | $@ | implicit_return.rb:52:3:52:11 | call to source | call to source |
+| implicit_return.rb:68:6:68:18 | call to m_else_ensure | implicit_return.rb:63:3:63:11 | call to source | implicit_return.rb:68:6:68:18 | call to m_else_ensure | $@ | implicit_return.rb:63:3:63:11 | call to source | call to source |

ruby/ql/test/library-tests/dataflow/implicit-return/implicit_return.rb

@@ -21,7 +21,7 @@ def m_body_ensure
   source(20)
 end
 
-sink(m_body_ensure) # $ MISSING: hasValueFlow=2
+sink(m_body_ensure) # $ hasValueFlow=2
 
 # Implicit return from a `rescue` clause.
 def m_rescue
@@ -30,7 +30,7 @@ def m_rescue
   source(3)
 end
 
-sink(m_rescue) # $ MISSING: hasValueFlow=3
+sink(m_rescue) # $ hasValueFlow=3
 
 # Implicit return from a `rescue` clause when an `ensure` clause is present.
 def m_rescue_ensure
@@ -41,7 +41,7 @@ def m_rescue_ensure
   source(40)
 end
 
-sink(m_rescue_ensure) # $ MISSING: hasValueFlow=4
+sink(m_rescue_ensure) # $ hasValueFlow=4
 
 # Implicit return from an `else` clause.
 def m_else
@@ -52,7 +52,7 @@ def m_else
   source(5)
 end
 
-sink(m_else) # $ MISSING: hasValueFlow=5
+sink(m_else) # $ hasValueFlow=5
 
 # Implicit return from an `else` clause when an `ensure` clause is present.
 def m_else_ensure
@@ -65,4 +65,4 @@ def m_else_ensure
   nil
 end
 
-sink(m_else_ensure) # $ MISSING: hasValueFlow=6
+sink(m_else_ensure) # $ hasValueFlow=6