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

     1  package grpc
     2  
     3  import (
     4  	"encoding/binary"
     5  	"fmt"
     6  	"io"
     7  )
     8  
     9  var (
    10  	maxMessageSize = 1024 * 1024 * 4
    11  	maxInt         = int(^uint(0) >> 1)
    12  )
    13  
    14  func decode(r io.Reader) (uint8, []byte, error) {
    15  	header := make([]byte, 5)
    16  
    17  	// read the header
    18  	if _, err := r.Read(header[:]); err != nil {
    19  		return uint8(0), nil, err
    20  	}
    21  
    22  	// get encoding format e.g compressed
    23  	cf := uint8(header[0])
    24  
    25  	// get message length
    26  	length := binary.BigEndian.Uint32(header[1:])
    27  
    28  	// no encoding format
    29  	if length == 0 {
    30  		return cf, nil, nil
    31  	}
    32  
    33  	//
    34  	if int64(length) > int64(maxInt) {
    35  		return cf, nil, fmt.Errorf("grpc: received message larger than max length allowed on current machine (%d vs. %d)", length, maxInt)
    36  	}
    37  	if int(length) > maxMessageSize {
    38  		return cf, nil, fmt.Errorf("grpc: received message larger than max (%d vs. %d)", length, maxMessageSize)
    39  	}
    40  
    41  	msg := make([]byte, int(length))
    42  
    43  	if _, err := r.Read(msg); err != nil {
    44  		if err == io.EOF {
    45  			err = io.ErrUnexpectedEOF
    46  		}
    47  		return cf, nil, err
    48  	}
    49  
    50  	return cf, msg, nil
    51  }
    52  
    53  func encode(cf uint8, buf []byte, w io.Writer) error {
    54  	header := make([]byte, 5)
    55  
    56  	// set compression
    57  	header[0] = byte(cf)
    58  
    59  	// write length as header
    60  	binary.BigEndian.PutUint32(header[1:], uint32(len(buf)))
    61  
    62  	// read the header
    63  	if _, err := w.Write(header[:]); err != nil {
    64  		return err
    65  	}
    66  
    67  	// write the buffer
    68  	_, err := w.Write(buf)
    69  	return err
    70  }