tinymq/conn/tcp2/tcp2.go

package tcp2

import (
	"crypto/rand"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"log"
	"net"
	"slices"
	"strings"
	"time"

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

const PROTO string = "tcp"
const VERSION uint8 = 2

// 第一级数据包的最大长度，如果超过则自动增加一个地址位到4个字节
const MAX_LENGTH = 0xFFFF

// 由于都是内部使用，感觉没有必要加这个限制
// const MAX2_LENGTH = 0x1FFFFFFF // 500 M，避免申请过大内存

type Tcp2 struct {
	cf       *config.Config
	conn     net.Conn
	cipher   *util.Cipher // 记录当前的加解密类
	compress bool         // 是否启用压缩
}

// 服务端
// hash 格式 encryptMethod:encryptKey
func Server(cf *config.Config, bind string, hash string, fn conn.ServerConnectFunc) (err error) {
	var ci *util.CipherInfo
	var encryptKey string
	if hash != "" {
		i := strings.Index(hash, ":")
		if i <= 0 {
			err = errors.New("hash is invalid")
			return
		}
		encryptMethod := hash[0:i]
		encryptKey = hash[i+1:]

		if c, ok := util.CipherMethod[encryptMethod]; ok {
			ci = c
		} else {
			return errors.New("Unsupported encryption method: " + encryptMethod)
		}
	}

	log.Printf("Listening and serving tcp on %s\n", bind)
	l, err := net.Listen("tcp", bind)
	if err != nil {
		log.Println("[tcp2 Server ERROR]", err)
		return
	}
	go func(l net.Listener) {
		defer l.Close()
		for {
			conn, err := l.Accept()
			if err != nil {
				log.Println("[accept ERROR]", err)
				return
			}
			go func(conn net.Conn) {
				if ci == nil {
					c := &Tcp2{
						cf:   cf,
						conn: conn,
					}
					fn(c)
					return
				}
				var eiv []byte
				var div []byte
				if ci.IvLen > 0 {
					// 服务端 IV
					eiv = pool.Get(ci.IvLen)
					defer pool.Put(eiv)
					_, err = rand.Read(eiv)
					if err != nil {
						log.Println("[tcp2 Server rand.Read ERROR]", err)
						return
					}
					// 发送 IV
					conn.SetWriteDeadline(time.Now().Add(time.Duration(cf.WriteWait) * time.Millisecond))
					if _, err := conn.Write(eiv); err != nil {
						log.Println("[tcp2 Server conn.Write ERROR]", err)
						return
					}

					// 读取 IV
					err = conn.SetReadDeadline(time.Now().Add(time.Millisecond * time.Duration(cf.ReadWait)))
					if err != nil {
						log.Println("[tcp2 Server SetReadDeadline ERROR]", err)
						return
					}
					div = pool.Get(ci.IvLen)
					defer pool.Put(div)
					_, err := io.ReadFull(conn, div)
					if err != nil {
						log.Println("[tcp2 Server ReadFull ERROR]", err)
						return
					}
				}
				cipher, err := util.NewCipher(ci, encryptKey, eiv, div)
				if err != nil {
					log.Println("[tcp2 NewCipher ERROR]", err)
					return
				}
				// 初始化
				c := &Tcp2{
					cf:     cf,
					conn:   conn,
					cipher: cipher,
				}
				fn(c)
			}(conn)
		}
	}(l)
	return
}

// 客户端，新建一个连接
func Dial(cf *config.Config, addr string, hash string) (conn.Connect, error) {
	// 没有加密的情况
	if hash == "" {
		conn, err := net.DialTimeout("tcp", addr, time.Duration(cf.ConnectTimeout)*time.Millisecond)
		if err != nil {
			return nil, err
		}
		c := &Tcp2{
			cf:   cf,
			conn: conn,
		}
		return c, nil
	}
	i := strings.Index(hash, ":")
	if i <= 0 {
		return nil, errors.New("hash is invalid")
	}
	encryptMethod := hash[0:i]
	encryptKey := hash[i+1:]
	ci, ok := util.CipherMethod[encryptMethod]
	if !ok {
		return nil, errors.New("Unsupported encryption method: " + encryptMethod)
	}
	conn, err := net.DialTimeout("tcp", addr, time.Duration(cf.ConnectTimeout)*time.Millisecond)
	if err != nil {
		return nil, err
	}
	var eiv []byte
	var div []byte
	if ci.IvLen > 0 {
		// 客户端 IV
		eiv = pool.Get(ci.IvLen)
		defer pool.Put(eiv)
		_, err = rand.Read(eiv)
		if err != nil {
			log.Println("[tcp2 Client rand.Read ERROR]", err)
			return nil, err
		}
		// 发送 IV
		conn.SetWriteDeadline(time.Now().Add(time.Duration(cf.WriteWait) * time.Millisecond))
		if _, err := conn.Write(eiv); err != nil {
			log.Println("[tcp2 Client conn.Write ERROR]", err)
			return nil, err
		}

		// 读取 IV
		err = conn.SetReadDeadline(time.Now().Add(time.Millisecond * time.Duration(cf.ReadWait)))
		if err != nil {
			log.Println("[tcp2 Client SetReadDeadline ERROR]", err)
			return nil, err
		}
		div = pool.Get(ci.IvLen)
		defer pool.Put(div)
		_, err := io.ReadFull(conn, div)
		if err != nil {
			log.Println("[tcp2 Client ReadFull ERROR]", err)
			return nil, err
		}
	}
	cipher, err := util.NewCipher(ci, encryptKey, eiv, div)
	if err != nil {
		log.Println("[tcp2 NewCipher ERROR]", err)
		return nil, err
	}
	// 初始化
	c := &Tcp2{
		cf:       cf,
		conn:     conn,
		cipher:   cipher,
		compress: true,
	}
	return c, nil
}

// 发送数据包到网络
// recycle 指示是否需要将raw放回内存池
func (c *Tcp2) WriteRawPackage(raw []byte, recycle bool) (err error) {
	var buf []byte
	var index int
	dlen := len(raw)
	if dlen >= MAX_LENGTH {
		buf = pool.Get(dlen + 2 + 4 + 1)
		index = 2 + 4
		binary.BigEndian.PutUint16(buf[:2], MAX_LENGTH)
		binary.BigEndian.PutUint32(buf[2:6], uint32(dlen))
	} else {
		buf = pool.Get(dlen + 2 + 1)
		index = 2
		binary.BigEndian.PutUint16(buf[:2], uint16(dlen))
	}
	copy(buf[index:index+dlen], raw)
	index += dlen
	buf[index] = util.CRC8(raw)

	if recycle {
		pool.Put(raw)
	}
	defer pool.Put(buf)

	if c.cipher != nil {
		c.cipher.Encrypt(buf, buf)
	}
	c.conn.SetWriteDeadline(time.Now().Add(time.Duration(c.cf.WriteWait) * time.Millisecond))
	for {
		n, err := c.conn.Write(buf)
		if err != nil {
			return err
		}
		if n < len(buf) {
			buf = buf[n:]
		} else {
			return nil
		}
	}
}

// 从连接中读取数据包，已经剥离了头部长度和尾部crc
// recycle 指示是否需要手动将缓存放回内存池
func (c *Tcp2) ReadRawPackage(deadline int) (buf []byte, recycle bool, err error) {
	err = c.conn.SetReadDeadline(time.Now().Add(time.Millisecond * time.Duration(deadline)))
	if err != nil {
		return nil, false, err
	}
	buf = pool.Get(2)
	// 读取数据流长度
	_, err = io.ReadFull(c.conn, buf)
	if err != nil {
		pool.Put(buf)
		return nil, false, err
	}
	// 将读出来的数据进行解密
	if c.cipher != nil {
		c.cipher.Decrypt(buf, buf)
	}
	dlen := uint32(binary.BigEndian.Uint16(buf))
	pool.Put(buf)
	if dlen < 2 {
		return nil, false, errors.New("length is less to 2")
	}
	if dlen >= MAX_LENGTH {
		// 数据包比较大，通过后面的4位长度来表示实际长度
		buf := pool.Get(4)
		_, err := io.ReadFull(c.conn, buf)
		if err != nil {
			pool.Put(buf)
			return nil, false, err
		}
		if c.cipher != nil {
			c.cipher.Decrypt(buf, buf)
		}
		dlen = binary.BigEndian.Uint32(buf)
		pool.Put(buf)
		if dlen < MAX_LENGTH {
			return nil, false, errors.New("wrong length in read message")
		}
	}
	// 读取指定长度的数据
	buf = pool.Get(int(dlen) + 1) // 最后一个是crc的值
	_, err = io.ReadFull(c.conn, buf)
	if err != nil {
		pool.Put(buf)
		return nil, false, err
	}
	if c.cipher != nil {
		c.cipher.Decrypt(buf, buf)
	}
	// 检查CRC8
	if util.CRC8(buf[:dlen]) != buf[dlen] {
		pool.Put(buf)
		return nil, false, errors.New("CRC error")
	}
	return buf[:dlen], true, nil
}

// 发送Auth信息
// 建立连接后第一个发送的消息
func (c *Tcp2) WriteAuthInfo(channel string, auth []byte) (err error) {
	protoLen := len(PROTO)
	if protoLen > 0xFF {
		return errors.New("length of protocol over")
	}

	channelLen := len(channel)
	if channelLen > 0xFFFF {
		return errors.New("length of channel over")
	}
	// 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 c.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 c.WriteRawPackage(buf, true)
}

// 获取Auth信息
// id(65502)+proto(string)+version(uint8)+option(byte)+channel(string)+auth([]byte)
func (c *Tcp2) ReadAuthInfo() (proto string, version uint8, channel string, auth []byte, err error) {
	defer func() {
		if r := recover(); r != nil {
			err = fmt.Errorf("recovered from panic: %v", r)
			return
		}
	}()
	msg, recycle, err := c.ReadRawPackage(c.cf.ReadWait)
	if err != nil {
		return
	}
	if recycle {
		defer pool.Put(msg)
	}
	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])
	if proto != PROTO {
		err = fmt.Errorf("wrong proto: %s", proto)
		return
	}
	index += protoLen

	version = msg[index]
	if version != VERSION {
		err = fmt.Errorf("require version %d, get version: %d", VERSION, version)
		return
	}
	index++

	c.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 (c *Tcp2) WriteRequest(id uint16, cmd string, data []byte) error {
	// 为了区分请求还是响应包，命令字符串不能超过127个字节
	cmdLen := len(cmd)
	if cmdLen > 0x7F {
		return errors.New("command length is more than 0x7F")
	}
	dlen := len(data)
	if c.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 c.cf.PrintMsg {
				log.Println("[CompressData]", len(data), "->", len(compressedData))
			}
		} else {
			buf[index] = 0
		}
		index++

		copy(buf[index:], compressedData)
		return c.WriteRawPackage(buf, true)
	} 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 c.WriteRawPackage(buf, true)
	}
}

