7474a92e3b
P5 阶段交付物： - services/msg: 消息通知服务（NestJS） - notifications: 发送通知 + ES 全文检索 + search - config/elasticsearch.ts: ES Client 单例 - package.json: 补充 @opentelemetry/sdk-node + exporter-trace-otlp-http - services/push-gateway: WebSocket 推送网关（Go Gin） - internal/hub/hub.go: WebSocket 连接池管理（Register/Unregister/SendToUser） - internal/ws/handler.go: JWT 鉴权 + WebSocket 升级 + 内部推送 API - services/ai: AI 辅助服务（Python FastAPI） - /chat + /chat/stream（SSE 流式） - /generate/question + /optimize/expression - config.py: OpenAI 兼容 API 配置 - packages/shared-proto/proto/msg.proto: NotificationService 契约（send/search） - packages/shared-proto/proto/ai.proto: AiService 契约（含 stream 方法）
100 lines
2.4 KiB
Go
100 lines
2.4 KiB
Go
package ws
|
||
|
||
import (
|
||
"net/http"
|
||
"strings"
|
||
|
||
"github.com/edu-cloud/push-gateway/internal/hub"
|
||
"github.com/gin-gonic/gin"
|
||
"github.com/golang-jwt/jwt/v5"
|
||
"github.com/gorilla/websocket"
|
||
)
|
||
|
||
var upgrader = websocket.Upgrader{
|
||
CheckOrigin: func(r *http.Request) bool {
|
||
return true // P5 骨架，生产环境需校验 origin
|
||
},
|
||
}
|
||
|
||
type Handler struct {
|
||
hub *hub.Hub
|
||
jwtSecret string
|
||
}
|
||
|
||
func NewHandler(h *hub.Hub, jwtSecret string) *Handler {
|
||
return &Handler{hub: h, jwtSecret: jwtSecret}
|
||
}
|
||
|
||
func (h *Handler) HandleWebSocket(c *gin.Context) {
|
||
// 从 query 参数获取 token（WebSocket 无法设置 Authorization 头）
|
||
tokenStr := c.Query("token")
|
||
if tokenStr == "" {
|
||
c.JSON(http.StatusUnauthorized, gin.H{"error": "missing token"})
|
||
return
|
||
}
|
||
|
||
token, err := jwt.Parse(tokenStr, func(t *jwt.Token) (interface{}, error) {
|
||
if _, ok := t.Method.(*jwt.SigningMethodHMAC); !ok {
|
||
return nil, jwt.ErrSignatureInvalid
|
||
}
|
||
return []byte(h.jwtSecret), nil
|
||
})
|
||
|
||
if err != nil || !token.Valid {
|
||
c.JSON(http.StatusUnauthorized, gin.H{"error": "invalid token"})
|
||
return
|
||
}
|
||
|
||
claims, ok := token.Claims.(jwt.MapClaims)
|
||
if !ok {
|
||
c.JSON(http.StatusUnauthorized, gin.H{"error": "invalid claims"})
|
||
return
|
||
}
|
||
|
||
userID, ok := claims["sub"].(string)
|
||
if !ok {
|
||
c.JSON(http.StatusUnauthorized, gin.H{"error": "missing user id"})
|
||
return
|
||
}
|
||
|
||
conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
|
||
if err != nil {
|
||
return
|
||
}
|
||
defer conn.Close()
|
||
|
||
h.hub.Register(userID, conn)
|
||
defer h.hub.Unregister(userID, conn)
|
||
|
||
// 读取循环（保持连接，处理心跳）
|
||
for {
|
||
_, msg, err := conn.ReadMessage()
|
||
if err != nil {
|
||
break
|
||
}
|
||
// 处理心跳 ping
|
||
if strings.ToLower(string(msg)) == "ping" {
|
||
conn.WriteMessage(websocket.TextMessage, []byte("pong"))
|
||
}
|
||
}
|
||
}
|
||
|
||
// PushHandler 接收来自 Msg 服务的推送请求
|
||
func (h *Handler) PushHandler(c *gin.Context) {
|
||
var req struct {
|
||
UserID string `json:"user_id"`
|
||
Message string `json:"message"`
|
||
}
|
||
if err := c.ShouldBindJSON(&req); err != nil {
|
||
c.JSON(http.StatusBadRequest, gin.H{"success": false, "error": gin.H{"code": "INVALID_REQUEST", "message": err.Error()}})
|
||
return
|
||
}
|
||
|
||
if err := h.hub.SendToUser(req.UserID, []byte(req.Message)); err != nil {
|
||
c.JSON(http.StatusInternalServerError, gin.H{"success": false, "error": gin.H{"code": "PUSH_FAILED", "message": err.Error()}})
|
||
return
|
||
}
|
||
|
||
c.JSON(http.StatusOK, gin.H{"success": true})
|
||
}
|