redeem_cache.go 1.8 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374
  1. package repository
  2. import (
  3. "context"
  4. "fmt"
  5. "time"
  6. "github.com/Wei-Shaw/sub2api/internal/service"
  7. "github.com/redis/go-redis/v9"
  8. )
  9. const (
  10. redeemRateLimitKeyPrefix = "redeem:ratelimit:v2:"
  11. redeemLockKeyPrefix = "redeem:lock:"
  12. redeemRateLimitWindow = 10 * time.Minute
  13. )
  14. var incrementRedeemAttemptScript = redis.NewScript(`
  15. local current = redis.call('INCR', KEYS[1])
  16. local ttl = redis.call('PTTL', KEYS[1])
  17. if current == 1 or ttl == -1 then
  18. redis.call('PEXPIRE', KEYS[1], ARGV[1])
  19. end
  20. return current
  21. `)
  22. // redeemRateLimitKey generates the Redis key for redeem attempt rate limiting.
  23. func redeemRateLimitKey(userID int64) string {
  24. return fmt.Sprintf("%s%d", redeemRateLimitKeyPrefix, userID)
  25. }
  26. // redeemLockKey generates the Redis key for redeem code locking.
  27. func redeemLockKey(code string) string {
  28. return redeemLockKeyPrefix + code
  29. }
  30. type redeemCache struct {
  31. rdb *redis.Client
  32. }
  33. func NewRedeemCache(rdb *redis.Client) service.RedeemCache {
  34. return &redeemCache{rdb: rdb}
  35. }
  36. func (c *redeemCache) GetRedeemAttemptCount(ctx context.Context, userID int64) (int, error) {
  37. key := redeemRateLimitKey(userID)
  38. count, err := c.rdb.Get(ctx, key).Int()
  39. if err == redis.Nil {
  40. return 0, nil
  41. }
  42. return count, err
  43. }
  44. func (c *redeemCache) IncrementRedeemAttemptCount(ctx context.Context, userID int64) error {
  45. key := redeemRateLimitKey(userID)
  46. return incrementRedeemAttemptScript.Run(
  47. ctx,
  48. c.rdb,
  49. []string{key},
  50. redeemRateLimitWindow.Milliseconds(),
  51. ).Err()
  52. }
  53. func (c *redeemCache) AcquireRedeemLock(ctx context.Context, code string, ttl time.Duration) (bool, error) {
  54. key := redeemLockKey(code)
  55. return c.rdb.SetNX(ctx, key, 1, ttl).Result()
  56. }
  57. func (c *redeemCache) ReleaseRedeemLock(ctx context.Context, code string) error {
  58. key := redeemLockKey(code)
  59. return c.rdb.Del(ctx, key).Err()
  60. }