tinymq/conn/ws2/ws2.go

package ws2

import (
	"crypto/rand"
	"encoding/binary"
	"errors"
	"fmt"
	"log"
	"net"
	"net/http"
	"net/url"
	"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"
	"github.com/gorilla/websocket"
)

const PROTO string = "ws"
const PROTO_STL string = "wss"
const VERSION uint8 = 2

type Ws2 struct {
	cf       *config.Config
	conn     *websocket.Conn
	cipher   *util.Cipher // 记录当前的加解密类，可以保证在没有ssl的情况下数据安全
	compress bool         // 是否启用压缩
}

var upgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool {
		return true // 允许任何Origin跨站访问
	},
} // use default options

// websocket 服务
// 如果有绑定参数，则进行绑定操作代码
func Server(cf *config.Config, bind string, path string, hash string, fn conn.ServerConnectFunc) (err error) {
	var ci *util.CipherInfo
	var encryptKey string
	if hash != "" {
		i := strings.Index(hash, ":")
		if i <= 0 {
			return errors.New("hash is invalid")
		}
		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)
		}
	}
	http.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) {
		conn, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			log.Println("[ws2 Server Upgrade ERROR]", err)
			return
		}
		if ci == nil {
			ws := &Ws2{
				cf:   cf,
				conn: conn,
			}
			fn(ws)
			return
		}
		var eiv []byte
		var div []byte
		if ci.IvLen > 0 {
			// 服务端 IV
			eiv = make([]byte, ci.IvLen)
			_, err = rand.Read(eiv)
			if err != nil {
				log.Println("[ws2 Server rand.Read ERROR]", err)
				return
			}
			// 发送 IV
			conn.SetWriteDeadline(time.Now().Add(time.Duration(cf.WriteWait) * time.Millisecond))
			if err := conn.WriteMessage(websocket.BinaryMessage, eiv); err != nil {
				log.Println("[ws2 Server conn.Write ERROR]", err)
				return
			}

			// 读取 IV
			err = conn.SetReadDeadline(time.Now().Add(time.Millisecond * time.Duration(cf.ReadWait)))
			if err != nil {
				log.Println("[ws2 Server SetReadDeadline ERROR]", err)
				return
			}
			_, div, err = conn.ReadMessage()
			if err != nil {
				log.Println("[ws2 Server ReadFull ERROR]", err)
				conn.Close()
				return
			}
		}
		cipher, err := util.NewCipher(ci, encryptKey, eiv, div)
		if err != nil {
			log.Println("[ws2 NewCipher ERROR]", err)
			return
		}
		ws := &Ws2{
			cf:     cf,
			conn:   conn,
			cipher: cipher,
		}
		fn(ws)
	})

	if bind != "" {
		go func() (err error) {
			defer func() {
				if err != nil {
					log.Fatal(err)
				}
			}()
			log.Printf("Listening and serving Websocket on %s\n", bind)
			// 暂时使用全局的方式，后面有需求再修改
			// 而且还没有 https 方式的绑定
			// 需要在前端增加其他的服务进行转换
			err = http.ListenAndServe(bind, nil)
			return
		}()
	}

	return
}

