From c33a0ef44072c56c6f94ce2cd03c587775cf7ba4 Mon Sep 17 00:00:00 2001
From: Sebastian Echeverria <sebastian.echeverria@gmail.com>
Date: Fri, 18 Oct 2024 12:56:17 -0400
Subject: [PATCH] add support for heredocs

---
 go.mod                   |  4 +--
 parse.go                 | 36 ++++++++++++++++++-----
 parse_test.go            | 60 ++++++++++++++++++++++++++++++++++++++
 pylib/main.go            | 50 +++++++++++++++++++++++++++++--
 pylib/support.c          | 49 ++++++++++++++++++++++++++-----
 tests/dockerfile_test.py | 63 ++++++++++++++++++++++++++++++++++++++++
 6 files changed, 242 insertions(+), 20 deletions(-)

diff --git a/go.mod b/go.mod
index 04d59b8..2c5f860 100644
--- a/go.mod
+++ b/go.mod
@@ -3,6 +3,6 @@ module github.com/asottile/dockerfile
 go 1.16
 
 require (
-	github.com/moby/buildkit v0.9.0
-	github.com/stretchr/testify v1.7.0
+	github.com/moby/buildkit v0.12.5
+	github.com/stretchr/testify v1.8.3
 )
diff --git a/parse.go b/parse.go
index 73107ea..e49136b 100644
--- a/parse.go
+++ b/parse.go
@@ -9,17 +9,25 @@ import (
 	"github.com/moby/buildkit/frontend/dockerfile/parser"
 )
 
+// Represents info about a heredoc.
+type Heredoc struct {
+	Name           string
+	FileDescriptor uint
+	Content        string
+}
+
 // Represents a single line (layer) in a Dockerfile.
 // For example `FROM ubuntu:xenial`
 type Command struct {
-	Cmd       string   // lowercased command name (ex: `from`)
-	SubCmd    string   // for ONBUILD only this holds the sub-command
-	Json      bool     // whether the value is written in json form
-	Original  string   // The original source line
-	StartLine int      // The original source line number which starts this command
-	EndLine   int      // The original source line number which ends this command
-	Flags     []string // Any flags such as `--from=...` for `COPY`.
-	Value     []string // The contents of the command (ex: `ubuntu:xenial`)
+	Cmd       string    // lowercased command name (ex: `from`)
+	SubCmd    string    // for ONBUILD only this holds the sub-command
+	Json      bool      // whether the value is written in json form
+	Original  string    // The original source line
+	StartLine int       // The original source line number which starts this command
+	EndLine   int       // The original source line number which ends this command
+	Flags     []string  // Any flags such as `--from=...` for `COPY`.
+	Value     []string  // The contents of the command (ex: `ubuntu:xenial`)
+	Heredocs  []Heredoc // Extra heredoc content attachments
 }
 
 // A failure in opening a file for reading.
@@ -78,6 +86,18 @@ func ParseReader(file io.Reader) ([]Command, error) {
 			cmd.Value = append(cmd.Value, n.Value)
 		}
 
+		if len(child.Heredocs) != 0 {
+			// For heredocs, add heredocs extra lines to Original,
+			// and to the heredocs list.
+			cmd.Original = cmd.Original + "\n"
+			for _, heredoc := range child.Heredocs {
+				cmd.Original = cmd.Original + heredoc.Content + heredoc.Name + "\n"
+				cmd.Heredocs = append(cmd.Heredocs, Heredoc{Name: heredoc.Name,
+					FileDescriptor: heredoc.FileDescriptor,
+					Content:        heredoc.Content})
+			}
+		}
+
 		ret = append(ret, cmd)
 	}
 	return ret, nil
diff --git a/parse_test.go b/parse_test.go
index bf3b18f..84c85b5 100644
--- a/parse_test.go
+++ b/parse_test.go
@@ -116,3 +116,63 @@ func TestParseFile(t *testing.T) {
 	}
 	assert.Equal(t, expected, cmds)
 }
