From 3573ac818f5ce53b9bdfb76db08314c57ffa131b Mon Sep 17 00:00:00 2001
From: magodo <wztdyl@sina.com>
Date: Sun, 9 Jun 2019 11:50:50 +0800
Subject: [PATCH] unmarshal json with `jsonpb` if accepter is pb

---
 codec/json/json.go      | 19 +++++++++++++------
 codec/json/marshaler.go |  6 ++++++
 2 files changed, 19 insertions(+), 6 deletions(-)

diff --git a/codec/json/json.go b/codec/json/json.go
index 42de4a4d..f17eef72 100644
--- a/codec/json/json.go
+++ b/codec/json/json.go
@@ -5,13 +5,16 @@ import (
 	"encoding/json"
 	"io"
 
+	"github.com/golang/protobuf/jsonpb"
+	"github.com/golang/protobuf/proto"
 	"github.com/micro/go-micro/codec"
 )
 
 type Codec struct {
-	Conn    io.ReadWriteCloser
-	Encoder *json.Encoder
-	Decoder *json.Decoder
+	Conn        io.ReadWriteCloser
+	Encoder     *json.Encoder
+	Decoder     *json.Decoder
+	Unmarshaler *jsonpb.Unmarshaler
 }
 
 func (c *Codec) ReadHeader(m *codec.Message, t codec.MessageType) error {
@@ -22,6 +25,9 @@ func (c *Codec) ReadBody(b interface{}) error {
 	if b == nil {
 		return nil
 	}
+	if pb, ok := b.(proto.Message); ok {
+		return c.Unmarshaler.UnmarshalNext(c.Decoder, pb)
+	}
 	return c.Decoder.Decode(b)
 }
 
@@ -42,8 +48,9 @@ func (c *Codec) String() string {
 
 func NewCodec(c io.ReadWriteCloser) codec.Codec {
 	return &Codec{
-		Conn:    c,
-		Decoder: json.NewDecoder(c),
-		Encoder: json.NewEncoder(c),
+		Conn:        c,
+		Decoder:     json.NewDecoder(c),
+		Encoder:     json.NewEncoder(c),
+		Unmarshaler: &jsonpb.Unmarshaler{},
 	}
 }
diff --git a/codec/json/marshaler.go b/codec/json/marshaler.go
index b9d0be28..a13e5f39 100644
--- a/codec/json/marshaler.go
+++ b/codec/json/marshaler.go
@@ -2,6 +2,9 @@ package json
 
 import (
 	"encoding/json"
+
+	"github.com/golang/protobuf/jsonpb"
+	"github.com/golang/protobuf/proto"
 )
 
 type Marshaler struct{}
@@ -11,6 +14,9 @@ func (j Marshaler) Marshal(v interface{}) ([]byte, error) {
 }
 
 func (j Marshaler) Unmarshal(d []byte, v interface{}) error {
+	if pb, ok := v.(proto.Message); ok {
+		return jsonpb.UnmarshalString(string(d), pb)
+	}
 	return json.Unmarshal(d, v)
 }