Fixing bug when parsing YAML documents with anchors

Signed-off-by: Dave Henderson <dhenderson@gmail.com>

2 files changed, 202 insertions, 9 deletions

diff --git a/data/data.go b/data/data.go
index b449b375..90639e02 100644
--- a/data/data.go
+++ b/data/data.go
@@ -8,6 +8,7 @@ import (
 	"bytes"
 	"encoding/csv"
 	"encoding/json"
+	"fmt"
 	"io"
 	"strings"
 
@@ -100,7 +101,9 @@ func YAML(in string) (map[string]interface{}, error) {
 			break
 		}
 	}
-	return obj, nil
+
+	err := stringifyYAMLMapMapKeys(obj)
+	return obj, err
 }
 
 // YAMLArray - Unmarshal a YAML Array
@@ -120,7 +123,64 @@ func YAMLArray(in string) ([]interface{}, error) {
 			break
 		}
 	}
-	return obj, nil
+	err := stringifyYAMLArrayMapKeys(obj)
+	return obj, err
+}
+
+// stringifyYAMLArrayMapKeys recurses into the input array and changes all
+// non-string map keys to string map keys. Modifies the input array.
+func stringifyYAMLArrayMapKeys(in []interface{}) error {
+	if _, changed := stringifyMapKeys(in); changed {
+		return fmt.Errorf("stringifyYAMLArrayMapKeys: output type did not match input type, this should be impossible")
+	}
+	return nil
+}
+
+// stringifyYAMLMapMapKeys recurses into the input map and changes all
+// non-string map keys to string map keys. Modifies the input map.
+func stringifyYAMLMapMapKeys(in map[string]interface{}) error {
+	if _, changed := stringifyMapKeys(in); changed {
+		return fmt.Errorf("stringifyYAMLMapMapKeys: output type did not match input type, this should be impossible")
+	}
+	return nil
+}
+
+// stringifyMapKeys recurses into in and changes all instances of
+// map[interface{}]interface{} to map[string]interface{}. This is useful to
+// work around the impedance mismatch between JSON and YAML unmarshaling that's
+// described here: https://github.com/go-yaml/yaml/issues/139
+//
+// Taken and modified from https://github.com/gohugoio/hugo/blob/cdfd1c99baa22d69e865294dfcd783811f96c880/parser/metadecoders/decoder.go#L257, Apache License 2.0
+// Originally inspired by https://github.com/stripe/stripe-mock/blob/24a2bb46a49b2a416cfea4150ab95781f69ee145/mapstr.go#L13, MIT License
+func stringifyMapKeys(in interface{}) (interface{}, bool) {
+	switch in := in.(type) {
+	case []interface{}:
+		for i, v := range in {
+			if vv, replaced := stringifyMapKeys(v); replaced {
+				in[i] = vv
+			}
+		}
+	case map[string]interface{}:
+		for k, v := range in {
+			if vv, changed := stringifyMapKeys(v); changed {
+				in[k] = vv
+			}
+		}
+	case map[interface{}]interface{}:
+		res := make(map[string]interface{})
+
+		for k, v := range in {
+			ks := conv.ToString(k)
+			if vv, replaced := stringifyMapKeys(v); replaced {
+				res[ks] = vv
+			} else {
+				res[ks] = v
+			}
+		}
+		return res, true
+	}
+
+	return nil, false
 }
 
 // TOML - Unmarshal a TOML Object
