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...) }