From 0cbf0942455f6a1d427afc912c170ae2fb7e1711 Mon Sep 17 00:00:00 2001
From: Tristan Cartledge <tristan@speakeasyapi.dev>
Date: Mon, 5 Aug 2024 14:23:41 +0100
Subject: [PATCH] fix: improved handling of complex yaml in YAMLNodeToJSON

---
 json/json.go      | 35 ++++++++++++++++++++------------
 json/json_test.go | 51 +++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 73 insertions(+), 13 deletions(-)

diff --git a/json/json.go b/json/json.go
index e73aa56..156a35b 100644
--- a/json/json.go
+++ b/json/json.go
@@ -3,14 +3,13 @@ package json
 import (
 	"encoding/json"
 	"fmt"
+	"reflect"
 
 	"github.com/pb33f/libopenapi/orderedmap"
 	"gopkg.in/yaml.v3"
 )
 
 // YAMLNodeToJSON converts yaml/json stored in a yaml.Node to json ordered matching the original yaml/json
-//
-// NOTE: The limitation is this won't work with YAML that is not compatible with JSON, ie yaml with anchors or complex map keys
 func YAMLNodeToJSON(node *yaml.Node, indentation string) ([]byte, error) {
 	v, err := handleYAMLNode(node)
 	if err != nil {
@@ -31,28 +30,38 @@ func handleYAMLNode(node *yaml.Node) (any, error) {
 	case yaml.ScalarNode:
 		return handleScalarNode(node)
 	case yaml.AliasNode:
-		panic("currently unsupported")
+		return handleYAMLNode(node.Alias)
 	default:
 		return nil, fmt.Errorf("unknown node kind: %v", node.Kind)
 	}
 }
 
 func handleMappingNode(node *yaml.Node) (any, error) {
-	m := orderedmap.New[string, yaml.Node]()
-
-	if err := node.Decode(m); err != nil {
-		return nil, err
-	}
-
 	v := orderedmap.New[string, any]()
-	for pair := orderedmap.First(m); pair != nil; pair = pair.Next() {
-		n := pair.Value()
-		vv, err := handleYAMLNode(&n)
+	for i, n := range node.Content {
+		if i%2 == 0 {
+			continue
+		}
+		keyNode := node.Content[i-1]
+		kv, err := handleYAMLNode(keyNode)
 		if err != nil {
 			return nil, err
 		}
 
-		v.Set(pair.Key(), vv)
+		if reflect.TypeOf(kv).Kind() != reflect.String {
+			keyData, err := json.Marshal(kv)
+			if err != nil {
+				return nil, err
+			}
+			kv = string(keyData)
+		}
+
+		vv, err := handleYAMLNode(n)
+		if err != nil {
+			return nil, err
+		}
+
+		v.Set(fmt.Sprintf("%v", kv), vv)
 	}
 
 	return v, nil
diff --git a/json/json_test.go b/json/json_test.go
index 211b3f8..8942c05 100644
--- a/json/json_test.go
+++ b/json/json_test.go
@@ -85,3 +85,54 @@ func TestYAMLNodeToJSON_FromJSON(t *testing.T) {
 
 	assert.Equal(t, j, string(o))
 }
+
+func TestYAMLNodeWithAnchorsToJSON(t *testing.T) {
+	y := `examples:
+  someExample: &someExample
+    key1: scalar1
+    key2: scalar2
+someValue: *someExample`
+
+	var v yaml.Node
+
+	err := yaml.Unmarshal([]byte(y), &v)
+	require.NoError(t, err)
+
+	j, err := json.YAMLNodeToJSON(&v, "  ")
+	require.NoError(t, err)
+
+	assert.Equal(t, `{
+  "examples": {
+    "someExample": {
+      "key1": "scalar1",
+      "key2": "scalar2"
+    }
+  },
+  "someValue": {
+    "key1": "scalar1",
+    "key2": "scalar2"
+  }
+}`, string(j))
+}
+
+func TestYAMLNodeWithComplexKeysToJSON(t *testing.T) {
+	y := `someMapWithComplexKeys:
+  {key1: scalar1, key2: scalar2}: {key1: scalar1, key2: scalar2}`
+
+	var v yaml.Node
+
+	err := yaml.Unmarshal([]byte(y), &v)
+	require.NoError(t, err)
+
+	j, err := json.YAMLNodeToJSON(&v, "  ")
+	require.NoError(t, err)
+
+	assert.Equal(t, `{
+  "someMapWithComplexKeys": {
+    "{\"key1\":\"scalar1\",\"key2\":\"scalar2\"}": {
+      "key1": "scalar1",
+      "key2": "scalar2"
+    }
+  }
+}`, string(j))
+}