From 919bb01b04324bd78af12609f281e255252cccfc Mon Sep 17 00:00:00 2001
From: AdoHe <coolhzb@163.com>
Date: Fri, 7 Oct 2016 11:22:44 +0800
Subject: [PATCH] fix yaml decode issue

---
 pkg/util/yaml/decoder.go      | 26 +++++++++++++++++++-------
 pkg/util/yaml/decoder_test.go |  8 ++++----
 2 files changed, 23 insertions(+), 11 deletions(-)

diff --git a/pkg/util/yaml/decoder.go b/pkg/util/yaml/decoder.go
index c65c2d6ba34..ed503d03905 100644
--- a/pkg/util/yaml/decoder.go
+++ b/pkg/util/yaml/decoder.go
@@ -137,7 +137,7 @@ func (d *YAMLDecoder) Close() error {
 }
 
 const yamlSeparator = "\n---"
-const separator = "---\n"
+const separator = "---"
 
 // splitYAMLDocument is a bufio.SplitFunc for splitting YAML streams into individual documents.
 func splitYAMLDocument(data []byte, atEOF bool) (advance int, token []byte, err error) {
@@ -246,16 +246,28 @@ func (r *YAMLReader) Read() ([]byte, error) {
 			return nil, err
 		}
 
-		if string(line) == separator || err == io.EOF {
+		sep := len([]byte(separator))
+		if i := bytes.Index(line, []byte(separator)); i == 0 {
+			// We have a potential document terminator
+			i += sep
+			after := line[i:]
+			if len(strings.TrimRightFunc(string(after), unicode.IsSpace)) == 0 {
+				if buffer.Len() != 0 {
+					return buffer.Bytes(), nil
+				}
+				if err == io.EOF {
+					return nil, err
+				}
+			}
+		}
+		if err == io.EOF {
 			if buffer.Len() != 0 {
+				// If we're at EOF, we have a final, non-terminated line. Return it.
 				return buffer.Bytes(), nil
 			}
-			if err == io.EOF {
-				return nil, err
-			}
-		} else {
-			buffer.Write(line)
+			return nil, err
 		}
+		buffer.Write(line)
 	}
 }
 
diff --git a/pkg/util/yaml/decoder_test.go b/pkg/util/yaml/decoder_test.go
index 8985907c456..b8c5e6987e3 100644
--- a/pkg/util/yaml/decoder_test.go
+++ b/pkg/util/yaml/decoder_test.go
@@ -106,7 +106,7 @@ func TestDecodeYAML(t *testing.T) {
 	s := NewYAMLToJSONDecoder(bytes.NewReader([]byte(`---
 stuff: 1
 
----
+---   
   `)))
 	obj := generic{}
 	if err := s.Decode(&obj); err != nil {
@@ -138,11 +138,11 @@ stuff: 1
 	obj := generic{}
 	err := s.Decode(&obj)
 	if err == nil {
-		t.Fatal("expected error with yaml: prefix, got no error")
+		t.Fatal("expected error with yaml: violate, got no error")
 	}
 	fmt.Printf("err: %s\n", err.Error())
-	if !strings.HasPrefix(err.Error(), "yaml: line 1:") {
-		t.Fatalf("expected %q to have 'yaml: line 1:' prefix", err.Error())
+	if !strings.HasPrefix(err.Error(), "yaml: line 2:") {
+		t.Fatalf("expected %q to have 'yaml: line 2:' found a tab character", err.Error())
 	}
 }