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

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