feat: 添加 Redis 客户端和 Token 黑名单模块

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
zerosaturation 2026-05-14 17:05:47 +08:00
parent 17674e3d54
commit 848908775a

View File

@ -0,0 +1,137 @@
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()
}