Fix JSON marshaling of Any containing object that implements JSONPBMarshaler (#361)

* Fix marshaling of Any containing object that implements JSONPBMarshaler.
* Add test for marshaling Any that implements JSONPBMarshaler.
* Add test for unmarshaling Any that implements JSONPBUnmarshaler.

Fixed issue #360.

jsonpb/jsonpb.go

@@ -124,6 +124,22 @@ func (m *Marshaler) marshalObject(out *errWriter, v proto.Message, indent, typeU
 		if err != nil {
 			return err
 		}
+		if typeURL != "" {
+			// we are marshaling this object to an Any type
+			var js map[string]*json.RawMessage
+			if err = json.Unmarshal(b, &js); err != nil {
+				return fmt.Errorf("type %T produced invalid JSON: %v", v, err)
+			}
+			turl, err := json.Marshal(typeURL)
+			if err != nil {
+				return fmt.Errorf("failed to marshal type URL %q to JSON: %v", typeURL, err)
+			}
+			js["@type"] = (*json.RawMessage)(&turl)
+			if b, err = json.Marshal(js); err != nil {
+				return err
+			}
+		}
+
 		out.write(string(b))
 		return out.err
 	}
@@ -441,9 +457,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":
@@ -629,9 +642,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",
@@ -651,13 +661,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)
 
@@ -678,7 +688,7 @@ func (u *Unmarshaler) unmarshalValue(target reflect.Value, inputValue json.RawMe
 				}
 
 				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 %T: %v", m, err)
 				}
 			} else {
 				delete(jsonFields, "@type")
@@ -688,13 +698,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 %T: %v", 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 %T into Any.Value: %v", m, err)
 			}
 			target.Field(1).SetBytes(b)
 

jsonpb/jsonpb_test.go

@@ -44,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"
@@ -461,7 +462,7 @@ func TestMarshaling(t *testing.T) {
 	}
 }
 
-func TestMarshalingWithJSONPBMarshaler(t *testing.T) {
+func TestMarshalJSONPBMarshaler(t *testing.T) {
 	rawJson := `{ "foo": "bar", "baz": [0, 1, 2, 3] }`
 	msg := dynamicMessage{rawJson: rawJson}
 	str, err := new(Marshaler).MarshalToString(&msg)
@@ -473,6 +474,24 @@ func TestMarshalingWithJSONPBMarshaler(t *testing.T) {
 	}
 }
 
+func TestMarshalAnyJSONPBMarshaler(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)
+	}
+}
+
 var unmarshalingTests = []struct {
 	desc        string
 	unmarshaler Unmarshaler
@@ -670,11 +689,10 @@ func TestUnmarshalingBadInput(t *testing.T) {
 	}
 }
 
-func TestUnmarshalWithJSONPBUnmarshaler(t *testing.T) {
+func TestUnmarshalJSONPBUnmarshaler(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 {
@@ -682,10 +700,39 @@ func TestUnmarshalWithJSONPBUnmarshaler(t *testing.T) {
 	}
 }
 
+func TestUnmarshalAnyJSONPBUnmarshaler(t *testing.T) {
+	rawJson := `{ "@type": "blah.com/` + dynamicMessageName + `", "foo": "bar", "baz": [0, 1, 2, 3] }`
+	var got anypb.Any
+	if err := Unmarshal(strings.NewReader(rawJson), &got); err != nil {
+		t.Errorf("an unexpected error occurred when parsing into JSONPBUnmarshaler: %v", err)
+	}
+
+	dm := &dynamicMessage{rawJson: `{"baz":[0,1,2,3],"foo":"bar"}`}
+	var want anypb.Any
+	if b, err := proto.Marshal(dm); err != nil {
+		t.Errorf("an unexpected error occurred when marshaling message: %v", err)
+	} else {
+		want.TypeUrl = "blah.com/" + dynamicMessageName
+		want.Value = b
+	}
+
+	if !proto.Equal(&got, &want) {
+		t.Errorf("message contents not set correctly after unmarshalling JSON: got %s, wanted %s", got, want)
+	}
+}
+
+const (
+	dynamicMessageName = "google.protobuf.jsonpb.testing.dynamicMessage"
+)
+func init() {
+	// we register the custom type below so that we can use it in Any types
+	proto.RegisterType((*dynamicMessage)(nil), dynamicMessageName)
+}
+
 // 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 {
-	rawJson string
+	rawJson string `protobuf:"bytes,1,opt,name=rawJson"`
 }
 
 func (m *dynamicMessage) Reset() {
@@ -703,7 +750,7 @@ 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
-}
+}