安全与可靠性: 修复配额并发、投票事务、密码授权与导出完整性

2026-11-05 12:17:05 +08:00
parent 7636456d4e
commit 25aee35030
18 changed files with 1002 additions and 194 deletions
@@ -0,0 +1,163 @@
+package db
+
+import (
+	"context"
+	"fmt"
+	"sync"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"gorm.io/gorm"
+)
+
+func createQuotaTestUser(t *testing.T, ctx context.Context, conn *gorm.DB, name string, dailyLimit int) *User {
+	t.Helper()
+	user := &User{Name: name, Domain: name, Email: name + "@example.test", Password: "test-only", DailyMessageLimit: dailyLimit}
+	require.NoError(t, conn.WithContext(ctx).Create(user).Error)
+	return user
+}
+
+func setFeatureQuota(t *testing.T, ctx context.Context, conn *gorm.DB, userID uint, feature string, limit int64) {
+	t.Helper()
+	require.NoError(t, NewFeatureControlsStore(conn).Save(ctx, userID, feature, UpdateFeatureControlOptions{Enabled: true, QuotaLimit: limit}))
+}
+
+func runConcurrently(n int, fn func(int) error) []error {
+	start := make(chan struct{})
+	errs := make([]error, n)
+	var wg sync.WaitGroup
+	for i := range n {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			<-start
+			errs[i] = fn(i)
+		}()
+	}
+	close(start)
+	wg.Wait()
+	return errs
+}
+
+func requireOnlyQuotaErrors(t *testing.T, errs []error, quotaErr error) int {
+	t.Helper()
+	successes := 0
+	for _, err := range errs {
+		if err == nil {
+			successes++
+			continue
+		}
+		require.ErrorIs(t, err, quotaErr)
+	}
+	return successes
+}
+
+func TestQuestionQuotaConcurrentLastSlot(t *testing.T) {
+	conn, _ := newTestDB(t)
+	ctx := context.Background()
+	owner := createQuotaTestUser(t, ctx, conn, "quota-question-owner", 0)
+	setFeatureQuota(t, ctx, conn, owner.ID, FeatureQuestion, 1)
+	questions := NewQuestionsStore(conn)
+	errs := runConcurrently(12, func(i int) error {
+		_, err := questions.Create(ctx, CreateQuestionOptions{UserID: owner.ID, Content: fmt.Sprintf("question-%d", i)})
+		return err
+	})
+	require.Equal(t, 1, requireOnlyQuotaErrors(t, errs, ErrFeatureQuotaReached))
+	var count int64
+	require.NoError(t, conn.Model(&Question{}).Where("user_id = ?", owner.ID).Count(&count).Error)
+	require.Equal(t, int64(1), count)
+	fc, err := NewFeatureControlsStore(conn).Get(ctx, owner.ID, FeatureQuestion)
+	require.NoError(t, err)
+	require.Equal(t, int64(1), fc.QuotaCount)
+}
+
+func TestVoteQuotasConcurrentLastSlot(t *testing.T) {
+	conn, _ := newTestDB(t)
+	ctx := context.Background()
+	owner := createQuotaTestUser(t, ctx, conn, "quota-vote-owner", 0)
+	setFeatureQuota(t, ctx, conn, owner.ID, FeatureVote, 1)
+	votes := NewVotesStore(conn)
+	options := NewVoteOptionsStore(conn)
+	submissions := NewVoteSubmissionsStore(conn)
+	ids := make([]uint, 2)
+	for i := range 2 {
+		vote, err := votes.Create(ctx, CreateVoteOptions{UserID: owner.ID, Title: fmt.Sprintf("vote-%d", i), Type: VoteTypeSingle, QuotaLimit: 1})
+		require.NoError(t, err)
+		require.NoError(t, votes.UpdateStatus(ctx, vote.ID, VoteStatusPublished))
+		option, err := options.Create(ctx, CreateVoteOptionOptions{VoteID: vote.ID, Content: "option"})
+		require.NoError(t, err)
+		ids[i] = option.ID
+	}
+	errs := runConcurrently(12, func(i int) error {
+		token := fmt.Sprintf("voter-%d", i)
+		_, err := submissions.Create(ctx, CreateVoteSubmissionOptions{VoteID: uint(i%2 + 1), OptionID: ids[i%2], VoterToken: &token, QuotaUserID: owner.ID})
+		return err
+	})
+	successes := 0
+	for _, err := range errs {
+		if err == nil {
+			successes++
+		} else {
+			require.Error(t, err)
+		}
+	}
+	require.Equal(t, 1, successes)
+	var count int64
+	require.NoError(t, conn.Model(&VoteSubmission{}).Count(&count).Error)
+	require.Equal(t, int64(1), count)
+	fc, err := NewFeatureControlsStore(conn).Get(ctx, owner.ID, FeatureVote)
+	require.NoError(t, err)
+	require.Equal(t, int64(1), fc.QuotaCount)
+	var ended int64
+	require.NoError(t, conn.Model(&Vote{}).Where("status = ?", VoteStatusEnded).Count(&ended).Error)
+	require.Equal(t, int64(1), ended)
+}
+
+func TestDailyMessageLimitAcrossConversationPaths(t *testing.T) {
+	conn, _ := newTestDB(t)
+	ctx := context.Background()
+	owner := createQuotaTestUser(t, ctx, conn, "quota-chat-owner", 1)
+	asker := createQuotaTestUser(t, ctx, conn, "quota-chat-asker", 0)
+	conversations := NewConversationStore(conn)
+	existing, err := conversations.Create(ctx, CreateConversationOptions{OwnerID: owner.ID, AskerID: asker.ID, MaxMessages: 20})
+	require.NoError(t, err)
+	var wg sync.WaitGroup
+	start := make(chan struct{})
+	errs := make([]error, 2)
+	wg.Add(2)
+	go func() {
+		defer wg.Done()
+		<-start
+		_, _, errs[0] = conversations.CreateWithFirstMessage(ctx, CreateConversationOptions{OwnerID: owner.ID, AskerID: asker.ID, MaxMessages: 20}, asker.ID, "first")
+	}()
+	go func() {
+		defer wg.Done()
+		<-start
+		_, errs[1] = conversations.SendMessage(ctx, SendConversationMessageOptions{ConversationID: existing.ID, SenderID: asker.ID, Content: "existing"})
+	}()
+	close(start)
+	wg.Wait()
+	require.Equal(t, 1, requireOnlyQuotaErrors(t, errs, ErrDailyMessageLimit))
+	var messages, convCount int64
+	require.NoError(t, conn.Model(&Message{}).Count(&messages).Error)
+	require.NoError(t, conn.Model(&Conversation{}).Count(&convCount).Error)
+	require.Equal(t, int64(1), messages)
+	require.Equal(t, int64(1), convCount)
+}
+
+func TestOwnerReplyDoesNotConsumeAskerDailyQuota(t *testing.T) {
+	conn, _ := newTestDB(t)
+	ctx := context.Background()
+	owner := createQuotaTestUser(t, ctx, conn, "quota-reply-owner", 1)
+	asker := createQuotaTestUser(t, ctx, conn, "quota-reply-asker", 0)
+	conversations := NewConversationStore(conn)
+	conv, _, err := conversations.CreateWithFirstMessage(ctx, CreateConversationOptions{OwnerID: owner.ID, AskerID: asker.ID, MaxMessages: 20}, asker.ID, "first")
+	require.NoError(t, err)
+	_, err = conversations.SendMessage(ctx, SendConversationMessageOptions{ConversationID: conv.ID, SenderID: owner.ID, Content: "reply"})
+	require.NoError(t, err)
+	_, err = conversations.SendMessage(ctx, SendConversationMessageOptions{ConversationID: conv.ID, SenderID: asker.ID, Content: "over"})
+	require.ErrorIs(t, err, ErrDailyMessageLimit)
+	var count int64
+	require.NoError(t, conn.Model(&Message{}).Count(&count).Error)
+	require.Equal(t, int64(2), count)
+}