Merge pull request #5720 from tamseno/fix/invitation-code-toctou-race

fix(auth): make invitation code consumption atomic with user creation

backend/internal/repository/user_repo.go

@@ -70,23 +70,30 @@ 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() }()
-		txClient = tx.Client()
-		txCtx = dbent.NewTxContext(ctx, tx)
+	var ownedTx *dbent.Tx
+	if existingTx := dbent.TxFromContext(ctx); existingTx != nil {
+		txClient = existingTx.Client()
 	} else {
-		// 已处于外部事务中（ErrTxStarted），复用当前事务 client 并由调用方负责提交/回滚。
-		if existingTx := dbent.TxFromContext(ctx); existingTx != nil {
-			txClient = existingTx.Client()
-		} else {
+		tx, err := r.client.Tx(ctx)
+		switch {
+		case errors.Is(err, dbent.ErrTxStarted):
+			// r.client 本身已是事务绑定 client（client 注入式事务，如集成测试
+			// 夹具 tx.Client()）：直接复用，提交/回滚由 client 的持有方负责。
 			txClient = r.client
+		case err != nil:
+			return err
+		default:
+			ownedTx = tx
+			defer func() { _ = ownedTx.Rollback() }()
+			txClient = tx.Client()
+			txCtx = dbent.NewTxContext(ctx, tx)
 		}
 	}
 
@@ -158,8 +165,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
 		}
 	}

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")
+	})
+}

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
+	}
+	defer func() { _ = tx.Rollback() }()
+	execCtx := dbent.NewTxContext(ctx, tx)
+	if err := commitUser(execCtx); err != nil {
+		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

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, "被拒绝的注册不应产生用户")
+}

backend/internal/service/auth_service_register_test.go

@@ -5,6 +5,7 @@ package service
 import (
 	"context"
 	"errors"
+	"sync"
 	"testing"
 	"time"
 
@@ -14,6 +15,7 @@ import (
 )
 
 type settingRepoStub struct {
+	mu               sync.Mutex
 	values           map[string]string
 	err              error
 	getValueCalls    int
@@ -25,6 +27,8 @@ func (s *settingRepoStub) Get(ctx context.Context, key string) (*Setting, error)
 }
 
 func (s *settingRepoStub) GetValue(ctx context.Context, key string) (string, error) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
 	s.getValueCalls++
 	if s.err != nil {
 		return "", s.err
@@ -40,6 +44,8 @@ func (s *settingRepoStub) Set(ctx context.Context, key, value string) error {
 }
 
 func (s *settingRepoStub) GetMultiple(ctx context.Context, keys []string) (map[string]string, error) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
 	s.getMultipleCalls++
 	if s.err != nil {
 		return nil, s.err