Jsonpb custom type - #411 (#491)

* added test for MarshalJSONPB not being called for custom type if it's an alias to a primitive (#411)

* removed work-around so test fails

* added comment thanking contributor

* bumped the go version

* regenerated after nil check

* regenerated after nil check

* #411 - check for UnmarshalJSON interface, regardless of reflect.Kind()

Makefile

@@ -131,6 +131,7 @@ regenerate:
 	make -C test/issue449 regenerate
 	make -C test/xxxfields regenerate
 	make -C test/issue435 regenerate
+	make -C test/issue411 regenerate
 
 	make gofmt
 

jsonpb/jsonpb.go

@@ -991,16 +991,16 @@ func (u *Unmarshaler) unmarshalValue(target reflect.Value, inputValue json.RawMe
 		return nil
 	}
 
-	// Handle nested messages.
-	if targetType.Kind() == reflect.Struct {
-		if prop != nil && len(prop.CustomType) > 0 && target.CanAddr() {
-			if m, ok := target.Addr().Interface().(interface {
-				UnmarshalJSON([]byte) error
-			}); ok {
-				return json.Unmarshal(inputValue, m)
-			}
+	if prop != nil && len(prop.CustomType) > 0 && target.CanAddr() {
+		if m, ok := target.Addr().Interface().(interface {
+			UnmarshalJSON([]byte) error
+		}); ok {
+			return json.Unmarshal(inputValue, m)
 		}
+	}
 
+	// Handle nested messages.
+	if targetType.Kind() == reflect.Struct {
 		var jsonFields map[string]json.RawMessage
 		if err := json.Unmarshal(inputValue, &jsonFields); err != nil {
 			return err

test/issue411/Makefile

@@ -0,0 +1,3 @@
+regenerate:
+	go install github.com/gogo/protobuf/protoc-gen-gogo
+	protoc-min-version --version="3.0.0" --gogo_out=Mgoogle/protobuf/timestamp.proto=github.com/gogo/protobuf/types:. --proto_path=../../../../../:../../protobuf/:. issue411.proto

test/issue411/ids.go

@@ -0,0 +1,216 @@
+package issue411
+
+import (
+	"encoding/base64"
+	"encoding/binary"
+	"fmt"
+	"strconv"
+)
+
+// TraceID is a random 128bit identifier for a trace
+type TraceID struct {
+	Low  uint64 `json:"lo"`
+	High uint64 `json:"hi"`
+}
+
+// SpanID is a random 64bit identifier for a span
+type SpanID uint64
+
+// ------- TraceID -------
+
+// NewTraceID creates a new TraceID from two 64bit unsigned ints.
+func NewTraceID(high, low uint64) TraceID {
+	return TraceID{High: high, Low: low}
+}
+
+func (t TraceID) String() string {
+	if t.High == 0 {
+		return fmt.Sprintf("%x", t.Low)
+	}
+	return fmt.Sprintf("%x%016x", t.High, t.Low)
+}
+
+// TraceIDFromString creates a TraceID from a hexadecimal string
+func TraceIDFromString(s string) (TraceID, error) {
+	var hi, lo uint64
+	var err error
+	if len(s) > 32 {
+		return TraceID{}, fmt.Errorf("TraceID cannot be longer than 32 hex characters: %s", s)
+	} else if len(s) > 16 {
+		hiLen := len(s) - 16
+		if hi, err = strconv.ParseUint(s[0:hiLen], 16, 64); err != nil {
+			return TraceID{}, err
+		}
+		if lo, err = strconv.ParseUint(s[hiLen:], 16, 64); err != nil {
+			return TraceID{}, err
+		}
+	} else {
+		if lo, err = strconv.ParseUint(s, 16, 64); err != nil {
+			return TraceID{}, err
+		}
+	}
+	return TraceID{High: hi, Low: lo}, nil
+}
+
+// MarshalText is called by encoding/json, which we do not want people to use.
+func (t TraceID) MarshalText() ([]byte, error) {
+	return nil, fmt.Errorf("unsupported method TraceID.MarshalText; please use github.com/gogo/protobuf/jsonpb for marshalling")
+}
+
+// UnmarshalText is called by encoding/json, which we do not want people to use.
+func (t *TraceID) UnmarshalText(text []byte) error {
+	return fmt.Errorf("unsupported method TraceID.UnmarshalText; please use github.com/gogo/protobuf/jsonpb for marshalling")
+}
+
+// Size returns the size of this datum in protobuf. It is always 16 bytes.
+func (t *TraceID) Size() int {
+	return 16
+}
+
+// Marshal converts trace ID into a binary representation. Called by protobuf serialization.
+func (t TraceID) Marshal() ([]byte, error) {
+	b := make([]byte, t.Size())
+	_, err := t.MarshalTo(b)
+	return b, err
+}
+
+// MarshalTo converts trace ID into a binary representation. Called by protobuf serialization.
+func (t *TraceID) MarshalTo(data []byte) (n int, err error) {
+	var b [16]byte
+	binary.BigEndian.PutUint64(b[:8], uint64(t.High))
+	binary.BigEndian.PutUint64(b[8:], uint64(t.Low))
+	return marshalBytes(data, b[:])
+}
+
+// Unmarshal inflates this trace ID from binary representation. Called by protobuf serialization.
+func (t *TraceID) Unmarshal(data []byte) error {
+	if len(data) < 16 {
+		return fmt.Errorf("buffer is too short")
+	}
+	t.High = binary.BigEndian.Uint64(data[:8])
+	t.Low = binary.BigEndian.Uint64(data[8:])
+	return nil
+}
+
+func marshalBytes(dst []byte, src []byte) (n int, err error) {
+	if len(dst) < len(src) {
+		return 0, fmt.Errorf("buffer is too short")
+	}
+	return copy(dst, src), nil
+}
+
+// MarshalJSON converts trace id into a base64 string enclosed in quotes.
+// Used by protobuf JSON serialization.
+// Example: {high:2, low:1} => "AAAAAAAAAAIAAAAAAAAAAQ==".
+func (t TraceID) MarshalJSON() ([]byte, error) {
+	var b [16]byte
+	t.MarshalTo(b[:]) // can only error on incorrect buffer size
+	s := make([]byte, 24+2)
+	base64.StdEncoding.Encode(s[1:25], b[:])
+	s[0], s[25] = '"', '"'
+	return s, nil
+}
+
+// UnmarshalJSON inflates trace id from base64 string, possibly enclosed in quotes.
+// User by protobuf JSON serialization.
+func (t *TraceID) UnmarshalJSON(data []byte) error {
+	s := string(data)
+	if l := len(s); l > 2 && s[0] == '"' && s[l-1] == '"' {
+		s = s[1 : l-1]
+	}
+	b, err := base64.StdEncoding.DecodeString(s)
+	if err != nil {
+		return fmt.Errorf("cannot unmarshal TraceID from string '%s': %v", string(data), err)
+	}
+	return t.Unmarshal(b)
+}
+
+// ------- SpanID -------
+
+// NewSpanID creates a new SpanID from a 64bit unsigned int.
+func NewSpanID(v uint64) SpanID {
+	return SpanID(v)
+}
+
+func (s SpanID) String() string {
+	return fmt.Sprintf("%x", uint64(s))
+}
+
+// SpanIDFromString creates a SpanID from a hexadecimal string
+func SpanIDFromString(s string) (SpanID, error) {
+	if len(s) > 16 {
+		return SpanID(0), fmt.Errorf("SpanID cannot be longer than 16 hex characters: %s", s)
+	}
+	id, err := strconv.ParseUint(s, 16, 64)
+	if err != nil {
+		return SpanID(0), err
+	}
+	return SpanID(id), nil
+}
+
+// MarshalText is called by encoding/json, which we do not want people to use.
+func (s SpanID) MarshalText() ([]byte, error) {
+	return nil, fmt.Errorf("unsupported method SpanID.MarshalText; please use github.com/gogo/protobuf/jsonpb for marshalling")
+}
+
+// UnmarshalText is called by encoding/json, which we do not want people to use.
+func (s *SpanID) UnmarshalText(text []byte) error {
+	return fmt.Errorf("unsupported method SpanID.UnmarshalText; please use github.com/gogo/protobuf/jsonpb for marshalling")
+}
+
+// Size returns the size of this datum in protobuf. It is always 8 bytes.
+func (s *SpanID) Size() int {
+	return 8
+}
+
+// Marshal converts span ID into a binary representation. Called by protobuf serialization.
+func (s SpanID) Marshal() ([]byte, error) {
+	b := make([]byte, s.Size())
+	_, err := s.MarshalTo(b)
+	return b, err
+}
+
+// MarshalTo converts span ID into a binary representation. Called by protobuf serialization.
+func (s *SpanID) MarshalTo(data []byte) (n int, err error) {
+	var b [8]byte
+	binary.BigEndian.PutUint64(b[:], uint64(*s))
+	return marshalBytes(data, b[:])
+}
+
+// Unmarshal inflates span ID from a binary representation. Called by protobuf serialization.
+func (s *SpanID) Unmarshal(data []byte) error {
+	if len(data) < 8 {
+		return fmt.Errorf("buffer is too short")
+	}
+	*s = NewSpanID(binary.BigEndian.Uint64(data))
+	return nil
+}
+
+// MarshalJSON converts span id into a base64 string enclosed in quotes.
+// Used by protobuf JSON serialization.
+// Example: {1} => "AAAAAAAAAAE=".
+func (s SpanID) MarshalJSON() ([]byte, error) {
+	var b [8]byte
+	s.MarshalTo(b[:]) // can only error on incorrect buffer size
+	v := make([]byte, 12+2)
+	base64.StdEncoding.Encode(v[1:13], b[:])
+	v[0], v[13] = '"', '"'
+	return v, nil
+}
+
+// UnmarshalJSON inflates span id from base64 string, possibly enclosed in quotes.
+// User by protobuf JSON serialization.
+//
+// There appears to be a bug in gogoproto, as this function is only called for numeric values.
+// https://github.com/gogo/protobuf/issues/411#issuecomment-393856837
+func (s *SpanID) UnmarshalJSON(data []byte) error {
+	str := string(data)
+	if l := len(str); l > 2 && str[0] == '"' && str[l-1] == '"' {
+		str = str[1 : l-1]
+	}
+	b, err := base64.StdEncoding.DecodeString(str)
+	if err != nil {
+		return fmt.Errorf("cannot unmarshal SpanID from string '%s': %v", string(data), err)
+	}
+	return s.Unmarshal(b)
+}

test/issue411/ids_test.go

@@ -0,0 +1,69 @@
+package issue411_test
+
+import (
+	"bytes"
+	"testing"
+
+	"github.com/gogo/protobuf/jsonpb"
+	"github.com/gogo/protobuf/proto"
+	"github.com/gogo/protobuf/test/issue411"
+)
+
+// Thanks to @yurishkuro for reporting this issue (#411) and providing this test case
+
+// TraceID/SpanID fields are defined as bytes in proto, backed by custom types in Go.
+// Unfortunately, that means they require manual implementations of proto & json serialization.
+// To ensure that it's the same as the default protobuf serialization, file jaeger_test.proto
+// contains a copy of SpanRef message without any gogo options. This test file is compiled with
+// plain protoc -go_out (without gogo). This test performs proto and JSON marshaling/unmarshaling
+// to ensure that the outputs of manual and standard serialization are identical.
+func TestTraceSpanIDMarshalProto(t *testing.T) {
+	testCases := []struct {
+		name      string
+		marshal   func(proto.Message) ([]byte, error)
+		unmarshal func([]byte, proto.Message) error
+		expected  string
+	}{
+		{
+			name:      "protobuf",
+			marshal:   proto.Marshal,
+			unmarshal: proto.Unmarshal,
+		},
+		{
+			name: "JSON",
+			marshal: func(m proto.Message) ([]byte, error) {
+				out := new(bytes.Buffer)
+				err := new(jsonpb.Marshaler).Marshal(out, m)
+				if err != nil {
+					return nil, err
+				}
+				return out.Bytes(), nil
+			},
+			unmarshal: func(in []byte, m proto.Message) error {
+				return jsonpb.Unmarshal(bytes.NewReader(in), m)
+			},
+			expected: `{"traceId":"AAAAAAAAAAIAAAAAAAAAAw==","spanId":"AAAAAAAAAAs="}`,
+		},
+	}
+	for _, testCase := range testCases {
+		t.Run(testCase.name, func(t *testing.T) {
+			o1 := issue411.Span{TraceID: issue411.NewTraceID(2, 3), SpanID: issue411.NewSpanID(11)}
+			d1, err := testCase.marshal(&o1)
+			if err != nil {
+				t.Fatalf("marshal error: %v", err)
+			}
+			// test unmarshal
+			var o2 issue411.Span
+			err = testCase.unmarshal(d1, &o2)
+			if err != nil {
+				t.Fatalf("umarshal error: %v", err)
+			}
+			if o1.TraceID != o2.TraceID {
+				t.Fatalf("TraceID: expected %v but got %v", o1, o2)
+			}
+			if o1.SpanID != o2.SpanID {
+				t.Fatalf("SpanID: expected %v but got %v", o1, o2)
+			}
+		})
+	}
+}