github.com/kaiya/goutils@v1.0.1-0.20230226104005-4ae4a4dc3688/tinyrpc/rpc/packet.go (about)

     1  package rpc
     2  
     3  import (
     4  	"encoding/binary"
     5  	"hash/adler32"
     6  	"io"
     7  
     8  	"github.com/pkg/errors"
     9  )
    10  
    11  var RPC_MAGIC = [4]byte{'p', 'y', 'x', 'i'}
    12  
    13  type Packet struct {
    14  	TotalSize uint32
    15  	Magic     [4]byte
    16  	Payload   []byte
    17  	Checksum  uint32
    18  }
    19  
    20  func EncodePacket(w io.Writer, payload []byte) error {
    21  	// len(Magic) + len(Checksum) == 8
    22  	totalsize := uint32(len(RPC_MAGIC) + len(payload) + 4)
    23  	// write totalsize
    24  	binary.Write(w, binary.BigEndian, totalsize)
    25  
    26  	sum := adler32.New()
    27  	ww := io.MultiWriter(sum, w)
    28  	// write magic bytes
    29  	binary.Write(ww, binary.BigEndian, RPC_MAGIC)
    30  
    31  	// write payload
    32  	ww.Write(payload)
    33  
    34  	// calc checksum
    35  	checksum := sum.Sum32()
    36  
    37  	// write checksum
    38  	return binary.Write(w, binary.BigEndian, checksum)
    39  }
    40  
    41  func DecodePacket(r io.Reader) ([]byte, error) {
    42  	var totalsize uint32
    43  	err := binary.Read(r, binary.BigEndian, &totalsize)
    44  	if err != nil {
    45  		return nil, errors.Wrap(err, "read totalsize")
    46  	}
    47  
    48  	if totalsize < 8 {
    49  		return nil, errors.New("totalsize too small")
    50  	}
    51  
    52  	sum := adler32.New()
    53  	rr := io.TeeReader(r, sum)
    54  
    55  	var magic [4]byte
    56  	binary.Read(rr, binary.BigEndian, &magic)
    57  	if magic != RPC_MAGIC {
    58  		return nil, errors.New("not tiny rpc byte")
    59  
    60  	}
    61  	payload := make([]byte, totalsize-8)
    62  	_, err = io.ReadFull(rr, payload)
    63  	if err != nil {
    64  		return nil, errors.Wrap(err, "io readfull payload")
    65  	}
    66  	var checksum uint32
    67  	// read from r, not rr!
    68  	err = binary.Read(r, binary.BigEndian, &checksum)
    69  	if err != nil {
    70  		return nil, errors.Wrap(err, "read checksum")
    71  	}
    72  	if checksum != sum.Sum32() {
    73  		return nil, errors.New("checksum mismatch")
    74  	}
    75  	return payload, nil
    76  }