2 files changed, 157 insertions, 16 deletions
diff --git a/vendor/github.com/golang/protobuf/jsonpb/jsonpb.go b/vendor/github.com/golang/protobuf/jsonpb/jsonpb.go
index 29bca020f..110ae1384 100644
--- a/vendor/github.com/golang/protobuf/jsonpb/jsonpb.go
+++ b/vendor/github.com/golang/protobuf/jsonpb/jsonpb.go
@@ -73,6 +73,31 @@ type Marshaler struct {
 
 	// Whether to use the original (.proto) name for fields.
 	OrigName bool
+
+	// A custom URL resolver to use when marshaling Any messages to JSON.
+	// If unset, the default resolution strategy is to extract the
+	// fully-qualified type name from the type URL and pass that to
+	// proto.MessageType(string).
+	AnyResolver AnyResolver
+}
+
+// AnyResolver takes a type URL, present in an Any message, and resolves it into
+// an instance of the associated message.
+type AnyResolver interface {
+	Resolve(typeUrl string) (proto.Message, error)
+}
+
+func defaultResolveAny(typeUrl string) (proto.Message, error) {
+	// Only the part of typeUrl after the last slash is relevant.
+	mname := typeUrl
+	if slash := strings.LastIndex(mname, "/"); slash >= 0 {
+		mname = mname[slash+1:]
+	}
+	mt := proto.MessageType(mname)
+	if mt == nil {
+		return nil, fmt.Errorf("unknown message type %q", mname)
+	}
+	return reflect.New(mt.Elem()).Interface().(proto.Message), nil
 }
 
 // JSONPBMarshaler is implemented by protobuf messages that customize the
@@ -344,16 +369,17 @@ func (m *Marshaler) marshalAny(out *errWriter, any proto.Message, indent string)
 	turl := v.Field(0).String()
 	val := v.Field(1).Bytes()
 
