github.com/tickoalcantara12/micro/v3@v3.0.0-20221007104245-9d75b9bcbab9/util/codec/grpc/util.go (about)

     1  // Licensed under the Apache License, Version 2.0 (the "License");
     2  // you may not use this file except in compliance with the License.
     3  // You may obtain a copy of the License at
     4  //
     5  //     https://www.apache.org/licenses/LICENSE-2.0
     6  //
     7  // Unless required by applicable law or agreed to in writing, software
     8  // distributed under the License is distributed on an "AS IS" BASIS,
     9  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    10  // See the License for the specific language governing permissions and
    11  // limitations under the License.
    12  //
    13  // Original source: github.com/micro/go-micro/v3/codec/grpc/util.go
    14  
    15  package grpc
    16  
    17  import (
    18  	"encoding/binary"
    19  	"fmt"
    20  	"io"
    21  )
    22  
    23  var (
    24  	MaxMessageSize = 1024 * 1024 * 4 // 4Mb
    25  	maxInt         = int(^uint(0) >> 1)
    26  )
    27  
    28  func decode(r io.Reader) (uint8, []byte, error) {
    29  	header := make([]byte, 5)
    30  
    31  	// read the header
    32  	if _, err := r.Read(header[:]); err != nil {
    33  		return uint8(0), nil, err
    34  	}
    35  
    36  	// get encoding format e.g compressed
    37  	cf := uint8(header[0])
    38  
    39  	// get message length
    40  	length := binary.BigEndian.Uint32(header[1:])
    41  
    42  	// no encoding format
    43  	if length == 0 {
    44  		return cf, nil, nil
    45  	}
    46  
    47  	//
    48  	if int64(length) > int64(maxInt) {
    49  		return cf, nil, fmt.Errorf("grpc: received message larger than max length allowed on current machine (%d vs. %d)", length, maxInt)
    50  	}
    51  	if int(length) > MaxMessageSize {
    52  		return cf, nil, fmt.Errorf("grpc: received message larger than max (%d vs. %d)", length, MaxMessageSize)
    53  	}
    54  
    55  	msg := make([]byte, int(length))
    56  
    57  	if _, err := r.Read(msg); err != nil {
    58  		if err == io.EOF {
    59  			err = io.ErrUnexpectedEOF
    60  		}
    61  		return cf, nil, err
    62  	}
    63  
    64  	return cf, msg, nil
    65  }
    66  
    67  func encode(cf uint8, buf []byte, w io.Writer) error {
    68  	header := make([]byte, 5)
    69  
    70  	// set compression
    71  	header[0] = byte(cf)
    72  
    73  	// write length as header
    74  	binary.BigEndian.PutUint32(header[1:], uint32(len(buf)))
    75  
    76  	// read the header
    77  	if _, err := w.Write(header[:]); err != nil {
    78  		return err
    79  	}
    80  
    81  	// write the buffer
    82  	_, err := w.Write(buf)
    83  	return err
    84  }