150 lines
3.8 KiB
Go
150 lines
3.8 KiB
Go
package cacherepo
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"github.com/go-redis/redis/v8"
|
|
"github.com/zeromicro/go-zero/core/logx"
|
|
"mingyang-admin-app-rpc/internal/svc"
|
|
"time"
|
|
)
|
|
|
|
type CacheRepository struct {
|
|
ctx context.Context
|
|
svcCtx *svc.ServiceContext
|
|
logx.Logger
|
|
}
|
|
|
|
func NewCacheRepository(ctx context.Context, svcCtx *svc.ServiceContext) *CacheRepository {
|
|
return &CacheRepository{
|
|
ctx: ctx,
|
|
svcCtx: svcCtx,
|
|
Logger: logx.WithContext(ctx),
|
|
}
|
|
}
|
|
|
|
// SetVerificationCode 设置验证码
|
|
func (r *CacheRepository) SetVerificationCode(ctx context.Context, key, code string, expiry time.Duration) error {
|
|
data := VerificationCodeData{
|
|
Code: code,
|
|
Attempts: 0,
|
|
CreatedAt: time.Now(),
|
|
}
|
|
|
|
jsonData, err := json.Marshal(data)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
return r.svcCtx.Redis.Set(ctx, key, jsonData, expiry).Err()
|
|
}
|
|
|
|
// GetVerificationCode 获取验证码
|
|
func (r *CacheRepository) GetVerificationCode(ctx context.Context, key string) (*VerificationCodeData, error) {
|
|
data, err := r.svcCtx.Redis.Get(ctx, key).Result()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
var codeData VerificationCodeData
|
|
if err := json.Unmarshal([]byte(data), &codeData); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return &codeData, nil
|
|
}
|
|
|
|
// IncrementVerificationAttempts 增加验证尝试次数
|
|
func (r *CacheRepository) IncrementVerificationAttempts(ctx context.Context, key string) error {
|
|
data, err := r.GetVerificationCode(ctx, key)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if data == nil {
|
|
return fmt.Errorf("verification code not found")
|
|
}
|
|
|
|
data.Attempts++
|
|
jsonData, err := json.Marshal(data)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
// 重新设置过期时间
|
|
ttl := r.svcCtx.Redis.TTL(ctx, key).Val()
|
|
return r.svcCtx.Redis.Set(ctx, key, jsonData, ttl).Err()
|
|
}
|
|
|
|
// DeleteVerificationCode 删除验证码
|
|
func (r *CacheRepository) DeleteVerificationCode(ctx context.Context, key string) error {
|
|
return r.svcCtx.Redis.Del(ctx, key).Err()
|
|
}
|
|
|
|
// SetPasswordResetToken 设置密码重置令牌
|
|
func (r *CacheRepository) SetPasswordResetToken(ctx context.Context, key, token string, expiry time.Duration) error {
|
|
return r.svcCtx.Redis.Set(ctx, key, token, expiry).Err()
|
|
}
|
|
|
|
// GetPasswordResetToken 获取密码重置令牌
|
|
func (r *CacheRepository) GetPasswordResetToken(ctx context.Context, key string) (string, error) {
|
|
return r.svcCtx.Redis.Get(ctx, key).Result()
|
|
}
|
|
|
|
// DeletePasswordResetToken 删除密码重置令牌
|
|
func (r *CacheRepository) DeletePasswordResetToken(ctx context.Context, key string) error {
|
|
return r.svcCtx.Redis.Del(ctx, key).Err()
|
|
}
|
|
|
|
// SetSession 设置会话
|
|
func (r *CacheRepository) SetSession(ctx context.Context, key string, data interface{}, expiry time.Duration) error {
|
|
jsonData, err := json.Marshal(data)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
return r.svcCtx.Redis.Set(ctx, key, jsonData, expiry).Err()
|
|
}
|
|
|
|
// GetSession 获取会话
|
|
func (r *CacheRepository) GetSession(ctx context.Context, key string, result interface{}) error {
|
|
data, err := r.svcCtx.Redis.Get(ctx, key).Result()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
return json.Unmarshal([]byte(data), result)
|
|
}
|
|
|
|
// DeleteSession 删除会话
|
|
func (r *CacheRepository) DeleteSession(ctx context.Context, key string) error {
|
|
return r.svcCtx.Redis.Del(ctx, key).Err()
|
|
}
|
|
|
|
// RateLimit 限流
|
|
func (r *CacheRepository) RateLimit(ctx context.Context, key string, limit int, window time.Duration) (bool, error) {
|
|
count, err := r.svcCtx.Redis.Get(ctx, key).Int()
|
|
if err != nil && !errors.Is(err, redis.Nil) {
|
|
return false, err
|
|
}
|
|
|
|
if count >= limit {
|
|
return false, nil
|
|
}
|
|
|
|
pipe := r.svcCtx.Redis.Pipeline()
|
|
pipe.Incr(ctx, key)
|
|
pipe.Expire(ctx, key, window)
|
|
|
|
_, err = pipe.Exec(ctx)
|
|
return err == nil, err
|
|
}
|
|
|
|
// VerificationCodeData 数据模型
|
|
type VerificationCodeData struct {
|
|
Code string `json:"code"`
|
|
Attempts int `json:"attempts"`
|
|
CreatedAt time.Time `json:"created_at"`
|
|
}
|