+
+func TestParseReaderHeredocs(t *testing.T) {
+	dockerfile := `RUN 3<<EOF
+source $HOME/.bashrc && echo $HOME
+echo "Hello" >> /hello
+echo "World!" >> /hello
+EOF
+`
+	cmds, err := ParseReader(bytes.NewBufferString(dockerfile))
+	assert.Nil(t, err)
+	expected := []Command{
+		Command{
+			Cmd:       "RUN",
+			Original:  dockerfile,
+			StartLine: 1,
+			EndLine:   5,
+			Flags:     []string{},
+			Value:     []string{"3<<EOF"},
+			Heredocs: []Heredoc{
+				Heredoc{
+					Name:           "EOF",
+					FileDescriptor: 3,
+					Content:        "source $HOME/.bashrc && echo $HOME\necho \"Hello\" >> /hello\necho \"World!\" >> /hello\n"},
+			},
+		},
+	}
+	assert.Equal(t, expected, cmds)
+}
+
+func TestParseReaderHeredocsMultiple(t *testing.T) {
+	dockerfile := `COPY <<FILE1 <<FILE2 /dest
+content 1
+FILE1
+content 2
+FILE2
+`
+	cmds, err := ParseReader(bytes.NewBufferString(dockerfile))
+	assert.Nil(t, err)
+	expected := []Command{
+		Command{
+			Cmd:       "COPY",
+			Original:  dockerfile,
+			StartLine: 1,
+			EndLine:   5,
+			Flags:     []string{},
+			Value:     []string{"<<FILE1", "<<FILE2", "/dest"},
+			Heredocs: []Heredoc{
+				Heredoc{
+					Name:           "FILE1",
+					FileDescriptor: 0,
+					Content:        "content 1\n"},
+				Heredoc{
+					Name:           "FILE2",
+					FileDescriptor: 0,
+					Content:        "content 2\n"},
+			},
+		},
+	}
+	assert.Equal(t, expected, cmds)
+}
diff --git a/pylib/main.go b/pylib/main.go
index 65a89da..14af386 100644
--- a/pylib/main.go
+++ b/pylib/main.go
@@ -10,7 +10,10 @@ package main
 // extern PyObject* PyDockerfile_GoParseError;
 // extern PyObject* PyDockerfile_NewCommand(
 //     PyObject*, PyObject*, PyObject*, PyObject*, PyObject*, PyObject*,