// 发送响应数据包到网络
// 网络格式：[id, stateCode, data]
func (c *Tcp2) WriteResponse(id uint16, state uint8, data []byte) error {
	dlen := len(data)
	if c.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 c.cf.PrintMsg {
				log.Println("[CompressData]", len(data), "->", len(compressedData))
			}
		} else {
			buf[index] = 0
		}
		index++

		copy(buf[index:], compressedData)
		return c.WriteRawPackage(buf, true)
	} 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 c.WriteRawPackage(buf, true)
	}
}

// 发送ping包
func (c *Tcp2) WritePing(id uint16) error {
	buf := pool.Get(2)
	index := 0
	binary.BigEndian.PutUint16(buf[index:index+2], id)
	return c.WriteRawPackage(buf, true)
}

// 获取信息
func (c *Tcp2) ReadMessage(deadline int) (msgType conn.MsgType, id uint16, cmd string, state uint8, data []byte, err error) {
	msg, recycle, err := c.ReadRawPackage(deadline)
	if err != nil {
		return
	}
	if recycle {
		defer pool.Put(msg)
	}
	msgLen := len(msg)
	id = binary.BigEndian.Uint16(msg[0:2])
	// ping信息
	if msgLen == 2 {
		msgType = conn.PingMsg
		return
	}

	if id > config.ID_MAX {
		err = fmt.Errorf("wrong message id: %d", id)
		return
	}

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

// 获取远程的地址
func (c *Tcp2) RemoteIP() net.IP {
	return util.AddrToIP(c.conn.RemoteAddr())
}

// 获取本地的地址
func (c *Tcp2) LocalIP() net.IP {
	return util.AddrToIP(c.conn.LocalAddr())
}

func (c *Tcp2) Close() error {
	return c.conn.Close()
}