diff --git a/data/data_test.go b/data/data_test.go
index ae2c31c7..08c49015 100644
--- a/data/data_test.go
+++ b/data/data_test.go
@@ -24,15 +24,22 @@ func TestUnmarshalObj(t *testing.T) {
 
 	test := func(actual map[string]interface{}, err error) {
 		assert.NoError(t, err)
-		assert.Equal(t, expected["foo"], actual["foo"])
-		assert.Equal(t, expected["one"], actual["one"])
-		assert.Equal(t, expected["true"], actual["true"])
+		assert.Equal(t, expected["foo"], actual["foo"], "foo")
+		assert.Equal(t, expected["one"], actual["one"], "one")
+		assert.Equal(t, expected["true"], actual["true"], "true")
 	}
 	test(JSON(`{"foo":{"bar":"baz"},"one":1.0,"true":true}`))
 	test(YAML(`foo:
   bar: baz
 one: 1.0
-true: true
+'true': true
+`))
+	test(YAML(`anchor: &anchor
+  bar: baz
+foo:
+  <<: *anchor
+one: 1.0
+'true': true
 `))
 	test(YAML(`# this comment marks an empty (nil!) document
 ---
@@ -41,7 +48,7 @@ true: true
 foo:
   bar: baz
 one: 1.0
-true: true
+'true': true
 `))
 
 	obj := make(map[string]interface{})
@@ -52,7 +59,6 @@ true: true
 }
 
 func TestUnmarshalArray(t *testing.T) {
-
 	expected := []interface{}{"foo", "bar",
 		map[string]interface{}{
 			"baz":   map[string]interface{}{"qux": true},
@@ -90,8 +96,43 @@ func TestUnmarshalArray(t *testing.T) {
 this shouldn't be reached
 `))
 
+	actual, err := YAMLArray(`---
+- foo: &foo
+    bar: baz
+- qux:
+    <<: *foo
+    quux: corge
+- baz:
+    qux: true
+    42: 18
+    false: blah
+`)
+	assert.NoError(t, err)
+	assert.EqualValues(t,
+		[]interface{}{
+			map[string]interface{}{
+				"foo": map[string]interface{}{
+					"bar": "baz",
+				},
+			},
+			map[string]interface{}{
+				"qux": map[string]interface{}{
+					"bar":  "baz",
+					"quux": "corge",
+				},
+			},
+			map[string]interface{}{
+				"baz": map[string]interface{}{
+					"qux":   true,
+					"42":    18,
+					"false": "blah",
+				},
+			},
+		},
+		actual)
+
 	obj := make([]interface{}, 1)
-	_, err := unmarshalArray(obj, "SOMETHING", func(in []byte, out interface{}) error {
+	_, err = unmarshalArray(obj, "SOMETHING", func(in []byte, out interface{}) error {
 		return errors.New("fail")
 	})
 	assert.EqualError(t, err, "Unable to unmarshal array SOMETHING: fail")
@@ -523,3 +564,95 @@ QUX='single quotes ignore $variables'
 	assert.NoError(t, err)
 	assert.EqualValues(t, expected, out)
 }
+
+func TestStringifyYAMLArrayMapKeys(t *testing.T) {
+	cases := []struct {
+		input    []interface{}
+		want     []interface{}
+		replaced bool
+	}{
+		{
+			[]interface{}{map[interface{}]interface{}{"a": 1, "b": 2}},
+			[]interface{}{map[string]interface{}{"a": 1, "b": 2}},
+			false,
+		},
+		{
+			[]interface{}{map[interface{}]interface{}{"a": []interface{}{1, map[interface{}]interface{}{"b": 2}}}},
+			[]interface{}{map[string]interface{}{"a": []interface{}{1, map[string]interface{}{"b": 2}}}},
+			false,
+		},
+		{
+			[]interface{}{map[interface{}]interface{}{true: 1, "b": false}},
+			[]interface{}{map[string]interface{}{"true": 1, "b": false}},
+			false,
+		},
+		{
+			[]interface{}{map[interface{}]interface{}{1: "a", 2: "b"}},
+			[]interface{}{map[string]interface{}{"1": "a", "2": "b"}},
+			false,
+		},
+		{
+			[]interface{}{map[interface{}]interface{}{"a": map[interface{}]interface{}{"b": 1}}},
+			[]interface{}{map[string]interface{}{"a": map[string]interface{}{"b": 1}}},
+			false,
+		},
+		{
+			[]interface{}{map[string]interface{}{"a": map[string]interface{}{"b": 1}}},
+			[]interface{}{map[string]interface{}{"a": map[string]interface{}{"b": 1}}},
+			false,
+		},
+		{
+			[]interface{}{map[interface{}]interface{}{1: "a", 2: "b"}},
+			[]interface{}{map[string]interface{}{"1": "a", "2": "b"}},
+			false,
+		},
+	}
+
+	for _, c := range cases {
+		err := stringifyYAMLArrayMapKeys(c.input)
+		assert.NoError(t, err)
+		assert.EqualValues(t, c.want, c.input)
+	}
+}
+
+func TestStringifyYAMLMapMapKeys(t *testing.T) {
+	cases := []struct {
+		input map[string]interface{}
+		want  map[string]interface{}
+	}{
+		{
+			map[string]interface{}{"root": map[interface{}]interface{}{"a": 1, "b": 2}},
+			map[string]interface{}{"root": map[string]interface{}{"a": 1, "b": 2}},
+		},
+		{
+			map[string]interface{}{"root": map[interface{}]interface{}{"a": []interface{}{1, map[interface{}]interface{}{"b": 2}}}},
+			map[string]interface{}{"root": map[string]interface{}{"a": []interface{}{1, map[string]interface{}{"b": 2}}}},
+		},
+		{
+			map[string]interface{}{"root": map[interface{}]interface{}{true: 1, "b": false}},
+			map[string]interface{}{"root": map[string]interface{}{"true": 1, "b": false}},
+		},
+		{
+			map[string]interface{}{"root": map[interface{}]interface{}{1: "a", 2: "b"}},
+			map[string]interface{}{"root": map[string]interface{}{"1": "a", "2": "b"}},
+		},
+		{
+			map[string]interface{}{"root": map[interface{}]interface{}{"a": map[interface{}]interface{}{"b": 1}}},
+			map[string]interface{}{"root": map[string]interface{}{"a": map[string]interface{}{"b": 1}}},
+		},
+		{
+			map[string]interface{}{"a": map[string]interface{}{"b": 1}},
+			map[string]interface{}{"a": map[string]interface{}{"b": 1}},
+		},
+		{
+			map[string]interface{}{"root": []interface{}{map[interface{}]interface{}{1: "a", 2: "b"}}},
+			map[string]interface{}{"root": []interface{}{map[string]interface{}{"1": "a", "2": "b"}}},
+		},
+	}
+
+	for _, c := range cases {
+		err := stringifyYAMLMapMapKeys(c.input)
+		assert.NoError(t, err)
+		assert.EqualValues(t, c.want, c.input)
+	}
+}