-	// Only the part of type_url after the last slash is relevant.
-	mname := turl
-	if slash := strings.LastIndex(mname, "/"); slash >= 0 {
-		mname = mname[slash+1:]
+	var msg proto.Message
+	var err error
+	if m.AnyResolver != nil {
+		msg, err = m.AnyResolver.Resolve(turl)
+	} else {
+		msg, err = defaultResolveAny(turl)
 	}
-	mt := proto.MessageType(mname)
-	if mt == nil {
-		return fmt.Errorf("unknown message type %q", mname)
+	if err != nil {
+		return err
 	}
-	msg := reflect.New(mt.Elem()).Interface().(proto.Message)
+
 	if err := proto.Unmarshal(val, msg); err != nil {
 		return err
 	}
@@ -590,6 +616,12 @@ type Unmarshaler struct {
 	// Whether to allow messages to contain unknown fields, as opposed to
 	// failing to unmarshal.
 	AllowUnknownFields bool
+
+	// A custom URL resolver to use when unmarshaling Any messages from JSON.
+	// If unset, the default resolution strategy is to extract the
+	// fully-qualified type name from the type URL and pass that to
+	// proto.MessageType(string).
+	AnyResolver AnyResolver
 }
 
 // UnmarshalNext unmarshals the next protocol buffer from a JSON object stream.
@@ -641,7 +673,8 @@ func (u *Unmarshaler) unmarshalValue(target reflect.Value, inputValue json.RawMe
 	if targetType.Kind() == reflect.Ptr {
 		// If input value is "null" and target is a pointer type, then the field should be treated as not set
 		// UNLESS the target is structpb.Value, in which case it should be set to structpb.NullValue.
-		if string(inputValue) == "null" && targetType != reflect.TypeOf(&stpb.Value{}) {
+		_, isJSONPBUnmarshaler := target.Interface().(JSONPBUnmarshaler)
+		if string(inputValue) == "null" && targetType != reflect.TypeOf(&stpb.Value{}) && !isJSONPBUnmarshaler {
 			return nil
 		}
 		target.Set(reflect.New(targetType.Elem()))
@@ -679,16 +712,17 @@ func (u *Unmarshaler) unmarshalValue(target reflect.Value, inputValue json.RawMe
 			}
 			target.Field(0).SetString(turl)
 
-			mname := turl
-			if slash := strings.LastIndex(mname, "/"); slash >= 0 {
-				mname = mname[slash+1:]
+			var m proto.Message
+			var err error
+			if u.AnyResolver != nil {
+				m, err = u.AnyResolver.Resolve(turl)
+			} else {
+				m, err = defaultResolveAny(turl)
 			}
-			mt := proto.MessageType(mname)
-			if mt == nil {
-				return fmt.Errorf("unknown message type %q", mname)
+			if err != nil {
+				return err
 			}
 
-			m := reflect.New(mt.Elem()).Interface().(proto.Message)
 			if _, ok := m.(wkt); ok {
 				val, ok := jsonFields["value"]
 				if !ok {
diff --git a/vendor/github.com/golang/protobuf/jsonpb/jsonpb_test.go b/vendor/github.com/golang/protobuf/jsonpb/jsonpb_test.go
index 254caa6c4..2428d0566 100644
--- a/vendor/github.com/golang/protobuf/jsonpb/jsonpb_test.go
+++ b/vendor/github.com/golang/protobuf/jsonpb/jsonpb_test.go
@@ -721,6 +721,65 @@ func TestUnmarshalingBadInput(t *testing.T) {
 	}
 }
 
+type funcResolver func(turl string) (proto.Message, error)
+
+func (fn funcResolver) Resolve(turl string) (proto.Message, error) {
+	return fn(turl)
+}
+
+func TestAnyWithCustomResolver(t *testing.T) {
+	var resolvedTypeUrls []string
+	resolver := funcResolver(func(turl string) (proto.Message, error) {
+		resolvedTypeUrls = append(resolvedTypeUrls, turl)
+		return new(pb.Simple), nil
+	})
+	msg := &pb.Simple{
+		OBytes:  []byte{1, 2, 3, 4},
+		OBool:   proto.Bool(true),
+		OString: proto.String("foobar"),
+		OInt64:  proto.Int64(1020304),
+	}
+	msgBytes, err := proto.Marshal(msg)
+	if err != nil {
+		t.Errorf("an unexpected error occurred when marshaling message: %v", err)
+	}
+	// make an Any with a type URL that won't resolve w/out custom resolver
+	any := &anypb.Any{
+		TypeUrl: "https://foobar.com/some.random.MessageKind",
+		Value:   msgBytes,
+	}
+
+	m := Marshaler{AnyResolver: resolver}
+	js, err := m.MarshalToString(any)
+	if err != nil {
+		t.Errorf("an unexpected error occurred when marshaling any to JSON: %v", err)
+	}
+	if len(resolvedTypeUrls) != 1 {
+		t.Errorf("custom resolver was not invoked during marshaling")
+	} else if resolvedTypeUrls[0] != "https://foobar.com/some.random.MessageKind" {
+		t.Errorf("custom resolver was invoked with wrong URL: got %q, wanted %q", resolvedTypeUrls[0], "https://foobar.com/some.random.MessageKind")
+	}
+	wanted := `{"@type":"https://foobar.com/some.random.MessageKind","oBool":true,"oInt64":"1020304","oString":"foobar","oBytes":"AQIDBA=="}`
+	if js != wanted {
+		t.Errorf("marshalling JSON produced incorrect output: got %s, wanted %s", js, wanted)
+	}
+
+	u := Unmarshaler{AnyResolver: resolver}
+	roundTrip := &anypb.Any{}
+	err = u.Unmarshal(bytes.NewReader([]byte(js)), roundTrip)
+	if err != nil {
+		t.Errorf("an unexpected error occurred when unmarshaling any from JSON: %v", err)
+	}
+	if len(resolvedTypeUrls) != 2 {
+		t.Errorf("custom resolver was not invoked during marshaling")
+	} else if resolvedTypeUrls[1] != "https://foobar.com/some.random.MessageKind" {
+		t.Errorf("custom resolver was invoked with wrong URL: got %q, wanted %q", resolvedTypeUrls[1], "https://foobar.com/some.random.MessageKind")
+	}
+	if !proto.Equal(any, roundTrip) {
+		t.Errorf("message contents not set correctly after unmarshalling JSON: got %s, wanted %s", roundTrip, any)
+	}
+}
+
 func TestUnmarshalJSONPBUnmarshaler(t *testing.T) {
 	rawJson := `{ "foo": "bar", "baz": [0, 1, 2, 3] }`
 	var msg dynamicMessage
@@ -732,6 +791,19 @@ func TestUnmarshalJSONPBUnmarshaler(t *testing.T) {
 	}
 }
 
+func TestUnmarshalNullWithJSONPBUnmarshaler(t *testing.T) {
+	rawJson := `{"stringField":null}`
+	var ptrFieldMsg ptrFieldMessage
+	if err := Unmarshal(strings.NewReader(rawJson), &ptrFieldMsg); err != nil {
+		t.Errorf("unmarshal error: %v", err)
+	}
+
+	want := ptrFieldMessage{StringField: &stringField{IsSet: true, StringValue: "null"}}
+	if !proto.Equal(&ptrFieldMsg, &want) {
+		t.Errorf("unmarshal result StringField: got %v, want %v", ptrFieldMsg, want)
+	}
+}
+
 func TestUnmarshalAnyJSONPBUnmarshaler(t *testing.T) {
 	rawJson := `{ "@type": "blah.com/` + dynamicMessageName + `", "foo": "bar", "baz": [0, 1, 2, 3] }`
 	var got anypb.Any
@@ -762,6 +834,41 @@ func init() {
 	proto.RegisterType((*dynamicMessage)(nil), dynamicMessageName)
 }
 
+type ptrFieldMessage struct {
+	StringField *stringField `protobuf:"bytes,1,opt,name=stringField"`
+}
+
+func (m *ptrFieldMessage) Reset() {
+}
+
+func (m *ptrFieldMessage) String() string {
+	return m.StringField.StringValue
+}
+
+func (m *ptrFieldMessage) ProtoMessage() {
+}
+
+type stringField struct {
+	IsSet       bool   `protobuf:"varint,1,opt,name=isSet"`
+	StringValue string `protobuf:"bytes,2,opt,name=stringValue"`
+}
+
+func (s *stringField) Reset() {
+}
+
+func (s *stringField) String() string {
+	return s.StringValue
+}
+
+func (s *stringField) ProtoMessage() {
+}
+
+func (s *stringField) UnmarshalJSONPB(jum *Unmarshaler, js []byte) error {
+	s.IsSet = true
+	s.StringValue = string(js)
+	return nil
+}
+
 // 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 {