codec.go

@@ -8,6 +8,11 @@ import (
 	"google.golang.org/grpc/encoding"
 )
 
+var (
+	_ encoding.Codec = &wrapMicroCodec{}
+	_ codec.Codec    = &wrapGrpcCodec{}
+)
+
 type wrapStream struct{ grpc.ClientStream }
 
 func (w *wrapStream) Write(d []byte) (int, error) {
@@ -29,20 +34,38 @@ func (w *wrapMicroCodec) Name() string {
 	return w.Codec.String()
 }
 
-type wrapGrpcCodec struct{ encoding.Codec }
-
-func (w *wrapGrpcCodec) String() string {
-	return w.Codec.Name()
-}
-
-func (w *wrapGrpcCodec) Marshal(v interface{}) ([]byte, error) {
+func (w *wrapMicroCodec) Marshal(v interface{}) ([]byte, error) {
 	if m, ok := v.(*codec.Frame); ok {
 		return m.Data, nil
 	}
 	return w.Codec.Marshal(v)
 }
 
-func (w *wrapGrpcCodec) Unmarshal(d []byte, v interface{}) error {
+func (w *wrapMicroCodec) Unmarshal(d []byte, v interface{}) error {
+	if d == nil || v == nil {
+		return nil
+	}
+	if m, ok := v.(*codec.Frame); ok {
+		m.Data = d
+		return nil
+	}
+	return w.Codec.Unmarshal(d, v)
+}
+
+type wrapGrpcCodec struct{ encoding.Codec }
+
+func (w *wrapGrpcCodec) String() string {
+	return w.Codec.Name()
+}
+
+func (w *wrapGrpcCodec) Marshal(v interface{}, opts ...codec.Option) ([]byte, error) {
+	if m, ok := v.(*codec.Frame); ok {
+		return m.Data, nil
+	}
+	return w.Codec.Marshal(v)
+}
+
+func (w *wrapGrpcCodec) Unmarshal(d []byte, v interface{}, opts ...codec.Option) error {
 	if d == nil || v == nil {
 		return nil
 	}