Zack Williams | e940c7a | 2019-08-21 14:25:39 -0700 | [diff] [blame] | 1 | /* |
| 2 | Copyright 2014 The Kubernetes Authors. |
| 3 | |
| 4 | Licensed under the Apache License, Version 2.0 (the "License"); |
| 5 | you may not use this file except in compliance with the License. |
| 6 | You may obtain a copy of the License at |
| 7 | |
| 8 | http://www.apache.org/licenses/LICENSE-2.0 |
| 9 | |
| 10 | Unless required by applicable law or agreed to in writing, software |
| 11 | distributed under the License is distributed on an "AS IS" BASIS, |
| 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 13 | See the License for the specific language governing permissions and |
| 14 | limitations under the License. |
| 15 | */ |
| 16 | |
| 17 | package yaml |
| 18 | |
| 19 | import ( |
| 20 | "bufio" |
| 21 | "bytes" |
| 22 | "encoding/json" |
| 23 | "fmt" |
| 24 | "io" |
| 25 | "io/ioutil" |
| 26 | "strings" |
| 27 | "unicode" |
| 28 | |
| 29 | "github.com/ghodss/yaml" |
| 30 | "github.com/golang/glog" |
| 31 | ) |
| 32 | |
| 33 | // ToJSON converts a single YAML document into a JSON document |
| 34 | // or returns an error. If the document appears to be JSON the |
| 35 | // YAML decoding path is not used (so that error messages are |
| 36 | // JSON specific). |
| 37 | func ToJSON(data []byte) ([]byte, error) { |
| 38 | if hasJSONPrefix(data) { |
| 39 | return data, nil |
| 40 | } |
| 41 | return yaml.YAMLToJSON(data) |
| 42 | } |
| 43 | |
| 44 | // YAMLToJSONDecoder decodes YAML documents from an io.Reader by |
| 45 | // separating individual documents. It first converts the YAML |
| 46 | // body to JSON, then unmarshals the JSON. |
| 47 | type YAMLToJSONDecoder struct { |
| 48 | reader Reader |
| 49 | } |
| 50 | |
| 51 | // NewYAMLToJSONDecoder decodes YAML documents from the provided |
| 52 | // stream in chunks by converting each document (as defined by |
| 53 | // the YAML spec) into its own chunk, converting it to JSON via |
| 54 | // yaml.YAMLToJSON, and then passing it to json.Decoder. |
| 55 | func NewYAMLToJSONDecoder(r io.Reader) *YAMLToJSONDecoder { |
| 56 | reader := bufio.NewReader(r) |
| 57 | return &YAMLToJSONDecoder{ |
| 58 | reader: NewYAMLReader(reader), |
| 59 | } |
| 60 | } |
| 61 | |
| 62 | // Decode reads a YAML document as JSON from the stream or returns |
| 63 | // an error. The decoding rules match json.Unmarshal, not |
| 64 | // yaml.Unmarshal. |
| 65 | func (d *YAMLToJSONDecoder) Decode(into interface{}) error { |
| 66 | bytes, err := d.reader.Read() |
| 67 | if err != nil && err != io.EOF { |
| 68 | return err |
| 69 | } |
| 70 | |
| 71 | if len(bytes) != 0 { |
| 72 | err := yaml.Unmarshal(bytes, into) |
| 73 | if err != nil { |
| 74 | return YAMLSyntaxError{err} |
| 75 | } |
| 76 | } |
| 77 | return err |
| 78 | } |
| 79 | |
| 80 | // YAMLDecoder reads chunks of objects and returns ErrShortBuffer if |
| 81 | // the data is not sufficient. |
| 82 | type YAMLDecoder struct { |
| 83 | r io.ReadCloser |
| 84 | scanner *bufio.Scanner |
| 85 | remaining []byte |
| 86 | } |
| 87 | |
| 88 | // NewDocumentDecoder decodes YAML documents from the provided |
| 89 | // stream in chunks by converting each document (as defined by |
| 90 | // the YAML spec) into its own chunk. io.ErrShortBuffer will be |
| 91 | // returned if the entire buffer could not be read to assist |
| 92 | // the caller in framing the chunk. |
| 93 | func NewDocumentDecoder(r io.ReadCloser) io.ReadCloser { |
| 94 | scanner := bufio.NewScanner(r) |
| 95 | scanner.Split(splitYAMLDocument) |
| 96 | return &YAMLDecoder{ |
| 97 | r: r, |
| 98 | scanner: scanner, |
| 99 | } |
| 100 | } |
| 101 | |
| 102 | // Read reads the previous slice into the buffer, or attempts to read |
| 103 | // the next chunk. |
| 104 | // TODO: switch to readline approach. |
| 105 | func (d *YAMLDecoder) Read(data []byte) (n int, err error) { |
| 106 | left := len(d.remaining) |
| 107 | if left == 0 { |
| 108 | // return the next chunk from the stream |
| 109 | if !d.scanner.Scan() { |
| 110 | err := d.scanner.Err() |
| 111 | if err == nil { |
| 112 | err = io.EOF |
| 113 | } |
| 114 | return 0, err |
| 115 | } |
| 116 | out := d.scanner.Bytes() |
| 117 | d.remaining = out |
| 118 | left = len(out) |
| 119 | } |
| 120 | |
| 121 | // fits within data |
| 122 | if left <= len(data) { |
| 123 | copy(data, d.remaining) |
| 124 | d.remaining = nil |
| 125 | return left, nil |
| 126 | } |
| 127 | |
| 128 | // caller will need to reread |
| 129 | copy(data, d.remaining[:len(data)]) |
| 130 | d.remaining = d.remaining[len(data):] |
| 131 | return len(data), io.ErrShortBuffer |
| 132 | } |
| 133 | |
| 134 | func (d *YAMLDecoder) Close() error { |
| 135 | return d.r.Close() |
| 136 | } |
| 137 | |
| 138 | const yamlSeparator = "\n---" |
| 139 | const separator = "---" |
| 140 | |
| 141 | // splitYAMLDocument is a bufio.SplitFunc for splitting YAML streams into individual documents. |
| 142 | func splitYAMLDocument(data []byte, atEOF bool) (advance int, token []byte, err error) { |
| 143 | if atEOF && len(data) == 0 { |
| 144 | return 0, nil, nil |
| 145 | } |
| 146 | sep := len([]byte(yamlSeparator)) |
| 147 | if i := bytes.Index(data, []byte(yamlSeparator)); i >= 0 { |
| 148 | // We have a potential document terminator |
| 149 | i += sep |
| 150 | after := data[i:] |
| 151 | if len(after) == 0 { |
| 152 | // we can't read any more characters |
| 153 | if atEOF { |
| 154 | return len(data), data[:len(data)-sep], nil |
| 155 | } |
| 156 | return 0, nil, nil |
| 157 | } |
| 158 | if j := bytes.IndexByte(after, '\n'); j >= 0 { |
| 159 | return i + j + 1, data[0 : i-sep], nil |
| 160 | } |
| 161 | return 0, nil, nil |
| 162 | } |
| 163 | // If we're at EOF, we have a final, non-terminated line. Return it. |
| 164 | if atEOF { |
| 165 | return len(data), data, nil |
| 166 | } |
| 167 | // Request more data. |
| 168 | return 0, nil, nil |
| 169 | } |
| 170 | |
| 171 | // decoder is a convenience interface for Decode. |
| 172 | type decoder interface { |
| 173 | Decode(into interface{}) error |
| 174 | } |
| 175 | |
| 176 | // YAMLOrJSONDecoder attempts to decode a stream of JSON documents or |
| 177 | // YAML documents by sniffing for a leading { character. |
| 178 | type YAMLOrJSONDecoder struct { |
| 179 | r io.Reader |
| 180 | bufferSize int |
| 181 | |
| 182 | decoder decoder |
| 183 | rawData []byte |
| 184 | } |
| 185 | |
| 186 | type JSONSyntaxError struct { |
| 187 | Line int |
| 188 | Err error |
| 189 | } |
| 190 | |
| 191 | func (e JSONSyntaxError) Error() string { |
| 192 | return fmt.Sprintf("json: line %d: %s", e.Line, e.Err.Error()) |
| 193 | } |
| 194 | |
| 195 | type YAMLSyntaxError struct { |
| 196 | err error |
| 197 | } |
| 198 | |
| 199 | func (e YAMLSyntaxError) Error() string { |
| 200 | return e.err.Error() |
| 201 | } |
| 202 | |
| 203 | // NewYAMLOrJSONDecoder returns a decoder that will process YAML documents |
| 204 | // or JSON documents from the given reader as a stream. bufferSize determines |
| 205 | // how far into the stream the decoder will look to figure out whether this |
| 206 | // is a JSON stream (has whitespace followed by an open brace). |
| 207 | func NewYAMLOrJSONDecoder(r io.Reader, bufferSize int) *YAMLOrJSONDecoder { |
| 208 | return &YAMLOrJSONDecoder{ |
| 209 | r: r, |
| 210 | bufferSize: bufferSize, |
| 211 | } |
| 212 | } |
| 213 | |
| 214 | // Decode unmarshals the next object from the underlying stream into the |
| 215 | // provide object, or returns an error. |
| 216 | func (d *YAMLOrJSONDecoder) Decode(into interface{}) error { |
| 217 | if d.decoder == nil { |
| 218 | buffer, origData, isJSON := GuessJSONStream(d.r, d.bufferSize) |
| 219 | if isJSON { |
| 220 | glog.V(4).Infof("decoding stream as JSON") |
| 221 | d.decoder = json.NewDecoder(buffer) |
| 222 | d.rawData = origData |
| 223 | } else { |
| 224 | glog.V(4).Infof("decoding stream as YAML") |
| 225 | d.decoder = NewYAMLToJSONDecoder(buffer) |
| 226 | } |
| 227 | } |
| 228 | err := d.decoder.Decode(into) |
| 229 | if jsonDecoder, ok := d.decoder.(*json.Decoder); ok { |
| 230 | if syntax, ok := err.(*json.SyntaxError); ok { |
| 231 | data, readErr := ioutil.ReadAll(jsonDecoder.Buffered()) |
| 232 | if readErr != nil { |
| 233 | glog.V(4).Infof("reading stream failed: %v", readErr) |
| 234 | } |
| 235 | js := string(data) |
| 236 | |
| 237 | // if contents from io.Reader are not complete, |
| 238 | // use the original raw data to prevent panic |
| 239 | if int64(len(js)) <= syntax.Offset { |
| 240 | js = string(d.rawData) |
| 241 | } |
| 242 | |
| 243 | start := strings.LastIndex(js[:syntax.Offset], "\n") + 1 |
| 244 | line := strings.Count(js[:start], "\n") |
| 245 | return JSONSyntaxError{ |
| 246 | Line: line, |
| 247 | Err: fmt.Errorf(syntax.Error()), |
| 248 | } |
| 249 | } |
| 250 | } |
| 251 | return err |
| 252 | } |
| 253 | |
| 254 | type Reader interface { |
| 255 | Read() ([]byte, error) |
| 256 | } |
| 257 | |
| 258 | type YAMLReader struct { |
| 259 | reader Reader |
| 260 | } |
| 261 | |
| 262 | func NewYAMLReader(r *bufio.Reader) *YAMLReader { |
| 263 | return &YAMLReader{ |
| 264 | reader: &LineReader{reader: r}, |
| 265 | } |
| 266 | } |
| 267 | |
| 268 | // Read returns a full YAML document. |
| 269 | func (r *YAMLReader) Read() ([]byte, error) { |
| 270 | var buffer bytes.Buffer |
| 271 | for { |
| 272 | line, err := r.reader.Read() |
| 273 | if err != nil && err != io.EOF { |
| 274 | return nil, err |
| 275 | } |
| 276 | |
| 277 | sep := len([]byte(separator)) |
| 278 | if i := bytes.Index(line, []byte(separator)); i == 0 { |
| 279 | // We have a potential document terminator |
| 280 | i += sep |
| 281 | after := line[i:] |
| 282 | if len(strings.TrimRightFunc(string(after), unicode.IsSpace)) == 0 { |
| 283 | if buffer.Len() != 0 { |
| 284 | return buffer.Bytes(), nil |
| 285 | } |
| 286 | if err == io.EOF { |
| 287 | return nil, err |
| 288 | } |
| 289 | } |
| 290 | } |
| 291 | if err == io.EOF { |
| 292 | if buffer.Len() != 0 { |
| 293 | // If we're at EOF, we have a final, non-terminated line. Return it. |
| 294 | return buffer.Bytes(), nil |
| 295 | } |
| 296 | return nil, err |
| 297 | } |
| 298 | buffer.Write(line) |
| 299 | } |
| 300 | } |
| 301 | |
| 302 | type LineReader struct { |
| 303 | reader *bufio.Reader |
| 304 | } |
| 305 | |
| 306 | // Read returns a single line (with '\n' ended) from the underlying reader. |
| 307 | // An error is returned iff there is an error with the underlying reader. |
| 308 | func (r *LineReader) Read() ([]byte, error) { |
| 309 | var ( |
| 310 | isPrefix bool = true |
| 311 | err error = nil |
| 312 | line []byte |
| 313 | buffer bytes.Buffer |
| 314 | ) |
| 315 | |
| 316 | for isPrefix && err == nil { |
| 317 | line, isPrefix, err = r.reader.ReadLine() |
| 318 | buffer.Write(line) |
| 319 | } |
| 320 | buffer.WriteByte('\n') |
| 321 | return buffer.Bytes(), err |
| 322 | } |
| 323 | |
| 324 | // GuessJSONStream scans the provided reader up to size, looking |
| 325 | // for an open brace indicating this is JSON. It will return the |
| 326 | // bufio.Reader it creates for the consumer. |
| 327 | func GuessJSONStream(r io.Reader, size int) (io.Reader, []byte, bool) { |
| 328 | buffer := bufio.NewReaderSize(r, size) |
| 329 | b, _ := buffer.Peek(size) |
| 330 | return buffer, b, hasJSONPrefix(b) |
| 331 | } |
| 332 | |
| 333 | var jsonPrefix = []byte("{") |
| 334 | |
| 335 | // hasJSONPrefix returns true if the provided buffer appears to start with |
| 336 | // a JSON open brace. |
| 337 | func hasJSONPrefix(buf []byte) bool { |
| 338 | return hasPrefix(buf, jsonPrefix) |
| 339 | } |
| 340 | |
| 341 | // Return true if the first non-whitespace bytes in buf is |
| 342 | // prefix. |
| 343 | func hasPrefix(buf []byte, prefix []byte) bool { |
| 344 | trim := bytes.TrimLeftFunc(buf, unicode.IsSpace) |
| 345 | return bytes.HasPrefix(trim, prefix) |
| 346 | } |