diff --git a/cmd/server/main.go b/cmd/server/main.go index 479fa9bb..e909a9fc 100644 --- a/cmd/server/main.go +++ b/cmd/server/main.go @@ -49,6 +49,8 @@ func main() { h.RegisterRoutes(r) // WebSocket route + // 建连不走中间件：浏览器先凭 JWT 与公司头调用 /api/auth/websocket-ticket 换票， + // 再以 /ws?ticket=... 握手；HandleWebSocket 消费一次性票据并实时复核授权。 r.GET("/ws", handler.HandleWebSocket) // 端口可通过 APP_PORT 覆盖，默认 8082，便于试用脚本探测空闲端口。 diff --git a/frontend/src/utils/companySocket.js b/frontend/src/utils/companySocket.js new file mode 100644 index 00000000..88c08019 --- /dev/null +++ b/frontend/src/utils/companySocket.js @@ -0,0 +1,127 @@ +// 公司级 WebSocket 连接：把“当前公司上下文”与“实时事件”绑定在一起，供库存刷新等 +// 页面复用。关键约束： +// 1. 建连不能用长期 JWT，也不能仅凭 URL 里的公司 ID 订阅；必须先用票据接口换取 +// 短期一次性票据，再以 /ws?ticket=... 握手（服务端消费票据并实时复核授权）。 +// 2. 公司切换或登出要关闭旧连接，并为当前公司重新取票重连。 +// 3. 消息处理核对事件公司标识与捕获的上下文版本；连接回调失效时不得更新页面。 +// 4. 断线后再次 connect 会重新取票，由服务端重新校验授权。 +// +// 刻意只依赖 organizationContext：与 axios 拦截器共用同一套“公司 ID + 版本”语义， +// 不引入 Pinia / Axios，避免模块成环。 +import { + captureOrganizationContext, + isCurrentOrganizationContext +} from './organizationContext' + +// 依据页面地址推导 WebSocket 端点；测试可覆盖 location。 +export function buildSocketUrl(location) { + const loc = location || globalThis.location + const protocol = loc?.protocol === 'https:' ? 'wss:' : 'ws:' + const host = loc?.host || 'localhost' + return `${protocol}//${host}/ws` +} + +// createCompanySocket 返回一个连接控制器。fetchTicket 由调用方注入（通常包一条 +// /auth/websocket-ticket 请求），onMessage 只会在“当前公司 + 上下文未失效”时被调用。 +export function createCompanySocket(options = {}) { + const { + fetchTicket, + onMessage = () => {}, + WebSocketImpl = globalThis.WebSocket, + location = globalThis.location, + log = () => {} + } = options + + let socket = null + // generation 每次 connect/close 都递增：任何一个回调只要发现自己的代次与当前不符， + // 就说明连接已被替换或已关闭，不得再更新页面。 + let generation = 0 + + function isCurrent(gen, snapshot) { + return gen === generation && isCurrentOrganizationContext(snapshot) + } + + function handleMessage(event, gen, snapshot) { + // 连接回调失效（公司切换/登出/已被新连接替换）时不更新页面。 + if (!isCurrent(gen, snapshot)) return + let payload + try { + payload = JSON.parse(event.data) + } catch { + return + } + // 只处理当前公司的事件：旧公司事件或公司标识不符的一律丢弃。 + if (payload?.organization_id !== snapshot.organizationId) return + onMessage(payload, snapshot) + } + + async function connect() { + // 先关闭旧连接，避免旧公司连接残留。 + close() + const snapshot = captureOrganizationContext() + // 没有当前公司（登出或尚未选择）时不建立连接。 + if (snapshot.organizationId == null) return + const gen = ++generation + if (typeof WebSocketImpl !== 'function') { + log('当前环境不支持 WebSocket，跳过建连') + return + } + + let response + try { + response = await fetchTicket() + } catch (error) { + log('取得 WebSocket 票据失败', error) + return + } + const data = response?.data ?? response + const ticket = data?.ticket + if (!ticket) { + log('WebSocket 票据为空，跳过建连') + return + } + // 取票期间公司可能已切换或登出：快照失效则放弃这次连接，交给新的 connect 处理。 + if (!isCurrent(gen, snapshot)) return + + const ws = new WebSocketImpl( + `${buildSocketUrl(location)}?ticket=${encodeURIComponent(ticket)}` + ) + socket = ws + ws.onopen = () => { + if (isCurrent(gen, snapshot)) log('WebSocket 已连接（公司 %s）', snapshot.organizationId) + } + ws.onmessage = event => handleMessage(event, gen, snapshot) + ws.onclose = () => { + // 仅当仍是本次连接时才清空引用，避免误清掉后来建立的新连接。 + if (gen === generation) socket = null + } + ws.onerror = () => {} + } + + function close() { + // 递增代次使旧连接的所有回调失效。 + generation += 1 + if (socket) { + const previous = socket + socket = null + try { + previous.close() + } catch { + // 已关闭的连接重复 close 会抛错，忽略即可。 + } + } + } + + return { + connect, + close, + // 暴露句柄与内部消息入口，供测试直接驱动；业务代码只用 connect/close。 + handleMessage, + get socket() { + return socket + }, + get generation() { + return generation + } + } +} diff --git a/frontend/src/utils/companySocket.test.js b/frontend/src/utils/companySocket.test.js new file mode 100644 index 00000000..94c8d413 --- /dev/null +++ b/frontend/src/utils/companySocket.test.js @@ -0,0 +1,212 @@ +import { beforeEach, describe, expect, test, vi } from 'vitest' +import { createCompanySocket, buildSocketUrl } from './companySocket' +import { + commitOrganizationId, + resetOrganizationContext +} from './organizationContext' + +// 记录每个被创建的连接，供断言 URL、close 与手工推送消息。 +class FakeWebSocket { + constructor(url) { + this.url = url + this.closed = false + FakeWebSocket.instances.push(this) + } + close() { + this.closed = true + } + // 模拟服务端推送：直接触发本次连接的 onmessage。 + emitMessage(payload) { + this.onmessage?.({ data: JSON.stringify(payload) }) + } +} +FakeWebSocket.instances = [] + +const LOCATION = { protocol: 'http:', host: 'wms.example' } + +function makeSocket({ fetchTicket, onMessage = vi.fn() } = {}) { + return createCompanySocket({ + fetchTicket, + onMessage, + WebSocketImpl: FakeWebSocket, + location: LOCATION + }) +} + +function ticketOk(ticket = 'ticket-1') { + return vi.fn().mockResolvedValue({ data: { ticket, expires_at: 0 } }) +} + +describe('companySocket 建连', () => { + beforeEach(() => { + FakeWebSocket.instances = [] + resetOrganizationContext() + }) + + test('用一次性票据以 /ws?ticket= 建连，URL 不含 JWT 与公司 ID', async () => { + commitOrganizationId(1) + const fetchTicket = ticketOk('abc123') + const socket = makeSocket({ fetchTicket }) + + await socket.connect() + + expect(fetchTicket).toHaveBeenCalledTimes(1) + expect(FakeWebSocket.instances).toHaveLength(1) + const url = FakeWebSocket.instances[0].url + expect(url).toBe('ws://wms.example/ws?ticket=abc123') + expect(url).not.toContain('Bearer') + expect(url).not.toContain('organization') + }) + + test('没有当前公司（登出）时不建连也不取票', async () => { + const fetchTicket = vi.fn() + const socket = makeSocket({ fetchTicket }) + + await socket.connect() + + expect(fetchTicket).not.toHaveBeenCalled() + expect(FakeWebSocket.instances).toHaveLength(0) + }) + + test('取票失败时不建连且不抛出', async () => { + commitOrganizationId(1) + const fetchTicket = vi.fn().mockRejectedValue(new Error('网络异常')) + const socket = makeSocket({ fetchTicket }) + + await expect(socket.connect()).resolves.toBeUndefined() + expect(FakeWebSocket.instances).toHaveLength(0) + }) + + test('https 页面推导 wss 端点', () => { + expect(buildSocketUrl({ protocol: 'https:', host: 'wms.example' })).toBe('wss://wms.example/ws') + }) +}) + +describe('companySocket 切换与登出', () => { + beforeEach(() => { + FakeWebSocket.instances = [] + resetOrganizationContext() + }) + + test('切换公司关闭旧连接并为当前公司重新取票重连', async () => { + commitOrganizationId(1) + const fetchTicket = ticketOk() + fetchTicket.mockResolvedValueOnce({ data: { ticket: 'old-ticket' } }) + fetchTicket.mockResolvedValueOnce({ data: { ticket: 'new-ticket' } }) + const socket = makeSocket({ fetchTicket }) + + await socket.connect() + const first = FakeWebSocket.instances[0] + + // 切到公司 2 后重连：旧连接必须关闭，并重新取票。 + commitOrganizationId(2) + await socket.connect() + + expect(first.closed).toBe(true) + expect(fetchTicket).toHaveBeenCalledTimes(2) + expect(FakeWebSocket.instances[1].url).toBe('ws://wms.example/ws?ticket=new-ticket') + expect(socket.socket).toBe(FakeWebSocket.instances[1]) + }) + + test('登出（close）关闭连接且回调失效', async () => { + commitOrganizationId(1) + const onMessage = vi.fn() + const socket = makeSocket({ fetchTicket: ticketOk(), onMessage }) + + await socket.connect() + const ws = FakeWebSocket.instances[0] + + socket.close() + + expect(ws.closed).toBe(true) + expect(socket.socket).toBeNull() + // 关闭后旧连接迟到的消息不得再更新页面。 + ws.emitMessage({ organization_id: 1, type: 'stock_update' }) + expect(onMessage).not.toHaveBeenCalled() + }) + + test('取票期间公司已切换则放弃这次连接', async () => { + commitOrganizationId(1) + let resolveTicket + const fetchTicket = vi.fn().mockReturnValue( + new Promise(resolve => { + resolveTicket = resolve + }) + ) + const socket = makeSocket({ fetchTicket }) + + const pending = socket.connect() + // 取票尚未返回时公司已切换。 + commitOrganizationId(2) + resolveTicket({ data: { ticket: 'stale' } }) + await pending + + expect(FakeWebSocket.instances).toHaveLength(0) + expect(socket.socket).toBeNull() + }) +}) + +describe('companySocket 消息隔离', () => { + beforeEach(() => { + FakeWebSocket.instances = [] + resetOrganizationContext() + }) + + test('只处理当前公司的事件', async () => { + commitOrganizationId(1) + const onMessage = vi.fn() + const socket = makeSocket({ fetchTicket: ticketOk(), onMessage }) + await socket.connect() + const ws = FakeWebSocket.instances[0] + + ws.emitMessage({ organization_id: 2, type: 'stock_update' }) + expect(onMessage).not.toHaveBeenCalled() + + ws.emitMessage({ organization_id: 1, type: 'stock_update', action: 'inbound' }) + expect(onMessage).toHaveBeenCalledTimes(1) + expect(onMessage.mock.calls[0][0].action).toBe('inbound') + }) + + test('上下文版本失效后的消息不更新页面', async () => { + commitOrganizationId(1) + const onMessage = vi.fn() + const socket = makeSocket({ fetchTicket: ticketOk(), onMessage }) + await socket.connect() + const ws = FakeWebSocket.instances[0] + + // 公司已切换：即使事件公司标识仍是 1，捕获的版本已失效，必须丢弃。 + commitOrganizationId(1) + ws.emitMessage({ organization_id: 1, type: 'stock_update' }) + + expect(onMessage).not.toHaveBeenCalled() + }) + + test('旧连接的回调不影响新连接', async () => { + commitOrganizationId(1) + const onMessage = vi.fn() + const socket = makeSocket({ fetchTicket: ticketOk(), onMessage }) + await socket.connect() + const first = FakeWebSocket.instances[0] + + await socket.connect() + const second = FakeWebSocket.instances[1] + + first.emitMessage({ organization_id: 1, type: 'stock_update' }) + expect(onMessage).not.toHaveBeenCalled() + + second.emitMessage({ organization_id: 1, type: 'stock_update' }) + expect(onMessage).toHaveBeenCalledTimes(1) + }) + + test('断线恢复后重新取票（重新校验授权）', async () => { + commitOrganizationId(1) + const fetchTicket = ticketOk() + const socket = makeSocket({ fetchTicket }) + await socket.connect() + + // 模拟断线后再次连接：必须重新取票，由服务端再次校验授权。 + await socket.connect() + + expect(fetchTicket).toHaveBeenCalledTimes(2) + }) +}) diff --git a/internal/handler/api.go b/internal/handler/api.go index b1233a5e..714f76ef 100644 --- a/internal/handler/api.go +++ b/internal/handler/api.go @@ -228,6 +228,8 @@ func (h *Handler) RegisterRoutes(r *gin.Engine) { me.POST("/organization-context", h.SwitchOrganizationContext) me.GET("/me", h.GetCurrentUser) me.POST("/logout", h.Logout) + // WebSocket 建连前的票据签发：需要登录与已校验的公司头，返回短期一次性票据。 + me.POST("/websocket-ticket", middleware.OrganizationMiddleware(), h.CreateWebSocketTicket) } } } @@ -549,8 +551,8 @@ func (h *Handler) CreateInbound(c *gin.Context) { response.InternalError(c, err.Error()) return } - // Broadcast stock update to all connected clients（带上公司归属，供按公司投递过滤） - GetHub().BroadcastStockUpdate("inbound", orgID, order.ID, order) + // 事务已提交后再广播，且带公司归属，仅同公司订阅者会收到。 + GetHub().BroadcastStockUpdate(orgID, "inbound", order.ID, order) response.Created(c, "Inbound order created successfully", order) } @@ -621,7 +623,7 @@ func (h *Handler) UpdateInbound(c *gin.Context) { response.InternalError(c, err.Error()) return } - GetHub().BroadcastStockUpdate("inbound", orgID, order.ID, order) + GetHub().BroadcastStockUpdate(orgID, "inbound", order.ID, order) response.Success(c, order) } @@ -746,8 +748,8 @@ func (h *Handler) CreateOutbound(c *gin.Context) { response.InternalError(c, err.Error()) return } - // Broadcast stock update to all connected clients（带上公司归属，供按公司投递过滤） - GetHub().BroadcastStockUpdate("outbound", orgID, order.ID, order) + // 事务已提交后再广播，且带公司归属，仅同公司订阅者会收到。 + GetHub().BroadcastStockUpdate(orgID, "outbound", order.ID, order) response.Created(c, "Outbound order created successfully", order) } @@ -826,7 +828,7 @@ func (h *Handler) UpdateOutbound(c *gin.Context) { response.InternalError(c, err.Error()) return } - GetHub().BroadcastStockUpdate("outbound", orgID, order.ID, order) + GetHub().BroadcastStockUpdate(orgID, "outbound", order.ID, order) response.Success(c, order) } @@ -1233,6 +1235,24 @@ func (h *Handler) GetCurrentUser(c *gin.Context) { response.Success(c, user) } +// CreateWebSocketTicket 为当前登录用户与已校验的公司签发一张短期一次性票据， +// 供 /ws 握手使用。JWT 与公司头仍走正常请求头，绝不放进 WebSocket 的 URL。 +func (h *Handler) CreateWebSocketTicket(c *gin.Context) { + userID := middleware.GetCurrentUserID(c) + organizationID := middleware.GetCurrentOrganizationID(c) + + ticket, err := service.GetWebSocketTicketService().Issue(userID, organizationID) + if err != nil { + response.InternalError(c, "生成 WebSocket 票据失败") + return + } + + response.Success(c, gin.H{ + "ticket": ticket.Ticket, + "expires_at": ticket.ExpiresAt.Unix(), + }) +} + // User management handlers (Admin only) func (h *Handler) ListUsers(c *gin.Context) { diff --git a/internal/handler/order_organization_test.go b/internal/handler/order_organization_test.go index 79ebcac0..bf375ff7 100644 --- a/internal/handler/order_organization_test.go +++ b/internal/handler/order_organization_test.go @@ -231,12 +231,12 @@ func TestCreateOrderWithoutWarehouseReturnsConflict(t *testing.T) { // TestBroadcastStockUpdateCarriesOrganizationID：库存事件带上公司 ID，供 Task 5D 按公司投递。 func TestBroadcastStockUpdateCarriesOrganizationID(t *testing.T) { - hub := &Hub{broadcast: make(chan []byte, 1)} - hub.BroadcastStockUpdate("inbound", 7, 3, nil) + hub := newHub(nil) + hub.BroadcastStockUpdate(7, "inbound", 3, nil) select { case msg := <-hub.broadcast: - if !strings.Contains(string(msg), `"organization_id":7`) { - t.Fatalf("事件未携带公司 ID: %s", string(msg)) + if msg.organizationID != 7 || !strings.Contains(string(msg.payload), `"organization_id":7`) { + t.Fatalf("事件未携带公司 ID: org=%d payload=%s", msg.organizationID, string(msg.payload)) } default: t.Fatal("没有广播任何事件") diff --git a/internal/handler/websocket.go b/internal/handler/websocket.go index 5f8d5e3a..9d2d3ab2 100644 --- a/internal/handler/websocket.go +++ b/internal/handler/websocket.go @@ -8,6 +8,8 @@ import ( "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{ @@ -18,22 +20,6 @@ var upgrader = websocket.Upgrader{ }, } -// Client represents a WebSocket client -type Client struct { - hub *Hub - conn *websocket.Conn - send chan []byte -} - -// Hub maintains the set of active clients and broadcasts messages to clients -type Hub struct { - clients map[*Client]bool - broadcast chan []byte - register chan *Client - unregister chan *Client - mu sync.RWMutex -} - // StockUpdateEvent represents a stock update notification type StockUpdateEvent struct { Type string `json:"type"` // "stock_update" @@ -44,18 +30,77 @@ type StockUpdateEvent struct { 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 = &Hub{ - clients: make(map[*Client]bool), - broadcast: make(chan []byte, 256), - register: make(chan *Client), - unregister: make(chan *Client), - } + hubInstance = newHub(organizationAccessVerifier{}) go hubInstance.run() }) return hubInstance @@ -65,37 +110,71 @@ func (h *Hub) run() { for { select { case client := <-h.register: - h.mu.Lock() h.clients[client] = true - h.mu.Unlock() - log.Printf("WebSocket client connected. Total clients: %d", len(h.clients)) + log.Printf("WebSocket 客户端已连接（公司 %d）。当前连接数：%d", client.organizationID, len(h.clients)) case client := <-h.unregister: - h.mu.Lock() if _, ok := h.clients[client]; ok { delete(h.clients, client) - close(client.send) + client.closeSend() + log.Printf("WebSocket 客户端已断开。当前连接数：%d", len(h.clients)) } - h.mu.Unlock() - log.Printf("WebSocket client disconnected. Total clients: %d", len(h.clients)) case message := <-h.broadcast: - h.mu.RLock() + // 只选择同公司订阅者；每个收件人单独起协程，事件循环本身不等待投递完成。 for client := range h.clients { - select { - case client.send <- message: - default: - close(client.send) - delete(h.clients, client) + if client.organizationID != message.organizationID { + continue } + go h.deliver(client, message.payload) } - h.mu.RUnlock() } } } -// BroadcastStockUpdate sends a stock update event to all connected clients -func (h *Hub) BroadcastStockUpdate(action string, organizationID, orderID int64, data interface{}) { +// 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, @@ -105,10 +184,10 @@ func (h *Hub) BroadcastStockUpdate(action string, organizationID, orderID int64, } msg, err := json.Marshal(event) if err != nil { - log.Printf("Error marshaling stock update event: %v", err) + log.Printf("序列化库存更新事件失败: %v", err) return } - h.broadcast <- msg + h.broadcast <- broadcastMessage{organizationID: organizationID, payload: msg} } func (c *Client) readPump() { @@ -144,17 +223,34 @@ func (c *Client) writePump() { } // 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 upgrade failed: %v", err) + log.Printf("WebSocket 升级失败: %v", err) return } client := &Client{ - hub: hub, - conn: conn, - send: make(chan []byte, 256), + hub: hub, + conn: conn, + send: make(chan []byte, 256), + userID: ticket.UserID, + organizationID: ticket.OrganizationID, + ticketID: ticket.Ticket, } hub.register <- client go client.writePump() diff --git a/internal/handler/websocket_test.go b/internal/handler/websocket_test.go new file mode 100644 index 00000000..60b3fca1 --- /dev/null +++ b/internal/handler/websocket_test.go @@ -0,0 +1,390 @@ +package handler + +import ( + "encoding/json" + "errors" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + + "github.com/DATA-DOG/go-sqlmock" + "github.com/gin-gonic/gin" + "github.com/wmms/wmms/internal/service" +) + +// withTicketService 替换进程内共享票据服务，返回可控 TTL 的实例供测试使用。 +func withTicketService(t *testing.T, ttl time.Duration) *service.WebSocketTicketService { + t.Helper() + svc := service.NewWebSocketTicketService(ttl) + restore := service.SetDefaultWebSocketTicketService(svc) + t.Cleanup(restore) + return svc +} + +// TestWebSocketTicketEndpointRequiresAuthAndOrganization：票据签发接口需要登录、公司头， +// 并通过公司授权校验；缺一不可。 +func TestWebSocketTicketEndpointRequiresAuthAndOrganization(t *testing.T) { + gin.SetMode(gin.TestMode) + + t.Run("未登录", func(t *testing.T) { + contextMock(t) + router := gin.New() + NewHandler().RegisterRoutes(router) + req := httptest.NewRequest("POST", "/api/auth/websocket-ticket", nil) + w := httptest.NewRecorder() + router.ServeHTTP(w, req) + if w.Code != 401 { + t.Fatalf("status=%d body=%s, want 401", w.Code, w.Body.String()) + } + }) + + t.Run("缺少公司头", func(t *testing.T) { + mock := contextMock(t) + expectContextPrincipal(mock, "warehouse_admin") + router := gin.New() + NewHandler().RegisterRoutes(router) + req := httptest.NewRequest("POST", "/api/auth/websocket-ticket", nil) + req.Header.Set("Authorization", "Bearer "+contextToken(t)) + w := httptest.NewRecorder() + router.ServeHTTP(w, req) + if w.Code != 400 || !strings.Contains(w.Body.String(), "X-Organization-ID") { + t.Fatalf("status=%d body=%s, want 400 含公司头提示", w.Code, w.Body.String()) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Fatal(err) + } + }) + + t.Run("公司不可访问", func(t *testing.T) { + mock := contextMock(t) + expectContextPrincipal(mock, "warehouse_admin") + mock.ExpectQuery("SELECT EXISTS").WithArgs(int64(2), int64(7), int64(7)). + WillReturnRows(sqlmock.NewRows([]string{"allowed"}).AddRow(false)) + router := gin.New() + NewHandler().RegisterRoutes(router) + req := httptest.NewRequest("POST", "/api/auth/websocket-ticket", nil) + req.Header.Set("Authorization", "Bearer "+contextToken(t)) + req.Header.Set("X-Organization-ID", "2") + w := httptest.NewRecorder() + router.ServeHTTP(w, req) + if w.Code != 403 { + t.Fatalf("status=%d body=%s, want 403", w.Code, w.Body.String()) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Fatal(err) + } + }) + + t.Run("成功签发短期一次性票据", func(t *testing.T) { + mock := contextMock(t) + expectContextPrincipal(mock, "warehouse_admin") + mock.ExpectQuery("SELECT EXISTS").WithArgs(int64(2), int64(7), int64(7)). + WillReturnRows(sqlmock.NewRows([]string{"allowed"}).AddRow(true)) + router := gin.New() + NewHandler().RegisterRoutes(router) + req := httptest.NewRequest("POST", "/api/auth/websocket-ticket", nil) + req.Header.Set("Authorization", "Bearer "+contextToken(t)) + req.Header.Set("X-Organization-ID", "2") + w := httptest.NewRecorder() + router.ServeHTTP(w, req) + if w.Code != 200 { + t.Fatalf("status=%d body=%s, want 200", w.Code, w.Body.String()) + } + var envelope struct { + Data struct { + Ticket string `json:"ticket"` + ExpiresAt int64 `json:"expires_at"` + } `json:"data"` + } + if err := json.Unmarshal(w.Body.Bytes(), &envelope); err != nil { + t.Fatal(err) + } + if len(envelope.Data.Ticket) < 32 { + t.Fatalf("票据应为足够长的随机串，实际 %q", envelope.Data.Ticket) + } + if envelope.Data.ExpiresAt <= time.Now().Unix() { + t.Fatalf("过期时间应在未来，实际 %d", envelope.Data.ExpiresAt) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Fatal(err) + } + }) +} + +// wsHandshake 用一次性票据发起 /ws 握手请求，返回响应记录器。 +// 未通过鉴权时在升级前即返回；本测试只覆盖“不能连接”的各种拒绝路径。 +func wsHandshake(ticket string) *httptest.ResponseRecorder { + gin.SetMode(gin.TestMode) + router := gin.New() + router.GET("/ws", HandleWebSocket) + path := "/ws" + if ticket != "" { + path += "?ticket=" + ticket + } + req := httptest.NewRequest("GET", path, nil) + w := httptest.NewRecorder() + router.ServeHTTP(w, req) + return w +} + +// TestWebSocketHandshakeRejectsInvalidTickets：票据缺失、未知、过期、已消费（重放） +// 都不能建立连接，一律 401。 +func TestWebSocketHandshakeRejectsInvalidTickets(t *testing.T) { + t.Run("缺少票据", func(t *testing.T) { + if w := wsHandshake(""); w.Code != 401 { + t.Fatalf("status=%d body=%s, want 401", w.Code, w.Body.String()) + } + }) + + t.Run("未知票据", func(t *testing.T) { + withTicketService(t, 30*time.Second) + if w := wsHandshake("0123456789abcdef"); w.Code != 401 { + t.Fatalf("status=%d body=%s, want 401", w.Code, w.Body.String()) + } + }) + + t.Run("过期票据", func(t *testing.T) { + svc := withTicketService(t, -time.Second) + ticket, err := svc.Issue(7, 2) + if err != nil { + t.Fatal(err) + } + if w := wsHandshake(ticket.Ticket); w.Code != 401 { + t.Fatalf("status=%d body=%s, want 401", w.Code, w.Body.String()) + } + }) + + t.Run("已消费票据重放", func(t *testing.T) { + svc := withTicketService(t, 30*time.Second) + ticket, err := svc.Issue(7, 2) + if err != nil { + t.Fatal(err) + } + // 模拟第一次握手已消费该票据。 + if _, err := svc.Consume(ticket.Ticket); err != nil { + t.Fatal(err) + } + if w := wsHandshake(ticket.Ticket); w.Code != 401 { + t.Fatalf("status=%d body=%s, want 401", w.Code, w.Body.String()) + } + }) +} + +// TestWebSocketHandshakeRevalidatesAuthorization：票据有效但公司授权已在数据库侧撤销时， +// 握手前复核失败，返回 403 且不建连。 +func TestWebSocketHandshakeRevalidatesAuthorization(t *testing.T) { + svc := withTicketService(t, 30*time.Second) + ticket, err := svc.Issue(7, 2) + if err != nil { + t.Fatal(err) + } + mock := contextMock(t) + expectContextPrincipal(mock, "warehouse_admin") + mock.ExpectQuery("SELECT EXISTS").WithArgs(int64(2), int64(7), int64(7)). + WillReturnRows(sqlmock.NewRows([]string{"allowed"}).AddRow(false)) + + w := wsHandshake(ticket.Ticket) + if w.Code != 403 { + t.Fatalf("status=%d body=%s, want 403", w.Code, w.Body.String()) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Fatal(err) + } +} + +// TestOrganizationAccessVerifierReloadsFromDatabase：复核器以数据库当前状态判定， +// 授权撤销或用户被删除都会失败，而不是相信 JWT 快照。 +func TestOrganizationAccessVerifierReloadsFromDatabase(t *testing.T) { + for _, tc := range []struct { + name string + canAccess bool + principalRows bool + wantErr bool + }{ + {"仍有授权", true, true, false}, + {"授权已撤销", false, true, true}, + {"用户已删除", false, false, true}, + } { + t.Run(tc.name, func(t *testing.T) { + mock := contextMock(t) + if tc.principalRows { + expectContextPrincipal(mock, "warehouse_admin") + mock.ExpectQuery("SELECT EXISTS").WithArgs(int64(2), int64(7), int64(7)). + WillReturnRows(sqlmock.NewRows([]string{"allowed"}).AddRow(tc.canAccess)) + } else { + mock.ExpectQuery("SELECT u.id, r.code FROM users u").WithArgs(int64(7)). + WillReturnRows(sqlmock.NewRows([]string{"id", "code"})) + } + + err := organizationAccessVerifier{}.Verify(7, 2) + if tc.wantErr != (err != nil) { + t.Fatalf("err=%v wantErr=%v", err, tc.wantErr) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Fatal(err) + } + }) + } +} + +// fakeVerifier 是 Hub 测试用的可注入复核器。 +type fakeVerifier struct { + mu sync.Mutex + revoked map[int64]bool + err error +} + +func newFakeVerifier() *fakeVerifier { + return &fakeVerifier{revoked: make(map[int64]bool)} +} + +func (f *fakeVerifier) Verify(userID, _ int64) error { + f.mu.Lock() + defer f.mu.Unlock() + if f.err != nil { + return f.err + } + if f.revoked[userID] { + return errors.New("授权已撤销") + } + return nil +} + +func (f *fakeVerifier) revoke(userID int64) { + f.mu.Lock() + f.revoked[userID] = true + f.mu.Unlock() +} + +func (f *fakeVerifier) fail(err error) { + f.mu.Lock() + f.err = err + f.mu.Unlock() +} + +func testClient(hub *Hub, userID, organizationID int64) *Client { + return &Client{ + hub: hub, + send: make(chan []byte, 8), + userID: userID, + organizationID: organizationID, + } +} + +func expectMessage(t *testing.T, c *Client, timeout time.Duration) []byte { + t.Helper() + select { + case msg, ok := <-c.send: + if !ok { + t.Fatal("连接在收到消息前已被关闭") + } + return msg + case <-time.After(timeout): + t.Fatal("超时未收到消息") + return nil + } +} + +func expectClosed(t *testing.T, c *Client, timeout time.Duration) { + t.Helper() + select { + case _, ok := <-c.send: + if !ok { + return + } + t.Fatal("连接上仍有未消费消息，未按预期关闭") + case <-time.After(timeout): + t.Fatal("连接未被关闭") + } +} + +func expectSilence(t *testing.T, c *Client, d time.Duration) { + t.Helper() + select { + case msg, ok := <-c.send: + if ok { + t.Fatalf("不应收到消息，却收到 %s", string(msg)) + } + t.Fatal("不应关闭连接") + case <-time.After(d): + } +} + +// TestHubDeliversOnlyToSameOrganization：A 公司的订阅者收不到 B 公司的入出库事件。 +func TestHubDeliversOnlyToSameOrganization(t *testing.T) { + hub := newHub(newFakeVerifier()) + go hub.run() + + a1 := testClient(hub, 11, 1) + a2 := testClient(hub, 12, 1) + b1 := testClient(hub, 21, 2) + for _, c := range []*Client{a1, a2, b1} { + hub.register <- c + } + + hub.BroadcastStockUpdate(1, "inbound", 100, nil) + + for _, c := range []*Client{a1, a2} { + msg := expectMessage(t, c, time.Second) + if !strings.Contains(string(msg), `"organization_id":1`) { + t.Fatalf("同公司事件应携带公司 1: %s", string(msg)) + } + } + // B 公司订阅者不得收到任何事件。 + expectSilence(t, b1, 150*time.Millisecond) +} + +// TestHubDropsRevokedSubscriberOnNextEvent：授权撤销后，订阅者收不到下一条事件， +// 连接被关闭；同公司其他订阅者不受影响。 +func TestHubDropsRevokedSubscriberOnNextEvent(t *testing.T) { + verifier := newFakeVerifier() + hub := newHub(verifier) + go hub.run() + + a1 := testClient(hub, 11, 1) + a2 := testClient(hub, 12, 1) + for _, c := range []*Client{a1, a2} { + hub.register <- c + } + + // 首个事件两人都收到。 + hub.BroadcastStockUpdate(1, "inbound", 1, nil) + expectMessage(t, a1, time.Second) + expectMessage(t, a2, time.Second) + + // 撤销 a1 的授权：下一条事件不再投递给它，并关闭其连接。 + verifier.revoke(11) + hub.BroadcastStockUpdate(1, "inbound", 2, nil) + expectClosed(t, a1, time.Second) + expectMessage(t, a2, time.Second) +} + +// TestHubClosesSubscribersWhenVerificationErrors：投递前复核遇到数据库异常时， +// 不投递并关闭连接。 +func TestHubClosesSubscribersWhenVerificationErrors(t *testing.T) { + verifier := newFakeVerifier() + hub := newHub(verifier) + go hub.run() + + c := testClient(hub, 11, 1) + hub.register <- c + + verifier.fail(errors.New("数据库异常")) + hub.BroadcastStockUpdate(1, "outbound", 1, nil) + + expectClosed(t, c, time.Second) +} + +// TestHubUnregisterClosesConnection：注销由事件循环串行处理，连接被关闭且不再收到事件。 +func TestHubUnregisterClosesConnection(t *testing.T) { + hub := newHub(nil) + go hub.run() + + c := testClient(hub, 11, 1) + hub.register <- c + hub.unregister <- c + + expectClosed(t, c, time.Second) +} diff --git a/internal/service/websocket_ticket.go b/internal/service/websocket_ticket.go new file mode 100644 index 00000000..d1645b99 --- /dev/null +++ b/internal/service/websocket_ticket.go @@ -0,0 +1,122 @@ +package service + +import ( + "crypto/rand" + "encoding/hex" + "errors" + "sync" + "time" +) + +// WebSocket 票据默认有效期：短期、一次性，把 /ws 握手从长期 JWT 中解耦。 +// 浏览器不能把 JWT 放 URL，也不能仅凭 URL 里的公司 ID 订阅，所以先用受保护的接口 +// 换取一张短暂票据，再凭它建连。 +const WebSocketTicketTTL = 30 * time.Second + +var ( + // ErrWebSocketTicketInvalid 表示票据不存在、格式非法或已被消费（重放）。 + ErrWebSocketTicketInvalid = errors.New("websocket ticket invalid or already used") + // ErrWebSocketTicketExpired 表示票据已超过有效期。 + ErrWebSocketTicketExpired = errors.New("websocket ticket expired") +) + +// WebSocketTicket 绑定一次握手所需的身份与公司：票据 ID 本身即一次性随机标识， +// 连同 userID、organizationID 与签发时间一起记录。本系统当前没有会话 ID， +// 因此用“随机票据 + 签发时间 + 数据库实时复核”共同承担“会话标识”的职责。 +type WebSocketTicket struct { + Ticket string + UserID int64 + OrganizationID int64 + IssuedAt time.Time + ExpiresAt time.Time +} + +// WebSocketTicketService 维护未消费的票据。票据只驻留内存，进程重启即全部失效。 +type WebSocketTicketService struct { + ttl time.Duration + now func() time.Time + mu sync.Mutex + tickets map[string]*WebSocketTicket +} + +// NewWebSocketTicketService 构造票据服务。ttl 可选，测试用它构造“已过期”票据， +// 不传则使用默认 30 秒。 +func NewWebSocketTicketService(ttl ...time.Duration) *WebSocketTicketService { + effective := WebSocketTicketTTL + if len(ttl) > 0 { + effective = ttl[0] + } + return &WebSocketTicketService{ + ttl: effective, + now: time.Now, + tickets: make(map[string]*WebSocketTicket), + } +} + +// Issue 为指定用户与公司签发一张一次性票据。 +func (s *WebSocketTicketService) Issue(userID, organizationID int64) (*WebSocketTicket, error) { + if userID <= 0 || organizationID <= 0 { + return nil, ErrInvalidOrganizationID + } + raw := make([]byte, 32) + if _, err := rand.Read(raw); err != nil { + return nil, err + } + now := s.now() + ticket := &WebSocketTicket{ + Ticket: hex.EncodeToString(raw), + UserID: userID, + OrganizationID: organizationID, + IssuedAt: now, + ExpiresAt: now.Add(s.ttl), + } + s.mu.Lock() + defer s.mu.Unlock() + // 顺手清理已过期票据，避免长期运行后 map 无限增长。 + s.purgeExpiredLocked(now) + s.tickets[ticket.Ticket] = ticket + return ticket, nil +} + +// Consume 校验并一次性消费票据。无论是否过期都先删除：过期票据、已消费票据 +// 和重放请求都无法再次连接。返回的票据已通过有效期检查。 +func (s *WebSocketTicketService) Consume(value string) (*WebSocketTicket, error) { + if value == "" { + return nil, ErrWebSocketTicketInvalid + } + s.mu.Lock() + defer s.mu.Unlock() + ticket, ok := s.tickets[value] + if !ok { + return nil, ErrWebSocketTicketInvalid + } + delete(s.tickets, value) + if s.now().After(ticket.ExpiresAt) { + return nil, ErrWebSocketTicketExpired + } + return ticket, nil +} + +func (s *WebSocketTicketService) purgeExpiredLocked(now time.Time) { + for key, ticket := range s.tickets { + if now.After(ticket.ExpiresAt) { + delete(s.tickets, key) + } + } +} + +// defaultTicketService 是进程内共享的票据服务，签发接口与握手共用同一实例。 +var defaultTicketService = NewWebSocketTicketService() + +// GetWebSocketTicketService 返回共享票据服务。 +func GetWebSocketTicketService() *WebSocketTicketService { + return defaultTicketService +} + +// SetDefaultWebSocketTicketService 仅供测试替换共享票据服务，返回恢复函数。 +// 生产代码请使用 GetWebSocketTicketService。 +func SetDefaultWebSocketTicketService(s *WebSocketTicketService) func() { + previous := defaultTicketService + defaultTicketService = s + return func() { defaultTicketService = previous } +} diff --git a/internal/service/websocket_ticket_test.go b/internal/service/websocket_ticket_test.go new file mode 100644 index 00000000..ed68b19d --- /dev/null +++ b/internal/service/websocket_ticket_test.go @@ -0,0 +1,125 @@ +package service + +import ( + "errors" + "testing" + "time" +) + +// TestWebSocketTicketIssueBindsIdentityAndExpiry：签发的票据绑定用户、公司、签发时间， +// 并在有效期内可被一次性消费。 +func TestWebSocketTicketIssueBindsIdentityAndExpiry(t *testing.T) { + base := time.Now() + svc := NewWebSocketTicketService(30 * time.Second) + svc.now = func() time.Time { return base } + + ticket, err := svc.Issue(7, 2) + if err != nil { + t.Fatalf("签发失败: %v", err) + } + if ticket.UserID != 7 || ticket.OrganizationID != 2 { + t.Fatalf("票据未绑定用户/公司: %+v", ticket) + } + if !ticket.IssuedAt.Equal(base) { + t.Fatalf("签发时间错误: %v", ticket.IssuedAt) + } + if !ticket.ExpiresAt.Equal(base.Add(30 * time.Second)) { + t.Fatalf("过期时间应为签发时间 + 30s: %v", ticket.ExpiresAt) + } + + // 有效期最后一刻仍可消费。 + svc.now = func() time.Time { return base.Add(29 * time.Second) } + consumed, err := svc.Consume(ticket.Ticket) + if err != nil { + t.Fatalf("有效期内消费失败: %v", err) + } + if consumed.UserID != 7 || consumed.OrganizationID != 2 { + t.Fatalf("消费后身份错误: %+v", consumed) + } +} + +// TestWebSocketTicketIsOneTimeOnly：票据一次性，重放或再次消费都被拒绝。 +func TestWebSocketTicketIsOneTimeOnly(t *testing.T) { + svc := NewWebSocketTicketService() + ticket, err := svc.Issue(7, 2) + if err != nil { + t.Fatal(err) + } + if _, err := svc.Consume(ticket.Ticket); err != nil { + t.Fatalf("首次消费应成功: %v", err) + } + if _, err := svc.Consume(ticket.Ticket); !errors.Is(err, ErrWebSocketTicketInvalid) { + t.Fatalf("重放应被拒绝，实际 err=%v", err) + } +} + +func TestWebSocketTicketRejectsUnknownAndEmpty(t *testing.T) { + for _, tc := range []struct { + name string + value string + }{ + {"empty", ""}, + {"unknown", "0123456789abcdef"}, + } { + t.Run(tc.name, func(t *testing.T) { + svc := NewWebSocketTicketService() + if _, err := svc.Consume(tc.value); !errors.Is(err, ErrWebSocketTicketInvalid) { + t.Fatalf("应返回无效票据错误，实际 %v", err) + } + }) + } +} + +// TestWebSocketTicketExpires：超过 30 秒有效期的票据不能再连接。 +func TestWebSocketTicketExpires(t *testing.T) { + base := time.Now() + svc := NewWebSocketTicketService(30 * time.Second) + svc.now = func() time.Time { return base } + ticket, err := svc.Issue(7, 2) + if err != nil { + t.Fatal(err) + } + + svc.now = func() time.Time { return base.Add(31 * time.Second) } + if _, err := svc.Consume(ticket.Ticket); !errors.Is(err, ErrWebSocketTicketExpired) { + t.Fatalf("过期票据应被拒绝，实际 %v", err) + } + // 过期票据已从集合中移除，再次消费按“未知票据”处理。 + if _, err := svc.Consume(ticket.Ticket); !errors.Is(err, ErrWebSocketTicketInvalid) { + t.Fatalf("过期后再次消费应为无效票据，实际 %v", err) + } +} + +func TestWebSocketTicketIssueRequiresValidIdentity(t *testing.T) { + svc := NewWebSocketTicketService() + for _, tc := range []struct { + name string + userID, orgID int64 + }{ + {"zero user", 0, 2}, + {"zero org", 7, 0}, + {"negative", -1, -2}, + } { + t.Run(tc.name, func(t *testing.T) { + if _, err := svc.Issue(tc.userID, tc.orgID); !errors.Is(err, ErrInvalidOrganizationID) { + t.Fatalf("应返回 ErrInvalidOrganizationID，实际 %v", err) + } + }) + } +} + +// TestWebSocketTicketValuesAreRandom：每张票据的随机标识互不相同。 +func TestWebSocketTicketValuesAreRandom(t *testing.T) { + svc := NewWebSocketTicketService() + first, err := svc.Issue(7, 2) + if err != nil { + t.Fatal(err) + } + second, err := svc.Issue(7, 2) + if err != nil { + t.Fatal(err) + } + if first.Ticket == second.Ticket || len(first.Ticket) < 32 { + t.Fatalf("票据应为足够长的随机串: %q / %q", first.Ticket, second.Ticket) + } +}