fox/ws/wsConn.go

package ws

import (
	"encoding/binary"
	"fmt"
	"github.com/fox/fox/log"
	"github.com/gorilla/websocket"
	"sync"
	"time"
)

// 连接id
var (
	wsMsgType  = websocket.BinaryMessage
	nextConnId uint32
)

// 客户端读写消息
type wsMessage struct {
	messageType int
	data        []byte
}

// 客户端连接
type wsConnect struct {
	wsConn       *websocket.Conn // 底层websocket
	inChan       chan *wsMessage // 读队列
	outChan      chan *wsMessage // 写队列
	mutex        sync.Mutex      // 避免重复关闭管道，加锁处理
	isClosed     bool
	closeCh      chan struct{} // 关闭通知
	id           uint32
	userId       int64
	onDisconnect func(IConn)
}

func newWsConnect(wsConn *websocket.Conn, onDisconnect func(IConn)) *wsConnect {
	return &wsConnect{
		wsConn:       wsConn,
		inChan:       make(chan *wsMessage, 1000),
		outChan:      make(chan *wsMessage, 1000),
		closeCh:      make(chan struct{}),
		isClosed:     false,
		id:           nextConnId,
		userId:       0,
		onDisconnect: onDisconnect,
	}
}

// 从读队列读取消息
func (c *wsConnect) readFromChan() (*wsMessage, error) {
	select {
	case msg := <-c.inChan:
		return msg, nil
	case <-c.closeCh:
		return nil, fmt.Errorf("连接已关闭")
	}
}

// 把消息放进写队列
func (c *wsConnect) sendMsg(msgType int, data []byte) error {
	select {
	case c.outChan <- &wsMessage{messageType: msgType, data: data}:
	case <-c.closeCh:
		return fmt.Errorf("连接已关闭")
	}
	return nil
}

// 把消息放进写队列
func (c *wsConnect) SendMsg(data []byte) error {
	return c.sendMsg(wsMsgType, data)
}

// 关闭链接
func (c *wsConnect) Close() {
	log.Debug(c.Log("关闭链接"))
	c.mutex.Lock()
	defer c.mutex.Unlock()
	if c.isClosed == false {
		c.isClosed = true
		wsMgr.Remove(c)
		close(c.closeCh)
		if c.onDisconnect != nil {
			c.onDisconnect(c)
		}
	}
}

// 循环从websocket中读取消息放入到读队列中
func (c *wsConnect) readWsLoop() {
	c.wsConn.SetReadLimit(maxMessageSize)
	_ = c.wsConn.SetReadDeadline(time.Now().Add(pongWait))
	for {
		// 读一个message
		msgType, data, err := c.wsConn.ReadMessage()
		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure, websocket.CloseNormalClosure) {
				log.Error(c.Log("消息读取出现错误:%v", err))
			}
			c.Close()
			return
		}
		switch msgType {
		case websocket.PingMessage:
			_ = c.sendMsg(websocket.PongMessage, []byte("pong"))
		case websocket.PongMessage:
		// _ = c.sendMsg(websocket.PingMessage, []byte("ping"))
		case websocket.CloseMessage:
			code := websocket.CloseNormalClosure
			reason := ""
			if len(data) >= 2 {
				code = int(binary.BigEndian.Uint16(data))
				reason = string(data[2:])
			}
			log.Debug(c.Log("关闭原因码:%d 描述:%s", code, reason))
			// 发送响应关闭帧（必须回传相同状态码）
			rspMsg := websocket.FormatCloseMessage(code, reason)
			_ = c.wsConn.WriteControl(websocket.CloseMessage, rspMsg, time.Now().Add(5*time.Second))
			c.Close()
		default:
			if msgType != wsMsgType {
				continue
			}
			msg := &wsMessage{messageType: msgType, data: data}
			select {
			case c.inChan <- msg:
			case <-c.closeCh:
				return
			}
		}
	}
}

func (c *wsConnect) writeWsLoop() {
	ticker := time.NewTicker(pingPeriod)
	defer ticker.Stop()
	for {
		select {
		// 取一个消息发送给客户端
		case msg := <-c.outChan:
			if err := c.wsConn.WriteMessage(msg.messageType, msg.data); err != nil {
				log.Error(c.Log("发送消息错误:%v", err))
				// 关闭连接
				c.Close()
				return
			}
		case <-c.closeCh:
			// 收到关闭通知
			return
		case <-ticker.C:
			_ = c.wsConn.SetWriteDeadline(time.Now().Add(writeWait))
			if err := c.wsConn.WriteMessage(websocket.PingMessage, []byte{}); err != nil {
				return
			}
		}
	}
}

func (c *wsConnect) handle(process func(IConn, []byte)) {
	for {
		msg, err := c.readFromChan()
		if err != nil {
			// log.Error(c.Log("获取消息错误:%v", err))
			break
		}
		// Log.Debug(c.Log("接收消息:%v", msg.data))
		process(c, msg.data)
	}
}

// 设置用户id
func (c *wsConnect) setUserId(uid int64) {
	c.userId = uid
}

// 获取连接id
func (c *wsConnect) Id() uint32 {
	return c.id
}

// 获取用户id
func (c *wsConnect) UserId() int64 {
	return c.userId
}

func (c *wsConnect) Name() string {
	return fmt.Sprintf("用户:%v, 地址:%v", c.userId, c.wsConn.RemoteAddr())
}

func (c *wsConnect) Log(format string, v ...interface{}) string {
	s := fmt.Sprintf("连接:%v, id:%v ", c.wsConn.RemoteAddr().String(), c.id)
	return s + fmt.Sprintf(format, v...)
}