tinymq/conn/util.go

package conn

import (
	"encoding/binary"
	"errors"
	"fmt"
	"log"
	"slices"

	"git.me9.top/git/tinymq/config"
	"git.me9.top/git/tinymq/util"
	"git.me9.top/git/tinymq/util/pool"
)

// 验证包编码，其中输出的数据需要回收
func AuthPackageEncode(proto string, version uint8, channel string, auth []byte, compress bool) (buf []byte) {
	protoLen := len(proto)
	protoLenMax := 0xFF
	if protoLen > protoLenMax {
		proto = proto[0:protoLenMax]
		protoLen = protoLenMax
	}

	channelLen := len(channel)
	channelLenMax := 0xFFFF
	if channelLen > channelLenMax {
		channel = channel[0:channelLenMax]
		channelLen = channelLenMax
	}
	// id(65502)+proto(string)+version(uint8)+option(byte)+channel(string)+auth([]byte)
	dlen := 2 + 1 + protoLen + 1 + 1 + 2 + channelLen + len(auth)
	buf = pool.Get(dlen)
	index := 0
	binary.BigEndian.PutUint16(buf[index:index+2], config.ID_AUTH)
	index += 2

	buf[index] = byte(protoLen)
	index++
	copy(buf[index:index+protoLen], []byte(proto))
	index += protoLen

	buf[index] = version
	index++

	// 指示是否支持压缩
	if compress {
		buf[index] = 0x01
	} else {
		buf[index] = 0
	}
	index++

	binary.BigEndian.PutUint16(buf[index:index+2], uint16(channelLen))
	index += 2

	copy(buf[index:index+channelLen], []byte(channel))
	index += channelLen
	copy(buf[index:], auth)
	return
}

// 验证包解码
// id(65502)+proto(string)+version(uint8)+option(byte)+channel(string)+auth([]byte)
func AuthPackageDecode(msg []byte, recycle bool) (proto string, version uint8, compress bool, channel string, auth []byte, err error) {
	// 避免错误的数据包，出现数组越界的错误
	defer func() {
		if recycle {
			pool.Put(msg)
		}
		if r := recover(); r != nil {
			err = fmt.Errorf("recovered from panic: %v", r)
			return
		}
	}()

	msgLen := len(msg)
	if msgLen < 9 {
		err = errors.New("wrong message length")
		return
	}
	index := 0
	id := binary.BigEndian.Uint16(msg[index : index+2])
	if id != config.ID_AUTH {
		err = fmt.Errorf("wrong message id: %d", id)
		return
	}
	index += 2

	protoLen := int(msg[index])
	if protoLen < 2 {
		err = errors.New("wrong proto length")
		return
	}
	index++
	proto = string(msg[index : index+protoLen])
	index += protoLen

	version = msg[index]
	index++

	compress = (msg[index] & 0x01) != 0
	index++

	channelLen := int(binary.BigEndian.Uint16(msg[index : index+2]))
	if channelLen < 2 {
		err = errors.New("wrong channel length")
		return
	}
	index += 2
	channel = string(msg[index : index+channelLen])
	index += channelLen

	if recycle {
		auth = slices.Clone(msg[index:])
	} else {
		auth = msg[index:]
	}
	return
}

// 请求包
func RequestPackageEncode(id uint16, cmd string, data []byte, compress bool, cf *config.Config) (buf []byte) {
	// 为了区分请求还是响应包，命令字符串不能超过127个字节
	cmdLenMax := 0x7F
	cmdLen := len(cmd)
	if cmdLen > cmdLenMax {
		cmd = cmd[0:cmdLenMax]
		cmdLen = cmdLenMax
	}
	dlen := len(data)
	if compress && dlen > 0 {
		compressedData, ok, _ := util.CompressData(data)
		// id(uint16)+cmdLen(byte)+len(cmd)+option(byte)+len(data)
		ddlen := 2 + 1 + cmdLen + 1 + len(compressedData)
		buf = pool.Get(ddlen)
		index := 0
		binary.BigEndian.PutUint16(buf[index:index+2], id)
		index += 2

		buf[index] = byte(cmdLen)
		index++
		copy(buf[index:index+cmdLen], cmd)
		index += cmdLen

		if ok {
			buf[index] = 0x01
			if cf.PrintMsg {
				log.Println("[CompressData]", len(data), "->", len(compressedData))
			}
		} else {
			buf[index] = 0
		}
		index++

		copy(buf[index:], compressedData)

	} else {
		ddlen := 2 + 1 + cmdLen + dlen
		buf = pool.Get(ddlen)
		index := 0
		binary.BigEndian.PutUint16(buf[index:index+2], id)
		index += 2
		buf[index] = byte(cmdLen)
		index++
		copy(buf[index:index+cmdLen], cmd)
		index += cmdLen
		copy(buf[index:], data)
	}
	return
}

