nnet/packet/packer.go

package packet

import (
	"bytes"
	"errors"
)

var _ Packer = (*DefaultPacker)(nil)

type DefaultPacker struct {
	buf  *bytes.Buffer
	size int  // 最近一次 长度
	typ  byte // 最近一次 数据帧类型
}

// Codec constants.
const (
	headLength    = 4
	maxPacketSize = 64 * 1024
)

var ErrPacketSizeExceed = errors.New("codec: packet size exceed")

func NewDefaultPacker() Packer {
	return &DefaultPacker{
		buf:  bytes.NewBuffer(nil),
		size: -1,
	}
}

func (d *DefaultPacker) Pack(typ Type, data []byte) ([]byte, error) {
	if typ < Handshake || typ > Kick {
		return nil, ErrWrongPacketType
	}

	p := &Packet{Type: typ, Length: uint32(len(data))}
	buf := make([]byte, p.Length+headLength)

	// header
	buf[0] = byte(p.Type)
	copy(buf[1:headLength], d.intToBytes(p.Length))

	// body
	copy(buf[headLength:], data)

	return buf, nil
}

// Encode packet data length to bytes(Big end)
func (d *DefaultPacker) intToBytes(n uint32) []byte {
	buf := make([]byte, 3)
	buf[0] = byte((n >> 16) & 0xFF)
	buf[1] = byte((n >> 8) & 0xFF)
	buf[2] = byte(n & 0xFF)
	return buf
}

func (d *DefaultPacker) Unpack(data []byte) ([]interface{}, error) {
	d.buf.Write(data) // copy

	var (
		packets []interface{}
		err     error
	)

	// 检查包长度
	if d.buf.Len() < headLength {
		return nil, err
	}

	// 第一次拆包
	if d.size < 0 {
		if err = d.readHeader(); err != nil {
			return nil, err
		}
	}

	for d.size <= d.buf.Len() {
		// 读取
		p := &Packet{
			Type:   Type(d.typ),
			Length: uint32(d.size),
			Data:   d.buf.Next(d.size),
		}
		packets = append(packets, p)

		// 剩余数据不满足至少一个数据帧，重置数据帧长度
		// 数据缓存内存 保留至 下一次进入本方法以继续拆包
		if d.buf.Len() < headLength {
			d.size = -1
			break
		}
		// 读取下一个包 next
		if err = d.readHeader(); err != nil {
			return packets, err
		}
	}
	return packets, nil
}

func (d *DefaultPacker) readHeader() error {
	header := d.buf.Next(headLength)
	d.typ = header[0]
	if d.typ < Handshake || d.typ > Kick {
		return ErrWrongPacketType
	}
	d.size = d.bytesToInt(header[1:])

	// 最大包限定
	if d.size > maxPacketSize {
		return ErrPacketSizeExceed
	}
	return nil
}

// Decode packet data length byte to int(Big end)
func (d *DefaultPacker) bytesToInt(b []byte) int {
	result := 0
	for _, v := range b {
		result = result<<8 + int(v)
	}
	return result
}