tinymq/conn/tcp2/tcp2.go

package tcp2

import (
	"crypto/rand"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"log"
	"net"
	"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) {
	buf := conn.AuthPackageEncode(PROTO, VERSION, channel, auth, c.compress)
	return c.WriteRawPackage(buf, true)
}

// 获取Auth信息
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
	}
	var compress bool
	proto, version, compress, channel, auth, err = conn.AuthPackageDecode(msg, recycle)
	if err != nil {
		return
	}
	if proto != PROTO {
		err = fmt.Errorf("wrong proto: %s", proto)
		return
	}
	if version != VERSION {
		err = fmt.Errorf("require version %d, get version: %d", VERSION, version)
		return
	}
	c.compress = compress
	return
}

// 发送请求数据包到网络
func (c *Tcp2) WriteRequest(id uint16, cmd string, data []byte) error {
	buf := conn.RequestPackageEncode(id, cmd, data, c.compress, c.cf)
	return c.WriteRawPackage(buf, true)
}

// 发送响应数据包到网络
// 网络格式：[id, stateCode, data]
func (c *Tcp2) WriteResponse(id uint16, state uint8, data []byte) error {
	buf := conn.ResponsePackageEncode(id, state, data, c.compress, c.cf)
	return c.WriteRawPackage(buf, true)
}

// 发送ping包
func (c *Tcp2) WritePing(id uint16) error {
	buf := conn.PingPackageEncode(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
	}
	return conn.PackageDecode(msg, recycle, c.compress, c.cf)
}

// 获取远程的地址
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()
}