diff --git a/codec.go b/codec.go index 65363d9..be2f9a2 100644 --- a/codec.go +++ b/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 }