117 lines
3.4 KiB
Go
117 lines
3.4 KiB
Go
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()
|
||
}
|
||
}
|