From 6b2565e44cb75c160337fd7a077cba6aff97ae5d Mon Sep 17 00:00:00 2001
From: Noah Fontes <noah.fontes@puppet.com>
Date: Fri, 21 Feb 2020 19:28:56 -0800
Subject: [PATCH] Update: Add support for a generic interface type that
 supports binary data in our transfer package

---
 encoding/transfer/json.go      | 90 ++++++++++++++++++++++++++++++++++
 encoding/transfer/json_test.go | 53 ++++++++++++++++++++
 2 files changed, 143 insertions(+)

diff --git a/encoding/transfer/json.go b/encoding/transfer/json.go
index e01e49e..2b6e772 100644
--- a/encoding/transfer/json.go
+++ b/encoding/transfer/json.go
@@ -49,3 +49,93 @@ func (tos *JSONOrStr) UnmarshalJSON(data []byte) error {
 
 	return json.Unmarshal(data, &tos.JSON)
 }
+
+// JSONInterface allows arbitrary embedding of encoded data within any JSON
+// type. The accepted interface values for the data correspond to those listed
+// in the Go documentation for encoding/json.Unmarshal.
+type JSONInterface struct {
+	Data interface{}
+}
+
+func (ji JSONInterface) MarshalJSON() ([]byte, error) {
+	switch dt := ji.Data.(type) {
+	case map[string]interface{}:
+		m := make(map[string]interface{}, len(dt))
+		for k, v := range dt {
+			m[k] = JSONInterface{Data: v}
+		}
+
+		return json.Marshal(m)
+	case []interface{}:
+		s := make([]interface{}, len(dt))
+		for i, v := range dt {
+			s[i] = JSONInterface{Data: v}
+		}
+
+		return json.Marshal(s)
+	case string:
+		jos, err := EncodeJSON([]byte(dt))
+		if err != nil {
+			return nil, err
+		}
+
+		return json.Marshal(jos)
+	default:
+		return json.Marshal(dt)
+	}
+}
+
+func (ji *JSONInterface) UnmarshalJSON(data []byte) error {
+	var v interface{}
+	if err := json.Unmarshal(data, &v); err != nil {
+		return err
+	}
+
+	var decode func(v interface{}) (interface{}, error)
+	decode = func(v interface{}) (interface{}, error) {
+		switch vt := v.(type) {
+		case map[string]interface{}:
+			ty, ok := vt["$encoding"].(string)
+			if ok {
+				d, _ := vt["data"].(string)
+
+				b, err := JSON{EncodingType: EncodingType(ty), Data: d}.Decode()
+				if err != nil {
+					return nil, err
+				}
+
+				// Plop this back into one of the interface types supported by
+				// json.Marshal and json.Unmarshal (string).
+				v = string(b)
+			}
+
+			for k, v := range vt {
+				v, err := decode(v)
+				if err != nil {
+					return nil, err
+				}
+
+				vt[k] = v
+			}
+		case []interface{}:
+			for i, v := range vt {
+				v, err := decode(v)
+				if err != nil {
+					return nil, err
+				}
+
+				vt[i] = v
+			}
+		}
+
+		return v, nil
+	}
+
+	v, err := decode(v)
+	if err != nil {
+		return err
+	}
+
+	ji.Data = v
+	return nil
+}
diff --git a/encoding/transfer/json_test.go b/encoding/transfer/json_test.go
index abb8176..1d35fe2 100644
--- a/encoding/transfer/json_test.go
+++ b/encoding/transfer/json_test.go
@@ -116,3 +116,56 @@ func TestJSONOrStrMarshalUnmarshal(t *testing.T) {
 		})
 	}
 }
+
+func TestJSONInterfaceMarshalUnmarshal(t *testing.T) {
+	cases := []struct {
+		description string
+		input       interface{}
+		expected    string
+	}{
+		{
+			description: "String at top level",
+			input:       "This is a normal string",
+			expected:    `"This is a normal string"`,
+		},
+		{
+			description: "JSON object",
+			input: map[string]interface{}{
+				"a": []interface{}{
+					"b",
+					nil,
+					true,
+					1.23,
+				},
+				"b": "This is a normal string",
+			},
+			expected: `{"a":["b",null,true,1.23],"b":"This is a normal string"}`,
+		},
+		{
+			description: "Non-UTF-8 string embedded in JSON structure",
+			input: map[string]interface{}{
+				"a": []interface{}{
+					"Hello, \x90\xA2\x8A\x45",
+					"This is a normal string",
+				},
+				"b": "Goodbye, \x90!",
+			},
+			expected: `{"a":[{"$encoding":"base64","data":"SGVsbG8sIJCiikU="},"This is a normal string"],"b":{"$encoding":"base64","data":"R29vZGJ5ZSwgkCE="}}`,
+		},
+	}
+
+	for _, c := range cases {
+		t.Run(c.description, func(t *testing.T) {
+			j := transfer.JSONInterface{Data: c.input}
+
+			js, err := json.Marshal(j)
+			require.NoError(t, err)
+			require.JSONEq(t, c.expected, string(js))
+
+			var ju transfer.JSONInterface
+			require.NoError(t, json.Unmarshal(js, &ju))
+
+			require.Equal(t, c.input, ju.Data)
+		})
+	}
+}