2019-06-03 18:44:43 +01:00
|
|
|
package grpc
|
|
|
|
|
|
|
|
import (
|
2020-11-26 01:17:21 +03:00
|
|
|
"io"
|
2019-06-03 18:44:43 +01:00
|
|
|
|
2020-09-20 16:08:45 +03:00
|
|
|
"github.com/unistack-org/micro/v3/codec"
|
2019-06-03 18:44:43 +01:00
|
|
|
"google.golang.org/grpc/encoding"
|
|
|
|
)
|
|
|
|
|
2020-11-26 01:17:21 +03:00
|
|
|
type wrapMicroCodec struct{ codec.Codec }
|
2019-06-17 20:05:58 +01:00
|
|
|
|
2020-11-26 01:17:21 +03:00
|
|
|
func (w *wrapMicroCodec) Name() string {
|
|
|
|
return w.Codec.String()
|
2019-06-03 18:44:43 +01:00
|
|
|
}
|
|
|
|
|
2020-11-26 01:17:21 +03:00
|
|
|
type wrapGrpcCodec struct{ encoding.Codec }
|
2019-06-03 18:44:43 +01:00
|
|
|
|
2020-11-26 01:17:21 +03:00
|
|
|
func (w *wrapGrpcCodec) String() string {
|
|
|
|
return w.Codec.Name()
|
2019-06-03 18:44:43 +01:00
|
|
|
}
|
|
|
|
|
2020-11-26 01:17:21 +03:00
|
|
|
func (w *wrapGrpcCodec) Marshal(v interface{}) ([]byte, error) {
|
2021-04-26 19:04:27 +03:00
|
|
|
if m, ok := v.(*codec.Frame); ok {
|
2020-11-26 01:17:21 +03:00
|
|
|
return m.Data, nil
|
2019-06-03 18:44:43 +01:00
|
|
|
}
|
2020-11-26 01:17:21 +03:00
|
|
|
return w.Codec.Marshal(v)
|
2019-06-03 18:44:43 +01:00
|
|
|
}
|
|
|
|
|
2021-04-26 19:04:27 +03:00
|
|
|
func (w *wrapGrpcCodec) Unmarshal(d []byte, v interface{}) error {
|
2020-11-26 01:17:21 +03:00
|
|
|
if d == nil || v == nil {
|
2020-09-20 16:08:45 +03:00
|
|
|
return nil
|
|
|
|
}
|
2021-04-26 19:04:27 +03:00
|
|
|
if m, ok := v.(*codec.Frame); ok {
|
2020-11-26 01:17:21 +03:00
|
|
|
m.Data = d
|
2020-09-20 16:08:45 +03:00
|
|
|
return nil
|
|
|
|
}
|
2020-11-26 01:17:21 +03:00
|
|
|
return w.Codec.Unmarshal(d, v)
|
2019-06-03 18:44:43 +01:00
|
|
|
}
|
2019-06-18 18:51:52 +01:00
|
|
|
|
2021-04-26 19:04:27 +03:00
|
|
|
func (w *wrapGrpcCodec) ReadHeader(conn io.Reader, m *codec.Message, mt codec.MessageType) error {
|
2019-06-18 18:51:52 +01:00
|
|
|
return nil
|
|
|
|
}
|
|
|
|
|
2021-04-26 19:04:27 +03:00
|
|
|
func (w *wrapGrpcCodec) ReadBody(conn io.Reader, v interface{}) error {
|
|
|
|
if m, ok := v.(*codec.Frame); ok {
|
2020-11-26 01:17:21 +03:00
|
|
|
_, err := conn.Read(m.Data)
|
|
|
|
return err
|
2019-06-18 18:51:52 +01:00
|
|
|
}
|
2020-11-26 01:17:21 +03:00
|
|
|
return codec.ErrInvalidMessage
|
2019-06-18 18:51:52 +01:00
|
|
|
}
|
|
|
|
|
2021-04-26 19:04:27 +03:00
|
|
|
func (w *wrapGrpcCodec) Write(conn io.Writer, m *codec.Message, v interface{}) error {
|
2019-06-18 18:51:52 +01:00
|
|
|
// if we don't have a body
|
|
|
|
if v != nil {
|
2021-04-26 19:04:27 +03:00
|
|
|
b, err := w.Marshal(v)
|
2019-06-18 18:51:52 +01:00
|
|
|
if err != nil {
|
|
|
|
return err
|
|
|
|
}
|
|
|
|
m.Body = b
|
|
|
|
}
|
|
|
|
// write the body using the framing codec
|
2020-11-26 01:17:21 +03:00
|
|
|
_, err := conn.Write(m.Body)
|
|
|
|
return err
|
2019-06-18 18:51:52 +01:00
|
|
|
}
|