hyapp-server/services/user-service/internal/service/auth/risk_cache.go

package auth

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"time"

	"github.com/redis/go-redis/v9"
	"hyapp/pkg/appcode"
	authdomain "hyapp/services/user-service/internal/domain/auth"
)

type RedisIPDecisionCache struct {
	client *redis.Client
}

var setRevokedSessionMaxTTLScript = redis.NewScript(`
local current_ttl = redis.call('PTTL', KEYS[1])
local requested_ttl = tonumber(ARGV[2])
if current_ttl == -2 then
  redis.call('PSETEX', KEYS[1], requested_ttl, ARGV[1])
  return requested_ttl
end
if current_ttl == -1 then
  return current_ttl
end
if current_ttl < requested_ttl then
  redis.call('PSETEX', KEYS[1], requested_ttl, ARGV[1])
  return requested_ttl
end
return current_ttl
`)

func NewRedisIPDecisionCache(client *redis.Client) *RedisIPDecisionCache {
	return &RedisIPDecisionCache{client: client}
}

func NewLoginRiskRedisClient(ctx context.Context, addr string, password string, db int) (*redis.Client, error) {
	client := redis.NewClient(&redis.Options{
		Addr:     strings.TrimSpace(addr),
		Password: password,
		DB:       db,
	})
	if err := client.Ping(ctx).Err(); err != nil {
		_ = client.Close()
		return nil, err
	}
	return client, nil
}

func (c *RedisIPDecisionCache) GetIPDecision(ctx context.Context, appCode string, clientIPHash string) (IPDecision, bool, error) {
	raw, err := c.client.Get(ctx, loginRiskIPDecisionKey(appCode, clientIPHash)).Result()
	if errors.Is(err, redis.Nil) {
		return IPDecision{}, false, nil
	}
	if err != nil {
		return IPDecision{}, false, err
	}
	var decision IPDecision
	if err := json.Unmarshal([]byte(raw), &decision); err != nil {
		return IPDecision{}, false, err
	}
	return decision, true, nil
}

func (c *RedisIPDecisionCache) SetIPDecision(ctx context.Context, appCode string, clientIPHash string, decision IPDecision, ttl time.Duration) error {
	raw, err := json.Marshal(decision)
	if err != nil {
		return err
	}
	return c.client.Set(ctx, loginRiskIPDecisionKey(appCode, clientIPHash), raw, ttl).Err()
}

func (c *RedisIPDecisionCache) SetRevokedSession(ctx context.Context, appCode string, sessionID string, reason string, ttl time.Duration) error {
	if c == nil || c.client == nil {
		return errors.New("session denylist redis client is unavailable")
	}
	ttlMs := ttl.Milliseconds()
	if ttlMs <= 0 {
		return fmt.Errorf("session denylist TTL must be positive")
	}
	// 多来源会并发撤销同一 sid（rotation/reuse/logout/risk）；普通 SET 会让较短旧任务覆盖较长 deny_until。
	// Lua 原子比较 PTTL，只允许延长绝对保护时间；已有永久 key 也绝不被改短。
	return setRevokedSessionMaxTTLScript.Run(ctx, c.client, []string{revokedSessionKey(appCode, sessionID)}, strings.TrimSpace(reason), ttlMs).Err()
}

func (c *RedisIPDecisionCache) SessionDenylistEpoch(ctx context.Context) (string, error) {
	if c == nil || c.client == nil {
		return "", errors.New("session denylist redis client is unavailable")
	}
	info, err := c.client.Info(ctx, "server").Result()
	if err != nil {
		return "", err
	}
	for _, line := range strings.Split(info, "\n") {
		key, value, found := strings.Cut(strings.TrimSpace(line), ":")
		if found && key == "run_id" && strings.TrimSpace(value) != "" {
			return strings.TrimSpace(value), nil
		}
	}
	return "", errors.New("redis run_id is unavailable")
}

func (c *RedisIPDecisionCache) GetIPWhitelist(ctx context.Context, appCode string) ([]string, bool, error) {
	raw, err := c.client.Get(ctx, loginRiskIPWhitelistKey(appCode)).Result()
	if errors.Is(err, redis.Nil) {
		return nil, false, nil
	}
	if err != nil {
		return nil, false, err
	}
	var ips []string
	if err := json.Unmarshal([]byte(raw), &ips); err != nil {
		return nil, false, err
	}
	return ips, true, nil
}

func (c *RedisIPDecisionCache) SetIPWhitelist(ctx context.Context, appCode string, ips []string, ttl time.Duration) error {
	raw, err := json.Marshal(ips)
	if err != nil {
		return err
	}
	return c.client.Set(ctx, loginRiskIPWhitelistKey(appCode), raw, ttl).Err()
}

func (c *RedisIPDecisionCache) GetUserWhitelist(ctx context.Context, appCode string) ([]int64, bool, error) {
	raw, err := c.client.Get(ctx, loginRiskUserWhitelistKey(appCode)).Result()
	if errors.Is(err, redis.Nil) {
		return nil, false, nil
	}
	if err != nil {
		return nil, false, err
	}
	var userIDs []int64
	if err := json.Unmarshal([]byte(raw), &userIDs); err != nil {
		return nil, false, err
	}
	return userIDs, true, nil
}

func (c *RedisIPDecisionCache) SetUserWhitelist(ctx context.Context, appCode string, userIDs []int64, ttl time.Duration) error {
	raw, err := json.Marshal(userIDs)
	if err != nil {
		return err
	}
	return c.client.Set(ctx, loginRiskUserWhitelistKey(appCode), raw, ttl).Err()
}

func (c *RedisIPDecisionCache) GetRegisterRiskConfig(ctx context.Context, appCode string) (authdomain.RegisterRiskConfig, bool, error) {
	raw, err := c.client.Get(ctx, registerRiskConfigKey(appCode)).Result()
	if errors.Is(err, redis.Nil) {
		return authdomain.RegisterRiskConfig{}, false, nil
	}
	if err != nil {
		return authdomain.RegisterRiskConfig{}, false, err
	}
	var config authdomain.RegisterRiskConfig
	if err := json.Unmarshal([]byte(raw), &config); err != nil {
		return authdomain.RegisterRiskConfig{}, false, err
	}
	return config, true, nil
}

func (c *RedisIPDecisionCache) SetRegisterRiskConfig(ctx context.Context, config authdomain.RegisterRiskConfig, ttl time.Duration) error {
	raw, err := json.Marshal(config)
	if err != nil {
		return err
	}
	return c.client.Set(ctx, registerRiskConfigKey(config.AppCode), raw, ttl).Err()
}

func loginRiskIPDecisionKey(appCode string, clientIPHash string) string {
	return "login_risk:ip:" + appcode.Normalize(appCode) + ":" + strings.TrimSpace(clientIPHash)
}

func loginRiskIPWhitelistKey(appCode string) string {
	return "login_risk:ip_whitelist:" + appcode.Normalize(appCode)
}

func loginRiskUserWhitelistKey(appCode string) string {
	return "login_risk:user_whitelist:" + appcode.Normalize(appCode)
}

func revokedSessionKey(appCode string, sessionID string) string {
	return "auth:revoked_session:" + appcode.Normalize(appCode) + ":" + strings.TrimSpace(sessionID)
}

func registerRiskConfigKey(appCode string) string {
	return "auth:risk_config:register:" + appcode.Normalize(appCode)
}