fix: split YAML documents with the Kubernetes YAML reader

approvals/kir_test.TestFailure.BadYAML.stderr.approved.txt

@@ -1 +1 @@
-error: error processing document: yaml: line 9: did not find expected ',' or ']'
+error: yaml: line 9: did not find expected ',' or ']'

processor/processor.go

@@ -1,45 +1,25 @@
 package processor
 
 import (
-	"bytes"
 	"fmt"
 	"io"
 	"os"
 
 	"github.com/mpv/kir/yamlparser"
 )
 
-// ProcessStdin reads a manifest stream from r and returns its images. Taking an
-// io.Reader (rather than reading os.Stdin directly) keeps it testable and lets
-// the CLI inject stdin.
+// ProcessStdin reads a (possibly multi-document) manifest stream from r and
+// returns its images. Taking an io.Reader keeps it testable and lets the CLI
+// inject stdin.
 func ProcessStdin(r io.Reader) ([]string, error) {
-	data, err := io.ReadAll(r)
-	if err != nil {
-		return nil, fmt.Errorf("error reading stdin: %v", err)
-	}
-	return processDocuments(data)
+	return yamlparser.ProcessReader(r)
 }
 
 func ProcessFile(filePath string) ([]string, error) {
-	data, err := os.ReadFile(filePath)
+	file, err := os.Open(filePath)
 	if err != nil {
 		return nil, fmt.Errorf("error reading file: %v", err)
 	}
-	return processDocuments(data)
-}
-
-// processDocuments splits a (possibly multi-document) YAML stream and collects
-// the images from every document. Both the file and stdin paths go through it
-// so they handle multi-document input identically.
-func processDocuments(data []byte) ([]string, error) {
-	var images []string
-	docs := bytes.Split(data, []byte("\n---\n"))
-	for _, doc := range docs {
-		imgs, err := yamlparser.ProcessData(doc)
-		if err != nil {
-			return nil, fmt.Errorf("error processing document: %v", err)
-		}
-		images = append(images, imgs...)
-	}
-	return images, nil
+	defer file.Close()
+	return yamlparser.ProcessReader(file)
 }

yamlparser/processreader_test.go

@@ -0,0 +1,68 @@
+package yamlparser
+
+import (
+	"strings"
+	"testing"
+)
+
+// ProcessReader must collect images from every document in a stream and split
+// documents robustly. A naive bytes.Split on "\n---\n" drops every document
+// after the first when the separator isn't exactly that literal — a separator
+// with trailing whitespace, or CRLF line endings. The leading-separator and
+// no-trailing-newline cases are general correctness checks: a plain bytes.Split
+// happens to handle those, but the YAML reader must too.
+func TestProcessReader(t *testing.T) {
+	pod := func(name, image string) string {
+		return "apiVersion: v1\nkind: Pod\nmetadata:\n  name: " + name +
+			"\nspec:\n  containers:\n  - name: c\n    image: " + image + "\n"
+	}
+
+	tests := []struct {
+		name string
+		data string
+		want []string
+	}{
+		{
+			name: "multiple documents",
+			data: pod("one", "image-one") + "---\n" + pod("two", "image-two"),
+			want: []string{"image-one", "image-two"},
+		},
+		{
+			name: "leading separator",
+			data: "---\n" + pod("one", "image-one"),
+			want: []string{"image-one"},
+		},
+		{
+			name: "separator with trailing whitespace",
+			data: pod("one", "image-one") + "---   \n" + pod("two", "image-two"),
+			want: []string{"image-one", "image-two"},
+		},
+		{
+			name: "crlf line endings",
+			data: strings.ReplaceAll(pod("one", "image-one")+"---\n"+pod("two", "image-two"), "\n", "\r\n"),
+			want: []string{"image-one", "image-two"},
+		},
+		{
+			name: "no trailing newline",
+			data: strings.TrimRight(pod("one", "image-one"), "\n"),
+			want: []string{"image-one"},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := ProcessReader(strings.NewReader(tt.data))
+			if err != nil {
+				t.Fatalf("ProcessReader() error = %v", err)
+			}
+			if len(got) != len(tt.want) {
+				t.Fatalf("expected %d images, got %d: %v", len(tt.want), len(got), got)
+			}
+			for i, img := range got {
+				if img != tt.want[i] {
+					t.Errorf("image %d: expected %q, got %q", i, tt.want[i], img)
+				}
+			}
+		})
+	}
+}

yamlparser/yamlparser.go

@@ -1,19 +1,47 @@
 package yamlparser
 
 import (
+	"bufio"
 	"fmt"
+	"io"
 	"slices"
 
 	"github.com/mpv/kir/k8s"
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/apimachinery/pkg/runtime/serializer"
+	utilyaml "k8s.io/apimachinery/pkg/util/yaml"
 	"k8s.io/client-go/kubernetes/scheme"
 )
 
 var supportedKinds = []string{"Pod", "Deployment", "DaemonSet", "ReplicaSet", "StatefulSet", "Job", "CronJob"}
 
+// ProcessReader reads a (possibly multi-document) YAML stream and returns the
+// container images of every supported workload it contains. Documents are
+// separated using the Kubernetes YAML reader, which correctly handles leading
+// and trailing "---" separators, separators followed by trailing whitespace,
+// CRLF line endings, and a final document without a trailing newline.
+func ProcessReader(r io.Reader) ([]string, error) {
+	var images []string
+	reader := utilyaml.NewYAMLReader(bufio.NewReader(r))
+	for {
+		doc, err := reader.Read()
+		if err == io.EOF {
+			break
+		}
+		if err != nil {
+			return nil, fmt.Errorf("error reading YAML document: %v", err)
+		}
+		imgs, err := ProcessData(doc)
+		if err != nil {
+			return nil, err
+		}
+		images = append(images, imgs...)
+	}
+	return images, nil
+}
+
 func ProcessData(data []byte) ([]string, error) {
 	// Decode the YAML file into a Kubernetes object
 	decode := serializer.NewCodecFactory(scheme.Scheme).UniversalDeserializer().Decode