// 响应数据包
// 网络格式：[id, stateCode, data]
func ResponsePackageEncode(id uint16, state uint8, data []byte, compress bool, cf *config.Config) (buf []byte) {
	dlen := len(data)
	if compress && dlen > 0 {
		compressedData, ok, _ := util.CompressData(data)
		ddlen := 2 + 1 + 1 + len(compressedData)
		buf = pool.Get(ddlen)
		index := 0
		binary.BigEndian.PutUint16(buf[index:index+2], id)
		index += 2

		buf[index] = state | 0x80
		index++

		if ok {
			buf[index] = 0x001
			if cf.PrintMsg {
				log.Println("[CompressData]", len(data), "->", len(compressedData))
			}
		} else {
			buf[index] = 0
		}
		index++

		copy(buf[index:], compressedData)
	} else {
		ddlen := 2 + 1 + len(data)
		buf = pool.Get(ddlen)
		index := 0
		binary.BigEndian.PutUint16(buf[index:index+2], id)
		index += 2
		buf[index] = state | 0x80
		index++
		copy(buf[index:], data)
	}
	return
}

func PingPackageEncode(id uint16) (buf []byte) {
	buf = pool.Get(2)
	index := 0
	binary.BigEndian.PutUint16(buf[index:index+2], id)
	return
}

// 解析收到的数据包
func PackageDecode(msg []byte, recycle bool, compress bool, cf *config.Config) (msgType MsgType, id uint16, cmd string, state uint8, data []byte, err error) {
	if recycle {
		defer pool.Put(msg)
	}
	id = binary.BigEndian.Uint16(msg[0:2])

	if id > config.ID_MAX {
		switch id {
		case config.ID_PROXY_CONNECT: // 代理连接请求
			msgType = MsgProxyConnect
			// 请求代理的地址
			if recycle {
				data = slices.Clone(msg[2:])
			} else {
				data = msg[2:]
			}
			return
		case config.ID_PROXY_RESULT: // 代理请求执行结果
			msgType = MsgProxyResult
			// 错误信息
			if recycle {
				data = slices.Clone(msg[2:])
			} else {
				data = msg[2:]
			}
			// 如果没有错误信息则表示成功
			if len(data) <= 0 {
				state = config.STATE_OK
			}
			return
		}
		err = fmt.Errorf("wrong message id: %d", id)
		return
	}

	// ping信息
	if len(msg) == 2 {
		msgType = MsgPing
		return
	}

	cmdx := msg[2]
	if (cmdx & 0x80) == 0 {
		// 请求包
		msgType = MsgRequest
		cmdLen := int(cmdx)
		cmd = string(msg[3 : cmdLen+3])
		data = msg[cmdLen+3:]
	} else {
		// 响应数据包
		msgType = MsgResponse
		state = cmdx & 0x7F
		data = msg[3:]
	}
	dataLen := len(data)
	if compress && dataLen > 1 {
		var ok bool
		data, ok, err = util.UncompressData(data)
		if !ok && recycle {
			data = slices.Clone(data)
		}
		if ok && cf.PrintMsg {
			log.Println("[UncompressData]", dataLen, "->", len(data))
		}
	} else if recycle && dataLen > 0 {
		data = slices.Clone(data)
	}
	return
}