github.com/btccom/go-micro/v2@v2.9.3/codec/grpc/grpc.go (about)

     1  // Package grpc provides a grpc codec
     2  package grpc
     3  
     4  import (
     5  	"encoding/json"
     6  	"errors"
     7  	"fmt"
     8  	"io"
     9  	"strings"
    10  
    11  	"github.com/golang/protobuf/proto"
    12  	"github.com/btccom/go-micro/v2/codec"
    13  )
    14  
    15  type Codec struct {
    16  	Conn        io.ReadWriteCloser
    17  	ContentType string
    18  }
    19  
    20  func (c *Codec) ReadHeader(m *codec.Message, t codec.MessageType) error {
    21  	if ct := m.Header["Content-Type"]; len(ct) > 0 {
    22  		c.ContentType = ct
    23  	}
    24  
    25  	if ct := m.Header["content-type"]; len(ct) > 0 {
    26  		c.ContentType = ct
    27  	}
    28  
    29  	// service method
    30  	path := m.Header[":path"]
    31  	if len(path) == 0 || path[0] != '/' {
    32  		m.Target = m.Header["Micro-Service"]
    33  		m.Endpoint = m.Header["Micro-Endpoint"]
    34  	} else {
    35  		// [ , a.package.Foo, Bar]
    36  		parts := strings.Split(path, "/")
    37  		if len(parts) != 3 {
    38  			return errors.New("Unknown request path")
    39  		}
    40  		service := strings.Split(parts[1], ".")
    41  		m.Endpoint = strings.Join([]string{service[len(service)-1], parts[2]}, ".")
    42  		m.Target = strings.Join(service[:len(service)-1], ".")
    43  	}
    44  
    45  	return nil
    46  }
    47  
    48  func (c *Codec) ReadBody(b interface{}) error {
    49  	// no body
    50  	if b == nil {
    51  		return nil
    52  	}
    53  
    54  	_, buf, err := decode(c.Conn)
    55  	if err != nil {
    56  		return err
    57  	}
    58  
    59  	switch c.ContentType {
    60  	case "application/grpc+json":
    61  		return json.Unmarshal(buf, b)
    62  	case "application/grpc+proto", "application/grpc":
    63  		return proto.Unmarshal(buf, b.(proto.Message))
    64  	}
    65  
    66  	return errors.New("Unsupported Content-Type")
    67  }
    68  
    69  func (c *Codec) Write(m *codec.Message, b interface{}) error {
    70  	var buf []byte
    71  	var err error
    72  
    73  	if ct := m.Header["Content-Type"]; len(ct) > 0 {
    74  		c.ContentType = ct
    75  	}
    76  
    77  	if ct := m.Header["content-type"]; len(ct) > 0 {
    78  		c.ContentType = ct
    79  	}
    80  
    81  	switch m.Type {
    82  	case codec.Request:
    83  		parts := strings.Split(m.Endpoint, ".")
    84  		m.Header[":method"] = "POST"
    85  		m.Header[":path"] = fmt.Sprintf("/%s.%s/%s", m.Target, parts[0], parts[1])
    86  		m.Header[":proto"] = "HTTP/2.0"
    87  		m.Header["te"] = "trailers"
    88  		m.Header["user-agent"] = "grpc-go/1.0.0"
    89  		m.Header[":authority"] = m.Target
    90  		m.Header["content-type"] = c.ContentType
    91  	case codec.Response:
    92  		m.Header["Trailer"] = "grpc-status" //, grpc-message"
    93  		m.Header["content-type"] = c.ContentType
    94  		m.Header[":status"] = "200"
    95  		m.Header["grpc-status"] = "0"
    96  		//		m.Header["grpc-message"] = ""
    97  	case codec.Error:
    98  		m.Header["Trailer"] = "grpc-status, grpc-message"
    99  		// micro end of stream
   100  		if m.Error == "EOS" {
   101  			m.Header["grpc-status"] = "0"
   102  		} else {
   103  			m.Header["grpc-message"] = m.Error
   104  			m.Header["grpc-status"] = "13"
   105  		}
   106  
   107  		return nil
   108  	}
   109  
   110  	// marshal content
   111  	switch c.ContentType {
   112  	case "application/grpc+json":
   113  		buf, err = json.Marshal(b)
   114  	case "application/grpc+proto", "application/grpc":
   115  		pb, ok := b.(proto.Message)
   116  		if ok {
   117  			buf, err = proto.Marshal(pb)
   118  		}
   119  	default:
   120  		err = errors.New("Unsupported Content-Type")
   121  	}
   122  	// check error
   123  	if err != nil {
   124  		m.Header["grpc-status"] = "8"
   125  		m.Header["grpc-message"] = err.Error()
   126  		return err
   127  	}
   128  
   129  	if len(buf) == 0 {
   130  		return nil
   131  	}
   132  
   133  	return encode(0, buf, c.Conn)
   134  }
   135  
   136  func (c *Codec) Close() error {
   137  	return c.Conn.Close()
   138  }
   139  
   140  func (c *Codec) String() string {
   141  	return "grpc"
   142  }
   143  
   144  func NewCodec(c io.ReadWriteCloser) codec.Codec {
   145  	return &Codec{
   146  		Conn:        c,
   147  		ContentType: "application/grpc",
   148  	}
   149  }