fix marshaling of any types to JSON when contained object has custom …

…marshaling/unmarshaling
golang · cybrcodr · May 31, 2017 · May 26, 2017 · May 27, 2017 · May 31, 2017
commit 03a02321e7c9e401cb002dd325f187bf45c728ff
diff --git a/jsonpb/jsonpb.go b/jsonpb/jsonpb.go
@@ -77,15 +77,26 @@ type Marshaler struct {
 // JSONPBMarshaler is implemented by protobuf messages that customize the
 // way they are marshaled to JSON. Messages that implement this should
 // also implement JSONPBUnmarshaler so that the custom format can be
-// parsed.
+// parsed. If this method produces JSON other than an object (for example,
+// produces an encoded JSON string), the type must also implement another
+// method:
+//    XXX_CustomJSON() bool
+// If this method is implemented and returns true when invoked, then the
+// message is allowed to marshal itself to JSON types other than objects.
 type JSONPBMarshaler interface {
 	MarshalJSONPB(*Marshaler) ([]byte, error)
 }
 
 // JSONPBUnmarshaler is implemented by protobuf messages that customize
 // the way they are unmarshaled from JSON. Messages that implement this
 // should also implement JSONPBMarshaler so that the custom format can be
-// produced.
+// produced. If this method expects to unmarshal itself from JSON other
+// than an object (for example, unmarshals from an encoded string), the
+// type must also implement another method:
+//    XXX_CustomJSON() bool
+// If this method is implemented and returns true when invoked, then the
+// message is allowed to unmarshal itself from JSON types other than
+// objects.
 type JSONPBUnmarshaler interface {
 	UnmarshalJSONPB(*Unmarshaler, []byte) error
 }
@@ -116,13 +127,43 @@ type wkt interface {
 	XXX_WellKnownType() string
 }
 
+// customJSON can be implemented by messages that have custom non-object JSON
+// representations. If the message implements this method and the method
+// returns true, the JSON representation will not be an object (for example,
+// the message marshals itself as a string).
+type customJSON interface {
+	XXX_CustomJSON() bool
+}
+
 // marshalObject writes a struct to the Writer.
 func (m *Marshaler) marshalObject(out *errWriter, v proto.Message, indent, typeURL string) error {
 	if jsm, ok := v.(JSONPBMarshaler); ok {
 		b, err := jsm.MarshalJSONPB(m)
 		if err != nil {
 			return err
 		}
+		if typeURL != "" {
+			// we are marshaling this object to an Any type
+			var js interface{}
+			if err = json.Unmarshal(b, &js); err != nil {
+				return fmt.Errorf("Type %v produced invalid JSON: %v", reflect.TypeOf(v), err)
+			}
+			m, ok := js.(map[string]interface{});
+			if ok {
+				m["@type"] = typeURL
+			} else {
+				if cj, ok := v.(customJSON); !ok || !cj.XXX_CustomJSON() {
+					return fmt.Errorf("Type %v illegally produced JSON that is not an object", reflect.TypeOf(v))
+				}
+				m = map[string]interface{}{}
+				m["@type"] = typeURL
+				m["value"] = js
+			}
+			if b, err = json.Marshal(m); err != nil {
+				return err
+			}
+		}
+
 		out.write(string(b))
 		return out.err
 	}
@@ -441,9 +482,6 @@ func (m *Marshaler) marshalValue(out *errWriter, prop *proto.Properties, v refle
 
 	// Handle well-known types.
 	// Most are handled up in marshalObject (because 99% are messages).
-	type wkt interface {
-		XXX_WellKnownType() string
-	}
 	if wkt, ok := v.Interface().(wkt); ok {
 		switch wkt.XXX_WellKnownType() {
 		case "NullValue":
@@ -611,9 +649,6 @@ func (u *Unmarshaler) unmarshalValue(target reflect.Value, inputValue json.RawMe
 	}
 
 	// Handle well-known types.
-	type wkt interface {
-		XXX_WellKnownType() string
-	}
 	if w, ok := target.Addr().Interface().(wkt); ok {
 		switch w.XXX_WellKnownType() {
 		case "DoubleValue", "FloatValue", "Int64Value", "UInt64Value",
@@ -633,13 +668,13 @@ func (u *Unmarshaler) unmarshalValue(target reflect.Value, inputValue json.RawMe
 			}
 
 			val, ok := jsonFields["@type"]
-			if !ok {
+			if !ok || val == nil {
 				return errors.New("Any JSON doesn't have '@type'")
 			}
 
 			var turl string
 			if err := json.Unmarshal([]byte(*val), &turl); err != nil {
-				return fmt.Errorf("can't unmarshal Any's '@type': %q", val)
+				return fmt.Errorf("can't unmarshal Any's '@type': %q", *val)
 			}
 			target.Field(0).SetString(turl)
 
@@ -653,14 +688,23 @@ func (u *Unmarshaler) unmarshalValue(target reflect.Value, inputValue json.RawMe
 			}
 
 			m := reflect.New(mt.Elem()).Interface().(proto.Message)
+			hasCustomJson := false
 			if _, ok := m.(wkt); ok {
+				hasCustomJson = true
+			} else if _, ok := m.(JSONPBUnmarshaler); ok {
+				// custom JSON unmarshallers may expect non-object JSON type
+				if cj, ok := m.(customJSON); ok && cj.XXX_CustomJSON() {
+					hasCustomJson = true
+				}
+			}
+			if hasCustomJson {
 				val, ok := jsonFields["value"]
 				if !ok {
 					return errors.New("Any JSON doesn't have 'value'")
 				}
 
 				if err := u.unmarshalValue(reflect.ValueOf(m).Elem(), *val, nil); err != nil {
-					return fmt.Errorf("can't unmarshal Any's WKT: %v", err)
+					return fmt.Errorf("can't unmarshal Any nested proto %v: %v", reflect.TypeOf(m), err)
 				}
 			} else {
 				delete(jsonFields, "@type")
@@ -670,13 +714,13 @@ func (u *Unmarshaler) unmarshalValue(target reflect.Value, inputValue json.RawMe
 				}
 
 				if err = u.unmarshalValue(reflect.ValueOf(m).Elem(), nestedProto, nil); err != nil {
-					return fmt.Errorf("can't unmarshal nested Any proto: %v", err)
+					return fmt.Errorf("can't unmarshal Any nested proto %v: %v", reflect.TypeOf(m), err)
 				}
 			}
 
 			b, err := proto.Marshal(m)
 			if err != nil {
-				return fmt.Errorf("can't marshal proto into Any.Value: %v", err)
+				return fmt.Errorf("can't marshal proto %v into Any.Value: %v", reflect.TypeOf(m), err)
 			}
 			target.Field(1).SetBytes(b)
 

diff --git a/jsonpb/jsonpb_test.go b/jsonpb/jsonpb_test.go
@@ -34,6 +34,7 @@ package jsonpb
 import (
 	"bytes"
 	"encoding/json"
+	"fmt"
 	"io"
 	"reflect"
 	"strings"
@@ -43,6 +44,7 @@ import (
 
 	pb "github.com/golang/protobuf/jsonpb/jsonpb_test_proto"
 	proto3pb "github.com/golang/protobuf/proto/proto3_proto"
+	"github.com/golang/protobuf/ptypes"
 	anypb "github.com/golang/protobuf/ptypes/any"
 	durpb "github.com/golang/protobuf/ptypes/duration"
 	stpb "github.com/golang/protobuf/ptypes/struct"
@@ -454,6 +456,55 @@ func TestMarshalingWithJSONPBMarshaler(t *testing.T) {
 	}
 }
 
+func TestMarshalingToAnyWithJSONPBMarshaler(t *testing.T) {
+	msg := dynamicMessage{rawJson: `{ "foo": "bar", "baz": [0, 1, 2, 3] }`}
+	a, err := ptypes.MarshalAny(&msg)
+	if err != nil {
+		t.Errorf("an unexpected error occurred when marshalling to Any: %v", err)
+	}
+	str, err := new(Marshaler).MarshalToString(a)
+	if err != nil {
+		t.Errorf("an unexpected error occurred when marshalling Any to JSON: %v", err)
+	}
+	// after custom marshaling, it's round-tripped through JSON decoding/encoding already,
+	// so the keys are sorted, whitespace is compacted, and "@type" key has been added
+	expected := `{"@type":"type.googleapis.com/` + dynamicMessageName +`","baz":[0,1,2,3],"foo":"bar"}`
+	if str != expected {
+		t.Errorf("marshalling JSON produced incorrect output: got %s, wanted %s", str, expected)
+	}
+}
+
+func TestMarshalingToAnyWithJSONPBMarshalerCustomJSON(t *testing.T) {
+	// first try will fail because we're producing invalid JSON object
+	var msg proto.Message = &dynamicMessage{rawJson: `"fubar"`}
+	a, err := ptypes.MarshalAny(msg)
+	if err != nil {
+		t.Errorf("an unexpected error occurred when marshalling to Any: %v", err)
+	}
+	str, err := new(Marshaler).MarshalToString(a)
+	if err == nil || !strings.Contains(err.Error(), "illegally produced JSON that is not an object") {
+		t.Errorf(`expected error message "illegally produced JSON that is not an object" but instead got %v`, err)
+	}
+
+	// now we try, but with custom marshaling that allows non-object JSON type
+	msg = &dynamicMessageCustomJson{dynamicMessage{rawJson: `"fubar"`}}
+	a, err = ptypes.MarshalAny(msg)
+	if err != nil {
+		t.Errorf("an unexpected error occurred when marshalling to Any: %v", err)
+	}
+	str, err = new(Marshaler).MarshalToString(a)
+	if err != nil {
+		t.Errorf("an unexpected error occurred when marshalling Any to JSON: %v", err)
+	}
+
+	// after custom marshaling, it's round-tripped through JSON decoding/encoding already,
+	// so the keys are sorted, whitespace is compacted, and "@type" key has been added
+	expected := `{"@type":"type.googleapis.com/` + dynamicMessageCustomJsonName +`","value":"fubar"}`
+	if str != expected {
+		t.Errorf("marshalling JSON produced incorrect output: got %s, wanted %s", str, expected)
+	}
+}
+
 var unmarshalingTests = []struct {
 	desc        string
 	unmarshaler Unmarshaler
@@ -654,15 +705,62 @@ func TestUnmarshalingBadInput(t *testing.T) {
 func TestUnmarshalWithJSONPBUnmarshaler(t *testing.T) {
 	rawJson := `{ "foo": "bar", "baz": [0, 1, 2, 3] }`
 	var msg dynamicMessage
-	err := Unmarshal(strings.NewReader(rawJson), &msg)
-	if err != nil {
+	if err := Unmarshal(strings.NewReader(rawJson), &msg); err != nil {
 		t.Errorf("an unexpected error occurred when parsing into JSONPBUnmarshaler: %v", err)
 	}
 	if msg.rawJson != rawJson {
 		t.Errorf("message contents not set correctly after unmarshalling JSON: got %s, wanted %s", msg.rawJson, rawJson)
 	}
 }
 
+func TestUnmarshalAnyWithJSONPBUnmarshaler(t *testing.T) {
+	rawJson := `{ "@type": "blah.com/` + dynamicMessageName + `", "foo": "bar", "baz": [0, 1, 2, 3] }`
+	var a anypb.Any
+	if err := Unmarshal(strings.NewReader(rawJson), &a); err != nil {
+		t.Errorf("an unexpected error occurred when parsing into JSONPBUnmarshaler: %v", err)
+	}
+	var msg ptypes.DynamicAny
+	if err := ptypes.UnmarshalAny(&a, &msg); err != nil {
+		t.Errorf("an unexpected error occurred when parsing from Any type: %v", err)
+	}
+	dm := msg.Message.(*dynamicMessage)
+	// when it gets to the custom unmarshal, it's round-tripped through JSON decoding/encoding
+	// already, so the keys are sorted, whitespace is compacted, and "@type" key has been removed
+	expected := `{"baz":[0,1,2,3],"foo":"bar"}`
+	if dm.rawJson != expected {
+		t.Errorf("message contents not set correctly after unmarshalling JSON: got %s, wanted %s", dm.rawJson, expected)
+	}
+}
+
+func TestUnmarshalAnyWithJSONPBUnmarshalerCustomJSON(t *testing.T) {
+	// with custom JSON, we unmarshal contents at "value" key instead of whole JSON object
+	rawJson := `{ "@type": "blah.com/` + dynamicMessageCustomJsonName + `", "value": "fubar"}`
+	var a anypb.Any
+	if err := Unmarshal(strings.NewReader(rawJson), &a); err != nil {
+		t.Errorf("an unexpected error occurred when parsing into JSONPBUnmarshaler: %v", err)
+	}
+	var msg ptypes.DynamicAny
+	if err := ptypes.UnmarshalAny(&a, &msg); err != nil {
+		t.Errorf("an unexpected error occurred when parsing from Any type: %v", err)
+	}
+	dm := msg.Message.(*dynamicMessageCustomJson)
+	// when it gets to the custom unmarshal, it's round-tripped through JSON decoding/encoding
+	// already, so the keys are sorted, whitespace is compacted, and "@type" key has been removed
+	if dm.rawJson != `"fubar"` {
+		t.Errorf("message contents not set correctly after unmarshalling JSON: got %s, wanted %s", dm.rawJson, `"fubar"`)
+	}
+}
+
+const (
+	dynamicMessageName = "google.protobuf.jsonpb.testing.dynamicMessage"
+	dynamicMessageCustomJsonName = "google.protobuf.jsonpb.testing.dynamicMessageCustomJson"
+)
+func init() {
+	// we register the custom types below so that we can use them in Any types
+	proto.RegisterType((*dynamicMessage)(nil), dynamicMessageName)
+	proto.RegisterType((*dynamicMessageCustomJson)(nil), dynamicMessageCustomJsonName)
+}
+
 // dynamicMessage implements protobuf.Message but is not a normal generated message type.
 // It provides implementations of JSONPBMarshaler and JSONPBUnmarshaler for JSON support.
 type dynamicMessage struct {
@@ -684,7 +782,45 @@ func (m *dynamicMessage) MarshalJSONPB(jm *Marshaler) ([]byte, error) {
 	return []byte(m.rawJson), nil
 }
 
-func (m *dynamicMessage) UnmarshalJSONPB(jum *Unmarshaler, json []byte) error {
-	m.rawJson = string(json)
+func (m *dynamicMessage) UnmarshalJSONPB(jum *Unmarshaler, js []byte) error {
+	m.rawJson = string(js)
+	return nil
+}
+
+func (m *dynamicMessage) Marshal() ([]byte, error) {
+	var b proto.Buffer
+	// tag #1, wire type length-delimited
+	if err := b.EncodeVarint(uint64((1 << 3) | proto.WireBytes)); err != nil {
+		return nil, err
+	}
+	if err := b.EncodeRawBytes([]byte(m.rawJson)); err != nil {
+		return nil, err
+	}
+	return b.Bytes(), nil
+}
+
+func (m *dynamicMessage) Unmarshal(b []byte) error {
+	// really dumb unmarshalling, just for test
+	buf := proto.NewBuffer(b)
+	if tagAndWire, err := buf.DecodeVarint(); err != nil {
+		return err
+	} else if tagAndWire != uint64((1 << 3) | proto.WireBytes) {
+		return fmt.Errorf("Unknown tag/wire type: %x", tagAndWire)
+	}
+	if bb, err := buf.DecodeRawBytes(true); err != nil {
+		return err
+	} else {
+		m.rawJson = string(bb)
+	}
 	return nil
 }
+
+// dynamicMessageCustomJson is the same as dynamicMessage except it also provides
+// XXX_CustomJSON method, for marshaling to non-object JSON types.
+type dynamicMessageCustomJson struct {
+	dynamicMessage
+}
+
+func (m *dynamicMessageCustomJson) XXX_CustomJSON() bool {
+	return true
+}