Introduce a new WebSocket messaging architecture that supports both standalone and cluster modes. This allows for horizontal scaling of WebSocket servers by using Redis Pub/Sub to synchronize messages across multiple instances. Key changes: - Added `ws.MessagePublisher` interface to abstract message distribution. - Implemented `ws.Bus` to handle cluster-mode messaging via Redis. - Added `ws.OnlineTracker` to manage user online status across the cluster. - Refactored multiple services (Chat, Group, Push, Call, etc.) to use the new `MessagePublisher` instead of a concrete `ws.Hub`. - Added WebSocket configuration options (mode, instance ID, channel, TTL, heartbeat) to `config.yaml` and `config.go`. - Updated dependency injection with Wire to support the new publisher and Redis client. - Improved logging by replacing standard `log` with `zap` in several service components.
299 lines
7.5 KiB
Go
299 lines
7.5 KiB
Go
package service
|
||
|
||
import (
|
||
"context"
|
||
"fmt"
|
||
"time"
|
||
|
||
"github.com/google/uuid"
|
||
|
||
"with_you/internal/model"
|
||
"with_you/internal/pkg/redis"
|
||
"with_you/internal/pkg/ws"
|
||
)
|
||
|
||
// Lua 脚本:原子性地检查状态并更新为 scanned
|
||
const scanScript = `
|
||
local key = KEYS[1]
|
||
local status = redis.call('HGET', key, 'status')
|
||
if status == 'pending' then
|
||
redis.call('HMSET', key, 'status', 'scanned', 'user_id', ARGV[1])
|
||
return 1
|
||
elseif status == 'scanned' then
|
||
return 2
|
||
else
|
||
return 0
|
||
end
|
||
`
|
||
|
||
const (
|
||
qrcodeSessionPrefix = "qrcode:session:"
|
||
qrcodeSessionTTL = 5 * time.Minute
|
||
)
|
||
|
||
// QRCodeStatus 二维码状态
|
||
type QRCodeStatus string
|
||
|
||
const (
|
||
QRCodeStatusPending QRCodeStatus = "pending"
|
||
QRCodeStatusScanned QRCodeStatus = "scanned"
|
||
QRCodeStatusConfirmed QRCodeStatus = "confirmed"
|
||
QRCodeStatusCancelled QRCodeStatus = "cancelled"
|
||
QRCodeStatusExpired QRCodeStatus = "expired"
|
||
)
|
||
|
||
// QRCodeSession 二维码会话
|
||
type QRCodeSession struct {
|
||
SessionID string `json:"session_id"`
|
||
Status QRCodeStatus `json:"status"`
|
||
UserID string `json:"user_id"`
|
||
CreatedAt int64 `json:"created_at"`
|
||
ExpiresAt int64 `json:"expires_at"`
|
||
}
|
||
|
||
// QRCodeLoginService 二维码登录服务
|
||
type QRCodeLoginService struct {
|
||
redis *redis.Client
|
||
wsHub ws.MessagePublisher
|
||
jwtService *JWTService
|
||
userService UserService
|
||
activityService UserActivityService
|
||
logService *LogService
|
||
}
|
||
|
||
// NewQRCodeLoginService 创建二维码登录服务
|
||
func NewQRCodeLoginService(redis *redis.Client, publisher ws.MessagePublisher, jwtService *JWTService, userService UserService, activityService UserActivityService, logService *LogService) *QRCodeLoginService {
|
||
return &QRCodeLoginService{
|
||
redis: redis,
|
||
wsHub: publisher,
|
||
jwtService: jwtService,
|
||
userService: userService,
|
||
activityService: activityService,
|
||
logService: logService,
|
||
}
|
||
}
|
||
|
||
// CreateSession 创建二维码会话
|
||
func (s *QRCodeLoginService) CreateSession(ctx context.Context) (*QRCodeSession, error) {
|
||
sessionID := uuid.New().String()
|
||
now := time.Now()
|
||
expiresAt := now.Add(qrcodeSessionTTL)
|
||
|
||
session := &QRCodeSession{
|
||
SessionID: sessionID,
|
||
Status: QRCodeStatusPending,
|
||
CreatedAt: now.Unix(),
|
||
ExpiresAt: expiresAt.Unix(),
|
||
}
|
||
|
||
key := qrcodeSessionPrefix + sessionID
|
||
data := map[string]any{
|
||
"status": string(session.Status),
|
||
"user_id": "",
|
||
"created_at": session.CreatedAt,
|
||
"expires_at": session.ExpiresAt,
|
||
}
|
||
|
||
if err := s.redis.HMSet(ctx, key, data); err != nil {
|
||
return nil, fmt.Errorf("failed to create session: %w", err)
|
||
}
|
||
|
||
if _, err := s.redis.Expire(ctx, key, qrcodeSessionTTL); err != nil {
|
||
return nil, fmt.Errorf("failed to set session ttl: %w", err)
|
||
}
|
||
|
||
return session, nil
|
||
}
|
||
|
||
// GetSession 获取会话
|
||
func (s *QRCodeLoginService) GetSession(ctx context.Context, sessionID string) (*QRCodeSession, error) {
|
||
key := qrcodeSessionPrefix + sessionID
|
||
data, err := s.redis.HGetAll(ctx, key)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("failed to get session: %w", err)
|
||
}
|
||
|
||
if len(data) == 0 {
|
||
return nil, fmt.Errorf("session not found or expired")
|
||
}
|
||
|
||
session := &QRCodeSession{
|
||
SessionID: sessionID,
|
||
Status: QRCodeStatus(data["status"]),
|
||
UserID: data["user_id"],
|
||
}
|
||
|
||
// 解析时间戳
|
||
if v := data["created_at"]; v != "" {
|
||
fmt.Sscanf(v, "%d", &session.CreatedAt)
|
||
}
|
||
if v := data["expires_at"]; v != "" {
|
||
fmt.Sscanf(v, "%d", &session.ExpiresAt)
|
||
}
|
||
|
||
// 检查是否过期
|
||
if time.Now().Unix() > session.ExpiresAt {
|
||
session.Status = QRCodeStatusExpired
|
||
}
|
||
|
||
return session, nil
|
||
}
|
||
|
||
// Scan 扫描二维码
|
||
func (s *QRCodeLoginService) Scan(ctx context.Context, sessionID, userID string) error {
|
||
key := qrcodeSessionPrefix + sessionID
|
||
|
||
// 使用 Lua 脚本实现原子性检查和更新,避免竞态条件
|
||
result, err := s.redis.GetClient().Eval(ctx, scanScript, []string{key}, userID).Int()
|
||
if err != nil {
|
||
return fmt.Errorf("failed to scan: %w", err)
|
||
}
|
||
|
||
switch result {
|
||
case 0:
|
||
return fmt.Errorf("session not found or expired")
|
||
case 2:
|
||
return fmt.Errorf("qrcode already scanned")
|
||
}
|
||
|
||
// 获取用户信息用于推送
|
||
user, err := s.userService.GetUserByID(ctx, userID)
|
||
if err != nil {
|
||
return fmt.Errorf("failed to get user info: %w", err)
|
||
}
|
||
|
||
// 推送扫码事件
|
||
payload := map[string]any{
|
||
"user": map[string]any{
|
||
"id": user.ID,
|
||
"nickname": user.Nickname,
|
||
"avatar": user.Avatar,
|
||
},
|
||
}
|
||
s.wsHub.PublishToUser(sessionID, "scanned", payload)
|
||
|
||
return nil
|
||
}
|
||
|
||
// Confirm 确认登录
|
||
func (s *QRCodeLoginService) Confirm(ctx context.Context, sessionID, userID string, ip, userAgent string) (*LoginResponse, error) {
|
||
session, err := s.GetSession(ctx, sessionID)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
if session.Status != QRCodeStatusScanned {
|
||
return nil, fmt.Errorf("invalid session status")
|
||
}
|
||
|
||
if session.UserID != userID {
|
||
return nil, fmt.Errorf("unauthorized")
|
||
}
|
||
|
||
// 获取用户信息
|
||
user, err := s.userService.GetUserByID(ctx, userID)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("user not found: %w", err)
|
||
}
|
||
|
||
// 复用JWT生成逻辑
|
||
accessToken, err := s.jwtService.GenerateAccessToken(user.ID, user.Username)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("failed to generate access token: %w", err)
|
||
}
|
||
|
||
refreshToken, err := s.jwtService.GenerateRefreshToken(user.ID, user.Username)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("failed to generate refresh token: %w", err)
|
||
}
|
||
|
||
// 更新状态
|
||
key := qrcodeSessionPrefix + sessionID
|
||
if err := s.redis.HSet(ctx, key, "status", string(QRCodeStatusConfirmed)); err != nil {
|
||
return nil, fmt.Errorf("failed to update status: %w", err)
|
||
}
|
||
|
||
// 复用活跃记录
|
||
if s.activityService != nil {
|
||
go func() {
|
||
s.activityService.RecordUserActive(context.Background(), userID, model.LoginTypeLogin, ip, userAgent)
|
||
}()
|
||
}
|
||
|
||
// 复用登录日志
|
||
if s.logService != nil {
|
||
go func() {
|
||
s.logService.LoginLog.RecordLogin(&model.LoginLog{
|
||
UserID: user.ID,
|
||
UserName: user.Username,
|
||
NickName: user.Nickname,
|
||
LoginType: string(model.LoginTypeQRCode),
|
||
Event: string(model.LoginEventLogin),
|
||
IP: ip,
|
||
UserAgent: userAgent,
|
||
Result: string(model.LoginResultSuccess),
|
||
})
|
||
}()
|
||
}
|
||
|
||
// 推送确认事件
|
||
response := &LoginResponse{
|
||
Token: accessToken,
|
||
RefreshToken: refreshToken,
|
||
User: user,
|
||
}
|
||
s.wsHub.PublishToUser(sessionID, "confirmed", response)
|
||
|
||
return response, nil
|
||
}
|
||
|
||
// LoginResponse 登录响应
|
||
type LoginResponse struct {
|
||
Token string `json:"token"`
|
||
RefreshToken string `json:"refresh_token"`
|
||
User *model.User `json:"user"`
|
||
}
|
||
|
||
// Cancel 取消登录
|
||
func (s *QRCodeLoginService) Cancel(ctx context.Context, sessionID, userID string) error {
|
||
session, err := s.GetSession(ctx, sessionID)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
|
||
if session.Status != QRCodeStatusScanned {
|
||
return fmt.Errorf("invalid session status")
|
||
}
|
||
|
||
if session.UserID != userID {
|
||
return fmt.Errorf("unauthorized")
|
||
}
|
||
|
||
key := qrcodeSessionPrefix + sessionID
|
||
if err := s.redis.HSet(ctx, key, "status", string(QRCodeStatusCancelled)); err != nil {
|
||
return fmt.Errorf("failed to update status: %w", err)
|
||
}
|
||
|
||
// 推送取消事件
|
||
s.wsHub.PublishToUser(sessionID, "cancelled", map[string]any{})
|
||
|
||
return nil
|
||
}
|
||
|
||
// GetWSHub 获取底层 WS Hub(用于 Subscribe 等本地操作)
|
||
func (s *QRCodeLoginService) GetWSHub() *ws.Hub {
|
||
switch p := s.wsHub.(type) {
|
||
case *ws.Bus:
|
||
return p.Hub()
|
||
case *ws.Hub:
|
||
return p
|
||
default:
|
||
return nil
|
||
}
|
||
}
|
||
|
||
// GetUserService 获取用户服务
|
||
func (s *QRCodeLoginService) GetUserService() UserService {
|
||
return s.userService
|
||
}
|