Fix issues with nested YAML (#148)

pkg/template/BUILD

@@ -15,6 +15,20 @@ go_library(
         "//third_party/go:evanphx_json-patch",
         "//third_party/go:ghodss_yaml",
         "//third_party/go:speps_go-hashids",
+        "//third_party/go:yaml_v3",
+    ],
+)
+
+go_test(
+    name = "template_test",
+    srcs = [
+        "load_test.go",
+    ],
+    resources = [":resources"],
+    deps = [
+        ":template",
+        "//third_party/go:heredoc",
+        "//third_party/go:stretchr_testify",
     ],
 )
 

pkg/template/load.go

@@ -5,6 +5,7 @@ import (
 	"embed"
 	"errors"
 	"fmt"
+	"io"
 	"io/ioutil"
 	"log"
 	"os"
@@ -15,6 +16,7 @@ import (
 
 	jsonpatch "github.com/evanphx/json-patch"
 	"github.com/ghodss/yaml"
+	goyaml "gopkg.in/yaml.v3"
 )
 
 // yamlPatches holds patch yaml files as static assets
@@ -51,7 +53,10 @@ func loadYAMLFile(path string) (ResourceDocs, error) {
 		return nil, fmt.Errorf("could not read file at path %s: %w", path, err)
 	}
 	resFileYamlDocs := ResourceDocs{}
-	yamlByteParts := bytes.Split(targetYAML, []byte(`---`))
+	yamlByteParts, err := splitYAML(targetYAML)
+	if err != nil {
+		return nil, fmt.Errorf("could not split YAML at path %s: %w", path, err)
+	}
 	yamlParts := ResourceDocs{}
 	for _, bytePart := range yamlByteParts {
 		yamlParts = append(yamlParts, ResourceDoc(bytePart))
@@ -77,6 +82,28 @@ func loadYAMLFile(path string) (ResourceDocs, error) {
 	return resFileYamlDocs, nil
 }
 
+func splitYAML(targetYAML []byte) ([][]byte, error) {
+	dec := goyaml.NewDecoder(bytes.NewReader(targetYAML))
+
+	var res [][]byte
+	for {
+		var value interface{}
+		err := dec.Decode(&value)
+		if err == io.EOF {
+			break
+		}
+		if err != nil {
+			return nil, err
+		}
+		valueBytes, err := goyaml.Marshal(value)
+		if err != nil {
+			return nil, err
+		}
+		res = append(res, valueBytes)
+	}
+	return res, nil
+}
+
 // PatchKindYAMLDocs stores all of the jsonpatch yaml docs found by type
 type PatchKindYAMLDocs map[string][]jsonpatch.Patch
 

pkg/template/load_test.go

@@ -0,0 +1,111 @@
+package template
+
+import (
+	"testing"
+
+	heredoc "github.com/makenowjust/heredoc/v2"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestSplitYAML(t *testing.T) {
+	var tests = []struct {
+		desc     string
+		inBytes  []byte
+		outBytes [][]byte
+		err      error
+	}{
+		{
+			desc: "a simple YAML file with no leading document separator",
+			inBytes: []byte(heredoc.Doc(`
+			a: Easy
+			b:
+			    c: 2
+			    d:
+			        - 3
+			        - 4
+			`)),
+			outBytes: [][]byte{[]byte(heredoc.Doc(`
+			a: Easy
+			b:
+			    c: 2
+			    d:
+			        - 3
+			        - 4
+			`))},
+			err: nil,
+		},
+		{
+			desc: "a simple YAML file with a leading document separator",
+			inBytes: []byte(heredoc.Doc(`
+			---
+			a: Easy
+			b:
+			    c: 2
+			    d:
+			        - 3
+			        - 4
+			`)),
+			outBytes: [][]byte{[]byte(heredoc.Doc(`
+			a: Easy
+			b:
+			    c: 2
+			    d:
+			        - 3
+			        - 4
+			`))},
+			err: nil,
+		},
+		{
+			desc: "a YAML file with no leading document separator but a string value containing ---",
+			inBytes: []byte(heredoc.Doc(`
+			a: Easy
+			b:
+			    c: 2 ---
+			    d:
+			        - 3
+			        - 4
+			`)),
+			outBytes: [][]byte{[]byte(heredoc.Doc(`
+			a: Easy
+			b:
+			    c: 2 ---
+			    d:
+			        - 3
+			        - 4
+			`))},
+			err: nil,
+		},
+		{
+			desc: "a YAML file with both a leading document separator and a string value containing ---",
+			inBytes: []byte(heredoc.Doc(`
+			---
+			a: Easy
+			b:
+			    c: 2 ---
+			    d:
+			        - 3
+			        - 4
+			`)),
+			outBytes: [][]byte{[]byte(heredoc.Doc(`
+			a: Easy
+			b:
+			    c: 2 ---
+			    d:
+			        - 3
+			        - 4
+			`))},
+			err: nil,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.desc, func(t *testing.T) {
+			outBytes, err := splitYAML(tt.inBytes)
+			if tt.err == nil {
+				assert.NoError(t, err)
+			} else {
+				assert.Error(t, err)
+			}
+			assert.Equal(t, tt.outBytes, outBytes)
+		})
+	}
+}

third_party/go/BUILD

@@ -754,3 +754,11 @@ go_module(
     module = "github.com/owenrumney/go-sarif/v2",
     version = "v2.1.2",
 )
+
+go_module(
+    name = "heredoc",
+    hashes = ["4411c83041d1ce5c59829c809861c95065fd687f83154e28441eea6f0ccec1f7"],
+    licences = ["MIT"],
+    module = "github.com/makenowjust/heredoc/v2",
+    version = "v2.0.1",
+)