tinymq/conn/ws2/ws2.go

package ws2

import (
	"crypto/rand"
	"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         // 是否启用压缩
	remoteIp net.IP
}

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) {
		// Get real remote address from proxy headers
		remoteIp := r.Header.Get("X-Real-IP")
		if remoteIp == "" {
			// X-Forwarded-For can contain a list of comma-separated IPs; the first is the original client
			forwarded := r.Header.Get("X-Forwarded-For")
			if forwarded != "" {
				remoteIp = strings.Split(forwarded, ",")[0]
				remoteIp = strings.TrimSpace(remoteIp)
			}
		}

		// Fallback to standard RemoteAddr if no proxy headers exist
		if remoteIp == "" {
			remoteIp = r.RemoteAddr
		}

		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,
			remoteIp: util.StrToIP(remoteIp),
		}
		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) (buf []byte, recycle bool, err error) {
	err = c.conn.SetReadDeadline(time.Now().Add(time.Millisecond * time.Duration(deadline)))
	if err != nil {
		return
	}
	_, buf, err = c.conn.ReadMessage()
	if err != nil {
		return
	}
	if c.cipher != nil {
		c.cipher.Decrypt(buf, buf)
	}
	return
}

// 发送Auth信息
// 建立连接后第一个发送的消息
func (c *Ws2) WriteAuthInfo(channel string, auth []byte) (err error) {
	buf := conn.AuthPackageEncode(PROTO, VERSION, channel, auth, c.compress)
	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) {
	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 && proto != PROTO_STL {
		err = fmt.Errorf("wrong proto: want %s, get %s", PROTO, proto)
		return
	}
	if version != VERSION {
		err = fmt.Errorf("require version %d, get version: %d", VERSION, version)
		return
	}
	c.compress = compress
	return
}

// 发送请求数据包到网络
func (c *Ws2) 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 *Ws2) 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 *Ws2) WritePing(id uint16) error {
	buf := conn.PingPackageEncode(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) {
	var msg []byte
	var recycle bool
	msg, recycle, err = c.ReadRawPackage(deadline)
	if err != nil {
		return
	}
	return conn.PackageDecode(msg, recycle, c.compress, c.cf)
}

// 获取远程的地址
func (c *Ws2) RemoteIP() net.IP {
	if c.remoteIp != nil {
		return c.remoteIp
	}
	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()
}