Ver código fonte

chore: optimize yaml unmarshaler (#1513)

chenquan 3 anos atrás
pai
commit
05cc62f5ff

+ 10 - 6
core/mapping/yamlunmarshaler.go

@@ -3,10 +3,8 @@ package mapping
 import (
 	"encoding/json"
 	"errors"
-	"io"
-	"io/ioutil"
-
 	"gopkg.in/yaml.v2"
+	"io"
 )
 
 // To make .json & .yaml consistent, we just use json as the tag key.
@@ -43,12 +41,18 @@ func unmarshalYamlBytes(content []byte, v interface{}, unmarshaler *Unmarshaler)
 }
 
 func unmarshalYamlReader(reader io.Reader, v interface{}, unmarshaler *Unmarshaler) error {
-	content, err := ioutil.ReadAll(reader)
-	if err != nil {
+	var res interface{}
+	if err := yaml.NewDecoder(reader).Decode(&res); err != nil {
 		return err
 	}
 
-	return unmarshalYamlBytes(content, v, unmarshaler)
+	out := cleanupMapValue(res)
+
+	if m, ok := out.(map[string]interface{}); ok {
+		return unmarshaler.Unmarshal(m, v)
+	}
+
+	return ErrUnsupportedType
 }
 
 // yamlUnmarshal YAML to map[string]interface{} instead of map[interface{}]interface{}.

+ 7 - 3
core/mapping/yamlunmarshaler_test.go

@@ -926,14 +926,18 @@ func TestUnmarshalYamlBytesError(t *testing.T) {
 }
 
 func TestUnmarshalYamlReaderError(t *testing.T) {
-	payload := `abcd: cdef`
-	reader := strings.NewReader(payload)
 	var v struct {
 		Any string
 	}
 
+	reader := strings.NewReader(`abcd: cdef`)
 	err := UnmarshalYamlReader(reader, &v)
 	assert.NotNil(t, err)
+
+	reader = strings.NewReader("chenquan")
+	err = UnmarshalYamlReader(reader, &v)
+	assert.ErrorIs(t, err, ErrUnsupportedType)
+
 }
 
 func TestUnmarshalYamlBadReader(t *testing.T) {
@@ -1011,6 +1015,6 @@ func TestUnmarshalYamlMapRune(t *testing.T) {
 
 type badReader struct{}
 
-func (b *badReader) Read(p []byte) (n int, err error) {
+func (b *badReader) Read(_ []byte) (n int, err error) {
 	return 0, io.ErrLimitReached
 }