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

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