From fa64578cb830d31637e87d502fff649086b7e305 Mon Sep 17 00:00:00 2001
From: Justin Chadwell <me@jedevc.com>
Date: Tue, 25 May 2021 17:16:11 +0100
Subject: [PATCH] Various cleanups and bugfixes for heredocs

Signed-off-by: Justin Chadwell <me@jedevc.com>
---
 frontend/dockerfile/instructions/parse.go | 15 ++++++++++-----
 frontend/dockerfile/parser/parser.go      | 10 +++++++---
 2 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/frontend/dockerfile/instructions/parse.go b/frontend/dockerfile/instructions/parse.go
index 757872d6a0d7..9f29e0a4bff5 100644
--- a/frontend/dockerfile/instructions/parse.go
+++ b/frontend/dockerfile/instructions/parse.go
@@ -248,13 +248,16 @@ func parseAdd(req parseRequest) (*AddCommand, error) {
 		return nil, err
 	}
 
-	if req.heredoc != nil && parser.IsHeredoc(req.args[0]) {
+	if req.heredoc != nil {
 		if len(req.args) != 2 {
 			return nil, errExactlyTwoArguments("ADD")
 		}
 		if parser.IsHeredoc(req.args[1]) {
 			return nil, errBadHeredoc("ADD", "a destination")
 		}
+		if !parser.IsHeredoc(req.args[0]) {
+			panic("First argument must be a heredoc, or somehow the parsing has gone wrong!")
+		}
 
 		return &AddCommand{
 			withNameAndCode: newWithNameAndCode(req),
@@ -285,13 +288,16 @@ func parseCopy(req parseRequest) (*CopyCommand, error) {
 		return nil, err
 	}
 
-	if req.heredoc != nil && parser.IsHeredoc(req.args[0]) {
+	if req.heredoc != nil {
 		if len(req.args) != 2 {
 			return nil, errExactlyTwoArguments("COPY")
 		}
 		if parser.IsHeredoc(req.args[1]) {
 			return nil, errBadHeredoc("COPY", "a destination")
 		}
+		if !parser.IsHeredoc(req.args[0]) {
+			panic("First argument must be a heredoc, or somehow the parsing has gone wrong!")
+		}
 
 		return &CopyCommand{
 			withNameAndCode: newWithNameAndCode(req),
@@ -394,9 +400,8 @@ func parseWorkdir(req parseRequest) (*WorkdirCommand, error) {
 
 func parseShellDependentCommand(req parseRequest, emptyAsNil bool, allowHeredoc bool) (ShellDependantCmdLine, error) {
 	if allowHeredoc && req.heredoc != nil {
-		if len(req.args) != 1 || !parser.IsHeredoc(req.args[0]) {
-			// FIXME: this is not a good error message
-			return ShellDependantCmdLine{}, errors.Errorf("Bad heredoc shell command")
+		if !(len(req.args) == 1 && parser.IsHeredoc(req.args[0])) {
+			return ShellDependantCmdLine{}, errExactlyOneArgument(req.command + " with heredoc")
 		}
 
 		cmds := make([]strslice.StrSlice, len(req.heredoc.Lines))
diff --git a/frontend/dockerfile/parser/parser.go b/frontend/dockerfile/parser/parser.go
index 6816042ebcbd..8a4607a1e60d 100644
--- a/frontend/dockerfile/parser/parser.go
+++ b/frontend/dockerfile/parser/parser.go
@@ -31,7 +31,7 @@ type Node struct {
 	Value       string          // actual content
 	Next        *Node           // the next item in the current sexp
 	Children    []*Node         // the children of this sexp
-	Heredoc     *Heredoc        // TODO
+	Heredoc     *Heredoc        // extra heredoc content attachment
 	Attributes  map[string]bool // special attributes for this node
 	Original    string          // original line used before parsing
 	Flags       []string        // only top Node should have this set
@@ -332,19 +332,23 @@ func Parse(rwc io.Reader) (*Result, error) {
 				heredocExpand = false
 			}
 
+			var heredocLine string
 			var heredocLines []string
 			for scanner.Scan() {
-				heredocLine := scanner.Text()
+				heredocLine = scanner.Text()
 				if heredocChomp {
 					heredocLine = strings.TrimLeftFunc(heredocLine, unicode.IsSpace)
 				}
+				currentLine++
 
 				if heredocLine == heredocName {
 					break
 				}
 				heredocLines = append(heredocLines, heredocLine)
 			}
-			// TODO: handle unclosed heredoc
+			if heredocLine != heredocName {
+				return nil, withLocation(errors.New("unterminated heredoc"), startLine, currentLine)
+			}
 
 			heredoc = &Heredoc{
 				Name:   heredocName,