412 lines
12 KiB
Go
412 lines
12 KiB
Go
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, contextToken(t))
|
||
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)
|
||
}
|
||
})
|
||
}
|
||
}
|
||
|
||
// TestOrganizationAccessVerifierDoesNotDropBusinessRoleRevocation：显式钉住复核边界。
|
||
// 复核的隔离边界是“用户仍存在 + 公司授权仍可访问”；仅撤销某个非管理员业务角色、
|
||
// 而公司授权仍在时，Verify 仍通过，订阅者不会被断开。角色级断流需要完整的会话/权限
|
||
// 体系，本期未实现，故用本用例固定该语义，避免日后被误认为已实现角色级撤销。
|
||
func TestOrganizationAccessVerifierDoesNotDropBusinessRoleRevocation(t *testing.T) {
|
||
mock := contextMock(t)
|
||
// 用户仍存在，有效角色只剩一个普通业务角色（原管理员角色已被撤销）。
|
||
mock.ExpectQuery("SELECT u.id, r.code FROM users u").WithArgs(int64(7)).
|
||
WillReturnRows(sqlmock.NewRows([]string{"id", "code"}).AddRow(7, "purchaser"))
|
||
// 公司授权仍然保留。
|
||
mock.ExpectQuery("SELECT EXISTS").WithArgs(int64(2), int64(7), int64(7)).
|
||
WillReturnRows(sqlmock.NewRows([]string{"allowed"}).AddRow(true))
|
||
|
||
if err := (organizationAccessVerifier{}).Verify(7, 2); err != nil {
|
||
t.Fatalf("仅业务角色变化、公司授权仍在时不应断流，实际 err=%v", err)
|
||
}
|
||
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)
|
||
}
|