Add unmarshal to struct provided by client

mxpv · crosbymichael · commit cd3ce7159eae · 2020-04-13T14:06:23.000-04:00
Signed-off-by: Maksym Pavlenko &lt;pavlenko.maksym@gmail.com&gt;
diff --git a/types.go b/types.go
@@ -47,14 +47,14 @@ func Register(v interface{}, args ...string) {
 	defer mu.Unlock()
 	if et, ok := registry[t]; ok {
 		if et != p {
-			panic(errors.Errorf("type registred with alternate path %q != %q", et, p))
+			panic(errors.Errorf("type registered with alternate path %q != %q", et, p))
 		}
 		return
 	}
 	registry[t] = p
 }
 
-// TypeURL returns the type url for a registred type.
+// TypeURL returns the type url for a registered type.
 func TypeURL(v interface{}) (string, error) {
 	mu.Lock()
 	u, ok := registry[tryDereference(v)]
@@ -120,16 +120,43 @@ func UnmarshalAny(any *types.Any) (interface{}, error) {
 }
 
 func UnmarshalByTypeURL(typeURL string, value []byte) (interface{}, error) {
+	return unmarshal(typeURL, value, nil)
+}
+
+func UnmarshalTo(any *types.Any, out interface{}) error {
+	return UnmarshalToByTypeURL(any.TypeUrl, any.Value, out)
+}
+
+func UnmarshalToByTypeURL(typeURL string, value []byte, out interface{}) error {
+	_, err := unmarshal(typeURL, value, out)
+	return err
+}
+
+func unmarshal(typeURL string, value []byte, v interface{}) (interface{}, error) {
 	t, err := getTypeByUrl(typeURL)
 	if err != nil {
 		return nil, err
 	}
-	v := reflect.New(t.t).Interface()
+
+	if v == nil {
+		v = reflect.New(t.t).Interface()
+	} else {
+		// Validate interface type provided by client
+		vURL, err := TypeURL(v)
+		if err != nil {
+			return nil, err
+		}
+		if typeURL != vURL {
+			return nil, errors.Errorf("can't unmarshal type %q to output %q", typeURL, vURL)
+		}
+	}
+
 	if t.isProto {
 		err = proto.Unmarshal(value, v.(proto.Message))
 	} else {
 		err = json.Unmarshal(value, v)
 	}
+
 	return v, err
 }
 
diff --git a/types_test.go b/types_test.go
@@ -101,6 +101,56 @@ func TestMarshalUnmarshal(t *testing.T) {
 	}
 }
 
+func TestMarshalUnmarshalTo(t *testing.T) {
+	clear()
+	Register(&test{}, "test")
+
+	in := &test{
+		Name: "koye",
+		Age:  6,
+	}
+	any, err := MarshalAny(in)
+	if err != nil {
+		t.Fatal(err)
+	}
+	out := &test{}
+	err = UnmarshalTo(any, out)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if out.Name != "koye" {
+		t.Fatal("invalid name")
+	}
+	if out.Age != 6 {
+		t.Fatal("invalid age")
+	}
+}
+
+type test2 struct {
+	Name string
+}
+
+func TestUnmarshalToInvalid(t *testing.T) {
+	clear()
+	Register(&test{}, "test1")
+	Register(&test2{}, "test2")
+
+	in := &test{
+		Name: "koye",
+		Age:  6,
+	}
+	any, err := MarshalAny(in)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	out := &test2{}
+	err = UnmarshalTo(any, out)
+	if err == nil || err.Error() != `can't unmarshal type "test1" to output "test2"` {
+		t.Fatalf("unexpected result: %+v", err)
+	}
+}
+
 func TestIs(t *testing.T) {
 	clear()
 	Register(&test{}, "test")
