package database import ( "context" "crypto/sha256" "encoding/json" "fmt" "time" "github.com/redis/go-redis/v9" ) const ( BlacklistKeyPrefix = "blacklist:token:" ) // RedisClient Redis 客户端单例 var RedisClient *redis.Client // Config Redis 配置 type RedisConfig struct { Host string Port int Password string DB int } // InitRedis 初始化 Redis 连接 func InitRedis(cfg RedisConfig) error { RedisClient = redis.NewClient(&redis.Options{ Addr: fmt.Sprintf("%s:%d", cfg.Host, cfg.Port), Password: cfg.Password, DB: cfg.DB, }) ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() if err := RedisClient.Ping(ctx).Err(); err != nil { return fmt.Errorf("failed to connect to redis: %w", err) } return nil } // CloseRedis 关闭 Redis 连接 func CloseRedis() error { if RedisClient != nil { return RedisClient.Close() } return nil } // GetRedis 获取 Redis 客户端实例 func GetRedis() *redis.Client { return RedisClient } // RedisHealthCheck 健康检查 func RedisHealthCheck() error { if RedisClient == nil { return fmt.Errorf("redis client is not initialized") } ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() return RedisClient.Ping(ctx).Err() } // BlacklistEntry 黑名单条目 type BlacklistEntry struct { UserID int64 `json:"user_id"` Reason string `json:"reason"` } // tokenToHash 将 Token 转换为 SHA256 哈希作为 Key func tokenToHash(token string) string { hash := sha256.Sum256([]byte(token)) return fmt.Sprintf("%x", hash) } // AddToBlacklist 添加 Token 到黑名单 func AddToBlacklist(ctx context.Context, token string, userID int64, banReason string, ttl time.Duration) error { if token == "" { return fmt.Errorf("token is empty") } if RedisClient == nil { return fmt.Errorf("redis client is not initialized") } key := BlacklistKeyPrefix + tokenToHash(token) entry := BlacklistEntry{UserID: userID, Reason: banReason} value, err := json.Marshal(entry) if err != nil { return fmt.Errorf("failed to marshal blacklist entry: %w", err) } return RedisClient.Set(ctx, key, value, ttl).Err() } // IsBlacklisted 检查 Token 是否在黑名单 func IsBlacklisted(ctx context.Context, token string) (bool, int64, string, error) { if token == "" { return false, 0, "", nil } if RedisClient == nil { return false, 0, "", fmt.Errorf("redis client is not initialized") } key := BlacklistKeyPrefix + tokenToHash(token) value, err := RedisClient.Get(ctx, key).Result() if err == redis.Nil { return false, 0, "", nil } if err != nil { return false, 0, "", err } var entry BlacklistEntry if err := json.Unmarshal([]byte(value), &entry); err != nil { return false, 0, "", fmt.Errorf("failed to unmarshal blacklist entry: %w", err) } return true, entry.UserID, entry.Reason, nil } // RemoveFromBlacklist 从黑名单移除 Token(用于解封) func RemoveFromBlacklist(ctx context.Context, token string) error { if token == "" { return nil } if RedisClient == nil { return fmt.Errorf("redis client is not initialized") } key := BlacklistKeyPrefix + tokenToHash(token) return RedisClient.Del(ctx, key).Err() }