5021b039c9
- 新增 POST /api/auth/websocket-ticket：需 JWT 与公司头，签发 30 秒有效的一次性随机票据 - /ws 握手消费票据并以数据库当前状态复核身份与公司授权，JWT 不进 URL - Hub 注册/注销/投递选人统一由事件循环串行管理，消除 RLock 内 delete 的并发缺陷 - 事件携带 organization_id，仅投递同公司订阅者，投递前复核失败即关闭连接 - 统一广播签名为 (organizationID, action, orderID, data) - 新增前端 companySocket 模块：按当前公司取票建连、切换/登出关旧连重连、消息校验组织与上下文版本
259 lines
7.6 KiB
Go
259 lines
7.6 KiB
Go
package handler
|
||
|
||
import (
|
||
"encoding/json"
|
||
"log"
|
||
"net/http"
|
||
"sync"
|
||
|
||
"github.com/gin-gonic/gin"
|
||
"github.com/gorilla/websocket"
|
||
"github.com/wmms/wmms/internal/service"
|
||
"github.com/wmms/wmms/pkg/response"
|
||
)
|
||
|
||
var upgrader = websocket.Upgrader{
|
||
ReadBufferSize: 1024,
|
||
WriteBufferSize: 1024,
|
||
CheckOrigin: func(r *http.Request) bool {
|
||
return true // Allow all origins for development
|
||
},
|
||
}
|
||
|
||
// StockUpdateEvent represents a stock update notification
|
||
type StockUpdateEvent struct {
|
||
Type string `json:"type"` // "stock_update"
|
||
Action string `json:"action"` // "inbound" or "outbound"
|
||
// OrganizationID 是本次库存变化所属公司，供按公司投递过滤（Task 5D 按它路由）。
|
||
OrganizationID int64 `json:"organization_id"`
|
||
OrderID int64 `json:"order_id"`
|
||
Data interface{} `json:"data,omitempty"`
|
||
}
|
||
|
||
// AuthorizationVerifier 在握手与每次投递前，以数据库当前状态复核订阅者是否仍可访问该公司。
|
||
// 实现必须实时查询数据库，不得依赖 JWT 中的角色快照；查询异常返回 error，调用方据此
|
||
// 拒绝连接或不投递。
|
||
type AuthorizationVerifier interface {
|
||
Verify(userID, organizationID int64) error
|
||
}
|
||
|
||
// organizationAccessVerifier 是生产用的复核器：先确认用户仍存在且有有效角色记录，
|
||
// 再确认其仍持有该公司的访问授权（超级管理员以启用的有效角色判定）。
|
||
type organizationAccessVerifier struct{}
|
||
|
||
func (organizationAccessVerifier) Verify(userID, organizationID int64) error {
|
||
access := service.NewOrganizationAccessService()
|
||
if _, err := access.EffectiveRoles(userID); err != nil {
|
||
return err
|
||
}
|
||
return access.Validate(userID, organizationID)
|
||
}
|
||
|
||
// Client 表示一个已完成票据握手、绑定到某用户与某公司的订阅者。
|
||
type Client struct {
|
||
hub *Hub
|
||
conn *websocket.Conn
|
||
send chan []byte
|
||
userID int64
|
||
organizationID int64
|
||
ticketID string
|
||
|
||
// sendMu 保护 send 通道的发送与关闭。投递协程与事件循环的注销流程可能并发，
|
||
// 用互斥锁加 closed 标记保证不会向已关闭通道发送，也不会重复关闭。
|
||
sendMu sync.Mutex
|
||
closed bool
|
||
}
|
||
|
||
// broadcastMessage 是送入事件循环的库存事件：携带公司维度，只发给同公司订阅者。
|
||
type broadcastMessage struct {
|
||
organizationID int64
|
||
payload []byte
|
||
}
|
||
|
||
// Hub 维护活跃订阅者并按公司投递事件。
|
||
//
|
||
// clients 只由事件循环 run 读写，不设互斥锁：注册、注销与投递选人都在同一协程里
|
||
// 串行完成，消除旧实现“在 RLock 内 delete(h.clients, client)”的并发缺陷。
|
||
// 投递前的数据库复核与网络写入放到独立协程，避免拖慢事件循环。
|
||
type Hub struct {
|
||
clients map[*Client]bool
|
||
broadcast chan broadcastMessage
|
||
register chan *Client
|
||
unregister chan *Client
|
||
verifier AuthorizationVerifier
|
||
}
|
||
|
||
var hubInstance *Hub
|
||
var hubOnce sync.Once
|
||
|
||
// newHub 构造 Hub；测试可注入假的 verifier，生产用 organizationAccessVerifier。
|
||
func newHub(verifier AuthorizationVerifier) *Hub {
|
||
return &Hub{
|
||
clients: make(map[*Client]bool),
|
||
broadcast: make(chan broadcastMessage, 256),
|
||
register: make(chan *Client),
|
||
unregister: make(chan *Client),
|
||
verifier: verifier,
|
||
}
|
||
}
|
||
|
||
// GetHub returns the singleton Hub instance
|
||
func GetHub() *Hub {
|
||
hubOnce.Do(func() {
|
||
hubInstance = newHub(organizationAccessVerifier{})
|
||
go hubInstance.run()
|
||
})
|
||
return hubInstance
|
||
}
|
||
|
||
func (h *Hub) run() {
|
||
for {
|
||
select {
|
||
case client := <-h.register:
|
||
h.clients[client] = true
|
||
log.Printf("WebSocket 客户端已连接（公司 %d）。当前连接数：%d", client.organizationID, len(h.clients))
|
||
|
||
case client := <-h.unregister:
|
||
if _, ok := h.clients[client]; ok {
|
||
delete(h.clients, client)
|
||
client.closeSend()
|
||
log.Printf("WebSocket 客户端已断开。当前连接数：%d", len(h.clients))
|
||
}
|
||
|
||
case message := <-h.broadcast:
|
||
// 只选择同公司订阅者；每个收件人单独起协程，事件循环本身不等待投递完成。
|
||
for client := range h.clients {
|
||
if client.organizationID != message.organizationID {
|
||
continue
|
||
}
|
||
go h.deliver(client, message.payload)
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// deliver 在发送前复核授权；校验失败或数据库异常则不投递并关闭连接。
|
||
func (h *Hub) deliver(client *Client, payload []byte) {
|
||
if h.verifier != nil {
|
||
if err := h.verifier.Verify(client.userID, client.organizationID); err != nil {
|
||
log.Printf("WebSocket 投递前复核失败，关闭连接（用户 %d，公司 %d）：%v", client.userID, client.organizationID, err)
|
||
h.unregister <- client
|
||
return
|
||
}
|
||
}
|
||
if !client.enqueue(payload) {
|
||
h.unregister <- client
|
||
}
|
||
}
|
||
|
||
// enqueue 尝试把消息写入发送缓冲，返回是否成功。缓冲已满（慢客户端）时关闭连接。
|
||
func (c *Client) enqueue(payload []byte) bool {
|
||
c.sendMu.Lock()
|
||
defer c.sendMu.Unlock()
|
||
if c.closed {
|
||
return false
|
||
}
|
||
select {
|
||
case c.send <- payload:
|
||
return true
|
||
default:
|
||
c.closed = true
|
||
close(c.send)
|
||
return false
|
||
}
|
||
}
|
||
|
||
// closeSend 幂等地关闭发送通道，唤醒 writePump 收尾并关闭连接。
|
||
func (c *Client) closeSend() {
|
||
c.sendMu.Lock()
|
||
defer c.sendMu.Unlock()
|
||
if !c.closed {
|
||
c.closed = true
|
||
close(c.send)
|
||
}
|
||
}
|
||
|
||
// BroadcastStockUpdate 把库存事件按公司投递：只有同公司订阅者会收到。
|
||
func (h *Hub) BroadcastStockUpdate(organizationID int64, action string, orderID int64, data interface{}) {
|
||
event := StockUpdateEvent{
|
||
Type: "stock_update",
|
||
Action: action,
|
||
OrganizationID: organizationID,
|
||
OrderID: orderID,
|
||
Data: data,
|
||
}
|
||
msg, err := json.Marshal(event)
|
||
if err != nil {
|
||
log.Printf("序列化库存更新事件失败: %v", err)
|
||
return
|
||
}
|
||
h.broadcast <- broadcastMessage{organizationID: organizationID, payload: msg}
|
||
}
|
||
|
||
func (c *Client) readPump() {
|
||
defer func() {
|
||
c.hub.unregister <- c
|
||
c.conn.Close()
|
||
}()
|
||
for {
|
||
_, _, err := c.conn.ReadMessage()
|
||
if err != nil {
|
||
if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
|
||
log.Printf("WebSocket error: %v", err)
|
||
}
|
||
break
|
||
}
|
||
}
|
||
}
|
||
|
||
func (c *Client) writePump() {
|
||
defer func() {
|
||
c.conn.Close()
|
||
}()
|
||
for {
|
||
message, ok := <-c.send
|
||
if !ok {
|
||
c.conn.WriteMessage(websocket.CloseMessage, []byte{})
|
||
return
|
||
}
|
||
if err := c.conn.WriteMessage(websocket.TextMessage, message); err != nil {
|
||
return
|
||
}
|
||
}
|
||
}
|
||
|
||
// HandleWebSocket handles WebSocket connections at /ws
|
||
//
|
||
// 浏览器以 /ws?ticket=... 建连：先消费一次性短期票据，再以数据库当前状态复核身份与
|
||
// 公司授权。JWT 不放 URL，仅凭 URL 里的公司 ID 也无法订阅。
|
||
func HandleWebSocket(c *gin.Context) {
|
||
ticket, err := service.GetWebSocketTicketService().Consume(c.Query("ticket"))
|
||
if err != nil {
|
||
// 票据缺失、非法、过期或已消费（重放）都不允许建立连接。
|
||
response.Unauthorized(c, "WebSocket 票据无效或已过期")
|
||
return
|
||
}
|
||
hub := GetHub()
|
||
if err := hub.verifier.Verify(ticket.UserID, ticket.OrganizationID); err != nil {
|
||
// 公司授权或角色已被撤销：与普通请求一致，拒绝建连，不等 JWT 过期。
|
||
response.Forbidden(c, "当前公司不可访问")
|
||
return
|
||
}
|
||
conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
|
||
if err != nil {
|
||
log.Printf("WebSocket 升级失败: %v", err)
|
||
return
|
||
}
|
||
client := &Client{
|
||
hub: hub,
|
||
conn: conn,
|
||
send: make(chan []byte, 256),
|
||
userID: ticket.UserID,
|
||
organizationID: ticket.OrganizationID,
|
||
ticketID: ticket.Ticket,
|
||
}
|
||
hub.register <- client
|
||
go client.writePump()
|
||
go client.readPump()
|
||
}
|