package repository import ( "context" "crypto/sha256" "encoding/hex" "errors" "fmt" "strconv" "strings" "time" "github.com/Wei-Shaw/sub2api/internal/service" "github.com/redis/go-redis/v9" ) const stickySessionPrefix = "sticky_session:" const liveCallPrefix = "live:call:" type gatewayCache struct { rdb *redis.Client } func NewGatewayCache(rdb *redis.Client) service.GatewayCache { return &gatewayCache{rdb: rdb} } // buildSessionKey 构建 session key,包含 groupID 实现分组隔离 // 格式: sticky_session:{groupID}:{sessionHash} func buildSessionKey(groupID int64, sessionHash string) string { return fmt.Sprintf("%s%d:%s", stickySessionPrefix, groupID, sessionHash) } func (c *gatewayCache) GetSessionAccountID(ctx context.Context, groupID int64, sessionHash string) (int64, error) { key := buildSessionKey(groupID, sessionHash) accountID, err := c.rdb.Get(ctx, key).Int64() if err != nil { if errors.Is(err, redis.Nil) { return 0, service.ErrStickySessionNotFound } return 0, err } return accountID, nil } func (c *gatewayCache) SetSessionAccountID(ctx context.Context, groupID int64, sessionHash string, accountID int64, ttl time.Duration) error { key := buildSessionKey(groupID, sessionHash) return c.rdb.Set(ctx, key, accountID, ttl).Err() } func (c *gatewayCache) RefreshSessionTTL(ctx context.Context, groupID int64, sessionHash string, ttl time.Duration) error { key := buildSessionKey(groupID, sessionHash) return c.rdb.Expire(ctx, key, ttl).Err() } // DeleteSessionAccountID 删除粘性会话与账号的绑定关系。 // 当检测到绑定的账号不可用(如状态错误、禁用、不可调度等)时调用, // 以便下次请求能够重新选择可用账号。 // // DeleteSessionAccountID removes the sticky session binding for the given session. // Called when the bound account becomes unavailable (e.g., error status, disabled, // or unschedulable), allowing subsequent requests to select a new available account. func (c *gatewayCache) DeleteSessionAccountID(ctx context.Context, groupID int64, sessionHash string) error { key := buildSessionKey(groupID, sessionHash) return c.rdb.Del(ctx, key).Err() } const ( grokVideoPendingBillingPrefix = "grok_video_pending:" grokVideoBilledPrefix = "grok_video_billed:" ) func (c *gatewayCache) SetGrokVideoPendingBilling(ctx context.Context, key string, payload []byte, ttl time.Duration) error { if c == nil || c.rdb == nil { return errors.New("gateway cache unavailable") } key = strings.TrimSpace(key) if key == "" || len(payload) == 0 { return errors.New("invalid grok video pending billing payload") } if ttl <= 0 { ttl = 24 * time.Hour } return c.rdb.Set(ctx, grokVideoPendingBillingPrefix+key, payload, ttl).Err() } func (c *gatewayCache) GetGrokVideoPendingBilling(ctx context.Context, key string) ([]byte, error) { if c == nil || c.rdb == nil { return nil, errors.New("gateway cache unavailable") } key = strings.TrimSpace(key) if key == "" { return nil, errors.New("invalid grok video pending billing key") } val, err := c.rdb.Get(ctx, grokVideoPendingBillingPrefix+key).Bytes() if err != nil { if errors.Is(err, redis.Nil) { return nil, nil } return nil, err } return val, nil } func (c *gatewayCache) ClaimGrokVideoBilled(ctx context.Context, key string, ttl time.Duration) (bool, error) { if c == nil || c.rdb == nil { return false, errors.New("gateway cache unavailable") } key = strings.TrimSpace(key) if key == "" { return false, errors.New("invalid grok video billed key") } if ttl <= 0 { ttl = 48 * time.Hour } return c.rdb.SetNX(ctx, grokVideoBilledPrefix+key, "1", ttl).Result() } func (c *gatewayCache) ReleaseGrokVideoBilled(ctx context.Context, key string) error { if c == nil || c.rdb == nil { return errors.New("gateway cache unavailable") } key = strings.TrimSpace(key) if key == "" { return errors.New("invalid grok video billed key") } return c.rdb.Del(ctx, grokVideoBilledPrefix+key).Err() } // Compile-time assertion: gatewayCache must implement CyberSessionBlockStore. var _ service.CyberSessionBlockStore = (*gatewayCache)(nil) var _ service.LiveCallStore = (*gatewayCache)(nil) const reasoningContentPrefix = "reasoning_content:" // reasoningContentDefaultTTL 是 reasoning 缓存的默认过期时间。Codex 会话可能 // 跨多天恢复,取 7 天;调用方传入非正 TTL 时兜底。 const reasoningContentDefaultTTL = 7 * 24 * time.Hour // SetReasoningContent 按 reasoning item id 缓存 reasoning 全文。 // itemID 或 content 为空时直接返回 nil(无可缓存内容,属正常情况而非错误)。 func (c *gatewayCache) SetReasoningContent(ctx context.Context, itemID string, content string, ttl time.Duration) error { if c == nil || c.rdb == nil { return errors.New("gateway cache unavailable") } itemID = strings.TrimSpace(itemID) if itemID == "" || content == "" { return nil } if ttl <= 0 { ttl = reasoningContentDefaultTTL } return c.rdb.Set(ctx, reasoningContentPrefix+itemID, content, ttl).Err() } // GetReasoningContent 返回缓存的 reasoning 全文;未命中返回 // service.ErrReasoningContentNotFound。 func (c *gatewayCache) GetReasoningContent(ctx context.Context, itemID string) (string, error) { if c == nil || c.rdb == nil { return "", errors.New("gateway cache unavailable") } itemID = strings.TrimSpace(itemID) if itemID == "" { return "", service.ErrReasoningContentNotFound } val, err := c.rdb.Get(ctx, reasoningContentPrefix+itemID).Result() if err != nil { if errors.Is(err, redis.Nil) { return "", service.ErrReasoningContentNotFound } return "", err } return val, nil } const cyberSessionBlockPrefix = "cyber_session_block:" // SetCyberSessionBlocked 把被 cyber_policy 命中的会话写入屏蔽表(TTL 自动过期)。 // 存储值 "1" 作为存在标记(IsCyberSessionBlocked 只检查 key 是否存在,不读值)。 func (c *gatewayCache) SetCyberSessionBlocked(ctx context.Context, key string, ttl time.Duration) error { return c.rdb.Set(ctx, cyberSessionBlockPrefix+key, "1", ttl).Err() } // IsCyberSessionBlocked 查询会话是否在屏蔽表中。 func (c *gatewayCache) IsCyberSessionBlocked(ctx context.Context, key string) (bool, error) { n, err := c.rdb.Exists(ctx, cyberSessionBlockPrefix+key).Result() if err != nil { return false, err } return n > 0, nil } var claimLiveControllerScript = redis.NewScript(` local key = KEYS[1] local target = ARGV[1] local owner = ARGV[2] local current = redis.call('HGET', key, 'controller') if current == false or current == 'closed' then return 0 end if target == 'observer' and current ~= 'pending' then return 0 end if target == 'proxy' and current ~= 'pending' and current ~= 'observer' and (current ~= 'proxy' or redis.call('HGET', key, 'controller_owner') ~= owner) then return 0 end redis.call('HSET', key, 'controller', target, 'controller_owner', owner) return 1 `) var markLiveCallClosedScript = redis.NewScript(` local key = KEYS[1] if redis.call('EXISTS', key) == 0 then return 0 end if redis.call('HGET', key, 'controller') == 'closed' then return 0 end redis.call('HSET', key, 'controller', 'closed', 'controller_owner', '') redis.call('EXPIRE', key, ARGV[1]) return 1 `) var releaseLiveControllerScript = redis.NewScript(` local key = KEYS[1] if redis.call('HGET', key, 'controller') ~= 'proxy' or redis.call('HGET', key, 'controller_owner') ~= ARGV[1] then return 0 end redis.call('HSET', key, 'controller', 'pending', 'controller_owner', '') return 1 `) func liveCallKey(callHash string) string { return liveCallPrefix + callHash } func HashLiveCallID(callID string) string { sum := sha256.Sum256([]byte(callID)) return hex.EncodeToString(sum[:]) } func (c *gatewayCache) SaveLiveCall(ctx context.Context, record *service.LiveCallRecord, ttl time.Duration) error { if record == nil || record.CallHash == "" || record.CallID == "" { return fmt.Errorf("invalid live call record") } values := map[string]any{ "call_id": record.CallID, "account_id": record.AccountID, "api_key_id": record.APIKeyID, "user_id": record.UserID, "group_id": record.GroupID, "subscription_id": record.SubscriptionID, "lease_id": record.LeaseID, "model": record.Model, "created_at": record.CreatedAt.UnixMilli(), "expires_at": record.ExpiresAt.UnixMilli(), "controller": record.Controller, "controller_owner": record.ControllerOwner, "user_agent": record.UserAgent, "ip_address": record.IPAddress, "inbound_endpoint": record.InboundEndpoint, "attestation": record.AttestationCiphertext, } key := liveCallKey(record.CallHash) pipe := c.rdb.TxPipeline() pipe.HSet(ctx, key, values) pipe.Expire(ctx, key, ttl) _, err := pipe.Exec(ctx) return err } func (c *gatewayCache) GetLiveCall(ctx context.Context, callHash string) (*service.LiveCallRecord, error) { values, err := c.rdb.HGetAll(ctx, liveCallKey(callHash)).Result() if err != nil { return nil, err } if len(values) == 0 { return nil, service.ErrLiveCallNotFound } parseInt := func(field string) int64 { value, _ := strconv.ParseInt(values[field], 10, 64) return value } createdAt := time.UnixMilli(parseInt("created_at")) expiresAt := time.UnixMilli(parseInt("expires_at")) return &service.LiveCallRecord{ CallID: values["call_id"], CallHash: callHash, AccountID: parseInt("account_id"), APIKeyID: parseInt("api_key_id"), UserID: parseInt("user_id"), GroupID: parseInt("group_id"), SubscriptionID: parseInt("subscription_id"), LeaseID: values["lease_id"], Model: values["model"], CreatedAt: createdAt, ExpiresAt: expiresAt, Controller: values["controller"], ControllerOwner: values["controller_owner"], UserAgent: values["user_agent"], IPAddress: values["ip_address"], InboundEndpoint: values["inbound_endpoint"], AttestationCiphertext: values["attestation"], }, nil } func (c *gatewayCache) ClaimLiveController(ctx context.Context, callHash, controller, owner string) (bool, error) { result, err := claimLiveControllerScript.Run(ctx, c.rdb, []string{liveCallKey(callHash)}, controller, owner).Int() return result == 1, err } func (c *gatewayCache) GetLiveController(ctx context.Context, callHash string) (string, error) { value, err := c.rdb.HGet(ctx, liveCallKey(callHash), "controller").Result() if err == redis.Nil { return "", service.ErrLiveCallNotFound } return value, err } func (c *gatewayCache) ReleaseLiveController(ctx context.Context, callHash, owner string) (bool, error) { result, err := releaseLiveControllerScript.Run(ctx, c.rdb, []string{liveCallKey(callHash)}, owner).Int() return result == 1, err } func (c *gatewayCache) MarkLiveCallClosed(ctx context.Context, callHash string, ttl time.Duration) (bool, error) { result, err := markLiveCallClosedScript.Run(ctx, c.rdb, []string{liveCallKey(callHash)}, int64(ttl.Seconds())).Int() return result == 1, err }