110 lines
2.9 KiB
Go
110 lines
2.9 KiB
Go
package websocket
|
||
|
||
import (
|
||
"encoding/json"
|
||
"time"
|
||
|
||
"github.com/flipped-aurora/gin-vue-admin/server/global"
|
||
"github.com/gin-gonic/gin"
|
||
"github.com/olahol/melody"
|
||
)
|
||
|
||
// HandleConnection 处理WebSocket连接
|
||
func HandleConnection(ctx *gin.Context) {
|
||
// 从中间件获取用户ID
|
||
userId, exists := ctx.Get("userId")
|
||
if !exists {
|
||
ctx.JSON(401, gin.H{"error": "用户认证失败"})
|
||
return
|
||
}
|
||
|
||
// 升级为WebSocket连接，并存储用户信息
|
||
global.MELODY.HandleRequestWithKeys(ctx.Writer, ctx.Request, map[string]interface{}{
|
||
"userId": userId,
|
||
})
|
||
}
|
||
|
||
// SendMessage 发送消息到WebSocket连接
|
||
func SendMessage(s *melody.Session, message interface{}) error {
|
||
data, err := json.Marshal(message)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
return s.Write(data)
|
||
}
|
||
|
||
// SendErrorMessage 发送错误消息
|
||
func SendErrorMessage(s *melody.Session, errorMsg string) {
|
||
SendMessage(s, NewErrorEvent(errorMsg))
|
||
}
|
||
|
||
// BroadcastToUser 向指定用户广播消息
|
||
func BroadcastToUser(userID uint, message interface{}) error {
|
||
messageBytes, err := json.Marshal(message)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
|
||
// 遍历所有连接，找到对应用户的连接
|
||
return global.MELODY.BroadcastFilter(messageBytes, func(s *melody.Session) bool {
|
||
if userIdInterface, exists := s.Get("userId"); exists {
|
||
if userId, ok := userIdInterface.(uint); ok {
|
||
return userId == userID
|
||
}
|
||
}
|
||
return false
|
||
})
|
||
}
|
||
|
||
// BroadcastMessageToUser 向指定用户广播流式消息
|
||
func BroadcastMessageToUser(userID uint, delta string) error {
|
||
return BroadcastToUser(userID, NewMessageEvent(delta))
|
||
}
|
||
|
||
// BroadcastErrorToUser 向指定用户广播错误消息
|
||
func BroadcastErrorToUser(userID uint, errorMsg string) error {
|
||
return BroadcastToUser(userID, NewErrorEvent(errorMsg))
|
||
}
|
||
|
||
// BroadcastDoneToUser 向指定用户广播完成消息
|
||
func BroadcastDoneToUser(userID uint, message, sessionId string) error {
|
||
return BroadcastToUser(userID, NewDoneEvent(message, sessionId))
|
||
}
|
||
|
||
// GetConnectedUserCount 获取当前连接的用户数量
|
||
func GetConnectedUserCount() int {
|
||
count := 0
|
||
global.MELODY.BroadcastFilter([]byte{}, func(s *melody.Session) bool {
|
||
if _, exists := s.Get("userId"); exists {
|
||
count++
|
||
}
|
||
return false // 不实际发送消息，只统计
|
||
})
|
||
return count
|
||
}
|
||
|
||
// SendPingToUser 向指定用户发送心跳消息
|
||
func SendPingToUser(userID uint) error {
|
||
pingMsg := Message{
|
||
Type: "ping",
|
||
Data: map[string]interface{}{
|
||
"timestamp": time.Now().Unix(),
|
||
},
|
||
}
|
||
return BroadcastToUser(userID, pingMsg)
|
||
}
|
||
|
||
// IsUserConnected 检查用户是否在线
|
||
func IsUserConnected(userID uint) bool {
|
||
connected := false
|
||
global.MELODY.BroadcastFilter([]byte{}, func(s *melody.Session) bool {
|
||
if userIdInterface, exists := s.Get("userId"); exists {
|
||
if userId, ok := userIdInterface.(uint); ok && userId == userID {
|
||
connected = true
|
||
}
|
||
}
|
||
return false // 不实际发送消息，只检查
|
||
})
|
||
return connected
|
||
}
|