Merge pull request #7676 from william-drakemond/fix/email-code-token-atomic

fix(email): 验证码尝试次数原子化 + 重置密码 token 哈希存储且原子单次消费
2026-10-15 23:22:59 +08:00 · 2026-10-02 14:32:42 +08:00
parent 43c4c882fd b1ce3c6d59
commit 2529371bab
9 changed files with 426 additions and 98 deletions
@@ -17,8 +17,42 @@ const (
 	passwordResetKeyPrefix       = "password_reset:"
 	passwordResetSentAtKeyPrefix = "password_reset_sent:"
 	notifyCodeUserRateKeyPrefix  = "notify_code_user_rate:"
+
+	// attemptsKeySuffix stores the failed-attempt counter next to a verification code.
+	// Kept in a separate key so it can be incremented atomically with INCR.
+	attemptsKeySuffix = ":attempts"
 )

+// incrAttemptsScript atomically increments the attempt counter for an existing
+// verification code and aligns the counter TTL with the code TTL.
+// KEYS[1] = code key, KEYS[2] = attempts key. Returns -1 when the code is missing.
+var incrAttemptsScript = redis.NewScript(`
+if redis.call('EXISTS', KEYS[1]) == 0 then
+  return -1
+end
+local n = redis.call('INCR', KEYS[2])
+local ttl = redis.call('PTTL', KEYS[1])
+if ttl > 0 then
+  redis.call('PEXPIRE', KEYS[2], ttl)
+end
+return n
+`)
+
+// consumeResetTokenScript atomically compares the stored token hash and deletes it.
+// KEYS[1] = reset key, ARGV[1] = expected token hash. Returns 1 on success, 0 otherwise.
+var consumeResetTokenScript = redis.NewScript(`
+local v = redis.call('GET', KEYS[1])
+if not v then
+  return 0
+end
+local ok, d = pcall(cjson.decode, v)
+if not ok or type(d) ~= 'table' or d['Token'] ~= ARGV[1] then
+  return 0
+end
+redis.call('DEL', KEYS[1])
+return 1
+`)
+
 // verifyCodeKey generates the Redis key for email verification code.
 // Email is lowercased for case-insensitive consistency.
 func verifyCodeKey(email string) string {
@@ -50,8 +84,7 @@ func NewEmailCache(rdb *redis.Client) service.EmailCache {
 	return &emailCache{rdb: rdb}
 }

-func (c *emailCache) GetVerificationCode(ctx context.Context, email string) (*service.VerificationCodeData, error) {
-	key := verifyCodeKey(email)
+func (c *emailCache) getCode(ctx context.Context, key string) (*service.VerificationCodeData, error) {
 	val, err := c.rdb.Get(ctx, key).Result()
 	if err != nil {
 		return nil, err
@@ -60,21 +93,56 @@ func (c *emailCache) GetVerificationCode(ctx context.Context, email string) (*se
 	if err := json.Unmarshal([]byte(val), &data); err != nil {
 		return nil, err
 	}
+	if n, err := c.rdb.Get(ctx, key+attemptsKeySuffix).Int(); err == nil && n > data.Attempts {
+		data.Attempts = n
+	}
 	return &data, nil
 }

-func (c *emailCache) SetVerificationCode(ctx context.Context, email string, data *service.VerificationCodeData, ttl time.Duration) error {
-	key := verifyCodeKey(email)
+func (c *emailCache) setCode(ctx context.Context, key string, data *service.VerificationCodeData, ttl time.Duration) error {
 	val, err := json.Marshal(data)
 	if err != nil {
 		return err
 	}
-	return c.rdb.Set(ctx, key, val, ttl).Err()
+	pipe := c.rdb.TxPipeline()
+	pipe.Set(ctx, key, val, ttl)
+	pipe.Del(ctx, key+attemptsKeySuffix)
+	if data.Attempts > 0 {
+		pipe.Set(ctx, key+attemptsKeySuffix, data.Attempts, ttl)
+	}
+	_, err = pipe.Exec(ctx)
+	return err
+}
+
+func (c *emailCache) incrCodeAttempts(ctx context.Context, key string) (int, error) {
+	n, err := incrAttemptsScript.Run(ctx, c.rdb, []string{key, key + attemptsKeySuffix}).Int()
+	if err != nil {
+		return 0, err
+	}
+	if n < 0 {
+		return 0, redis.Nil
+	}
+	return n, nil
+}
+
+func (c *emailCache) deleteCode(ctx context.Context, key string) error {
+	return c.rdb.Del(ctx, key, key+attemptsKeySuffix).Err()
+}
+
+func (c *emailCache) GetVerificationCode(ctx context.Context, email string) (*service.VerificationCodeData, error) {
+	return c.getCode(ctx, verifyCodeKey(email))
+}
+
+func (c *emailCache) SetVerificationCode(ctx context.Context, email string, data *service.VerificationCodeData, ttl time.Duration) error {
+	return c.setCode(ctx, verifyCodeKey(email), data, ttl)
+}
+
+func (c *emailCache) IncrVerificationCodeAttempts(ctx context.Context, email string) (int, error) {
+	return c.incrCodeAttempts(ctx, verifyCodeKey(email))
 }

 func (c *emailCache) DeleteVerificationCode(ctx context.Context, email string) error {
-	key := verifyCodeKey(email)
-	return c.rdb.Del(ctx, key).Err()
+	return c.deleteCode(ctx, verifyCodeKey(email))
 }

 // Password reset token methods
@@ -101,6 +169,16 @@ func (c *emailCache) SetPasswordResetToken(ctx context.Context, email string, da
 	return c.rdb.Set(ctx, key, val, ttl).Err()
 }

+// ConsumePasswordResetToken atomically deletes the stored reset token when its
+// stored hash equals tokenHash. Returns true only for the single winning caller.
+func (c *emailCache) ConsumePasswordResetToken(ctx context.Context, email, tokenHash string) (bool, error) {
+	n, err := consumeResetTokenScript.Run(ctx, c.rdb, []string{passwordResetKey(email)}, tokenHash).Int()
+	if err != nil {
+		return false, err
+	}
+	return n == 1, nil
+}
+
 func (c *emailCache) DeletePasswordResetToken(ctx context.Context, email string) error {
 	key := passwordResetKey(email)
 	return c.rdb.Del(ctx, key).Err()
@@ -122,30 +200,19 @@ func (c *emailCache) SetPasswordResetEmailCooldown(ctx context.Context, email st
 // Notify email verification code methods

 func (c *emailCache) GetNotifyVerifyCode(ctx context.Context, email string) (*service.VerificationCodeData, error) {
-	key := notifyVerifyKey(email)
-	val, err := c.rdb.Get(ctx, key).Result()
-	if err != nil {
-		return nil, err
-	}
-	var data service.VerificationCodeData
-	if err := json.Unmarshal([]byte(val), &data); err != nil {
-		return nil, err
-	}
-	return &data, nil
+	return c.getCode(ctx, notifyVerifyKey(email))
 }

 func (c *emailCache) SetNotifyVerifyCode(ctx context.Context, email string, data *service.VerificationCodeData, ttl time.Duration) error {
-	key := notifyVerifyKey(email)
-	val, err := json.Marshal(data)
-	if err != nil {
-		return err
-	}
-	return c.rdb.Set(ctx, key, val, ttl).Err()
+	return c.setCode(ctx, notifyVerifyKey(email), data, ttl)
+}
+
+func (c *emailCache) IncrNotifyVerifyCodeAttempts(ctx context.Context, email string) (int, error) {
+	return c.incrCodeAttempts(ctx, notifyVerifyKey(email))
 }

 func (c *emailCache) DeleteNotifyVerifyCode(ctx context.Context, email string) error {
-	key := notifyVerifyKey(email)
-	return c.rdb.Del(ctx, key).Err()
+	return c.deleteCode(ctx, notifyVerifyKey(email))
 }

 // User-level rate limiting for notify email verification codes
@@ -0,0 +1,150 @@
+package repository
+
+import (
+	"context"
+	"crypto/sha256"
+	"encoding/hex"
+	"errors"
+	"strings"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/Wei-Shaw/sub2api/internal/service"
+	"github.com/alicebob/miniredis/v2"
+	"github.com/redis/go-redis/v9"
+	"github.com/stretchr/testify/require"
+)
+
+func newMiniredisEmailCache(t *testing.T) (service.EmailCache, *miniredis.Miniredis, *redis.Client) {
+	t.Helper()
+	mr := miniredis.RunT(t)
+	rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()})
+	t.Cleanup(func() { _ = rdb.Close() })
+	return NewEmailCache(rdb), mr, rdb
+}
+
+func TestEmailCache_ConcurrentWrongCodesCannotExceedAttemptCap(t *testing.T) {
+	cache, _, _ := newMiniredisEmailCache(t)
+	ctx := context.Background()
+	email := "user@example.com"
+
+	svc := service.NewEmailService(nil, cache)
+	require.NoError(t, cache.SetVerificationCode(ctx, email, &service.VerificationCodeData{
+		Code:      "123456",
+		CreatedAt: time.Now(),
+		ExpiresAt: time.Now().Add(15 * time.Minute),
+	}, 15*time.Minute))
+
+	const workers = 50
+	var invalid, maxed atomic.Int32
+	var wg sync.WaitGroup
+	for i := 0; i < workers; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			err := svc.VerifyCode(ctx, email, "000000")
+			switch {
+			case errors.Is(err, service.ErrInvalidVerifyCode):
+				invalid.Add(1)
+			case errors.Is(err, service.ErrVerifyCodeMaxAttempts):
+				maxed.Add(1)
+			default:
+				t.Errorf("unexpected result: %v", err)
+			}
+		}()
+	}
+	wg.Wait()
+
+	// Only attempts 1..4 may return "invalid"; every other guess is rejected by the cap.
+	require.LessOrEqual(t, int(invalid.Load()), 4)
+	require.Equal(t, workers, int(invalid.Load()+maxed.Load()))
+
+	// Even the correct code is now rejected.
+	require.ErrorIs(t, svc.VerifyCode(ctx, email, "123456"), service.ErrVerifyCodeMaxAttempts)
+
+	data, err := cache.GetVerificationCode(ctx, email)
+	require.NoError(t, err)
+	require.GreaterOrEqual(t, data.Attempts, 5)
+}
+
+func TestEmailCache_AttemptsResetOnNewCodeAndTTLFollowsCode(t *testing.T) {
+	cache, mr, _ := newMiniredisEmailCache(t)
+	ctx := context.Background()
+	email := "User@Example.com"
+
+	require.NoError(t, cache.SetVerificationCode(ctx, email, &service.VerificationCodeData{Code: "1"}, time.Minute))
+	n, err := cache.IncrVerificationCodeAttempts(ctx, email)
+	require.NoError(t, err)
+	require.Equal(t, 1, n)
+	require.Greater(t, mr.TTL(verifyCodeKey(email)+attemptsKeySuffix), time.Duration(0))
+
+	require.NoError(t, cache.SetVerificationCode(ctx, email, &service.VerificationCodeData{Code: "2"}, time.Minute))
+	data, err := cache.GetVerificationCode(ctx, email)
+	require.NoError(t, err)
+	require.Equal(t, 0, data.Attempts)
+
+	require.NoError(t, cache.DeleteVerificationCode(ctx, email))
+	_, err = cache.IncrVerificationCodeAttempts(ctx, email)
+	require.Error(t, err)
+	require.False(t, mr.Exists(verifyCodeKey(email)+attemptsKeySuffix))
+}
+
+func TestEmailCache_PasswordResetTokenHashedAndSingleUse(t *testing.T) {
+	cache, mr, _ := newMiniredisEmailCache(t)
+	ctx := context.Background()
+	email := "reset@example.com"
+
+	svc := service.NewEmailService(nil, cache)
+
+	// Seed the token the same way SendPasswordResetEmail does (hash only).
+	token, err := svc.GeneratePasswordResetToken()
+	require.NoError(t, err)
+	sum := sha256.Sum256([]byte(token))
+	require.NoError(t, cache.SetPasswordResetToken(ctx, email, &service.PasswordResetTokenData{
+		Token: hex.EncodeToString(sum[:]), CreatedAt: time.Now(),
+	}, 30*time.Minute))
+
+	raw, err := mr.Get(passwordResetKey(email))
+	require.NoError(t, err)
+	require.False(t, strings.Contains(raw, token), "plaintext token must not be stored")
+
+	require.NoError(t, svc.VerifyPasswordResetToken(ctx, email, token))
+	require.ErrorIs(t, svc.ConsumePasswordResetToken(ctx, email, "wrong"), service.ErrInvalidResetToken)
+
+	const workers = 30
+	var ok atomic.Int32
+	var wg sync.WaitGroup
+	for i := 0; i < workers; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			if svc.ConsumePasswordResetToken(ctx, email, token) == nil {
+				ok.Add(1)
+			}
+		}()
+	}
+	wg.Wait()
+	require.Equal(t, int32(1), ok.Load())
+	require.False(t, mr.Exists(passwordResetKey(email)))
+}
+
+func TestEmailCache_ConsumePasswordResetTokenMismatchKeepsToken(t *testing.T) {
+	cache, mr, _ := newMiniredisEmailCache(t)
+	ctx := context.Background()
+	email := "keep@example.com"
+	require.NoError(t, cache.SetPasswordResetToken(ctx, email, &service.PasswordResetTokenData{Token: "abc"}, time.Minute))
+
+	ok, err := cache.ConsumePasswordResetToken(ctx, email, "xyz")
+	require.NoError(t, err)
+	require.False(t, ok)
+	require.True(t, mr.Exists(passwordResetKey(email)))
+
+	ok, err = cache.ConsumePasswordResetToken(ctx, email, "abc")
+	require.NoError(t, err)
+	require.True(t, ok)
+	ok, err = cache.ConsumePasswordResetToken(ctx, email, "abc")
+	require.NoError(t, err)
+	require.False(t, ok)
+}