github.com/micro/go-micro/v2@v2.9.1/server/grpc/codec.go (about)

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