From ef6ce6878bd8f9bdf807fa4d1c606118a267e9ef Mon Sep 17 00:00:00 2001
From: Matthieu Vachon <matt@streamingfast.io>
Date: Tue, 17 Dec 2024 11:05:16 -0500
Subject: [PATCH] Fixed wrong handling of `bytes-encoding` in some situations

---
 cmd/tools/compare/tools_compare_blocks.go | 12 +++----
 cmd/tools/print/printer_json.go           |  8 ++---
 json/marshallers.go                       | 44 +++++++++++++++++++----
 3 files changed, 45 insertions(+), 19 deletions(-)

diff --git a/cmd/tools/compare/tools_compare_blocks.go b/cmd/tools/compare/tools_compare_blocks.go
index 59b99cc..f57e57b 100644
--- a/cmd/tools/compare/tools_compare_blocks.go
+++ b/cmd/tools/compare/tools_compare_blocks.go
@@ -30,7 +30,7 @@ import (
 	"github.com/streamingfast/dstore"
 	firecore "github.com/streamingfast/firehose-core"
 	"github.com/streamingfast/firehose-core/cmd/tools/check"
-	"github.com/streamingfast/firehose-core/json"
+	fcjson "github.com/streamingfast/firehose-core/json"
 	fcproto "github.com/streamingfast/firehose-core/proto"
 	"github.com/streamingfast/firehose-core/types"
 	"go.uber.org/multierr"
@@ -361,16 +361,16 @@ func Compare(reference proto.Message, current proto.Message, includeUnknownField
 
 	//todo: check if there is a equals that do not compare unknown fields
 	if !proto.Equal(reference, current) {
-		var opts []json.MarshallerOption
+		var opts []fcjson.MarshallerOption
 		if !includeUnknownFields {
-			opts = append(opts, json.WithoutUnknownFields())
+			opts = append(opts, fcjson.WithoutUnknownFields())
 		}
 
-		if bytesEncoding == "base58" {
-			opts = append(opts, json.WithBytesEncoderFunc(json.ToBase58))
+		if bytesEncoding != "" {
+			opts = append(opts, fcjson.WithBytesEncoding(bytesEncoding))
 		}
 
-		encoder := json.NewMarshaller(registry, opts...)
+		encoder := fcjson.NewMarshaller(registry, opts...)
 
 		referenceAsJSON, err := encoder.MarshalToString(reference)
 		cli.NoError(err, "marshal JSON reference")
diff --git a/cmd/tools/print/printer_json.go b/cmd/tools/print/printer_json.go
index cd05e93..be614bd 100644
--- a/cmd/tools/print/printer_json.go
+++ b/cmd/tools/print/printer_json.go
@@ -34,12 +34,8 @@ type JSONOutputPrinter struct {
 func NewJSONOutputPrinter(bytesEncoding string, singleLine bool, registry *fcproto.Registry) (OutputPrinter, error) {
 	var options []fcjson.MarshallerOption
 
-	if bytesEncoding == "base58" {
-		options = append(options, fcjson.WithBytesEncoderFunc(fcjson.ToBase58))
-	}
-
-	if bytesEncoding == "base64" {
-		options = append(options, fcjson.WithBytesEncoderFunc(fcjson.ToBase64))
+	if bytesEncoding != "" {
+		options = append(options, fcjson.WithBytesEncoding(bytesEncoding))
 	}
 
 	return &JSONOutputPrinter{
diff --git a/json/marshallers.go b/json/marshallers.go
index 73e5d6c..4f72de6 100644
--- a/json/marshallers.go
+++ b/json/marshallers.go
@@ -26,11 +26,13 @@ type kv struct {
 	value any
 }
 
+type CustomEncoderFunc func(encoder *jsontext.Encoder, t []byte, options json.Options) error
+
 type Marshaller struct {
 	marshallers          *json.Marshalers
 	includeUnknownFields bool
 	registry             *proto.Registry
-	bytesEncoderFunc     func(encoder *jsontext.Encoder, t []byte, options json.Options) error
+	bytesEncoderFunc     CustomEncoderFunc
 }
 
 type MarshallerOption func(*Marshaller)
@@ -40,7 +42,15 @@ func WithoutUnknownFields() MarshallerOption {
 		e.includeUnknownFields = false
 	}
 }
-func WithBytesEncoderFunc(f func(encoder *jsontext.Encoder, t []byte, options json.Options) error) MarshallerOption {
+
+func WithBytesEncoding(bytesEncoding string) MarshallerOption {
+	return withBytesEncoderFunc(jsonEncoder(bytesEncoding))
+}
+
+// Deprecated: Use WithBytesEncoding instead passing the encoding directly.
+var WithBytesEncoderFunc = withBytesEncoderFunc
+
+func withBytesEncoderFunc(f CustomEncoderFunc) MarshallerOption {
 	return func(e *Marshaller) {
 		e.bytesEncoderFunc = f
 	}
@@ -125,8 +135,7 @@ func (m *Marshaller) dynamicpbMessage(encoder *jsontext.Encoder, msg *dynamicpb.
 		x := msg.GetUnknown()
 		fieldNumber, ofType, l := protowire.ConsumeField(x)
 		if l > 0 {
-			var unknownValue []byte
-			unknownValue = x[:l]
+			unknownValue := x[:l]
 			kvl = append(kvl, &kv{
 				key:   fmt.Sprintf("__unknown_fields_%d_with_type_%d__", fieldNumber, ofType),
 				value: hex.EncodeToString(unknownValue),
@@ -213,15 +222,36 @@ func EncodeHex(bytes []byte) string {
 
 // Encode encodes the given bytes using the specified encoding.
 func Encode(bytesEncoding string, bytes []byte) string {
+	return Encoder(bytesEncoding)(bytes)
+}
+
+// Encoder returns the encoder function that will converts received bytes
+// into a string of the specified encoding.
+func Encoder(bytesEncoding string) func([]byte) string {
+	switch {
+	case strings.EqualFold(bytesEncoding, "base58"):
+		return base58.Encode
+
+	case strings.EqualFold(bytesEncoding, "base64"):
+		return base64.StdEncoding.EncodeToString
+
+	case strings.EqualFold(bytesEncoding, "hex"):
+		return hex.EncodeToString
+	}
+
+	panic(fmt.Errorf("unsupported bytes encoding: %s", bytesEncoding))
+}
+
+func jsonEncoder(bytesEncoding string) jsonEncoderFunc {
 	switch {
 	case strings.EqualFold(bytesEncoding, "base58"):
-		return base58.Encode(bytes)
+		return ToBase58
 
 	case strings.EqualFold(bytesEncoding, "base64"):
-		return base64.StdEncoding.EncodeToString(bytes)
+		return ToBase64
 
 	case strings.EqualFold(bytesEncoding, "hex"):
-		return hex.EncodeToString(bytes)
+		return ToHex
 	}
 
 	panic(fmt.Errorf("unsupported bytes encoding: %s", bytesEncoding))