Browse Source

Harden yaml stream decoding in `fluctl save`

Use go-yaml's Decoder instead of our own
Alfonso Acosta 9 months ago
parent
commit
8edd13f0f6
2 changed files with 12 additions and 47 deletions
  1. 3
    3
      cluster/kubernetes/resource/load.go
  2. 9
    44
      cmd/fluxctl/save_cmd.go

+ 3
- 3
cluster/kubernetes/resource/load.go View File

@@ -134,9 +134,9 @@ func ParseMultidoc(multidoc []byte, source string) (map[string]KubeManifest, err
134 134
 	var err error
135 135
 	for {
136 136
 		// In order to use the decoder to extract raw documents
137
-		// from the stream, we decode generically and encode again
138
-		// The result it the raw document (pretty-printed and
139
-		// without comments though)
137
+		// from the stream, we decode generically and encode again.
138
+		// The result is the raw document from the stream
139
+		// (pretty-printed and without comments)
140 140
 		var val interface{}
141 141
 		if err := decoder.Decode(&val); err != nil {
142 142
 			break

+ 9
- 44
cmd/fluxctl/save_cmd.go View File

@@ -1,7 +1,6 @@
1 1
 package main
2 2
 
3 3
 import (
4
-	"bufio"
5 4
 	"bytes"
6 5
 	"context"
7 6
 	"fmt"
@@ -11,7 +10,7 @@ import (
11 10
 
12 11
 	"github.com/pkg/errors"
13 12
 	"github.com/spf13/cobra"
14
-	yaml "gopkg.in/yaml.v2"
13
+	"gopkg.in/yaml.v2"
15 14
 )
16 15
 
17 16
 type saveOpts struct {
@@ -63,9 +62,6 @@ func (opts *saveOpts) RunE(cmd *cobra.Command, args []string) error {
63 62
 		return errors.Wrap(err, "exporting config")
64 63
 	}
65 64
 
66
-	yamls := bufio.NewScanner(bytes.NewReader(config))
67
-	yamls.Split(splitYAMLDocument)
68
-
69 65
 	if opts.path != "-" {
70 66
 		// check supplied path is a directory
71 67
 		if info, err := os.Stat(opts.path); err != nil {
@@ -75,11 +71,14 @@ func (opts *saveOpts) RunE(cmd *cobra.Command, args []string) error {
75 71
 		}
76 72
 	}
77 73
 
78
-	for yamls.Scan() {
74
+	decoder := yaml.NewDecoder(bytes.NewReader(config))
75
+
76
+	var decoderErr error
77
+	for {
79 78
 		var object saveObject
80 79
 		// Most unwanted fields are ignored at this point
81
-		if err := yaml.Unmarshal(yamls.Bytes(), &object); err != nil {
82
-			return errors.Wrap(err, "unmarshalling exported yaml")
80
+		if decoderErr = decoder.Decode(&object); decoderErr != nil {
81
+			break
83 82
 		}
84 83
 
85 84
 		// Filter out remaining unwanted keys from unstructured fields
@@ -91,8 +90,8 @@ func (opts *saveOpts) RunE(cmd *cobra.Command, args []string) error {
91 90
 		}
92 91
 	}
93 92
 
94
-	if yamls.Err() != nil {
95
-		return errors.Wrap(yamls.Err(), "splitting exported yaml")
93
+	if decoderErr != io.EOF {
94
+		return errors.Wrap(err, "unmarshalling exported yaml")
96 95
 	}
97 96
 
98 97
 	return nil
@@ -211,37 +210,3 @@ func abbreviateKind(kind string) string {
211 210
 		return kind
212 211
 	}
213 212
 }
214
-
215
-// Copied from k8s.io/client-go/1.5/pkg/util/yaml/decoder.go
216
-
217
-const yamlSeparator = "\n---"
218
-
219
-// splitYAMLDocument is a bufio.SplitFunc for splitting YAML streams into individual documents.
220
-func splitYAMLDocument(data []byte, atEOF bool) (advance int, token []byte, err error) {
221
-	if atEOF && len(data) == 0 {
222
-		return 0, nil, nil
223
-	}
224
-	sep := len([]byte(yamlSeparator))
225
-	if i := bytes.Index(data, []byte(yamlSeparator)); i >= 0 {
226
-		// We have a potential document terminator
227
-		i += sep
228
-		after := data[i:]
229
-		if len(after) == 0 {
230
-			// we can't read any more characters
231
-			if atEOF {
232
-				return len(data), data[:len(data)-sep], nil
233
-			}
234
-			return 0, nil, nil
235
-		}
236
-		if j := bytes.IndexByte(after, '\n'); j >= 0 {
237
-			return i + j + 1, data[0 : i-sep], nil
238
-		}
239
-		return 0, nil, nil
240
-	}
241
-	// If we're at EOF, we have a final, non-terminated line. Return it.
242
-	if atEOF {
243
-		return len(data), data, nil
244
-	}
245
-	// Request more data.
246
-	return 0, nil, nil
247
-}

Loading…
Cancel
Save