2022-12-07 16:06:36 +08:00
|
|
|
package channel
|
|
|
|
|
|
|
|
import (
|
|
|
|
"errors"
|
2023-05-05 10:43:56 +08:00
|
|
|
"fmt"
|
2022-12-07 16:06:36 +08:00
|
|
|
"github.com/gorilla/websocket"
|
|
|
|
"message-pusher/common"
|
|
|
|
"message-pusher/model"
|
|
|
|
"sync"
|
2022-12-08 14:50:22 +08:00
|
|
|
"time"
|
2022-12-07 16:06:36 +08:00
|
|
|
)
|
|
|
|
|
2022-12-08 14:50:22 +08:00
|
|
|
const (
|
|
|
|
writeWait = 10 * time.Second
|
|
|
|
pongWait = 60 * time.Second
|
|
|
|
pingPeriod = (pongWait * 9) / 10
|
|
|
|
maxMessageSize = 512
|
|
|
|
)
|
2022-12-07 16:06:36 +08:00
|
|
|
|
2022-12-08 14:50:22 +08:00
|
|
|
type webSocketClient struct {
|
2023-05-05 10:43:56 +08:00
|
|
|
key string
|
2022-12-08 14:50:22 +08:00
|
|
|
conn *websocket.Conn
|
2022-12-22 17:59:12 +08:00
|
|
|
message chan *model.Message
|
2022-12-08 14:50:22 +08:00
|
|
|
pong chan bool
|
|
|
|
stop chan bool
|
|
|
|
timestamp int64
|
|
|
|
}
|
|
|
|
|
|
|
|
func (c *webSocketClient) handleDataReading() {
|
|
|
|
c.conn.SetReadLimit(maxMessageSize)
|
|
|
|
_ = c.conn.SetReadDeadline(time.Now().Add(pongWait))
|
|
|
|
c.conn.SetPongHandler(func(string) error {
|
|
|
|
return c.conn.SetReadDeadline(time.Now().Add(pongWait))
|
|
|
|
})
|
|
|
|
for {
|
|
|
|
messageType, _, err := c.conn.ReadMessage()
|
|
|
|
if err != nil {
|
|
|
|
if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseNoStatusReceived, websocket.CloseAbnormalClosure) {
|
|
|
|
common.SysError("error read WebSocket client: " + err.Error())
|
|
|
|
}
|
|
|
|
c.close()
|
|
|
|
break
|
|
|
|
}
|
|
|
|
switch messageType {
|
|
|
|
case websocket.PingMessage:
|
|
|
|
c.pong <- true
|
|
|
|
case websocket.CloseMessage:
|
|
|
|
c.close()
|
|
|
|
break
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
func (c *webSocketClient) handleDataWriting() {
|
|
|
|
pingTicker := time.NewTicker(pingPeriod)
|
|
|
|
defer func() {
|
|
|
|
pingTicker.Stop()
|
|
|
|
clientConnMapMutex.Lock()
|
2023-05-05 10:43:56 +08:00
|
|
|
client, ok := clientMap[c.key]
|
2022-12-08 14:50:22 +08:00
|
|
|
// otherwise we may delete the new added client!
|
|
|
|
if ok && client.timestamp == c.timestamp {
|
2023-05-05 10:43:56 +08:00
|
|
|
delete(clientMap, c.key)
|
2022-12-08 14:50:22 +08:00
|
|
|
}
|
|
|
|
clientConnMapMutex.Unlock()
|
|
|
|
err := c.conn.Close()
|
|
|
|
if err != nil {
|
|
|
|
common.SysError("error close WebSocket client: " + err.Error())
|
|
|
|
}
|
|
|
|
}()
|
|
|
|
for {
|
|
|
|
select {
|
|
|
|
case message := <-c.message:
|
|
|
|
_ = c.conn.SetWriteDeadline(time.Now().Add(writeWait))
|
|
|
|
err := c.conn.WriteJSON(message)
|
|
|
|
if err != nil {
|
|
|
|
common.SysError("error write data to WebSocket client: " + err.Error())
|
|
|
|
return
|
|
|
|
}
|
|
|
|
case <-c.pong:
|
|
|
|
err := c.conn.WriteMessage(websocket.PongMessage, nil)
|
|
|
|
if err != nil {
|
|
|
|
common.SysError("error send pong to WebSocket client: " + err.Error())
|
|
|
|
return
|
|
|
|
}
|
|
|
|
case <-pingTicker.C:
|
|
|
|
_ = c.conn.SetWriteDeadline(time.Now().Add(writeWait))
|
|
|
|
err := c.conn.WriteMessage(websocket.PingMessage, nil)
|
|
|
|
if err != nil {
|
|
|
|
common.SysError("error write data to WebSocket client: " + err.Error())
|
|
|
|
return
|
|
|
|
}
|
|
|
|
case <-c.stop:
|
|
|
|
err := c.conn.WriteMessage(websocket.CloseMessage, nil)
|
|
|
|
if err != nil {
|
|
|
|
common.SysError("error write data to WebSocket client: " + err.Error())
|
|
|
|
}
|
|
|
|
return
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2022-12-22 17:59:12 +08:00
|
|
|
func (c *webSocketClient) sendMessage(message *model.Message) {
|
2022-12-08 14:50:22 +08:00
|
|
|
c.message <- message
|
2022-12-07 16:06:36 +08:00
|
|
|
}
|
|
|
|
|
2022-12-08 14:50:22 +08:00
|
|
|
func (c *webSocketClient) close() {
|
|
|
|
// should only be called once
|
|
|
|
c.stop <- true
|
|
|
|
// the defer function in handleDataWriting will do the cleanup
|
2022-12-07 16:06:36 +08:00
|
|
|
}
|
|
|
|
|
2023-05-05 10:43:56 +08:00
|
|
|
var clientMap map[string]*webSocketClient
|
2022-12-08 14:50:22 +08:00
|
|
|
var clientConnMapMutex sync.Mutex
|
|
|
|
|
|
|
|
func init() {
|
2022-12-07 16:06:36 +08:00
|
|
|
clientConnMapMutex.Lock()
|
2023-05-05 10:43:56 +08:00
|
|
|
clientMap = make(map[string]*webSocketClient)
|
2022-12-07 16:06:36 +08:00
|
|
|
clientConnMapMutex.Unlock()
|
|
|
|
}
|
|
|
|
|
2023-05-05 10:43:56 +08:00
|
|
|
func RegisterClient(channelName string, userId int, conn *websocket.Conn) {
|
|
|
|
key := fmt.Sprintf("%s:%d", channelName, userId)
|
2022-12-07 16:06:36 +08:00
|
|
|
clientConnMapMutex.Lock()
|
2023-05-05 10:43:56 +08:00
|
|
|
oldClient, existed := clientMap[key]
|
2022-12-07 16:06:36 +08:00
|
|
|
clientConnMapMutex.Unlock()
|
|
|
|
if existed {
|
2022-12-22 17:59:12 +08:00
|
|
|
byeMessage := &model.Message{
|
2022-12-07 16:06:36 +08:00
|
|
|
Title: common.SystemName,
|
|
|
|
Description: "其他客户端已连接服务器,本客户端已被挤下线!",
|
|
|
|
}
|
2022-12-08 14:50:22 +08:00
|
|
|
oldClient.sendMessage(byeMessage)
|
|
|
|
oldClient.close()
|
2022-12-07 16:06:36 +08:00
|
|
|
}
|
2022-12-22 17:59:12 +08:00
|
|
|
helloMessage := &model.Message{
|
2022-12-07 16:06:36 +08:00
|
|
|
Title: common.SystemName,
|
|
|
|
Description: "客户端连接成功!",
|
|
|
|
}
|
2022-12-08 14:50:22 +08:00
|
|
|
newClient := &webSocketClient{
|
2023-05-05 10:43:56 +08:00
|
|
|
key: key,
|
2022-12-08 14:50:22 +08:00
|
|
|
conn: conn,
|
2022-12-22 17:59:12 +08:00
|
|
|
message: make(chan *model.Message),
|
2022-12-08 14:50:22 +08:00
|
|
|
pong: make(chan bool),
|
|
|
|
stop: make(chan bool),
|
|
|
|
timestamp: time.Now().UnixMilli(),
|
2022-12-07 16:06:36 +08:00
|
|
|
}
|
2022-12-08 14:50:22 +08:00
|
|
|
go newClient.handleDataWriting()
|
|
|
|
go newClient.handleDataReading()
|
|
|
|
defer newClient.sendMessage(helloMessage)
|
|
|
|
clientConnMapMutex.Lock()
|
2023-05-05 10:43:56 +08:00
|
|
|
clientMap[key] = newClient
|
2022-12-08 14:50:22 +08:00
|
|
|
clientConnMapMutex.Unlock()
|
2022-12-07 16:06:36 +08:00
|
|
|
}
|
|
|
|
|
2023-05-05 10:43:56 +08:00
|
|
|
func SendClientMessage(message *model.Message, user *model.User, channel_ *model.Channel) error {
|
|
|
|
key := fmt.Sprintf("%s:%d", channel_.Name, user.Id)
|
2022-12-07 16:06:36 +08:00
|
|
|
clientConnMapMutex.Lock()
|
2023-05-05 10:43:56 +08:00
|
|
|
client, existed := clientMap[key]
|
2022-12-07 16:06:36 +08:00
|
|
|
clientConnMapMutex.Unlock()
|
|
|
|
if !existed {
|
|
|
|
return errors.New("客户端未连接")
|
|
|
|
}
|
2022-12-08 14:50:22 +08:00
|
|
|
client.sendMessage(message)
|
|
|
|
return nil
|
2022-12-07 16:06:36 +08:00
|
|
|
}
|