-//     PyObject*, PyObject*
+//     PyObject*, PyObject*, PyObject*
+// );
+// extern PyObject* PyDockerfile_NewHeredoc(
+//     PyObject*, PyObject*, PyObject*
 // );
 import "C"
 import (
@@ -65,6 +68,40 @@ func sliceToTuple(strs []string) *C.PyObject {
 	return ret
 }
 
+func heredocsToPy(heredocs []dockerfile.Heredoc) *C.PyObject {
+	var pyName, pyFileDescriptor, pyContent *C.PyObject
+	var ret *C.PyObject
+	decrefAll := func() {
+		C.Py_DecRef(pyName)
+		C.Py_DecRef(pyFileDescriptor)
+		C.Py_DecRef(pyContent)
+		C.Py_DecRef(ret)
+	}
+
+	ret = C.PyTuple_New(C.Py_ssize_t(len(heredocs)))
+	for i, heredoc := range heredocs {
+		pyName := stringToPy(heredoc.Name)
+		if pyName == nil {
+			decrefAll()
+			return nil
+		}
+
+		pyFileDescriptor := C.PyLong_FromLong(C.long(heredoc.FileDescriptor))
+
+		pyContent := stringToPy(heredoc.Content)
+		if pyContent == nil {
+			decrefAll()
+			return nil
+		}
+
+		pyHeredoc := C.PyDockerfile_NewHeredoc(
+			pyName, pyFileDescriptor, pyContent,
+		)
+		C.PyTuple_SetItem(ret, C.Py_ssize_t(i), pyHeredoc)
+	}
+	return ret
+}
+
 func boolToInt(b bool) int {
 	if b {
 		return 1
@@ -74,7 +111,7 @@ func boolToInt(b bool) int {
 }
 
 func cmdsToPy(cmds []dockerfile.Command) *C.PyObject {
-	var pyCmd, pySubCmd, pyJson, pyOriginal, pyStartLine, pyEndLine, pyValue *C.PyObject
+	var pyCmd, pySubCmd, pyJson, pyOriginal, pyStartLine, pyEndLine, pyValue, pyHeredocs *C.PyObject
 	var pyFlags *C.PyObject
 	var ret *C.PyObject
 	decrefAll := func() {
@@ -86,6 +123,7 @@ func cmdsToPy(cmds []dockerfile.Command) *C.PyObject {
 		C.Py_DecRef(pyEndLine)
 		C.Py_DecRef(pyFlags)
 		C.Py_DecRef(pyValue)
+		C.Py_DecRef(pyHeredocs)
 		C.Py_DecRef(ret)
 	}
 
@@ -126,8 +164,14 @@ func cmdsToPy(cmds []dockerfile.Command) *C.PyObject {
 			return nil
 		}
 
+		pyHeredocs = heredocsToPy(cmd.Heredocs)
+		if pyHeredocs == nil {
+			decrefAll()
+			return nil
+		}
+
 		pyCmd := C.PyDockerfile_NewCommand(
-			pyCmd, pySubCmd, pyJson, pyOriginal, pyStartLine, pyEndLine, pyFlags, pyValue,
+			pyCmd, pySubCmd, pyJson, pyOriginal, pyStartLine, pyEndLine, pyFlags, pyValue, pyHeredocs,
 		)
 		C.PyTuple_SetItem(ret, C.Py_ssize_t(i), pyCmd)
 	}
diff --git a/pylib/support.c b/pylib/support.c
index 5aa2237..b6b387a 100644
--- a/pylib/support.c
+++ b/pylib/support.c
@@ -30,11 +30,26 @@ PyObject* PyDockerfile_NewCommand(
     PyObject* start_line,
     PyObject* end_line,
     PyObject* flags,
-    PyObject* value
+    PyObject* value,
+    PyObject* heredocs
 ) {
     return PyObject_CallFunction(
-        PyDockerfile_Command, "OOOOOOOO",
-        cmd, sub_cmd, json, original, start_line, end_line, flags, value
+        PyDockerfile_Command, "OOOOOOOOO",
+        cmd, sub_cmd, json, original, start_line, end_line, flags, value, heredocs
+    );
+}
+
+/* Heredoc namedtuple */
+PyObject* PyDockerfile_Heredoc;
+
+PyObject* PyDockerfile_NewHeredoc(
+    PyObject* name,
+    PyObject* file_descriptor,
+    PyObject* content
+) {
+    return PyObject_CallFunction(
+        PyDockerfile_Heredoc, "OOO",
+        name, file_descriptor, content
     );
 }
 
@@ -53,15 +68,35 @@ static PyObject* _setup_module(PyObject* module) {
         PyModule_AddObject(module, "GoParseError", PyDockerfile_GoParseError);
 
         PyObject* collections = PyImport_ImportModule("collections");
-        PyDockerfile_Command = PyObject_CallMethod(
-            collections, "namedtuple", "ss",
-            "Command", "cmd sub_cmd json original start_line end_line flags value"
-        );
+
+        // Set up a Command namedtuple object, with empty default for heredocs substructure.
+        PyObject *args = Py_BuildValue("ss", "Command",  "cmd sub_cmd json original start_line end_line flags value heredocs");
+        PyObject *kwargs = PyDict_New();
+        PyObject* defaults = Py_BuildValue("(())");
+        PyDict_SetItemString(kwargs, "defaults", defaults);
+        PyObject *namedtuple_method = PyObject_GetAttrString(collections, "namedtuple");
+        PyDockerfile_Command = PyObject_Call(namedtuple_method, args, kwargs);
+        Py_DECREF(args);
+        Py_DECREF(kwargs);
+        Py_DECREF(defaults);
+        Py_DECREF(namedtuple_method);
         PyObject_SetAttrString(
             PyDockerfile_Command, "__module__",
             PyObject_GetAttrString(module, "__name__")
         );
         PyModule_AddObject(module, "Command", PyDockerfile_Command);
+
+        // Set up a Heredoc namedtuple object.
+        PyDockerfile_Heredoc = PyObject_CallMethod(
+            collections, "namedtuple", "ss",
+            "Heredoc", "name file_descriptor content"
+        );
+        PyObject_SetAttrString(
+            PyDockerfile_Heredoc, "__module__",
+            PyObject_GetAttrString(module, "__name__")
+        );
+        PyModule_AddObject(module, "Heredoc", PyDockerfile_Heredoc);
+
         Py_XDECREF(collections);
     }
     return module;
diff --git a/tests/dockerfile_test.py b/tests/dockerfile_test.py
index 8457d93..e35e87c 100644
--- a/tests/dockerfile_test.py
+++ b/tests/dockerfile_test.py
@@ -101,3 +101,66 @@ def test_parse_file_success():
             start_line=2, end_line=2, original='CMD ["echo", "hi"]',
         ),
     )
+
+
+def test_heredoc_string_success():
+    test_string = (
+        'RUN 3<<EOF\n'
+        'source $HOME/.bashrc && echo $HOME\n'
+        'echo "Hello" >> /hello\n'
+        'echo "World!" >> /hello\n'
+        'EOF\n'
+    )
+    ret = dockerfile.parse_string(test_string)
+    assert ret == (
+        dockerfile.Command(
+            cmd='RUN', sub_cmd=None, json=False, flags=(),
+            value=(
+                '3<<EOF',
+            ),
+            start_line=1, end_line=5, original=test_string,
+            heredocs=(
+                dockerfile.Heredoc(
+                    name='EOF',
+                    content='source $HOME/.bashrc && echo $HOME\n'
+                            'echo "Hello" >> /hello\n'
+                            'echo "World!" >> /hello\n',
+                    file_descriptor=3,
+                ),
+            ),
+        ),
+    )
+
+
+def test_heredoc_string_multiple_success():
+    test_string = (
+        'COPY <<FILE1 <<FILE2 /dest\n'
+        'content 1\n'
+        'FILE1\n'
+        'content 2\n'
+        'FILE2\n'
+    )
+    ret = dockerfile.parse_string(test_string)
+    assert ret == (
+        dockerfile.Command(
+            cmd='COPY', sub_cmd=None, json=False, flags=(),
+            value=(
+                '<<FILE1',
+                '<<FILE2',
+                '/dest',
+            ),
+            start_line=1, end_line=5, original=test_string,
+            heredocs=(
+                dockerfile.Heredoc(
+                    name='FILE1',
+                    content='content 1\n',
+                    file_descriptor=0,
+                ),
+                dockerfile.Heredoc(
+                    name='FILE2',
+                    content='content 2\n',
+                    file_descriptor=0,
+                ),
+            ),
+        ),
+    )