// 客户端，新建一个连接
func Dial(cf *config.Config, scheme string, addr string, path string, hash string) (conn.Connect, error) {
	u := url.URL{Scheme: scheme, Host: addr, Path: path}
	// 没有加密的情况
	if hash == "" {
		conn, _, err := (&websocket.Dialer{
			HandshakeTimeout: time.Duration(time.Millisecond * time.Duration(cf.ConnectTimeout)),
		}).Dial(u.String(), nil)
		if err != nil {
			return nil, err
		}
		ws := &Ws2{
			cf:   cf,
			conn: conn,
		}
		return ws, 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 := (&websocket.Dialer{
		HandshakeTimeout: time.Duration(time.Millisecond * time.Duration(cf.ConnectTimeout)),
	}).Dial(u.String(), nil)
	if err != nil {
		return nil, err
	}
	var eiv []byte
	var div []byte
	if ci.IvLen > 0 {
		// 客户端 IV
		eiv = make([]byte, ci.IvLen)
		_, err = rand.Read(eiv)
		if err != nil {
			log.Println("[ws2 Client rand.Read ERROR]", err)
			return nil, err
		}
		// 发送 IV
		conn.SetWriteDeadline(time.Now().Add(time.Duration(cf.WriteWait) * time.Millisecond))
		if err := conn.WriteMessage(websocket.BinaryMessage, eiv); err != nil {
			log.Println("[ws2 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("[ws2 Client SetReadDeadline ERROR]", err)
			return nil, err
		}
		_, div, err = conn.ReadMessage()
		if err != nil {
			log.Println("[ws2 Client ReadFull ERROR]", err)
			return nil, err
		}
	}
	cipher, err := util.NewCipher(ci, encryptKey, eiv, div)
	if err != nil {
		log.Println("[ws2 NewCipher ERROR]", err)
		return nil, err
	}
	ws := &Ws2{
		cf:       cf,
		conn:     conn,
		cipher:   cipher,
		compress: true,
	}
	return ws, nil
}

// 发送数据到网络
// 如果有加密函数的话会直接修改源数据
func (c *Ws2) WriteRawPackage(buf []byte, recycle bool) (err error) {
	if recycle {
		defer pool.Put(buf)
	}
	if c.cipher != nil {
		c.cipher.Encrypt(buf, buf)
	}
	c.conn.SetWriteDeadline(time.Now().Add(time.Millisecond * time.Duration(c.cf.WriteWait)))
	return c.conn.WriteMessage(websocket.BinaryMessage, buf)
}

// 从连接中读取数据包
// recycle 指示是否需要手动将缓存放回内存池
func (c *Ws2) ReadRawPackage(deadline int) (msg []byte, recycle bool, err error) {
	err = c.conn.SetReadDeadline(time.Now().Add(time.Millisecond * time.Duration(deadline)))
	if err != nil {
		return
	}
	_, msg, err = c.conn.ReadMessage()
	return
}

// 发送Auth信息
// 建立连接后第一个发送的消息
func (c *Ws2) 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)
	index := 0

	// buf := make([]byte, dlen)
	buf := pool.Get(dlen)
	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 *Ws2) 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, _, err := c.ReadRawPackage(c.cf.ReadWait)
	// err = c.conn.SetReadDeadline(time.Now().Add(time.Millisecond * time.Duration(c.cf.ReadWait)))
	// if err != nil {
	// 	return
	// }
	// _, msg, err := c.conn.ReadMessage()
	if err != nil {
		return
	}
	msgLen := len(msg)
	if msgLen < 9 {
		err = errors.New("wrong message length")
		return
	}
	// 将读出来的数据进行解密
	if c.cipher != nil {
		c.cipher.Decrypt(msg, msg)
	}
	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

	auth = msg[index:]
	return
}

// 发送请求数据包到网络
func (c *Ws2) WriteRequest(id uint16, cmd string, data []byte) error {
	// 为了区分请求还是响应包，命令字符串不能超过127个字节，如果超过则报错
	cmdLen := len(cmd)
	if cmdLen > 0x7F {
		return errors.New("length of command more than 0x7F")
	}
	dlen := len(data)
	if c.compress && dlen > 0 {
		compressedData, ok, _ := util.CompressData(data)
		ddlen := 2 + 1 + cmdLen + 1 + len(compressedData)
		// buf := make([]byte, ddlen) // 申请内存
		buf := pool.Get(ddlen)
		index := 0
		binary.BigEndian.PutUint16(buf[index:index+2], id)
		index += 2

		buf[index] = byte(cmdLen)
		index++
		copy(buf[index:], 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)
		binary.BigEndian.PutUint16(buf[0:2], id)
		buf[2] = byte(cmdLen)
		copy(buf[3:], cmd)
		copy(buf[3+cmdLen:], data)
		return c.WriteRawPackage(buf, true)
	}
}

// 发送响应数据包到网络
// 网络格式：[id, stateCode, data]
func (c *Ws2) 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)
		binary.BigEndian.PutUint16(buf[0:2], id)
		buf[2] = state | 0x80
		if ok {
			buf[3] = 0x01
			if c.cf.PrintMsg {
				log.Println("[CompressData]", len(data), "->", len(compressedData))
			}
		} else {
			buf[3] = 0
		}
		copy(buf[4:], compressedData)
		return c.WriteRawPackage(buf, true)
	} else {
		ddlen := 2 + 1 + dlen
		buf := pool.Get(ddlen)
		binary.BigEndian.PutUint16(buf[0:2], id)
		buf[2] = state | 0x80
		copy(buf[3:], data)
		return c.WriteRawPackage(buf, true)
	}
}

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

// 获取信息
func (c *Ws2) ReadMessage(deadline int) (msgType conn.MsgType, id uint16, cmd string, state uint8, data []byte, err error) {
	msg, _, err := c.ReadRawPackage(deadline)
	// err = c.conn.SetReadDeadline(time.Now().Add(time.Millisecond * time.Duration(deadline)))
	// if err != nil {
	// 	return
	// }
	// _, msg, err := c.conn.ReadMessage()
	if err != nil {
		return
	}
	msgLen := len(msg)
	if msgLen < 2 {
		err = errors.New("message length less than 2")
		return
	}
	// 将读出来的数据进行解密
	if c.cipher != nil {
		c.cipher.Decrypt(msg, 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 && c.cf.PrintMsg {
			log.Println("[UncompressData]", dataLen, "->", len(data))
		}
	}
	return
}

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

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

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