fox/ws/wsConn.go

package ws

import (
	"encoding/binary"
	"fmt"
	"github.com/fox/fox/ipb"
	"github.com/fox/fox/log"
	"github.com/fox/fox/safeChan"
	"github.com/golang/protobuf/proto"
	"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       *safeChan.SafeChan[*wsMessage] // 读队列
	outChan      *safeChan.SafeChan[*wsMessage] // 写队列
	id           uint32
	userId       int64
	onDisconnect func(IConn)
	once         sync.Once
}

func newWsConnect(wsConn *websocket.Conn, onDisconnect func(IConn)) *wsConnect {
	c := &wsConnect{
		wsConn:       wsConn,
		inChan:       safeChan.NewSafeChan[*wsMessage](1000),
		outChan:      safeChan.NewSafeChan[*wsMessage](1000),
		id:           nextConnId,
		userId:       0,
		onDisconnect: onDisconnect,
	}
	return c
}

// 把消息放进写队列
func (c *wsConnect) SendMsg(data []byte) error {
	return c.outChan.Write(&wsMessage{messageType: wsMsgType, data: data})
}

// 关闭链接
func (c *wsConnect) Close() {
	c.once.Do(func() {
		log.Debug(c.Log("关闭链接"))
		_ = c.wsConn.WriteMessage(websocket.CloseMessage, []byte{})
		c.inChan.Close()
		c.outChan.Close()
		_ = c.wsConn.Close()
		if c.onDisconnect != nil {
			c.onDisconnect(c)
		}
		wsMgr.Remove(c)
	})
}

// 循环从websocket中读取消息放入到读队列中
func (c *wsConnect) readWsLoop() {
	c.wsConn.SetReadLimit(maxMessageSize)
	_ = c.wsConn.SetReadDeadline(time.Now().Add(pongWait))
	c.wsConn.SetPongHandler(func(string) error {
		log.Debug(c.Log("received pong. from client"))
		_ = c.wsConn.SetReadDeadline(time.Now().Add(pongWait))
		return nil
	})

	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))
			}
			log.Debug(c.Log("关闭连接:%v", err))
			c.Close()
			return
		}
		switch msgType {
		//case websocket.PingMessage:
		//	log.Debug(c.Log("received ping message"))
		//	_ = c.sendMsg(websocket.PongMessage, []byte("pong"))
		//case websocket.PongMessage:
		//	log.Debug(c.Log("received pong from client"))
		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}
			_ = c.inChan.Write(msg)
		}
	}
}

func (c *wsConnect) writeWsLoop() {
	ticker := time.NewTicker(pingPeriod)
	defer ticker.Stop()
	for {
		select {
		// 取一个消息发送给客户端
		case msg := <-c.outChan.Reader():
			if err := c.wsConn.WriteMessage(msg.messageType, msg.data); err != nil {
				iMsg := &ipb.InternalMsg{}
				_ = proto.Unmarshal(msg.data, iMsg)
				log.Error(c.Log("发送消息错误:%v 消息长度:%v 消息内容:%v", err, len(msg.data), iMsg))
				// 关闭连接
				c.Close()
				return
			}
		case <-ticker.C:
			if err := c.wsConn.WriteMessage(websocket.PingMessage, []byte("ping")); err != nil {
				log.Error(c.Log("发送心跳失败:%v", err))
				c.Close()
				return
			} else {
				log.Debug(c.Log("发送心跳"))
			}
		}
	}
}

func (c *wsConnect) handle(process func(IConn, []byte)) {
	for {
		select {
		case msg, ok := <-c.inChan.Reader():
			if ok {
				process(c, msg.data)
			} else {
				c.Close()
			}
		}
	}
}

// 设置用户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) Addr() string {
	return c.wsConn.RemoteAddr().String()
}

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