88 lines
1.9 KiB
Go
88 lines
1.9 KiB
Go
package chathub
|
|
|
|
import (
|
|
"fmt"
|
|
"time"
|
|
|
|
"hfb_sys/backend/internal/middleware"
|
|
"hfb_sys/backend/pkg/response"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
)
|
|
|
|
const heartbeatInterval = 30 * time.Second
|
|
|
|
type Handler struct {
|
|
hub *Hub
|
|
}
|
|
|
|
func NewHandler(hub *Hub) *Handler {
|
|
return &Handler{hub: hub}
|
|
}
|
|
|
|
// UserEvents 处理用户端 SSE 连接: GET /api/chats/events
|
|
func (h *Handler) UserEvents(c *gin.Context) {
|
|
value, ok := c.Get(middleware.ContextUserID)
|
|
if !ok {
|
|
response.Unauthorized(c, "缺少用户上下文")
|
|
return
|
|
}
|
|
userID, ok := value.(uint64)
|
|
if !ok {
|
|
response.Unauthorized(c, "用户上下文无效")
|
|
return
|
|
}
|
|
h.serveSSE(c, "user", userID)
|
|
}
|
|
|
|
// AdminEvents 处理管理端 SSE 连接: GET /api/admin/chats/events
|
|
func (h *Handler) AdminEvents(c *gin.Context) {
|
|
value, ok := c.Get(middleware.ContextAdminID)
|
|
if !ok {
|
|
response.Unauthorized(c, "缺少管理员上下文")
|
|
return
|
|
}
|
|
adminID, ok := value.(uint64)
|
|
if !ok {
|
|
response.Unauthorized(c, "管理员上下文无效")
|
|
return
|
|
}
|
|
h.serveSSE(c, "admin", adminID)
|
|
}
|
|
|
|
func (h *Handler) serveSSE(c *gin.Context, pType string, pID uint64) {
|
|
ch := h.hub.Subscribe(pType, pID)
|
|
defer h.hub.Unsubscribe(pType, pID, ch)
|
|
|
|
c.Header("Content-Type", "text/event-stream")
|
|
c.Header("Cache-Control", "no-cache")
|
|
c.Header("Connection", "keep-alive")
|
|
c.Header("X-Accel-Buffering", "no")
|
|
|
|
// 发送初始连接确认
|
|
fmt.Fprintf(c.Writer, "event: connected\ndata: {\"ok\":true}\n\n")
|
|
c.Writer.Flush()
|
|
|
|
heartbeat := time.NewTicker(heartbeatInterval)
|
|
defer heartbeat.Stop()
|
|
|
|
clientGone := c.Request.Context().Done()
|
|
|
|
for {
|
|
select {
|
|
case <-clientGone:
|
|
return
|
|
case <-heartbeat.C:
|
|
fmt.Fprintf(c.Writer, ":heartbeat\n\n")
|
|
c.Writer.Flush()
|
|
case event, ok := <-ch:
|
|
if !ok {
|
|
return
|
|
}
|
|
data := MarshalEvent(event)
|
|
fmt.Fprintf(c.Writer, "event: %s\ndata: %s\n\n", event.Type, data)
|
|
c.Writer.Flush()
|
|
}
|
|
}
|
|
}
|