From 6c579ff608b57f87b9bce5ec56674294c6f770e2 Mon Sep 17 00:00:00 2001
From: Arthur Baars <aibaars@github.com>
Date: Thu, 3 Dec 2020 13:30:43 +0100
Subject: [PATCH] CFG: link heredoc start to its body

---
 .../internal/ControlFlowGraphImpl.qll         | 21 +++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/ql/src/codeql_ruby/controlflow/internal/ControlFlowGraphImpl.qll b/ql/src/codeql_ruby/controlflow/internal/ControlFlowGraphImpl.qll
index 0d18b7e894e..95aa779c9ef 100644
--- a/ql/src/codeql_ruby/controlflow/internal/ControlFlowGraphImpl.qll
+++ b/ql/src/codeql_ruby/controlflow/internal/ControlFlowGraphImpl.qll
@@ -431,6 +431,27 @@ private module Trees {
     }
   }
 
+  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.getValue().regexpCapture("^<<[-~]?[`']?(.*)[`']?$", 1) and
+          end = unique(HeredocEnd x | x = doc.getChild(_)) 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()
+        )
+    }
+  }
+
   private class IdentifierTree extends LeafTree, Identifier { }
 
   private class IfElsifTree extends PreOrderTree, IfElsifAstNode {