From dd954ea943a2a51646a80810504aea80aed15d3c Mon Sep 17 00:00:00 2001
From: Arthur Baars <aibaars@github.com>
Date: Wed, 16 Dec 2020 11:17:09 +0100
Subject: [PATCH] CFG: correct flow for lambda bodies

Lambda bodies are parsed as nested do-blocks or normal blocks.
This is actually incorrect, as the body of a lambda can't have
parameters. However, we can "inline" such blocks to get the
desired control flow.
---
 .../controlflow/ControlFlowGraph.qll          | 19 +++++++++++++++++--
 .../internal/ControlFlowGraphImpl.qll         | 18 +++++++++++++-----
 .../controlflow/graph/Cfg.expected            | 15 ++-------------
 3 files changed, 32 insertions(+), 20 deletions(-)

diff --git a/ql/src/codeql_ruby/controlflow/ControlFlowGraph.qll b/ql/src/codeql_ruby/controlflow/ControlFlowGraph.qll
index 7d0253410f3..98a2398db09 100644
--- a/ql/src/codeql_ruby/controlflow/ControlFlowGraph.qll
+++ b/ql/src/codeql_ruby/controlflow/ControlFlowGraph.qll
@@ -76,6 +76,8 @@ private module CfgScope {
   }
 
   private class DoBlockScope extends Range, Generated::DoBlock {
+    DoBlockScope() { not this.getParent() instanceof Generated::Lambda }
+
     final override string getName() { result = "do block" }
 
     final override predicate entry(Generated::AstNode first) {
@@ -88,6 +90,8 @@ private module CfgScope {
   }
 
   private class BlockScope extends Range, Generated::Block {
+    BlockScope() { not this.getParent() instanceof Generated::Lambda }
+
     final override string getName() { result = "block" }
 
     final override predicate entry(Generated::AstNode first) {
@@ -103,11 +107,22 @@ private module CfgScope {
     final override string getName() { result = "lambda" }
 
     final override predicate entry(Generated::AstNode first) {
-      first(this.(Trees::LambdaTree).getFirstChildNode(), first)
+      first(this.getParameters(), first)
+      or
+      not exists(this.getParameters()) and
+      (
+        first(this.getBody().(Trees::DoBlockTree).firstBody(), first)
+        or
+        first(this.getBody().(Trees::BlockTree).getFirstChildNode(), first)
+      )
     }
 
     final override predicate exit(Generated::AstNode last, Completion c) {
-      last(this.(Trees::LambdaTree).getLastChildNode(), last, c)
+      last(this.getBody().(Trees::BlockTree).getLastChildNode(), last, c)
+      or
+      this.getBody().(Trees::RescueEnsureBlockTree).lastBody(last, c)
+      or
+      not exists(this.getBody()) and last(this.getParameters(), last, c)
     }
   }
 }
diff --git a/ql/src/codeql_ruby/controlflow/internal/ControlFlowGraphImpl.qll b/ql/src/codeql_ruby/controlflow/internal/ControlFlowGraphImpl.qll
index d0f0d255251..4474165b4a0 100644
--- a/ql/src/codeql_ruby/controlflow/internal/ControlFlowGraphImpl.qll
+++ b/ql/src/codeql_ruby/controlflow/internal/ControlFlowGraphImpl.qll
@@ -431,7 +431,7 @@ module Trees {
     override predicate isHidden() { any() }
   }
 
-  private class DoBlockTree extends RescueEnsureBlockTree, PostOrderTree, DoBlock {
+  class DoBlockTree extends RescueEnsureBlockTree, PostOrderTree, DoBlock {
     final override predicate first(AstNode first) { first = this }
 
     final override AstNode getChildNode(int i, boolean rescuable) {
@@ -679,11 +679,19 @@ module Trees {
     final override AstNode getDefaultValue() { result = this.getValue() }
   }
 
-  class LambdaTree extends StandardNode, PreOrderTree, PostOrderTree, Lambda {
-    final override AstNode getChildNode(int i) {
-      result = this.getParameters() and i = 0
+  class LambdaTree extends PreOrderTree, PostOrderTree, Lambda {
+    final override predicate propagatesAbnormal(AstNode child) { child = this.getParameters() }
+
+    final override predicate succ(AstNode pred, AstNode succ, Completion c) {
+      this.getBody().(ControlFlowTree).succ(pred, succ, c)
       or
-      result = this.getBody() and i = 1
+      last(this.getParameters(), pred, c) and
+      c instanceof NormalCompletion and
+      (
+        first(this.getBody().(DoBlockTree).firstBody(), succ)
+        or
+        first(this.getBody().(BlockTree).getFirstChildNode(), succ)
+      )
     }
   }
 
diff --git a/ql/test/library-tests/controlflow/graph/Cfg.expected b/ql/test/library-tests/controlflow/graph/Cfg.expected
index a5bd861ba6b..a2fd71aa2eb 100644
--- a/ql/test/library-tests/controlflow/graph/Cfg.expected
+++ b/ql/test/library-tests/controlflow/graph/Cfg.expected
@@ -49,9 +49,6 @@ cfg.rb:
 #  120| enter lambda
 #-----|  -> x
 
-#  120| enter block
-#-----|  -> y
-
 #  142| enter print
 #-----|  -> puts
 
@@ -1375,7 +1372,7 @@ cfg.rb:
 #-----|  -> swap
 
 #  120| DestructuredParameter
-#-----|  -> Block
+#-----|  -> y
 
 #  120| x
 #-----|  -> y
@@ -1383,11 +1380,8 @@ cfg.rb:
 #  120| y
 #-----|  -> DestructuredParameter
 
-#  120| Block
-#-----|  -> exit lambda (normal)
-
 #  120| Array
-#-----|  -> exit block (normal)
+#-----|  -> exit lambda (normal)
 
 #  120| y
 #-----|  -> x
@@ -3690,8 +3684,6 @@ cfg.rb:
 
 #  120| exit lambda
 
-#  120| exit block
-
 #  142| exit print
 
 #  149| exit method
@@ -3821,9 +3813,6 @@ cfg.rb:
 #  120| exit lambda (normal)
 #-----|  -> exit lambda
 
-#  120| exit block (normal)
-#-----|  -> exit block
-
 #  142| exit print (normal)
 #-----|  -> exit print