From 8501e30b6a1f37d82a9185669def6085dc54dfd1 Mon Sep 17 00:00:00 2001
From: Arthur Baars <aibaars@github.com>
Date: Tue, 15 Dec 2020 12:44:16 +0100
Subject: [PATCH] CFG: fix linking heredoc start to heredoc body

---
 .../internal/ControlFlowGraphImpl.qll         | 39 +++++++++----------
 .../controlflow/graph/Cfg.expected            | 30 ++++++++++++++
 .../controlflow/graph/heredoc.rb              |  7 ++++
 3 files changed, 56 insertions(+), 20 deletions(-)
 create mode 100644 ql/test/library-tests/controlflow/graph/heredoc.rb

diff --git a/ql/src/codeql_ruby/controlflow/internal/ControlFlowGraphImpl.qll b/ql/src/codeql_ruby/controlflow/internal/ControlFlowGraphImpl.qll
index adff80357b8..cc4f3d06811 100644
--- a/ql/src/codeql_ruby/controlflow/internal/ControlFlowGraphImpl.qll
+++ b/ql/src/codeql_ruby/controlflow/internal/ControlFlowGraphImpl.qll
@@ -37,6 +37,7 @@ private import codeql_ruby.controlflow.ControlFlowGraph
 private import Completion
 private import SuccessorTypes
 private import Splitting
+private import codeql.files.FileSystem
 
 private AstNode parent(AstNode n) {
   result.getAFieldOrChild() = n and
@@ -606,29 +607,27 @@ module Trees {
 
   private class HashSplatParameterTree extends LeafTree, HashSplatParameter { }
 
-  private class HeredocBeginningTree extends StandardPreOrderTree, HeredocBeginning {
-    pragma[noinline]
-    private string getName() {
-      result = this.getValue().regexpCapture("^<<[-~]?[`']?(.*)[`']?$", 1)
-    }
+  private HeredocBody heredoc(HeredocBeginning start) {
+    exists(int i, File f |
+      start =
+        rank[i](HeredocBeginning b |
+          f = b.getLocation().getFile()
+        |
+          b order by b.getLocation().getStartLine(), b.getLocation().getStartColumn()
+        ) and
+      result =
+        rank[i](HeredocBody b |
+          f = b.getLocation().getFile()
+        |
+          b order by b.getLocation().getStartLine(), b.getLocation().getStartColumn()
+        )
+    )
+  }
 
+  private class HeredocBeginningTree extends StandardPreOrderTree, HeredocBeginning {
     final override AstNode getChildNode(int i) {
       i = 0 and
-      result =
-        min(string name, HeredocBody doc, HeredocEnd end |
-          name = this.getName() and
-          end = unique(HeredocEnd x | x = doc.getChild(_) | x) and
-          end.getValue() = name and
-          doc.getLocation().getFile() = this.getLocation().getFile() and
-          (
-            doc.getLocation().getStartLine() > this.getLocation().getStartLine()
-            or
-            doc.getLocation().getStartLine() = this.getLocation().getStartLine() and
-            doc.getLocation().getStartColumn() > this.getLocation().getStartColumn()
-          )
-        |
-          doc order by doc.getLocation().getStartLine(), doc.getLocation().getStartColumn()
-        )
+      result = heredoc(this)
     }
   }
 
diff --git a/ql/test/library-tests/controlflow/graph/Cfg.expected b/ql/test/library-tests/controlflow/graph/Cfg.expected
index ca1694f09fc..88146edbf0c 100644
--- a/ql/test/library-tests/controlflow/graph/Cfg.expected
+++ b/ql/test/library-tests/controlflow/graph/Cfg.expected
@@ -68,6 +68,10 @@ exit.rb:
 #    8| enter m2
 #-----|  -> x
 
+heredoc.rb:
+#    1| enter double_heredoc
+#-----|  -> <<A
+
 ifs.rb:
 #    1| enter m1
 #-----|  -> x
@@ -1950,6 +1954,25 @@ exit.rb:
 #   12| String
 #-----|  -> puts
 
+heredoc.rb:
+#    2| MethodCall
+#-----|  -> exit double_heredoc (normal)
+
+#    2| puts
+#-----|  -> MethodCall
+
+#    2| <<A
+#-----|  -> HeredocBody
+
+#    2| <<A
+#-----|  -> HeredocBody
+
+#    2| HeredocBody
+#-----|  -> <<A
+
+#    4| HeredocBody
+#-----|  -> puts
+
 ifs.rb:
 #    1| x
 #-----|  -> x
@@ -3380,6 +3403,9 @@ exit.rb:
 
 #    8| exit m2
 
+heredoc.rb:
+#    1| exit double_heredoc
+
 ifs.rb:
 #    1| exit m1
 
@@ -3505,6 +3531,10 @@ exit.rb:
 #    8| exit m2 (normal)
 #-----|  -> exit m2
 
+heredoc.rb:
+#    1| exit double_heredoc (normal)
+#-----|  -> exit double_heredoc
+
 ifs.rb:
 #    1| exit m1 (normal)
 #-----|  -> exit m1
diff --git a/ql/test/library-tests/controlflow/graph/heredoc.rb b/ql/test/library-tests/controlflow/graph/heredoc.rb
new file mode 100644
index 00000000000..09e1b771ea8
--- /dev/null
+++ b/ql/test/library-tests/controlflow/graph/heredoc.rb
@@ -0,0 +1,7 @@
+def double_heredoc
+  puts(<<A, <<A)
+    hello
+A
+   world
+A
+end