140 lines
4.3 KiB
Go
140 lines
4.3 KiB
Go
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 本身即一次性随机标识，
|
||
// 连同会话指纹和令牌到期时间记录；不保存原始 JWT。
|
||
type WebSocketTicket struct {
|
||
Ticket string
|
||
UserID int64
|
||
OrganizationID int64
|
||
IssuedAt time.Time
|
||
ExpiresAt time.Time
|
||
SessionID string
|
||
SessionExpiresAt int64
|
||
}
|
||
|
||
// 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, sessionToken ...string) (*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),
|
||
}
|
||
if len(sessionToken) > 0 {
|
||
principal, err := NewAuthService().ValidateToken(sessionToken[0])
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
if principal.UserID != userID {
|
||
return nil, ErrWebSocketTicketInvalid
|
||
}
|
||
ticket.SessionID = TokenSessionID(sessionToken[0])
|
||
ticket.SessionExpiresAt = principal.Exp
|
||
}
|
||
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
|
||
}
|
||
if ticket.SessionID != "" {
|
||
if err := ValidateSession(ticket.SessionID, ticket.SessionExpiresAt); err != nil {
|
||
return nil, err
|
||
}
|
||
}
|
||
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 }
|
||
}
|