diff --git a/backend/internal/repository/user_repo.go b/backend/internal/repository/user_repo.go index def7462c1..4b8841bb6 100644 --- a/backend/internal/repository/user_repo.go +++ b/backend/internal/repository/user_repo.go @@ -70,24 +70,25 @@ func (r *userRepository) create(ctx context.Context, userIn *service.User, guard // 统一使用 ent 的事务：保证用户与允许分组的更新原子化， // 并避免基于 *sql.Tx 手动构造 ent client 导致的 ExecQuerier 断言错误。 - tx, err := r.client.Tx(ctx) - if err != nil && !errors.Is(err, dbent.ErrTxStarted) { - return err - } - + // + // 注意：ent 的 Client.Tx 不感知上下文中的事务（只检查 driver 类型）， + // 因此必须显式检查 TxFromContext：当调用方已开启外部事务（如注册时的 + // “建用户 + 占用邀请码”原子事务），直接复用其 client，由调用方统一提交/回滚， + // 否则用户写入会落入独立事务并自行提交，导致外层事务无法回滚（孤儿用户）。 var txClient *dbent.Client txCtx := ctx - if err == nil { - defer func() { _ = tx.Rollback() }() + var ownedTx *dbent.Tx + if existingTx := dbent.TxFromContext(ctx); existingTx != nil { + txClient = existingTx.Client() + } else { + tx, err := r.client.Tx(ctx) + if err != nil { + return err + } + ownedTx = tx + defer func() { _ = ownedTx.Rollback() }() txClient = tx.Client() txCtx = dbent.NewTxContext(ctx, tx) - } else { - // 已处于外部事务中（ErrTxStarted），复用当前事务 client 并由调用方负责提交/回滚。 - if existingTx := dbent.TxFromContext(ctx); existingTx != nil { - txClient = existingTx.Client() - } else { - txClient = r.client - } } lockKeys := []string{normalizedEmailUniquenessLockKey(userIn.Email)} @@ -158,8 +159,8 @@ func (r *userRepository) create(ctx context.Context, userIn *service.User, guard return err } - if tx != nil { - if err := tx.Commit(); err != nil { + if ownedTx != nil { + if err := ownedTx.Commit(); err != nil { return err } } diff --git a/backend/internal/repository/user_repo_invitation_claim_test.go b/backend/internal/repository/user_repo_invitation_claim_test.go new file mode 100644 index 000000000..ba24c365b --- /dev/null +++ b/backend/internal/repository/user_repo_invitation_claim_test.go @@ -0,0 +1,113 @@ +//go:build integration + +package repository + +import ( + "context" + "testing" + + dbent "github.com/Wei-Shaw/sub2api/ent" + "github.com/Wei-Shaw/sub2api/ent/redeemcode" + "github.com/Wei-Shaw/sub2api/ent/user" + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/stretchr/testify/require" +) + +// TestCreateWithEmailAliasGuardJoinsOuterTransaction 验证用户创建会加入调用方开启的 +// 外部 ent 事务（注册流程“建用户 + 占用邀请码”原子性的基础）： +// - 外层事务回滚后，用户与邀请码占用必须一并撤销（不得残留孤儿账号）； +// - 外层事务提交后，用户与邀请码占用同时生效。 +// +// 回归背景：此前 create() 通过 r.client.Tx(ctx) 自开事务且自行 Commit —— ent 的 +// Client.Tx 不感知上下文事务（只检查 driver 类型），ErrTxStarted 分支实际是死代码， +// 导致外层事务包不住用户写入；并发注册同一邀请码时，败者用户的写入已自行提交， +// 即使注册被拒绝也会留下可登录的孤儿账号（1 个邀请码仍可生成任意数量账号）。 +func TestCreateWithEmailAliasGuardJoinsOuterTransaction(t *testing.T) { + client := testEntClient(t) + userRepo := NewUserRepository(client, integrationDB) + redeemRepo := NewRedeemCodeRepository(client) + + ctx := context.Background() + + // 清理：本测试会真实提交少量数据，确保不影响同包其它集成测试。 + var committedUserEmails []string + var seededCodeIDs []int64 + t.Cleanup(func() { + if len(committedUserEmails) > 0 { + _, _ = client.User.Delete().Where(user.EmailIn(committedUserEmails...)).Exec(ctx) + } + if len(seededCodeIDs) > 0 { + _, _ = client.RedeemCode.Delete().Where(redeemcode.IDIn(seededCodeIDs...)).Exec(ctx) + } + }) + + seedCode := func(code string) int64 { + _, err := client.RedeemCode.Create(). + SetCode(code). + SetType(service.RedeemTypeInvitation). + SetStatus(service.StatusUnused). + SetValue(0). + Save(ctx) + require.NoError(t, err, "seed redeem code") + c, err := client.RedeemCode.Query().Where(redeemcode.CodeEQ(code)).Only(ctx) + require.NoError(t, err) + seededCodeIDs = append(seededCodeIDs, c.ID) + return c.ID + } + + t.Run("rollback removes user and releases claim", func(t *testing.T) { + codeID := seedCode("ITX-RACE-ROLLBACK-001") + tx, err := client.Tx(ctx) + require.NoError(t, err) + txCtx := dbent.NewTxContext(ctx, tx) + + u := &service.User{ + Email: "itx-rollback@example.com", + PasswordHash: "test-password-hash", + Role: service.RoleUser, + Status: service.StatusActive, + Balance: 0, + Concurrency: 1, + } + require.NoError(t, userRepo.CreateWithEmailAliasGuard(txCtx, u)) + require.Greater(t, u.ID, int64(0), "create 应回填用户 ID") + require.NoError(t, redeemRepo.Use(txCtx, codeID, u.ID)) + require.NoError(t, tx.Rollback()) + + exists, err := userRepo.ExistsByEmail(ctx, "itx-rollback@example.com") + require.NoError(t, err) + require.False(t, exists, "回滚后不得残留孤儿用户") + + after, err := client.RedeemCode.Get(ctx, codeID) + require.NoError(t, err) + require.Equal(t, service.StatusUnused, after.Status, "回滚后邀请码应保持 unused") + }) + + t.Run("commit persists user and claim together", func(t *testing.T) { + codeID := seedCode("ITX-RACE-COMMIT-001") + tx, err := client.Tx(ctx) + require.NoError(t, err) + txCtx := dbent.NewTxContext(ctx, tx) + + u := &service.User{ + Email: "itx-commit@example.com", + PasswordHash: "test-password-hash", + Role: service.RoleUser, + Status: service.StatusActive, + Balance: 0, + Concurrency: 1, + } + require.NoError(t, userRepo.CreateWithEmailAliasGuard(txCtx, u)) + require.NoError(t, redeemRepo.Use(txCtx, codeID, u.ID)) + require.NoError(t, tx.Commit()) + committedUserEmails = append(committedUserEmails, u.Email) + + exists, err := userRepo.ExistsByEmail(ctx, "itx-commit@example.com") + require.NoError(t, err) + require.True(t, exists, "提交后用户应存在") + + after, err := client.RedeemCode.Get(ctx, codeID) + require.NoError(t, err) + require.Equal(t, service.StatusUsed, after.Status, "提交后邀请码应为 used") + }) +} diff --git a/backend/internal/service/auth_service.go b/backend/internal/service/auth_service.go index 2beba0d40..72adbe890 100644 --- a/backend/internal/service/auth_service.go +++ b/backend/internal/service/auth_service.go @@ -245,13 +245,15 @@ func (s *AuthService) RegisterWithVerification(ctx context.Context, email, passw Status: StatusActive, } - if err := s.createUserWithRegistrationEmailGuard(ctx, user); err != nil { + if err := s.createUserAndClaimInvitation(ctx, user, invitationRedeemCode); err != nil { // 优先检查邮箱冲突错误（竞态条件下可能发生） switch { case errors.Is(err, ErrEmailExists): return "", nil, ErrEmailExists case errors.Is(err, ErrEmailDomainRegistrationLimit): return "", nil, ErrEmailDomainRegistrationLimit + case errors.Is(err, ErrInvitationCodeInvalid): + return "", nil, ErrInvitationCodeInvalid default: logger.LegacyPrintf("service.auth", "[Auth] Database error creating user: %v", err) return "", nil, ErrServiceUnavailable @@ -273,13 +275,8 @@ func (s *AuthService) RegisterWithVerification(ctx context.Context, email, passw } } - // 标记邀请码为已使用（如果使用了邀请码） - if invitationRedeemCode != nil { - if err := s.redeemRepo.Use(ctx, invitationRedeemCode.ID, user.ID); err != nil { - // 邀请码标记失败不影响注册，只记录日志 - logger.LegacyPrintf("service.auth", "[Auth] Failed to mark invitation code as used for user %d: %v", user.ID, err) - } - } + // 邀请码占用已由 createUserAndClaimInvitation 在“用户创建 + 邀请码占用”的 + // 同一个数据库事务内原子完成（一次性约束，见函数注释），此处不再单独标记。 // 应用优惠码（如果提供且功能已启用） if promoCode != "" && s.promoService != nil && s.settingService != nil && s.settingService.IsPromoCodeEnabled(ctx) { if err := s.promoService.ApplyPromoCode(ctx, user.ID, promoCode); err != nil { @@ -1274,6 +1271,64 @@ func (s *AuthService) createUserWithRegistrationEmailGuard(ctx context.Context, return quotaRepo.CreateWithEmailAliasGuardAndDomainLimit(ctx, user, domain) } +// createUserAndClaimInvitation 原子化完成“用户创建 + 邀请码占用”。 +// +// 背景：邀请码属于一次性凭证，必须保证“一个邀请码最多注册一个账号”。旧实现先检查 +// CanUse()、再创建用户、最后才 redeemRepo.Use()（且失败仅记日志），检查与消耗分离且 +// 不在同一事务，并发注册可在同一邀请码上同时通过检查并各自创建账号（TOCTOU 竞态）。 +// +// 本实现把两者放入同一个数据库事务： +// - 占用走 redeemRepo.Use 的条件更新（WHERE status='unused'，乐观锁）； +// - 并发下只有一个事务能占用成功，其余事务回滚——既不产生多余账号，也不让码被烧掉； +// - 事务回滚同时撤销用户创建，避免“账号已建、码被占用”的中间态。 +// +// 无邀请码时保持原单次创建路径（不开事务）；entClient 缺失的异常配置下退化为顺序执行， +// 并发正确性仍由 Use 的条件更新兜底（可能产生孤儿用户，但不会放行第二个注册）。 +func (s *AuthService) createUserAndClaimInvitation(ctx context.Context, user *User, invitation *RedeemCode) error { + commitUser := func(execCtx context.Context) error { + if err := s.createUserWithRegistrationEmailGuard(execCtx, user); err != nil { + return err + } + if invitation == nil { + return nil + } + // createUserWithRegistrationEmailGuard 会回填 user.ID（applyUserEntityToService）， + // 直接以其原子占用邀请码；占用失败即整体回滚（含用户创建，见 user_repo.create + // 对外部事务的复用）。 + if err := s.redeemRepo.Use(execCtx, invitation.ID, user.ID); err != nil { + // 并发下唯一的合法失败路径：另一个注册已占用该码 + logger.LegacyPrintf("service.auth", + "[Auth] Rejected registration: invitation code %s already claimed (user_id=%d err=%v)", + invitation.Code, user.ID, err) + return ErrInvitationCodeInvalid + } + return nil + } + + if invitation == nil { + return commitUser(ctx) + } + if s.entClient == nil { + return commitUser(ctx) + } + + tx, err := s.entClient.Tx(ctx) + if err != nil { + logger.LegacyPrintf("service.auth", "[Auth] Failed to start registration transaction: %v", err) + return ErrServiceUnavailable + } + execCtx := dbent.NewTxContext(ctx, tx) + if err := commitUser(execCtx); err != nil { + _ = tx.Rollback() + return err + } + if err := tx.Commit(); err != nil { + logger.LegacyPrintf("service.auth", "[Auth] Failed to commit registration transaction: %v", err) + return ErrServiceUnavailable + } + return nil +} + func buildEmailSuffixNotAllowedError(whitelist []string) error { if len(whitelist) == 0 { return ErrEmailSuffixNotAllowed diff --git a/backend/internal/service/auth_service_invitation_race_test.go b/backend/internal/service/auth_service_invitation_race_test.go new file mode 100644 index 000000000..e1e7c82ae --- /dev/null +++ b/backend/internal/service/auth_service_invitation_race_test.go @@ -0,0 +1,236 @@ +//go:build unit + +package service + +import ( + "context" + "errors" + "fmt" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +// raceSafeUserRepo 是仅覆盖注册路径的并发安全用户仓储桩。 +// 未用到的方法走嵌入接口（调用即 panic，注册路径不会触发）。 +type raceSafeUserRepo struct { + UserRepository + + mu sync.Mutex + nextID int64 + byEmail map[string]*User + byID map[int64]*User +} + +func newRaceSafeUserRepo() *raceSafeUserRepo { + return &raceSafeUserRepo{nextID: 1, byEmail: map[string]*User{}, byID: map[int64]*User{}} +} + +func (s *raceSafeUserRepo) ExistsByEmail(_ context.Context, email string) (bool, error) { + s.mu.Lock() + defer s.mu.Unlock() + _, ok := s.byEmail[email] + return ok, nil +} + +func (s *raceSafeUserRepo) ExistsByEmailAlias(ctx context.Context, email string) (bool, error) { + return s.ExistsByEmail(ctx, email) +} + +func (s *raceSafeUserRepo) CreateWithEmailAliasGuard(_ context.Context, user *User) error { + s.mu.Lock() + defer s.mu.Unlock() + if _, ok := s.byEmail[user.Email]; ok { + return ErrEmailExists + } + user.ID = s.nextID + s.nextID++ + clone := *user + s.byEmail[user.Email] = &clone + s.byID[user.ID] = &clone + return nil +} + +func (s *raceSafeUserRepo) GetByEmail(_ context.Context, email string) (*User, error) { + s.mu.Lock() + defer s.mu.Unlock() + u, ok := s.byEmail[email] + if !ok { + return nil, ErrUserNotFound + } + clone := *u + return &clone, nil +} + +func (s *raceSafeUserRepo) GetByID(_ context.Context, id int64) (*User, error) { + s.mu.Lock() + defer s.mu.Unlock() + u, ok := s.byID[id] + if !ok { + return nil, ErrUserNotFound + } + clone := *u + return &clone, nil +} + +func (s *raceSafeUserRepo) Update(context.Context, *User, UserUpdateFields) error { + return nil +} + +// raceSafeRedeemRepo 是并发安全的兑换码仓储桩：Use 以互斥锁 + 状态条件 +// 模拟数据库的条件更新（WHERE status='unused'），语义与线上实现一致。 +type raceSafeRedeemRepo struct { + RedeemCodeRepository + + mu sync.Mutex + codes map[string]*RedeemCode +} + +func (s *raceSafeRedeemRepo) GetByCode(_ context.Context, code string) (*RedeemCode, error) { + s.mu.Lock() + defer s.mu.Unlock() + c, ok := s.codes[code] + if !ok { + return nil, ErrRedeemCodeNotFound + } + clone := *c + return &clone, nil +} + +func (s *raceSafeRedeemRepo) Use(_ context.Context, id, userID int64) error { + s.mu.Lock() + defer s.mu.Unlock() + for _, c := range s.codes { + if c.ID != id { + continue + } + if c.Status != StatusUnused { + return ErrRedeemCodeUsed + } + now := time.Now().UTC() + c.Status = StatusUsed + c.UsedBy = &userID + c.UsedAt = &now + return nil + } + return ErrRedeemCodeNotFound +} + +// TestAuthService_Register_InvitationCodeSingleUseUnderConcurrency 回归测试： +// 同一邀请码并发注册必须恰好成功 1 次，其余请求以 INVITATION_CODE_INVALID 拒绝。 +// +// 修复前：邀请码“检查(CanUse) 与 标记已用(Use)”分离且不在同一事务，Use 失败被吞， +// 并发请求全部注册成功（一个邀请码可创建任意数量账号）。此测试在该实现下必然失败。 +// 修复后：用户创建与邀请码占用在同一事务内原子完成（或退化路径下由 Use 条件更新 +// 兜底），并发下仅最先占码的注册成功。 +func TestAuthService_Register_InvitationCodeSingleUseUnderConcurrency(t *testing.T) { + const code = "INV-RACE-001" + userRepo := newRaceSafeUserRepo() + redeemRepo := &raceSafeRedeemRepo{codes: map[string]*RedeemCode{ + code: {ID: 1, Code: code, Type: RedeemTypeInvitation, Status: StatusUnused}, + }} + settings := map[string]string{ + "registration_enabled": "true", + "invitation_code_enabled": "true", + } + svc := newOAuthEmailFlowAuthService( + userRepo, + redeemRepo, + &refreshTokenCacheStub{}, + settings, + nil, // emailCache：注册不要求邮箱验证，保持关闭 + &userPlatformQuotaRepoStub{}, + ) + + const n = 8 + ctx := context.Background() + start := make(chan struct{}) + results := make(chan error, n) + var wg sync.WaitGroup + for i := 0; i < n; i++ { + wg.Add(1) + go func(i int) { + defer wg.Done() + <-start + email := fmt.Sprintf("race-%d@example.com", i) + _, _, err := svc.RegisterWithVerification(ctx, email, "Password123!", "", "", code, "") + results <- err + }(i) + } + close(start) + wg.Wait() + close(results) + + successes := 0 + rejected := 0 + for err := range results { + switch { + case err == nil: + successes++ + case errors.Is(err, ErrInvitationCodeInvalid): + rejected++ + default: + t.Fatalf("unexpected registration error: %v", err) + } + } + require.Equal(t, 1, successes, "同一邀请码并发注册必须恰好成功 1 次") + require.Equal(t, n-1, rejected, "其余并发请求必须以 INVITATION_CODE_INVALID 拒绝") + + claimed, err := redeemRepo.GetByCode(ctx, code) + require.NoError(t, err) + require.Equal(t, StatusUsed, claimed.Status, "邀请码最终必须处于 used 状态") + require.NotNil(t, claimed.UsedBy, "used_by 必须记录实际注册用户") +} + +// TestAuthService_Register_InvitationCodeRejectedWhenAlreadyUsed 顺序路径回归： +// 已使用过的邀请码再次注册（即便换邮箱）必须被拒绝。 +func TestAuthService_Register_InvitationCodeRejectedWhenAlreadyUsed(t *testing.T) { + const code = "INV-RACE-002" + userRepo := newRaceSafeUserRepo() + redeemRepo := &raceSafeRedeemRepo{codes: map[string]*RedeemCode{ + code: {ID: 2, Code: code, Type: RedeemTypeInvitation, Status: StatusUsed}, + }} + settings := map[string]string{ + "registration_enabled": "true", + "invitation_code_enabled": "true", + } + svc := newOAuthEmailFlowAuthService( + userRepo, + redeemRepo, + &refreshTokenCacheStub{}, + settings, + nil, + &userPlatformQuotaRepoStub{}, + ) + + _, _, err := svc.RegisterWithVerification(context.Background(), "later@example.com", "Password123!", "", "", code, "") + require.ErrorIs(t, err, ErrInvitationCodeInvalid) +} + +// TestAuthService_Register_InvitationCodeMissingWhenEnabled 门控回归： +// 邀请码开启时，不带邀请码的注册必须被拒绝（不产生用户）。 +func TestAuthService_Register_InvitationCodeMissingWhenEnabled(t *testing.T) { + userRepo := newRaceSafeUserRepo() + redeemRepo := &raceSafeRedeemRepo{codes: map[string]*RedeemCode{}} + settings := map[string]string{ + "registration_enabled": "true", + "invitation_code_enabled": "true", + } + svc := newOAuthEmailFlowAuthService( + userRepo, + redeemRepo, + &refreshTokenCacheStub{}, + settings, + nil, + &userPlatformQuotaRepoStub{}, + ) + + _, _, err := svc.RegisterWithVerification(context.Background(), "no-invite@example.com", "Password123!", "", "", "", "") + require.ErrorIs(t, err, ErrInvitationCodeRequired) + + ok, err := userRepo.ExistsByEmail(context.Background(), "no-invite@example.com") + require.NoError(t, err) + require.False(t, ok, "被拒绝的注册不应产生用户") +}