56 lines
1.4 KiB
Go
56 lines
1.4 KiB
Go
package websocket
|
||
|
||
import (
|
||
"context"
|
||
"fmt"
|
||
"net/http"
|
||
|
||
"github.com/gin-gonic/gin"
|
||
"github.com/go-admin-team/go-admin-core/logger"
|
||
"github.com/go-admin-team/go-admin-core/sdk/pkg/jwtauth"
|
||
"github.com/gorilla/websocket"
|
||
)
|
||
|
||
var upgrader = websocket.Upgrader{
|
||
CheckOrigin: func(r *http.Request) bool {
|
||
return true // 允许跨域
|
||
},
|
||
}
|
||
|
||
func ServeWS(hub *Hub) gin.HandlerFunc {
|
||
return func(c *gin.Context) {
|
||
// 这里假设 JWT 中间件已验证，且用户 ID 在 Context
|
||
claims := jwtauth.ExtractClaims(c)
|
||
userID, ok := claims["identity"].(float64)
|
||
if !ok {
|
||
c.JSON(http.StatusUnauthorized, gin.H{"error": "invalid user ID in token"})
|
||
return
|
||
}
|
||
|
||
// 设置 Subprotocols 支持客户端传来的协议（token等）
|
||
upgrader.Subprotocols = []string{c.GetHeader("Sec-WebSocket-Protocol")}
|
||
|
||
conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
|
||
if err != nil {
|
||
logger.Errorf("WebSocket upgrade failed: %v", err)
|
||
c.JSON(http.StatusBadRequest, gin.H{"error": "upgrade failed"})
|
||
return
|
||
}
|
||
|
||
ctx, cancel := context.WithCancel(context.Background())
|
||
|
||
client := &Client{
|
||
ID: fmt.Sprintf("%v", userID),
|
||
Conn: conn,
|
||
Send: make(chan []byte, 1024),
|
||
Context: ctx,
|
||
CancelFunc: cancel,
|
||
}
|
||
|
||
hub.Register <- client
|
||
|
||
go client.Read(hub)
|
||
go client.Write()
|
||
}
|
||
}
|