origin/network/tcp_msg.go

package network

import (
	"encoding/binary"
	"errors"
	"github.com/duanhf2012/origin/v2/util/bytespool"
	"io"
	"math"
)

// MsgParser --------------
// | len | data |
// --------------
type MsgParser struct {
	LenMsgLen    int
	MinMsgLen    uint32
	MaxReadMsgLen    uint32
	MaxWriteMsgLen   uint32
	LittleEndian bool

	bytespool.IBytesMemPool
}

func (p *MsgParser) getMaxMsgLen() uint32 {
	switch p.LenMsgLen {
	case 1:
		return math.MaxUint8
	case 2:
		return math.MaxUint16
	case 4:
		return math.MaxUint32
	default:
		panic("LenMsgLen value must be 1 or 2 or 4")
	}
}

func (p *MsgParser) Init() {
	p.IBytesMemPool = bytespool.NewMemAreaPool()
}

// goroutine safe
func (p *MsgParser) Read(r io.Reader) ([]byte, error) {
	var b [4]byte
	bufMsgLen := b[:p.LenMsgLen]

	// read len
	if _, err := io.ReadFull(r, bufMsgLen); err != nil {
		return nil, err
	}

	// parse len
	var msgLen uint32
	switch p.LenMsgLen {
	case 1:
		msgLen = uint32(bufMsgLen[0])
	case 2:
		if p.LittleEndian {
			msgLen = uint32(binary.LittleEndian.Uint16(bufMsgLen))
		} else {
			msgLen = uint32(binary.BigEndian.Uint16(bufMsgLen))
		}
	case 4:
		if p.LittleEndian {
			msgLen = binary.LittleEndian.Uint32(bufMsgLen)
		} else {
			msgLen = binary.BigEndian.Uint32(bufMsgLen)
		}
	}

	// check len
	if msgLen > p.MaxReadMsgLen {
		return nil, errors.New("message too long")
	} else if msgLen < p.MinMsgLen {
		return nil, errors.New("message too short")
	}

	// data
	msgData := p.MakeBytes(int(msgLen))
	if _, err := io.ReadFull(r, msgData[:msgLen]); err != nil {
		p.ReleaseBytes(msgData)
		return nil, err
	}

	return msgData[:msgLen], nil
}

// goroutine safe
func (p *MsgParser) Write(conn io.Writer, args ...[]byte) error {
	// get len
	var msgLen uint32
	for i := 0; i < len(args); i++ {
		msgLen += uint32(len(args[i]))
	}

	// check len
	if p.MaxWriteMsgLen > 0 && msgLen > p.MaxWriteMsgLen {
		return errors.New("message too long")
	} else if msgLen < p.MinMsgLen {
		return errors.New("message too short")
	}

	//msg := make([]byte, uint32(p.lenMsgLen)+msgLen)
	msg := p.MakeBytes(p.LenMsgLen + int(msgLen))
	// write len
	switch p.LenMsgLen {
	case 1:
		msg[0] = byte(msgLen)
	case 2:
		if p.LittleEndian {
			binary.LittleEndian.PutUint16(msg, uint16(msgLen))
		} else {
			binary.BigEndian.PutUint16(msg, uint16(msgLen))
		}
	case 4:
		if p.LittleEndian {
			binary.LittleEndian.PutUint32(msg, msgLen)
		} else {
			binary.BigEndian.PutUint32(msg, msgLen)
		}
	}

	// write data
	l := p.LenMsgLen
	for i := 0; i < len(args); i++ {
		copy(msg[l:], args[i])
		l += len(args[i])
	}

	conn.Write(msg)

	return nil
}

func (p *MsgParser) GetRecyclerReaderBytes() func(data []byte) {
	return func(data []byte) {
		p.IBytesMemPool.ReleaseBytes(data)
	}
}