email_cache.go 5.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177
  1. package repository
  2. import (
  3. "context"
  4. "encoding/json"
  5. "fmt"
  6. "strings"
  7. "time"
  8. "github.com/Wei-Shaw/sub2api/internal/service"
  9. "github.com/redis/go-redis/v9"
  10. )
  11. const (
  12. verifyCodeKeyPrefix = "verify_code:"
  13. notifyVerifyKeyPrefix = "notify_verify:"
  14. passwordResetKeyPrefix = "password_reset:"
  15. passwordResetSentAtKeyPrefix = "password_reset_sent:"
  16. notifyCodeUserRateKeyPrefix = "notify_code_user_rate:"
  17. )
  18. // verifyCodeKey generates the Redis key for email verification code.
  19. // Email is lowercased for case-insensitive consistency.
  20. func verifyCodeKey(email string) string {
  21. return verifyCodeKeyPrefix + strings.ToLower(email)
  22. }
  23. // notifyVerifyKey generates the Redis key for notify email verification code.
  24. // Email is lowercased to prevent case-sensitive key mismatch (the business layer
  25. // uses strings.EqualFold for comparison).
  26. func notifyVerifyKey(email string) string {
  27. return notifyVerifyKeyPrefix + strings.ToLower(email)
  28. }
  29. // passwordResetKey generates the Redis key for password reset token.
  30. func passwordResetKey(email string) string {
  31. return passwordResetKeyPrefix + strings.ToLower(email)
  32. }
  33. // passwordResetSentAtKey generates the Redis key for password reset email sent timestamp.
  34. func passwordResetSentAtKey(email string) string {
  35. return passwordResetSentAtKeyPrefix + strings.ToLower(email)
  36. }
  37. type emailCache struct {
  38. rdb *redis.Client
  39. }
  40. func NewEmailCache(rdb *redis.Client) service.EmailCache {
  41. return &emailCache{rdb: rdb}
  42. }
  43. func (c *emailCache) GetVerificationCode(ctx context.Context, email string) (*service.VerificationCodeData, error) {
  44. key := verifyCodeKey(email)
  45. val, err := c.rdb.Get(ctx, key).Result()
  46. if err != nil {
  47. return nil, err
  48. }
  49. var data service.VerificationCodeData
  50. if err := json.Unmarshal([]byte(val), &data); err != nil {
  51. return nil, err
  52. }
  53. return &data, nil
  54. }
  55. func (c *emailCache) SetVerificationCode(ctx context.Context, email string, data *service.VerificationCodeData, ttl time.Duration) error {
  56. key := verifyCodeKey(email)
  57. val, err := json.Marshal(data)
  58. if err != nil {
  59. return err
  60. }
  61. return c.rdb.Set(ctx, key, val, ttl).Err()
  62. }
  63. func (c *emailCache) DeleteVerificationCode(ctx context.Context, email string) error {
  64. key := verifyCodeKey(email)
  65. return c.rdb.Del(ctx, key).Err()
  66. }
  67. // Password reset token methods
  68. func (c *emailCache) GetPasswordResetToken(ctx context.Context, email string) (*service.PasswordResetTokenData, error) {
  69. key := passwordResetKey(email)
  70. val, err := c.rdb.Get(ctx, key).Result()
  71. if err != nil {
  72. return nil, err
  73. }
  74. var data service.PasswordResetTokenData
  75. if err := json.Unmarshal([]byte(val), &data); err != nil {
  76. return nil, err
  77. }
  78. return &data, nil
  79. }
  80. func (c *emailCache) SetPasswordResetToken(ctx context.Context, email string, data *service.PasswordResetTokenData, ttl time.Duration) error {
  81. key := passwordResetKey(email)
  82. val, err := json.Marshal(data)
  83. if err != nil {
  84. return err
  85. }
  86. return c.rdb.Set(ctx, key, val, ttl).Err()
  87. }
  88. func (c *emailCache) DeletePasswordResetToken(ctx context.Context, email string) error {
  89. key := passwordResetKey(email)
  90. return c.rdb.Del(ctx, key).Err()
  91. }
  92. // Password reset email cooldown methods
  93. func (c *emailCache) IsPasswordResetEmailInCooldown(ctx context.Context, email string) bool {
  94. key := passwordResetSentAtKey(email)
  95. exists, err := c.rdb.Exists(ctx, key).Result()
  96. return err == nil && exists > 0
  97. }
  98. func (c *emailCache) SetPasswordResetEmailCooldown(ctx context.Context, email string, ttl time.Duration) error {
  99. key := passwordResetSentAtKey(email)
  100. return c.rdb.Set(ctx, key, "1", ttl).Err()
  101. }
  102. // Notify email verification code methods
  103. func (c *emailCache) GetNotifyVerifyCode(ctx context.Context, email string) (*service.VerificationCodeData, error) {
  104. key := notifyVerifyKey(email)
  105. val, err := c.rdb.Get(ctx, key).Result()
  106. if err != nil {
  107. return nil, err
  108. }
  109. var data service.VerificationCodeData
  110. if err := json.Unmarshal([]byte(val), &data); err != nil {
  111. return nil, err
  112. }
  113. return &data, nil
  114. }
  115. func (c *emailCache) SetNotifyVerifyCode(ctx context.Context, email string, data *service.VerificationCodeData, ttl time.Duration) error {
  116. key := notifyVerifyKey(email)
  117. val, err := json.Marshal(data)
  118. if err != nil {
  119. return err
  120. }
  121. return c.rdb.Set(ctx, key, val, ttl).Err()
  122. }
  123. func (c *emailCache) DeleteNotifyVerifyCode(ctx context.Context, email string) error {
  124. key := notifyVerifyKey(email)
  125. return c.rdb.Del(ctx, key).Err()
  126. }
  127. // User-level rate limiting for notify email verification codes
  128. func notifyCodeUserRateKey(userID int64) string {
  129. return notifyCodeUserRateKeyPrefix + fmt.Sprintf("%d", userID)
  130. }
  131. func (c *emailCache) IncrNotifyCodeUserRate(ctx context.Context, userID int64, window time.Duration) (int64, error) {
  132. key := notifyCodeUserRateKey(userID)
  133. count, err := c.rdb.Incr(ctx, key).Result()
  134. if err != nil {
  135. return 0, err
  136. }
  137. // Always set TTL (idempotent) to avoid orphan keys if process crashes between INCR and EXPIRE.
  138. if err := c.rdb.Expire(ctx, key, window).Err(); err != nil {
  139. return count, fmt.Errorf("expire notify code rate key: %w", err)
  140. }
  141. return count, nil
  142. }
  143. func (c *emailCache) GetNotifyCodeUserRate(ctx context.Context, userID int64) (int64, error) {
  144. key := notifyCodeUserRateKey(userID)
  145. count, err := c.rdb.Get(ctx, key).Int64()
  146. if err != nil {
  147. return 0, err
  148. }
  149. return count, nil
  150. }