From 20f848508000601120b811e216c3aec63803bca4 Mon Sep 17 00:00:00 2001
From: "M. J. Fromberger" <fromberger@interchain.io>
Date: Fri, 22 Apr 2022 18:02:41 -0700
Subject: [PATCH] Make encoding of HexBytes values more robust. (#8398)

The HexBytes wrapper type handles decoding byte strings from JSON.  In the RPC
API, hashes are encoded as hex digits rather than the standard base64.

Simplify the implementation of this wrapper using the TextMarshaler interface,
which the encoding/json package uses for values (like these) that are meant to
be wrapped in JSON strings.

In addition, allow HexBytes values to be decoded from either hex OR base64
input. This preserves all existing use, but will allow us to remove some
reflection special cases in the RPC decoder plumbing.

Update tests to correctly tolerate empty/nil.
---
 libs/bytes/bytes.go      | 58 +++++++++++++++-------------------------
 libs/bytes/bytes_test.go |  2 +-
 2 files changed, 22 insertions(+), 38 deletions(-)

diff --git a/libs/bytes/bytes.go b/libs/bytes/bytes.go
index dd8e39737..a6be32f05 100644
--- a/libs/bytes/bytes.go
+++ b/libs/bytes/bytes.go
@@ -1,21 +1,16 @@
 package bytes
 
 import (
-	"bytes"
+	"encoding/base64"
 	"encoding/hex"
-	"encoding/json"
 	"fmt"
 	"strings"
 )
 
-// HexBytes enables HEX-encoding for json/encoding.
+// HexBytes is a wrapper around []byte that encodes data as hexadecimal strings
+// for use in JSON.
 type HexBytes []byte
 
-var (
-	_ json.Marshaler   = HexBytes{}
-	_ json.Unmarshaler = &HexBytes{}
-)
-
 // Marshal needed for protobuf compatibility
 func (bz HexBytes) Marshal() ([]byte, error) {
 	return bz, nil
@@ -27,41 +22,30 @@ func (bz *HexBytes) Unmarshal(data []byte) error {
 	return nil
 }
 
-// MarshalJSON implements the json.Marshaler interface. The encoding is a JSON
-// quoted string of hexadecimal digits.
-func (bz HexBytes) MarshalJSON() ([]byte, error) {
-	size := hex.EncodedLen(len(bz)) + 2 // +2 for quotation marks
-	buf := make([]byte, size)
-	hex.Encode(buf[1:], []byte(bz))
-	buf[0] = '"'
-	buf[size-1] = '"'
-
-	// Ensure letter digits are capitalized.
-	for i := 1; i < size-1; i++ {
-		if buf[i] >= 'a' && buf[i] <= 'f' {
-			buf[i] = 'A' + (buf[i] - 'a')
-		}
-	}
-	return buf, nil
+// MarshalText encodes a HexBytes value as hexadecimal digits.
+// This method is used by json.Marshal.
+func (bz HexBytes) MarshalText() ([]byte, error) {
+	enc := hex.EncodeToString([]byte(bz))
+	return []byte(strings.ToUpper(enc)), nil
 }
 
-// UnmarshalJSON implements the json.Umarshaler interface.
-func (bz *HexBytes) UnmarshalJSON(data []byte) error {
-	if bytes.Equal(data, []byte("null")) {
+// UnmarshalText handles decoding of HexBytes from JSON strings.
+// This method is used by json.Unmarshal.
+// It allows decoding of both hex and base64-encoded byte arrays.
+func (bz *HexBytes) UnmarshalText(data []byte) error {
+	input := string(data)
+	if input == "" || input == "null" {
 		return nil
 	}
-
-	if len(data) < 2 || data[0] != '"' || data[len(data)-1] != '"' {
-		return fmt.Errorf("invalid hex string: %s", data)
-	}
-
-	bz2, err := hex.DecodeString(string(data[1 : len(data)-1]))
+	dec, err := hex.DecodeString(input)
 	if err != nil {
-		return err
+		dec, err = base64.StdEncoding.DecodeString(input)
+
+		if err != nil {
+			return err
+		}
 	}
-
-	*bz = bz2
-
+	*bz = HexBytes(dec)
 	return nil
 }
 
diff --git a/libs/bytes/bytes_test.go b/libs/bytes/bytes_test.go
index ce39935e8..fb1200a04 100644
--- a/libs/bytes/bytes_test.go
+++ b/libs/bytes/bytes_test.go
@@ -61,7 +61,7 @@ func TestJSONMarshal(t *testing.T) {
 				t.Fatal(err)
 			}
 			assert.Equal(t, ts2.B1, tc.input)
-			assert.Equal(t, ts2.B2, HexBytes(tc.input))
+			assert.Equal(t, string(ts2.B2), string(tc.input))
 		})
 	}
 }