feat: 按公司隔离 WebSocket 订阅与事件投递
- 新增 POST /api/auth/websocket-ticket：需 JWT 与公司头，签发 30 秒有效的一次性随机票据 - /ws 握手消费票据并以数据库当前状态复核身份与公司授权，JWT 不进 URL - Hub 注册/注销/投递选人统一由事件循环串行管理，消除 RLock 内 delete 的并发缺陷 - 事件携带 organization_id，仅投递同公司订阅者，投递前复核失败即关闭连接 - 统一广播签名为 (organizationID, action, orderID, data) - 新增前端 companySocket 模块：按当前公司取票建连、切换/登出关旧连重连、消息校验组织与上下文版本
This commit is contained in:
文睿 张
parent
a1bc0539dc
commit
5021b039c9
|
|
@ -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，便于试用脚本探测空闲端口。
|
||||
|
|
|
|||
127
frontend/src/utils/companySocket.js
Normal file
127
frontend/src/utils/companySocket.js
Normal file
|
|
@ -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
|
||||
}
|
||||
}
|
||||
}
|
||||
212
frontend/src/utils/companySocket.test.js
Normal file
212
frontend/src/utils/companySocket.test.js
Normal file
|
|
@ -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)
|
||||
})
|
||||
})
|
||||
|
|
@ -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) {
|
||||
|
|
|
|||
|
|
@ -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("没有广播任何事件")
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
390
internal/handler/websocket_test.go
Normal file
390
internal/handler/websocket_test.go
Normal file
|
|
@ -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)
|
||||
}
|
||||
122
internal/service/websocket_ticket.go
Normal file
122
internal/service/websocket_ticket.go
Normal file
|
|
@ -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 }
|
||||
}
|
||||
125
internal/service/websocket_ticket_test.go
Normal file
125
internal/service/websocket_ticket_test.go
Normal file
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
Loading…
Reference in New Issue
Block a user