This commit is contained in:
@@ -0,0 +1,109 @@
|
||||
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
|
||||
}
|
||||
@@ -0,0 +1,54 @@
|
||||
package websocket
|
||||
|
||||
// Message WebSocket下发消息结构体
|
||||
type Message struct {
|
||||
Type string `json:"type"` // 消息类型：message, error, done
|
||||
Data interface{} `json:"data"` // 消息数据
|
||||
}
|
||||
|
||||
// MessageData 流式消息数据结构体
|
||||
type MessageData struct {
|
||||
Delta string `json:"delta"` // 增量消息内容
|
||||
}
|
||||
|
||||
// ErrorData 错误数据结构体
|
||||
type ErrorData struct {
|
||||
Error string `json:"error"` // 错误信息
|
||||
}
|
||||
|
||||
// DoneData 完成数据结构体
|
||||
type DoneData struct {
|
||||
Message string `json:"message"` // 完整消息内容
|
||||
SessionId string `json:"sessionId"` // 会话ID
|
||||
}
|
||||
|
||||
// NewMessageEvent 创建消息事件
|
||||
func NewMessageEvent(delta string) Message {
|
||||
return Message{
|
||||
Type: "message",
|
||||
Data: MessageData{
|
||||
Delta: delta,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// NewErrorEvent 创建错误事件
|
||||
func NewErrorEvent(errorMsg string) Message {
|
||||
return Message{
|
||||
Type: "error",
|
||||
Data: ErrorData{
|
||||
Error: errorMsg,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// NewDoneEvent 创建完成事件
|
||||
func NewDoneEvent(message, sessionId string) Message {
|
||||
return Message{
|
||||
Type: "done",
|
||||
Data: DoneData{
|
||||
Message: message,
|
||||
SessionId: sessionId,
|
||||
},
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,116 @@
|
||||
package websocket
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/flipped-aurora/gin-vue-admin/server/global"
|
||||
"github.com/flipped-aurora/gin-vue-admin/server/utils"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/olahol/melody"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// InitWebSocketServer 初始化WebSocket服务器
|
||||
func InitWebSocketServer() {
|
||||
// 创建melody实例
|
||||
global.MELODY = melody.New()
|
||||
|
||||
// 设置WebSocket配置
|
||||
global.MELODY.Config.MaxMessageSize = 1024 * 1024 // 1MB
|
||||
global.MELODY.Config.MessageBufferSize = 256
|
||||
global.MELODY.Config.PongWait = 60 * time.Second // 等待pong消息的时间
|
||||
global.MELODY.Config.PingPeriod = 54 * time.Second // 发送ping消息的间隔
|
||||
global.MELODY.Config.WriteWait = 10 * time.Second // 写入超时时间
|
||||
global.MELODY.Config.ConcurrentMessageHandling = false // 禁用并发消息处理
|
||||
|
||||
// 连接建立时的处理
|
||||
global.MELODY.HandleConnect(func(s *melody.Session) {
|
||||
if userIdInterface, exists := s.Get("userId"); exists {
|
||||
if userId, ok := userIdInterface.(uint); ok {
|
||||
global.GVA_LOG.Info("WebSocket连接建立",
|
||||
zap.Uint("userId", userId),
|
||||
zap.String("remoteAddr", s.Request.RemoteAddr))
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
// 连接断开时的处理
|
||||
global.MELODY.HandleDisconnect(func(s *melody.Session) {
|
||||
if userIdInterface, exists := s.Get("userId"); exists {
|
||||
if userId, ok := userIdInterface.(uint); ok {
|
||||
global.GVA_LOG.Info("WebSocket连接断开",
|
||||
zap.Uint("userId", userId),
|
||||
zap.String("remoteAddr", s.Request.RemoteAddr))
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
// 错误处理
|
||||
global.MELODY.HandleError(func(s *melody.Session, err error) {
|
||||
// 检查是否是正常关闭或客户端主动断开
|
||||
errMsg := err.Error()
|
||||
if strings.Contains(errMsg, "close 1000") ||
|
||||
strings.Contains(errMsg, "close 1001") ||
|
||||
strings.Contains(errMsg, "use of closed network connection") {
|
||||
// 正常关闭，只记录debug级别日志
|
||||
if userIdInterface, exists := s.Get("userId"); exists {
|
||||
if userId, ok := userIdInterface.(uint); ok {
|
||||
global.GVA_LOG.Debug("WebSocket正常关闭",
|
||||
zap.String("reason", errMsg),
|
||||
zap.Uint("userId", userId))
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// 异常错误，记录error级别日志
|
||||
if userIdInterface, exists := s.Get("userId"); exists {
|
||||
if userId, ok := userIdInterface.(uint); ok {
|
||||
global.GVA_LOG.Error("WebSocket异常错误",
|
||||
zap.Error(err),
|
||||
zap.Uint("userId", userId))
|
||||
}
|
||||
} else {
|
||||
global.GVA_LOG.Error("WebSocket异常错误", zap.Error(err))
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
// 由于只下发数据，不处理客户端消息，所以不设置消息处理器
|
||||
global.GVA_LOG.Info("WebSocket服务器初始化完成")
|
||||
}
|
||||
|
||||
// WebSocketAuthMiddleware WebSocket鉴权中间件
|
||||
func WebSocketAuthMiddleware() gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
// 从查询参数获取token
|
||||
token := c.Query("token")
|
||||
if token == "" {
|
||||
// 从头部获取token
|
||||
token = c.GetHeader("Authorization")
|
||||
if token != "" && len(token) > 7 && token[:7] == "Bearer " {
|
||||
token = token[7:]
|
||||
}
|
||||
}
|
||||
|
||||
if token == "" {
|
||||
c.JSON(http.StatusUnauthorized, gin.H{"error": "缺少认证token"})
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
|
||||
// 解析token获取用户ID
|
||||
j := utils.NewJWT()
|
||||
claims, err := j.ParseAppUserToken(token)
|
||||
if err != nil {
|
||||
global.GVA_LOG.Error("WebSocket token解析失败", zap.Error(err))
|
||||
c.JSON(http.StatusUnauthorized, gin.H{"error": "token无效"})
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
|
||||
// 将用户ID存储到上下文中
|
||||
c.Set("userId", claims.AppBaseClaims.ID)
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user