micro/codec/grpc/util.go

71 lines
1.3 KiB
Go
Raw Normal View History

2019-01-02 15:55:06 +03:00
package grpc
import (
"encoding/binary"
"fmt"
"io"
)
var (
maxMessageSize = 1024 * 1024 * 4
maxInt = int(^uint(0) >> 1)
)
func decode(r io.Reader) (uint8, []byte, error) {
header := make([]byte, 5)
// read the header
if _, err := r.Read(header[:]); err != nil {
return uint8(0), nil, err
}
// get encoding format e.g compressed
cf := uint8(header[0])
// get message length
length := binary.BigEndian.Uint32(header[1:])
// no encoding format
if length == 0 {
return cf, nil, nil
}
//
if int64(length) > int64(maxInt) {
return cf, nil, fmt.Errorf("grpc: received message larger than max length allowed on current machine (%d vs. %d)", length, maxInt)
}
if int(length) > maxMessageSize {
return cf, nil, fmt.Errorf("grpc: received message larger than max (%d vs. %d)", length, maxMessageSize)
}
msg := make([]byte, int(length))
if _, err := r.Read(msg); err != nil {
if err == io.EOF {
err = io.ErrUnexpectedEOF
}
return cf, nil, err
}
return cf, msg, nil
}
func encode(cf uint8, buf []byte, w io.Writer) error {
header := make([]byte, 5)
// set compression
header[0] = byte(cf)
// write length as header
binary.BigEndian.PutUint32(header[1:], uint32(len(buf)))
// read the header
if _, err := w.Write(header[:]); err != nil {
return err
}
// write the buffer
_, err := w.Write(buf)
return err
}