gitee.com/liuxuezhan/go-micro-v1.18.0@v1.0.0/client/grpc/codec.go (about)

     1  package grpc
     2  
     3  import (
     4  	"fmt"
     5  	"strings"
     6  
     7  	b "bytes"
     8  
     9  	"github.com/golang/protobuf/jsonpb"
    10  	"github.com/golang/protobuf/proto"
    11  	jsoniter "github.com/json-iterator/go"
    12  	"gitee.com/liuxuezhan/go-micro-v1.18.0/codec"
    13  	"gitee.com/liuxuezhan/go-micro-v1.18.0/codec/bytes"
    14  	"google.golang.org/grpc"
    15  	"google.golang.org/grpc/encoding"
    16  )
    17  
    18  type jsonCodec struct{}
    19  type protoCodec struct{}
    20  type bytesCodec struct{}
    21  type wrapCodec struct{ encoding.Codec }
    22  
    23  var jsonpbMarshaler = &jsonpb.Marshaler{}
    24  
    25  var (
    26  	defaultGRPCCodecs = map[string]encoding.Codec{
    27  		"application/json":         jsonCodec{},
    28  		"application/proto":        protoCodec{},
    29  		"application/protobuf":     protoCodec{},
    30  		"application/octet-stream": protoCodec{},
    31  		"application/grpc":         protoCodec{},
    32  		"application/grpc+json":    jsonCodec{},
    33  		"application/grpc+proto":   protoCodec{},
    34  		"application/grpc+bytes":   bytesCodec{},
    35  	}
    36  
    37  	json = jsoniter.ConfigCompatibleWithStandardLibrary
    38  )
    39  
    40  // UseNumber fix unmarshal Number(8234567890123456789) to interface(8.234567890123457e+18)
    41  func UseNumber() {
    42  	json = jsoniter.Config{
    43  		UseNumber:              true,
    44  		EscapeHTML:             true,
    45  		SortMapKeys:            true,
    46  		ValidateJsonRawMessage: true,
    47  	}.Froze()
    48  }
    49  
    50  func (w wrapCodec) String() string {
    51  	return w.Codec.Name()
    52  }
    53  
    54  func (w wrapCodec) Marshal(v interface{}) ([]byte, error) {
    55  	b, ok := v.(*bytes.Frame)
    56  	if ok {
    57  		return b.Data, nil
    58  	}
    59  	return w.Codec.Marshal(v)
    60  }
    61  
    62  func (w wrapCodec) Unmarshal(data []byte, v interface{}) error {
    63  	b, ok := v.(*bytes.Frame)
    64  	if ok {
    65  		b.Data = data
    66  		return nil
    67  	}
    68  	return w.Codec.Unmarshal(data, v)
    69  }
    70  
    71  func (protoCodec) Marshal(v interface{}) ([]byte, error) {
    72  	b, ok := v.(*bytes.Frame)
    73  	if ok {
    74  		return b.Data, nil
    75  	}
    76  	return proto.Marshal(v.(proto.Message))
    77  }
    78  
    79  func (protoCodec) Unmarshal(data []byte, v interface{}) error {
    80  	return proto.Unmarshal(data, v.(proto.Message))
    81  }
    82  
    83  func (protoCodec) Name() string {
    84  	return "proto"
    85  }
    86  
    87  func (bytesCodec) Marshal(v interface{}) ([]byte, error) {
    88  	b, ok := v.(*[]byte)
    89  	if !ok {
    90  		return nil, fmt.Errorf("failed to marshal: %v is not type of *[]byte", v)
    91  	}
    92  	return *b, nil
    93  }
    94  
    95  func (bytesCodec) Unmarshal(data []byte, v interface{}) error {
    96  	b, ok := v.(*[]byte)
    97  	if !ok {
    98  		return fmt.Errorf("failed to unmarshal: %v is not type of *[]byte", v)
    99  	}
   100  	*b = data
   101  	return nil
   102  }
   103  
   104  func (bytesCodec) Name() string {
   105  	return "bytes"
   106  }
   107  
   108  func (jsonCodec) Marshal(v interface{}) ([]byte, error) {
   109  	if pb, ok := v.(proto.Message); ok {
   110  		s, err := jsonpbMarshaler.MarshalToString(pb)
   111  
   112  		return []byte(s), err
   113  	}
   114  
   115  	return json.Marshal(v)
   116  }
   117  
   118  func (jsonCodec) Unmarshal(data []byte, v interface{}) error {
   119  	if pb, ok := v.(proto.Message); ok {
   120  		return jsonpb.Unmarshal(b.NewReader(data), pb)
   121  	}
   122  
   123  	return json.Unmarshal(data, v)
   124  }
   125  
   126  func (jsonCodec) Name() string {
   127  	return "json"
   128  }
   129  
   130  type grpcCodec struct {
   131  	// headers
   132  	id       string
   133  	target   string
   134  	method   string
   135  	endpoint string
   136  
   137  	s grpc.ClientStream
   138  	c encoding.Codec
   139  }
   140  
   141  func (g *grpcCodec) ReadHeader(m *codec.Message, mt codec.MessageType) error {
   142  	md, err := g.s.Header()
   143  	if err != nil {
   144  		return err
   145  	}
   146  	if m == nil {
   147  		m = new(codec.Message)
   148  	}
   149  	if m.Header == nil {
   150  		m.Header = make(map[string]string)
   151  	}
   152  	for k, v := range md {
   153  		m.Header[k] = strings.Join(v, ",")
   154  	}
   155  	m.Id = g.id
   156  	m.Target = g.target
   157  	m.Method = g.method
   158  	m.Endpoint = g.endpoint
   159  	return nil
   160  }
   161  
   162  func (g *grpcCodec) ReadBody(v interface{}) error {
   163  	if f, ok := v.(*bytes.Frame); ok {
   164  		return g.s.RecvMsg(f)
   165  	}
   166  	return g.s.RecvMsg(v)
   167  }
   168  
   169  func (g *grpcCodec) Write(m *codec.Message, v interface{}) error {
   170  	// if we don't have a body
   171  	if v != nil {
   172  		return g.s.SendMsg(v)
   173  	}
   174  	// write the body using the framing codec
   175  	return g.s.SendMsg(&bytes.Frame{Data: m.Body})
   176  }
   177  
   178  func (g *grpcCodec) Close() error {
   179  	return g.s.CloseSend()
   180  }
   181  
   182  func (g *grpcCodec) String() string {
   183  	return g.c.Name()
   184  }