From 78737f5a3844dfe463c81b1f9b32092423915eb0 Mon Sep 17 00:00:00 2001
From: brianpursley <bpursley@cinlogic.com>
Date: Thu, 24 Jun 2021 21:32:39 -0400
Subject: [PATCH] Updated ByteReader to allow white space and comments on the
 same line after --- and throw an error if anything else is detected

---
 kyaml/kio/byteio_reader.go      | 40 +++++++++++++++++++++--
 kyaml/kio/byteio_reader_test.go | 56 +++++++++++++++++++++++++++++++++
 2 files changed, 94 insertions(+), 2 deletions(-)

diff --git a/kyaml/kio/byteio_reader.go b/kyaml/kio/byteio_reader.go
index 7469698f4..271b26d22 100644
--- a/kyaml/kio/byteio_reader.go
+++ b/kyaml/kio/byteio_reader.go
@@ -7,6 +7,7 @@ import (
 	"bytes"
 	"fmt"
 	"io"
+	"regexp"
 	"sort"
 	"strings"
 
@@ -133,6 +134,37 @@ type ByteReader struct {
 
 var _ Reader = &ByteReader{}
 
+// splitDocuments returns a slice of all documents contained in a YAML string. Multiple documents can be divided by the
+// YAML document separator (---). It allows for white space and comments to be after the separator on the same line,
+// but will return an error if anything else is on the line.
+func splitDocuments(s string) ([]string, error) {
+	docs := make([]string, 0)
+	if len(s) > 0 {
+		// The YAML document separator is any line that starts with ---
+		yamlSeparatorRegexp := regexp.MustCompile(`\n---.*\n`)
+
+		// Find all separators, check them for invalid content, and append each document to docs
+		separatorLocations := yamlSeparatorRegexp.FindAllStringIndex(s, -1)
+		prev := 0
+		for i := range separatorLocations {
+			loc := separatorLocations[i]
+			separator := s[loc[0]:loc[1]]
+
+			// If the next non-whitespace character on the line following the separator is not a comment, return an error
+			trimmedContentAfterSeparator := strings.TrimSpace(separator[4:])
+			if len(trimmedContentAfterSeparator) > 0 && trimmedContentAfterSeparator[0] != '#' {
+				return nil, errors.Errorf("invalid document separator: %s", strings.TrimSpace(separator))
+			}
+
+			docs = append(docs, s[prev:loc[0]])
+			prev = loc[1]
+		}
+		docs = append(docs, s[prev:])
+	}
+
+	return docs, nil
+}
+
 func (r *ByteReader) Read() ([]*yaml.RNode, error) {
 	output := ResourceNodeSlice{}
 
@@ -144,8 +176,12 @@ func (r *ByteReader) Read() ([]*yaml.RNode, error) {
 		return nil, errors.Wrap(err)
 	}
 
-	// replace the ending \r\n (line ending used in windows) with \n and then separate by \n---\n
-	values := strings.Split(strings.ReplaceAll(input.String(), "\r\n", "\n"), "\n---\n")
+	// Replace the ending \r\n (line ending used in windows) with \n and then split it into multiple YAML documents
+	// if it contains document separators (---)
+	values, err := splitDocuments(strings.ReplaceAll(input.String(), "\r\n", "\n"))
+	if err != nil {
+		return nil, errors.Wrap(err)
+	}
 
 	index := 0
 	for i := range values {
diff --git a/kyaml/kio/byteio_reader_test.go b/kyaml/kio/byteio_reader_test.go
index 41c43182b..2d04e0a6b 100644
--- a/kyaml/kio/byteio_reader_test.go
+++ b/kyaml/kio/byteio_reader_test.go
@@ -372,6 +372,62 @@ metadata:
 			},
 			instance: ByteReader{},
 		},
+
+		//
+		//
+		//
+		{
+			name: "white_space_after_document_separator_should_be_ignored",
+			input: `
+a: b
+---         
+c: d
+`,
+			expectedItems: []string{
+				`
+a: b
+`,
+				`
+c: d
+`,
+			},
+			instance: ByteReader{OmitReaderAnnotations: true},
+		},
+
+		//
+		//
+		//
+		{
+			name: "comment_after_document_separator_should_be_ignored",
+			input: `
+a: b
+--- #foo
+c: d
+`,
+			expectedItems: []string{
+				`
+a: b
+`,
+				`
+c: d
+`,
+			},
+			instance: ByteReader{OmitReaderAnnotations: true},
+		},
+
+		//
+		//
+		//
+		{
+			name: "anything_after_document_separator_other_than_white_space_or_comment_is_an_error",
+			input: `
+a: b
+--- foo
+c: d
+`,
+			err:      "invalid document separator: --- foo",
+			instance: ByteReader{OmitReaderAnnotations: true},
+		},
 	}
 
 	for i := range testCases {