89 lines
2.0 KiB
Go
89 lines
2.0 KiB
Go
package modules
|
|
|
|
import (
|
|
"encoding/json"
|
|
"log"
|
|
"net/http"
|
|
"time"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/gorilla/websocket"
|
|
"github.com/lineup/app-server/internal"
|
|
"github.com/lineup/app-server/internal/wkapi"
|
|
)
|
|
|
|
var upgrader = websocket.Upgrader{
|
|
CheckOrigin: func(r *http.Request) bool { return true },
|
|
}
|
|
|
|
// ClientModule manages user WebSocket connections for real-time push.
|
|
type ClientModule struct {
|
|
hub *internal.Hub
|
|
wk *wkapi.Client
|
|
}
|
|
|
|
// NewClientModule creates a new ClientModule.
|
|
func NewClientModule(hub *internal.Hub, wk *wkapi.Client) *ClientModule {
|
|
return &ClientModule{hub: hub, wk: wk}
|
|
}
|
|
|
|
// RegisterRoutes registers WebSocket routes.
|
|
func (m *ClientModule) RegisterRoutes(r *gin.Engine) {
|
|
r.GET("/ws", m.handleWS)
|
|
}
|
|
|
|
func (m *ClientModule) handleWS(c *gin.Context) {
|
|
token := c.Query("token")
|
|
uid := c.Query("uid")
|
|
if token == "" || uid == "" {
|
|
c.JSON(http.StatusBadRequest, gin.H{"error": "参数错误: token和uid必填"})
|
|
return
|
|
}
|
|
|
|
// Verify token with WuKongIM (register if not exists, or validate)
|
|
if err := m.wk.RegisterToken(uid, token); err != nil {
|
|
c.JSON(http.StatusUnauthorized, gin.H{"error": "Token验证失败: " + err.Error()})
|
|
return
|
|
}
|
|
|
|
conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
|
|
if err != nil {
|
|
log.Printf("WebSocket升级失败: %v", err)
|
|
return
|
|
}
|
|
|
|
m.hub.Register(uid, conn)
|
|
log.Printf("客户端连接: uid=%s", uid)
|
|
|
|
// Send a welcome message
|
|
welcome, _ := json.Marshal(map[string]interface{}{
|
|
"event": "connected",
|
|
"uid": uid,
|
|
})
|
|
conn.WriteMessage(websocket.TextMessage, welcome)
|
|
|
|
// Read loop (keep connection alive, handle pings)
|
|
go m.readLoop(uid, conn)
|
|
}
|
|
|
|
func (m *ClientModule) readLoop(uid string, conn *websocket.Conn) {
|
|
defer func() {
|
|
m.hub.Unregister(uid)
|
|
conn.Close()
|
|
log.Printf("客户端断开: uid=%s", uid)
|
|
}()
|
|
|
|
conn.SetReadDeadline(time.Now().Add(60 * time.Second))
|
|
conn.SetPongHandler(func(string) error {
|
|
conn.SetReadDeadline(time.Now().Add(60 * time.Second))
|
|
return nil
|
|
})
|
|
|
|
for {
|
|
_, _, err := conn.ReadMessage()
|
|
if err != nil {
|
|
break
|
|
}
|
|
}
|
|
}
|