Files
feeeei 26b31a83e4 fix(gemini): 显式缓存修改串行化、失败回滚,确定金额费用按 Key / 订阅预留
同一缓存的并发 PATCH 可能以与上游提交不同的顺序写回本地到期时间;
确定金额的存储费只预留用户余额,并发续期可越过 API Key 总额度与订阅额度;
上游已修改而本地写回失败时没有可恢复的状态。

- 新增迁移 245:缓存记录增加修改意图(mutation_token / mutation_lease_until)。
  PATCH 在调用上游前登记意图并取得 2 分钟租约,条件为没有进行中的修改且到期时间
  未变;上游修改、本地写回与计费均由持有者完成,并发的修改返回 409 ABORTED
- 上游结果未知、本地写回失败或进程退出时修改意图保留且不计费;租约过期后由下一次
  修改或后台每分钟的恢复流程把上游有效期回滚到修改前的到期时间(已到期则删除上游
  缓存与本地记录)。上游明确拒绝或请求未发出时直接清除意图
- 修改 / 删除上游缓存的单次调用超时改为 30 秒;请求发出前的失败包装为
  ErrGeminiCachedContentRequestNotSent
- CheckChargeAffordable 与 ReserveInflightCharge 合并为 ReserveFixedCharge:在途预留
  新增订阅、API Key 两个维度(Redis billing:inflight_scope:*),先登记本笔预留再读取
  库中权威用量(Key quota_used 与限速窗口、订阅日 / 周 / 月用量)叠加同维度在途合计
  校验;预留随计费任务交接,落库后释放
- GetRateLimitData 一并读取 quota_used
2026-10-12 17:14:56 +08:00

1219 lines
40 KiB
Go
Raw Permalink Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
package service
import (
"context"
"crypto/rand"
"encoding/hex"
"fmt"
"html"
"math"
"sort"
"strconv"
"strings"
"sync"
"sync/atomic"
"time"
"github.com/Wei-Shaw/sub2api/internal/config"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
"github.com/Wei-Shaw/sub2api/internal/pkg/ip"
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
"github.com/Wei-Shaw/sub2api/internal/pkg/timezone"
"github.com/dgraph-io/ristretto"
"golang.org/x/sync/singleflight"
)
var (
ErrAPIKeyNotFound = infraerrors.NotFound("API_KEY_NOT_FOUND", "api key not found")
ErrGroupNotAllowed = infraerrors.Forbidden("GROUP_NOT_ALLOWED", "user is not allowed to bind this group")
ErrAPIKeyExists = infraerrors.Conflict("API_KEY_EXISTS", "api key already exists")
ErrAPIKeyTooShort = infraerrors.BadRequest("API_KEY_TOO_SHORT", "api key must be at least 16 characters")
ErrAPIKeyInvalidChars = infraerrors.BadRequest("API_KEY_INVALID_CHARS", "api key can only contain letters, numbers, underscores, and hyphens")
ErrAPIKeyRateLimited = infraerrors.TooManyRequests("API_KEY_RATE_LIMITED", "too many failed attempts, please try again later")
ErrAPIKeyCreateLimited = infraerrors.TooManyRequests("API_KEY_CREATE_RATE_LIMITED", "too many api keys created recently, please try again later")
ErrAPIKeyCountExceeded = infraerrors.Forbidden("API_KEY_COUNT_EXCEEDED", "api key count limit reached, please delete unused keys first")
ErrAPIKeyAuthOverloaded = infraerrors.ServiceUnavailable("API_KEY_AUTH_OVERLOADED", "api key authentication is temporarily overloaded")
ErrInvalidIPPattern = infraerrors.BadRequest("INVALID_IP_PATTERN", "invalid IP or CIDR pattern")
// ErrAPIKeyExpired = infraerrors.Forbidden("API_KEY_EXPIRED", "api key has expired")
ErrAPIKeyExpired = infraerrors.Forbidden("API_KEY_EXPIRED", "api key 已过期")
// ErrAPIKeyQuotaExhausted = infraerrors.TooManyRequests("API_KEY_QUOTA_EXHAUSTED", "api key quota exhausted")
ErrAPIKeyQuotaExhausted = infraerrors.TooManyRequests("API_KEY_QUOTA_EXHAUSTED", "api key 额度已用完")
// Rate limit errors
ErrAPIKeyRateLimit5hExceeded = infraerrors.TooManyRequests("API_KEY_RATE_5H_EXCEEDED", "api key 5小时限额已用完")
ErrAPIKeyRateLimit1dExceeded = infraerrors.TooManyRequests("API_KEY_RATE_1D_EXCEEDED", "api key 日限额已用完")
ErrAPIKeyRateLimit7dExceeded = infraerrors.TooManyRequests("API_KEY_RATE_7D_EXCEEDED", "api key 7天限额已用完")
)
const (
MaxAPIKeyCredentialBytes = 128
defaultAuthLookupConcurrency = 64
defaultNegativeAuthCacheSize = 16384
apiKeyMaxErrorsPerHour = 20
apiKeyCreateCountWindow = time.Hour
apiKeyLastUsedMinTouch = 30 * time.Second
apiKeySortCurrentConcurrency = "current_concurrency"
// DB 写失败后的短退避,避免请求路径持续同步重试造成写风暴与高延迟。
apiKeyLastUsedFailBackoff = 5 * time.Second
)
// APIKeyUpdateFields 声明 APIKeyRepository.Update 允许写回的列。
//
// 与 UserUpdateFields 同理:api_keys 的用量列由计费热路径原子递增
// (IncrementQuotaUsed / IncrementRateLimitUsage 的 quota_used、usage_5h/1d/7d),
// 若编辑 Key 时无条件整行回写,并发累计的配额与限流计数就会被旧快照覆盖。
// 因此调用方必须显式声明要改的列。
type APIKeyUpdateFields struct {
Name bool
Status bool
Quota bool
GroupID bool
ExpiresAt bool
// QuotaUsed 仅供"重置配额用量"路径声明;常规计费走 IncrementQuotaUsed。
QuotaUsed bool
// RateLimits 覆盖 rate_limit_5h / _1d / _7d 三个阈值。
RateLimits bool
// RateLimitUsage 覆盖 usage_5h/_1d/_7d 与三个窗口起点,
// 仅供"重置限流用量"路径声明;常规计费走 IncrementRateLimitUsage。
RateLimitUsage bool
// IPRules 覆盖 ip_whitelist 与 ip_blacklist。
IPRules bool
}
// IsEmpty 报告该次 Update 是否不写任何列。
func (f APIKeyUpdateFields) IsEmpty() bool {
return f == APIKeyUpdateFields{}
}
type APIKeyRepository interface {
Create(ctx context.Context, key *APIKey) error
GetByID(ctx context.Context, id int64) (*APIKey, error)
// GetKeyAndOwnerID 仅获取 API Key 的 key 与所有者 ID,用于删除等轻量场景
GetKeyAndOwnerID(ctx context.Context, id int64) (string, int64, error)
GetByKey(ctx context.Context, key string) (*APIKey, error)
// GetByKeyForAuth 认证专用查询,返回最小字段集
GetByKeyForAuth(ctx context.Context, key string) (*APIKey, error)
// Update 只写 fields 中显式声明的列,其余列保持库中当前值。
Update(ctx context.Context, key *APIKey, fields APIKeyUpdateFields) error
Delete(ctx context.Context, id int64) error
// DeleteWithAudit keeps the legacy interface name for rolling-upgrade compatibility.
// Implementations must tombstone the key and soft-delete it atomically without
// retaining the deleted credential material.
DeleteWithAudit(ctx context.Context, id int64) error
ListByUserID(ctx context.Context, userID int64, params pagination.PaginationParams, filters APIKeyListFilters) ([]APIKey, *pagination.PaginationResult, error)
VerifyOwnership(ctx context.Context, userID int64, apiKeyIDs []int64) ([]int64, error)
CountByUserID(ctx context.Context, userID int64) (int64, error)
ExistsByKey(ctx context.Context, key string) (bool, error)
ListByGroupID(ctx context.Context, groupID int64, params pagination.PaginationParams) ([]APIKey, *pagination.PaginationResult, error)
SearchAPIKeys(ctx context.Context, userID int64, keyword string, limit int) ([]APIKey, error)
ClearGroupIDByGroupID(ctx context.Context, groupID int64) (int64, error)
// UpdateGroupIDByUserAndGroup 将用户下绑定 oldGroupID 的所有 Key 迁移到 newGroupID
UpdateGroupIDByUserAndGroup(ctx context.Context, userID, oldGroupID, newGroupID int64) (int64, error)
CountByGroupID(ctx context.Context, groupID int64) (int64, error)
ListKeysByUserID(ctx context.Context, userID int64) ([]string, error)
ListKeysByGroupID(ctx context.Context, groupID int64) ([]string, error)
// Quota methods
IncrementQuotaUsed(ctx context.Context, id int64, amount float64) (float64, error)
UpdateLastUsed(ctx context.Context, id int64, usedAt time.Time) error
// Rate limit methods
IncrementRateLimitUsage(ctx context.Context, id int64, cost float64) error
ResetRateLimitWindows(ctx context.Context, id int64) error
GetRateLimitData(ctx context.Context, id int64) (*APIKeyRateLimitData, error)
}
type apiKeyAllByUserIDLister interface {
ListAllByUserID(ctx context.Context, userID int64, filters APIKeyListFilters) ([]APIKey, error)
}
// APIKeyRateLimitData holds rate limit usage and window state for an API key.
type APIKeyRateLimitData struct {
// QuotaUsed 是总额度已用量(quota_used)。
QuotaUsed float64
Usage5h float64
Usage1d float64
Usage7d float64
Window5hStart *time.Time
Window1dStart *time.Time
Window7dStart *time.Time
}
// EffectiveUsage5h returns the 5h window usage, or 0 if the window has expired.
func (d *APIKeyRateLimitData) EffectiveUsage5h() float64 {
if IsWindowExpired(d.Window5hStart, RateLimitWindow5h) {
return 0
}
return d.Usage5h
}
// EffectiveUsage1d returns the 1d window usage, or 0 if the window has expired.
func (d *APIKeyRateLimitData) EffectiveUsage1d() float64 {
if IsWindowExpired(d.Window1dStart, RateLimitWindow1d) {
return 0
}
return d.Usage1d
}
// EffectiveUsage7d returns the 7d window usage, or 0 if the window has expired.
func (d *APIKeyRateLimitData) EffectiveUsage7d() float64 {
if IsWindowExpired(d.Window7dStart, RateLimitWindow7d) {
return 0
}
return d.Usage7d
}
// APIKeyQuotaUsageState captures the latest quota fields after an atomic quota update.
// It is intentionally small so repositories can return it from a single SQL statement.
type APIKeyQuotaUsageState struct {
QuotaUsed float64
Quota float64
Key string
Status string
}
// APIKeyCache defines cache operations for API key service
type APIKeyCache interface {
GetCreateAttemptCount(ctx context.Context, userID int64) (int, error)
IncrementCreateAttemptCount(ctx context.Context, userID int64) error
IncrementCreateCount(ctx context.Context, userID int64, window time.Duration) (int64, error)
IncrementDailyUsage(ctx context.Context, apiKey string) error
SetDailyUsageExpiry(ctx context.Context, apiKey string, ttl time.Duration) error
GetAuthCache(ctx context.Context, key string) (*APIKeyAuthCacheEntry, error)
SetAuthCache(ctx context.Context, key string, entry *APIKeyAuthCacheEntry, ttl time.Duration) error
DeleteAuthCache(ctx context.Context, key string) error
// Pub/Sub for L1 cache invalidation across instances
PublishAuthCacheInvalidation(ctx context.Context, cacheKey string) error
SubscribeAuthCacheInvalidation(ctx context.Context, handler func(cacheKey string)) error
}
type authCacheSubscriptionReadyKey struct{}
func withAuthCacheSubscriptionReady(ctx context.Context, ready func()) context.Context {
return context.WithValue(ctx, authCacheSubscriptionReadyKey{}, ready)
}
// NotifyAuthCacheSubscriptionReady lets cache implementations report that the
// server acknowledged the subscription without widening the public cache API.
func NotifyAuthCacheSubscriptionReady(ctx context.Context) {
if ready, ok := ctx.Value(authCacheSubscriptionReadyKey{}).(func()); ok && ready != nil {
ready()
}
}
// APIKeyAuthCacheInvalidator 提供认证缓存失效能力
type APIKeyAuthCacheInvalidator interface {
InvalidateAuthCacheByKey(ctx context.Context, key string)
InvalidateAuthCacheByUserID(ctx context.Context, userID int64)
InvalidateAuthCacheByGroupID(ctx context.Context, groupID int64)
}
// CreateAPIKeyRequest 创建API Key请求
type CreateAPIKeyRequest struct {
Name string `json:"name"`
GroupID *int64 `json:"group_id"`
CustomKey *string `json:"custom_key"` // 可选的自定义key
IPWhitelist []string `json:"ip_whitelist"` // IP 白名单
IPBlacklist []string `json:"ip_blacklist"` // IP 黑名单
// Quota fields
Quota float64 `json:"quota"` // Quota limit in USD (0 = unlimited)
ExpiresInDays *int `json:"expires_in_days"` // Days until expiry (nil = never expires)
// Rate limit fields (0 = unlimited)
RateLimit5h float64 `json:"rate_limit_5h"`
RateLimit1d float64 `json:"rate_limit_1d"`
RateLimit7d float64 `json:"rate_limit_7d"`
}
// UpdateAPIKeyRequest 更新API Key请求
type UpdateAPIKeyRequest struct {
Name *string `json:"name"`
GroupID *int64 `json:"group_id"`
Status *string `json:"status"`
IPWhitelist *[]string `json:"ip_whitelist"` // IP 白名单(nil 不修改,空数组清空)
IPBlacklist *[]string `json:"ip_blacklist"` // IP 黑名单(nil 不修改,空数组清空)
// Quota fields
Quota *float64 `json:"quota"` // Quota limit in USD (nil = no change, 0 = unlimited)
ExpiresAt *time.Time `json:"expires_at"` // Expiration time (nil = no change)
ClearExpiration bool `json:"-"` // Clear expiration (internal use)
ResetQuota *bool `json:"reset_quota"` // Reset quota_used to 0
// Rate limit fields (nil = no change, 0 = unlimited)
RateLimit5h *float64 `json:"rate_limit_5h"`
RateLimit1d *float64 `json:"rate_limit_1d"`
RateLimit7d *float64 `json:"rate_limit_7d"`
ResetRateLimitUsage *bool `json:"reset_rate_limit_usage"` // Reset all usage counters to 0
}
func validateAPIKeyLimit(v float64) error {
if math.IsNaN(v) || math.IsInf(v, 0) || v < 0 {
return infraerrors.BadRequest("API_KEY_LIMIT_INVALID", "API key limits must be finite and non-negative")
}
return nil
}
func validateCreateAPIKeyRequest(req CreateAPIKeyRequest) error {
for _, v := range []float64{req.Quota, req.RateLimit5h, req.RateLimit1d, req.RateLimit7d} {
if err := validateAPIKeyLimit(v); err != nil {
return err
}
}
if req.ExpiresInDays != nil && *req.ExpiresInDays <= 0 {
return infraerrors.BadRequest("API_KEY_EXPIRY_INVALID", "expires_in_days must be greater than zero")
}
return nil
}
func validateUpdateAPIKeyRequest(req UpdateAPIKeyRequest) error {
for _, v := range []*float64{req.Quota, req.RateLimit5h, req.RateLimit1d, req.RateLimit7d} {
if v != nil {
if err := validateAPIKeyLimit(*v); err != nil {
return err
}
}
}
return nil
}
// APIKeyService API Key服务
// RateLimitCacheInvalidator invalidates rate limit cache entries on manual reset.
type RateLimitCacheInvalidator interface {
InvalidateAPIKeyRateLimit(ctx context.Context, keyID int64) error
}
type APIKeyService struct {
apiKeyRepo APIKeyRepository
userRepo UserRepository
groupRepo GroupRepository
userSubRepo UserSubscriptionRepository
userGroupRateRepo UserGroupRateRepository
cache APIKeyCache
rateLimitCacheInvalid RateLimitCacheInvalidator // optional: invalidate Redis rate limit cache
concurrencyService *ConcurrencyService
cfg *config.Config
authCacheL1 *ristretto.Cache
authNegativeCacheL1 *ristretto.Cache
authCfg apiKeyAuthCacheConfig
authGroup singleflight.Group
authLookupSlots chan struct{}
authLookupTotal atomic.Uint64
authLookupRejected atomic.Uint64
authLookupInFlight atomic.Int64
invalidAuthAbuse *invalidAuthAbuseLimiter
authInvalidationStart sync.Once
authInvalidationStop sync.Once
authInvalidationCancel context.CancelFunc
authInvalidationWG sync.WaitGroup
authInvalidationConnected atomic.Bool
authInvalidationFailures atomic.Uint64
lastUsedTouchL1 sync.Map // keyID -> nextAllowedAt(time.Time)
lastUsedTouchSF singleflight.Group
}
type APIKeyAuthLookupMetrics struct {
Total uint64 `json:"total"`
Rejected uint64 `json:"rejected"`
InFlight int64 `json:"in_flight"`
Capacity int `json:"capacity"`
}
func (s *APIKeyService) AuthLookupMetrics() APIKeyAuthLookupMetrics {
if s == nil {
return APIKeyAuthLookupMetrics{}
}
return APIKeyAuthLookupMetrics{
Total: s.authLookupTotal.Load(),
Rejected: s.authLookupRejected.Load(),
InFlight: s.authLookupInFlight.Load(),
Capacity: cap(s.authLookupSlots),
}
}
// NewAPIKeyService 创建API Key服务实例
func NewAPIKeyService(
apiKeyRepo APIKeyRepository,
userRepo UserRepository,
groupRepo GroupRepository,
userSubRepo UserSubscriptionRepository,
userGroupRateRepo UserGroupRateRepository,
cache APIKeyCache,
cfg *config.Config,
) *APIKeyService {
svc := &APIKeyService{
apiKeyRepo: apiKeyRepo,
userRepo: userRepo,
groupRepo: groupRepo,
userSubRepo: userSubRepo,
userGroupRateRepo: userGroupRateRepo,
cache: cache,
cfg: cfg,
}
svc.initAuthCache(cfg)
lookupConcurrency := defaultAuthLookupConcurrency
if cfg != nil && cfg.APIKeyAuth.LookupConcurrency > 0 {
lookupConcurrency = cfg.APIKeyAuth.LookupConcurrency
}
svc.authLookupSlots = make(chan struct{}, lookupConcurrency)
svc.invalidAuthAbuse = newInvalidAuthAbuseLimiter(cfg)
return svc
}
// SetRateLimitCacheInvalidator sets the optional rate limit cache invalidator.
// Called after construction (e.g. in wire) to avoid circular dependencies.
func (s *APIKeyService) SetRateLimitCacheInvalidator(inv RateLimitCacheInvalidator) {
s.rateLimitCacheInvalid = inv
}
func (s *APIKeyService) SetConcurrencyService(concurrencyService *ConcurrencyService) {
s.concurrencyService = concurrencyService
}
func (s *APIKeyService) compileAPIKeyIPRules(apiKey *APIKey) {
if apiKey == nil {
return
}
apiKey.CompiledIPWhitelist = ip.CompileIPRules(apiKey.IPWhitelist)
apiKey.CompiledIPBlacklist = ip.CompileIPRules(apiKey.IPBlacklist)
}
// GenerateKey 生成随机API Key
func (s *APIKeyService) GenerateKey() (string, error) {
// 生成32字节随机数据
bytes := make([]byte, 32)
if _, err := rand.Read(bytes); err != nil {
return "", fmt.Errorf("generate random bytes: %w", err)
}
// 转换为十六进制字符串并添加前缀
prefix := s.cfg.Default.APIKeyPrefix
if prefix == "" {
prefix = "sk-"
}
key := prefix + hex.EncodeToString(bytes)
return key, nil
}
// ValidateCustomKey 验证自定义API Key格式
func (s *APIKeyService) ValidateCustomKey(key string) error {
// 检查长度
if len(key) < 16 {
return ErrAPIKeyTooShort
}
// 检查字符:只允许字母、数字、下划线、连字符
for _, c := range key {
if (c >= 'a' && c <= 'z') ||
(c >= 'A' && c <= 'Z') ||
(c >= '0' && c <= '9') ||
c == '_' || c == '-' {
continue
}
return ErrAPIKeyInvalidChars
}
return nil
}
// checkAPIKeyRateLimit 检查用户创建自定义Key的错误次数是否超限
func (s *APIKeyService) checkAPIKeyRateLimit(ctx context.Context, userID int64) error {
if s.cache == nil {
return nil
}
count, err := s.cache.GetCreateAttemptCount(ctx, userID)
if err != nil {
// Redis 出错时不阻止用户操作
return nil
}
if count >= apiKeyMaxErrorsPerHour {
return ErrAPIKeyRateLimited
}
return nil
}
// checkAPIKeyCreateLimits 校验创建 API Key 的防滥用限制(对自定义与自动生成的 Key 一视同仁)。
// 数量上限按未删除的 Key 计;创建次数按固定窗口累计,删除 Key 不返还次数,
// 以阻断"删除后反复新建"的循环。Redis 出错时与自定义 Key 限流一致,不阻止用户操作。
func (s *APIKeyService) checkAPIKeyCreateLimits(ctx context.Context, userID int64) error {
if s.cfg == nil {
return nil
}
if maxActive := s.cfg.APIKeyCreate.MaxActivePerUser; maxActive > 0 {
count, err := s.apiKeyRepo.CountByUserID(ctx, userID)
if err != nil {
return fmt.Errorf("count api keys: %w", err)
}
if count >= int64(maxActive) {
return ErrAPIKeyCountExceeded
}
}
if maxPerHour := s.cfg.APIKeyCreate.MaxPerUserPerHour; maxPerHour > 0 && s.cache != nil {
count, err := s.cache.IncrementCreateCount(ctx, userID, apiKeyCreateCountWindow)
if err != nil {
return nil
}
if count > int64(maxPerHour) {
return ErrAPIKeyCreateLimited
}
}
return nil
}
// incrementAPIKeyErrorCount 增加用户创建自定义Key的错误计数
func (s *APIKeyService) incrementAPIKeyErrorCount(ctx context.Context, userID int64) {
if s.cache == nil {
return
}
_ = s.cache.IncrementCreateAttemptCount(ctx, userID)
}
// canUserBindGroup 检查用户是否可以绑定指定分组
// 对于订阅类型分组:检查用户是否有有效订阅
// 对于标准类型分组:使用原有的 AllowedGroups 和 IsExclusive 逻辑
func (s *APIKeyService) canUserBindGroup(ctx context.Context, user *User, group *Group) bool {
// 订阅类型分组:需要有效订阅
if group.IsSubscriptionType() {
_, err := s.userSubRepo.GetActiveByUserIDAndGroupID(ctx, user.ID, group.ID)
return err == nil // 有有效订阅则允许
}
// 标准类型分组:使用原有逻辑
return user.CanBindGroup(group.ID, group.IsExclusive)
}
// Create 创建API Key
func (s *APIKeyService) Create(ctx context.Context, userID int64, req CreateAPIKeyRequest) (*APIKey, error) {
if err := validateCreateAPIKeyRequest(req); err != nil {
return nil, err
}
// 验证用户存在
user, err := s.userRepo.GetByID(ctx, userID)
if err != nil {
return nil, fmt.Errorf("get user: %w", err)
}
// 验证 IP 白名单格式
if len(req.IPWhitelist) > 0 {
if invalid := ip.ValidateIPPatterns(req.IPWhitelist); len(invalid) > 0 {
return nil, fmt.Errorf("%w: %v", ErrInvalidIPPattern, invalid)
}
}
// 验证 IP 黑名单格式
if len(req.IPBlacklist) > 0 {
if invalid := ip.ValidateIPPatterns(req.IPBlacklist); len(invalid) > 0 {
return nil, fmt.Errorf("%w: %v", ErrInvalidIPPattern, invalid)
}
}
// 验证分组权限(如果指定了分组)
if req.GroupID != nil {
group, err := s.groupRepo.GetByID(ctx, *req.GroupID)
if err != nil {
return nil, fmt.Errorf("get group: %w", err)
}
// 检查用户是否可以绑定该分组
if !s.canUserBindGroup(ctx, user, group) {
return nil, ErrGroupNotAllowed
}
}
var key string
// 判断是否使用自定义Key
if req.CustomKey != nil && *req.CustomKey != "" {
// 检查限流(仅对自定义key进行限流)
if err := s.checkAPIKeyRateLimit(ctx, userID); err != nil {
return nil, err
}
// 验证自定义Key格式
if err := s.ValidateCustomKey(*req.CustomKey); err != nil {
return nil, err
}
// 检查Key是否已存在
exists, err := s.apiKeyRepo.ExistsByKey(ctx, *req.CustomKey)
if err != nil {
return nil, fmt.Errorf("check key exists: %w", err)
}
if exists {
// Key已存在,增加错误计数
s.incrementAPIKeyErrorCount(ctx, userID)
return nil, ErrAPIKeyExists
}
key = *req.CustomKey
} else {
// 生成随机API Key
var err error
key, err = s.GenerateKey()
if err != nil {
return nil, fmt.Errorf("generate key: %w", err)
}
}
if err := s.checkAPIKeyCreateLimits(ctx, userID); err != nil {
return nil, err
}
// 创建API Key记录
apiKey := &APIKey{
UserID: userID,
Key: key,
Name: html.EscapeString(req.Name),
GroupID: req.GroupID,
Status: StatusActive,
IPWhitelist: req.IPWhitelist,
IPBlacklist: req.IPBlacklist,
Quota: req.Quota,
QuotaUsed: 0,
RateLimit5h: req.RateLimit5h,
RateLimit1d: req.RateLimit1d,
RateLimit7d: req.RateLimit7d,
}
// Set expiration time if specified
if req.ExpiresInDays != nil && *req.ExpiresInDays > 0 {
expiresAt := time.Now().AddDate(0, 0, *req.ExpiresInDays)
apiKey.ExpiresAt = &expiresAt
}
if err := s.apiKeyRepo.Create(ctx, apiKey); err != nil {
return nil, fmt.Errorf("create api key: %w", err)
}
s.InvalidateAuthCacheByKey(ctx, apiKey.Key)
s.compileAPIKeyIPRules(apiKey)
return apiKey, nil
}
// List 获取用户的API Key列表
func (s *APIKeyService) List(ctx context.Context, userID int64, params pagination.PaginationParams, filters APIKeyListFilters) ([]APIKey, *pagination.PaginationResult, error) {
if normalizedAPIKeySortBy(params.SortBy) == apiKeySortCurrentConcurrency {
return s.listByCurrentConcurrency(ctx, userID, params, filters)
}
keys, pagination, err := s.apiKeyRepo.ListByUserID(ctx, userID, params, filters)
if err != nil {
return nil, nil, fmt.Errorf("list api keys: %w", err)
}
s.fillCurrentConcurrency(ctx, keys)
return keys, pagination, nil
}
func (s *APIKeyService) listByCurrentConcurrency(ctx context.Context, userID int64, params pagination.PaginationParams, filters APIKeyListFilters) ([]APIKey, *pagination.PaginationResult, error) {
repo, ok := s.apiKeyRepo.(apiKeyAllByUserIDLister)
if !ok {
return nil, nil, fmt.Errorf("list api keys by current concurrency: repository does not support unpaginated API key listing")
}
keys, err := repo.ListAllByUserID(ctx, userID, filters)
if err != nil {
return nil, nil, fmt.Errorf("list api keys: %w", err)
}
s.fillCurrentConcurrency(ctx, keys)
sortAPIKeysByCurrentConcurrency(keys, params.NormalizedSortOrder(pagination.SortOrderDesc))
return paginateAPIKeys(keys, params), apiKeyPaginationResult(int64(len(keys)), params), nil
}
func normalizedAPIKeySortBy(sortBy string) string {
return strings.ToLower(strings.TrimSpace(sortBy))
}
func sortAPIKeysByCurrentConcurrency(keys []APIKey, sortOrder string) {
desc := sortOrder != pagination.SortOrderAsc
sort.SliceStable(keys, func(i, j int) bool {
if keys[i].CurrentConcurrency == keys[j].CurrentConcurrency {
if desc {
return keys[i].ID > keys[j].ID
}
return keys[i].ID < keys[j].ID
}
if desc {
return keys[i].CurrentConcurrency > keys[j].CurrentConcurrency
}
return keys[i].CurrentConcurrency < keys[j].CurrentConcurrency
})
}
func paginateAPIKeys(keys []APIKey, params pagination.PaginationParams) []APIKey {
if len(keys) == 0 {
return []APIKey{}
}
limit := params.Limit()
page := params.Page
if page < 1 {
page = 1
}
offset := (page - 1) * limit
if offset >= len(keys) {
return []APIKey{}
}
end := offset + limit
if end > len(keys) {
end = len(keys)
}
return keys[offset:end]
}
func apiKeyPaginationResult(total int64, params pagination.PaginationParams) *pagination.PaginationResult {
limit := params.Limit()
pages := int(total) / limit
if int(total)%limit > 0 {
pages++
}
return &pagination.PaginationResult{
Total: total,
Page: params.Page,
PageSize: limit,
Pages: pages,
}
}
func (s *APIKeyService) fillCurrentConcurrency(ctx context.Context, keys []APIKey) {
if s == nil || s.concurrencyService == nil || len(keys) == 0 {
return
}
ids := make([]int64, 0, len(keys))
for i := range keys {
if keys[i].ID > 0 {
ids = append(ids, keys[i].ID)
}
}
counts, err := s.concurrencyService.GetAPIKeyConcurrencyBatch(ctx, ids)
if err != nil {
return
}
for i := range keys {
keys[i].CurrentConcurrency = counts[keys[i].ID]
}
}
func (s *APIKeyService) currentConcurrencyForAPIKey(ctx context.Context, apiKeyID int64) int {
if s == nil || s.concurrencyService == nil || apiKeyID <= 0 {
return 0
}
counts, err := s.concurrencyService.GetAPIKeyConcurrencyBatch(ctx, []int64{apiKeyID})
if err != nil {
return 0
}
return counts[apiKeyID]
}
func (s *APIKeyService) VerifyOwnership(ctx context.Context, userID int64, apiKeyIDs []int64) ([]int64, error) {
if len(apiKeyIDs) == 0 {
return []int64{}, nil
}
validIDs, err := s.apiKeyRepo.VerifyOwnership(ctx, userID, apiKeyIDs)
if err != nil {
return nil, fmt.Errorf("verify api key ownership: %w", err)
}
return validIDs, nil
}
// GetByID 根据ID获取API Key
func (s *APIKeyService) GetByID(ctx context.Context, id int64) (*APIKey, error) {
apiKey, err := s.apiKeyRepo.GetByID(ctx, id)
if err != nil {
return nil, fmt.Errorf("get api key: %w", err)
}
s.compileAPIKeyIPRules(apiKey)
if apiKey != nil {
apiKey.CurrentConcurrency = s.currentConcurrencyForAPIKey(ctx, apiKey.ID)
}
return apiKey, nil
}
// GetByKey 根据Key字符串获取API Key(用于认证)
func (s *APIKeyService) GetByKey(ctx context.Context, key string) (*APIKey, error) {
if len(key) == 0 || len(key) > MaxAPIKeyCredentialBytes {
return nil, ErrAPIKeyNotFound
}
cacheKey := s.authCacheKey(key)
if entry, ok := s.getAuthCacheEntry(ctx, cacheKey); ok {
if apiKey, used, err := s.applyAuthCacheEntry(key, entry); used {
if err != nil {
return nil, fmt.Errorf("get api key: %w", err)
}
s.compileAPIKeyIPRules(apiKey)
return apiKey, nil
}
}
if s.authCfg.singleflight {
value, err, _ := s.authGroup.Do(cacheKey, func() (any, error) {
return s.loadAuthCacheEntry(ctx, key, cacheKey)
})
if err != nil {
return nil, err
}
entry, _ := value.(*APIKeyAuthCacheEntry)
if apiKey, used, err := s.applyAuthCacheEntry(key, entry); used {
if err != nil {
return nil, fmt.Errorf("get api key: %w", err)
}
s.compileAPIKeyIPRules(apiKey)
return apiKey, nil
}
} else {
entry, err := s.loadAuthCacheEntry(ctx, key, cacheKey)
if err != nil {
return nil, err
}
if apiKey, used, err := s.applyAuthCacheEntry(key, entry); used {
if err != nil {
return nil, fmt.Errorf("get api key: %w", err)
}
s.compileAPIKeyIPRules(apiKey)
return apiKey, nil
}
}
apiKey, err := s.lookupAPIKeyForAuth(ctx, key)
if err != nil {
return nil, fmt.Errorf("get api key: %w", err)
}
apiKey.Key = key
s.compileAPIKeyIPRules(apiKey)
return apiKey, nil
}
// Update 更新API Key
func (s *APIKeyService) Update(ctx context.Context, id int64, userID int64, req UpdateAPIKeyRequest) (*APIKey, error) {
if err := validateUpdateAPIKeyRequest(req); err != nil {
return nil, err
}
apiKey, err := s.apiKeyRepo.GetByID(ctx, id)
if err != nil {
return nil, fmt.Errorf("get api key: %w", err)
}
// 验证所有权
if apiKey.UserID != userID {
return nil, ErrInsufficientPerms
}
// 验证 IP 白名单格式
if req.IPWhitelist != nil && len(*req.IPWhitelist) > 0 {
if invalid := ip.ValidateIPPatterns(*req.IPWhitelist); len(invalid) > 0 {
return nil, fmt.Errorf("%w: %v", ErrInvalidIPPattern, invalid)
}
}
// 验证 IP 黑名单格式
if req.IPBlacklist != nil && len(*req.IPBlacklist) > 0 {
if invalid := ip.ValidateIPPatterns(*req.IPBlacklist); len(invalid) > 0 {
return nil, fmt.Errorf("%w: %v", ErrInvalidIPPattern, invalid)
}
}
// fields 只登记本次请求真正要改的列。quota_used 与 usage_5h/1d/7d 由计费热路径
// 原子递增,除非用户显式点了"重置",否则这里不用快照把它们写回去。
var fields APIKeyUpdateFields
// 下面若干分支会顺带把 Status 改回 active(配额扩容、清除过期等),
// 所以用原始值比对来决定是否写 status,而不是只看 req.Status。
originalStatus := apiKey.Status
// 更新字段
if req.Name != nil {
apiKey.Name = html.EscapeString(*req.Name)
fields.Name = true
}
if req.GroupID != nil {
// 验证分组权限
user, err := s.userRepo.GetByID(ctx, userID)
if err != nil {
return nil, fmt.Errorf("get user: %w", err)
}
group, err := s.groupRepo.GetByID(ctx, *req.GroupID)
if err != nil {
return nil, fmt.Errorf("get group: %w", err)
}
if !s.canUserBindGroup(ctx, user, group) {
return nil, ErrGroupNotAllowed
}
apiKey.GroupID = req.GroupID
fields.GroupID = true
}
if req.Status != nil {
apiKey.Status = *req.Status
fields.Status = true
}
// Update quota fields
if req.Quota != nil {
apiKey.Quota = *req.Quota
fields.Quota = true
// If quota now has room, or is changed to unlimited, reactivate exhausted keys.
if apiKey.Status == StatusAPIKeyQuotaExhausted && (*req.Quota <= 0 || *req.Quota > apiKey.QuotaUsed) {
apiKey.Status = StatusActive
}
}
if req.ResetQuota != nil && *req.ResetQuota {
apiKey.QuotaUsed = 0
fields.QuotaUsed = true
// If resetting quota and status was quota_exhausted, reactivate
if apiKey.Status == StatusAPIKeyQuotaExhausted {
apiKey.Status = StatusActive
}
}
if req.ClearExpiration {
apiKey.ExpiresAt = nil
fields.ExpiresAt = true
// If clearing expiry and status was expired, reactivate
if apiKey.Status == StatusAPIKeyExpired {
apiKey.Status = StatusActive
}
} else if req.ExpiresAt != nil {
apiKey.ExpiresAt = req.ExpiresAt
fields.ExpiresAt = true
// If extending expiry and status was expired, reactivate
if apiKey.Status == StatusAPIKeyExpired && time.Now().Before(*req.ExpiresAt) {
apiKey.Status = StatusActive
}
}
// 更新 IP 限制(nil 不修改,空数组清空设置)
if req.IPWhitelist != nil {
apiKey.IPWhitelist = *req.IPWhitelist
fields.IPRules = true
}
if req.IPBlacklist != nil {
apiKey.IPBlacklist = *req.IPBlacklist
fields.IPRules = true
}
// Update rate limit configuration
if req.RateLimit5h != nil {
apiKey.RateLimit5h = *req.RateLimit5h
fields.RateLimits = true
}
if req.RateLimit1d != nil {
apiKey.RateLimit1d = *req.RateLimit1d
fields.RateLimits = true
}
if req.RateLimit7d != nil {
apiKey.RateLimit7d = *req.RateLimit7d
fields.RateLimits = true
}
resetRateLimit := req.ResetRateLimitUsage != nil && *req.ResetRateLimitUsage
if resetRateLimit {
apiKey.Usage5h = 0
apiKey.Usage1d = 0
apiKey.Usage7d = 0
apiKey.Window5hStart = nil
apiKey.Window1dStart = nil
apiKey.Window7dStart = nil
fields.RateLimitUsage = true
}
// 上面的自动复活分支可能改了 status,这里统一登记。
if apiKey.Status != originalStatus {
fields.Status = true
}
if err := s.apiKeyRepo.Update(ctx, apiKey, fields); err != nil {
return nil, fmt.Errorf("update api key: %w", err)
}
s.InvalidateAuthCacheByKey(ctx, apiKey.Key)
s.compileAPIKeyIPRules(apiKey)
// Invalidate Redis rate limit cache so reset takes effect immediately
if resetRateLimit && s.rateLimitCacheInvalid != nil {
_ = s.rateLimitCacheInvalid.InvalidateAPIKeyRateLimit(ctx, apiKey.ID)
}
return apiKey, nil
}
// Delete 删除API Key
func (s *APIKeyService) Delete(ctx context.Context, id int64, userID int64) error {
key, ownerID, err := s.apiKeyRepo.GetKeyAndOwnerID(ctx, id)
if err != nil {
return fmt.Errorf("get api key: %w", err)
}
// 验证当前用户是否为该 API Key 的所有者
if ownerID != userID {
return ErrInsufficientPerms
}
// 事务内:写审计 + 软删除(tombstone)。
if err := s.apiKeyRepo.DeleteWithAudit(ctx, id); err != nil {
return fmt.Errorf("delete api key: %w", err)
}
// 删除成功后再清理缓存,避免"缓存已清但删除失败"的竞态。
// 注意:不清零创建相关计数,否则"删除后反复新建"即可绕过创建限流。
s.InvalidateAuthCacheByKey(ctx, key)
s.lastUsedTouchL1.Delete(id)
return nil
}
// ValidateKey 验证API Key是否有效(用于认证中间件)
func (s *APIKeyService) ValidateKey(ctx context.Context, key string) (*APIKey, *User, error) {
// 获取API Key
apiKey, err := s.GetByKey(ctx, key)
if err != nil {
return nil, nil, err
}
// 检查API Key状态
if !apiKey.IsActive() {
return nil, nil, infraerrors.Unauthorized("API_KEY_INACTIVE", "api key is not active")
}
// 获取用户信息
user, err := s.userRepo.GetByID(ctx, apiKey.UserID)
if err != nil {
return nil, nil, fmt.Errorf("get user: %w", err)
}
// 检查用户状态
if !user.IsActive() {
return nil, nil, ErrUserNotActive
}
return apiKey, user, nil
}
// TouchLastUsed 通过防抖更新 api_keys.last_used_at,减少高频写放大。
// 该操作为尽力而为,不应阻塞主请求链路。
func (s *APIKeyService) TouchLastUsed(ctx context.Context, keyID int64) error {
if keyID <= 0 {
return nil
}
now := time.Now()
if v, ok := s.lastUsedTouchL1.Load(keyID); ok {
if nextAllowedAt, ok := v.(time.Time); ok && now.Before(nextAllowedAt) {
return nil
}
}
_, err, _ := s.lastUsedTouchSF.Do(strconv.FormatInt(keyID, 10), func() (any, error) {
latest := time.Now()
if v, ok := s.lastUsedTouchL1.Load(keyID); ok {
if nextAllowedAt, ok := v.(time.Time); ok && latest.Before(nextAllowedAt) {
return nil, nil
}
}
if err := s.apiKeyRepo.UpdateLastUsed(ctx, keyID, latest); err != nil {
s.lastUsedTouchL1.Store(keyID, latest.Add(apiKeyLastUsedFailBackoff))
return nil, fmt.Errorf("touch api key last used: %w", err)
}
s.lastUsedTouchL1.Store(keyID, latest.Add(apiKeyLastUsedMinTouch))
return nil, nil
})
return err
}
// IncrementUsage 增加API Key使用次数(可选:用于统计)
func (s *APIKeyService) IncrementUsage(ctx context.Context, keyID int64) error {
// 使用Redis计数器
if s.cache != nil {
cacheKey := fmt.Sprintf("apikey:usage:%d:%s", keyID, timezone.Now().Format("2006-01-02"))
if err := s.cache.IncrementDailyUsage(ctx, cacheKey); err != nil {
return fmt.Errorf("increment usage: %w", err)
}
// 设置24小时过期
_ = s.cache.SetDailyUsageExpiry(ctx, cacheKey, 24*time.Hour)
}
return nil
}
// GetAvailableGroups 获取用户有权限绑定的分组列表
// 返回用户可以选择的分组:
// - 标准类型分组:公开的(非专属)或用户被明确允许的
// - 订阅类型分组:用户有有效订阅的
func (s *APIKeyService) GetAvailableGroups(ctx context.Context, userID int64) ([]Group, error) {
// 获取用户信息
user, err := s.userRepo.GetByID(ctx, userID)
if err != nil {
return nil, fmt.Errorf("get user: %w", err)
}
// 获取所有活跃分组
allGroups, err := s.groupRepo.ListActive(ctx)
if err != nil {
return nil, fmt.Errorf("list active groups: %w", err)
}
// 获取用户的所有有效订阅
activeSubscriptions, err := s.userSubRepo.ListActiveByUserID(ctx, userID)
if err != nil {
return nil, fmt.Errorf("list active subscriptions: %w", err)
}
// 构建订阅分组 ID 集合
subscribedGroupIDs := make(map[int64]bool)
for _, sub := range activeSubscriptions {
subscribedGroupIDs[sub.GroupID] = true
}
// 过滤出用户有权限的分组
availableGroups := make([]Group, 0)
for _, group := range allGroups {
if s.canUserBindGroupInternal(user, &group, subscribedGroupIDs) {
availableGroups = append(availableGroups, group)
}
}
return availableGroups, nil
}
// canUserBindGroupInternal 内部方法,检查用户是否可以绑定分组(使用预加载的订阅数据)
func (s *APIKeyService) canUserBindGroupInternal(user *User, group *Group, subscribedGroupIDs map[int64]bool) bool {
// 订阅类型分组:需要有效订阅
if group.IsSubscriptionType() {
return subscribedGroupIDs[group.ID]
}
// 标准类型分组:使用原有逻辑
return user.CanBindGroup(group.ID, group.IsExclusive)
}
func (s *APIKeyService) SearchAPIKeys(ctx context.Context, userID int64, keyword string, limit int) ([]APIKey, error) {
keys, err := s.apiKeyRepo.SearchAPIKeys(ctx, userID, keyword, limit)
if err != nil {
return nil, fmt.Errorf("search api keys: %w", err)
}
return keys, nil
}
// GetUserGroupVisibility 返回 user_allowed_groups 授权及有效订阅的分组 ID 集合,
// 以及该用户是否开启了公开分组限制。开启时公开分组的可见性也要落在该集合内。
//
// 与 GetAvailableGroups 的区别:这里保留普通授权分组的「橱窗」语义,不检查
// 分组是否活跃;有效订阅也授予对应分组的可见性。返回值恒非 nil。
func (s *APIKeyService) GetUserGroupVisibility(ctx context.Context, userID int64) (map[int64]struct{}, bool, error) {
user, err := s.userRepo.GetByID(ctx, userID)
if err != nil {
return nil, false, fmt.Errorf("get user: %w", err)
}
allowed := make(map[int64]struct{}, len(user.AllowedGroups))
for _, id := range user.AllowedGroups {
allowed[id] = struct{}{}
}
subscriptions, err := s.userSubRepo.ListActiveByUserID(ctx, userID)
if err != nil {
return nil, false, fmt.Errorf("list active subscriptions: %w", err)
}
for _, sub := range subscriptions {
allowed[sub.GroupID] = struct{}{}
}
return allowed, user.RestrictPublicGroups, nil
}
// GetUserGroupRates 获取用户的专属分组倍率配置
// 返回 map[groupID]rateMultiplier
func (s *APIKeyService) GetUserGroupRates(ctx context.Context, userID int64) (map[int64]float64, error) {
if s.userGroupRateRepo == nil {
return nil, nil
}
rates, err := s.userGroupRateRepo.GetByUserID(ctx, userID)
if err != nil {
return nil, fmt.Errorf("get user group rates: %w", err)
}
return rates, nil
}
// CheckAPIKeyQuotaAndExpiry checks if the API key is valid for use (not expired, quota not exhausted)
// Returns nil if valid, error if invalid
func (s *APIKeyService) CheckAPIKeyQuotaAndExpiry(apiKey *APIKey) error {
// Check expiration
if apiKey.IsExpired() {
return ErrAPIKeyExpired
}
// Check quota
if apiKey.IsQuotaExhausted() {
return ErrAPIKeyQuotaExhausted
}
return nil
}
// UpdateQuotaUsed updates the quota_used field after a request
// Also checks if quota is exhausted and updates status accordingly
func (s *APIKeyService) UpdateQuotaUsed(ctx context.Context, apiKeyID int64, cost float64) error {
if cost <= 0 {
return nil
}
type quotaStateReader interface {
IncrementQuotaUsedAndGetState(ctx context.Context, id int64, amount float64) (*APIKeyQuotaUsageState, error)
}
if repo, ok := s.apiKeyRepo.(quotaStateReader); ok {
state, err := repo.IncrementQuotaUsedAndGetState(ctx, apiKeyID, cost)
if err != nil {
return fmt.Errorf("increment quota used: %w", err)
}
if state != nil && state.Status == StatusAPIKeyQuotaExhausted && strings.TrimSpace(state.Key) != "" {
s.InvalidateAuthCacheByKey(ctx, state.Key)
}
return nil
}
// Use repository to atomically increment quota_used
newQuotaUsed, err := s.apiKeyRepo.IncrementQuotaUsed(ctx, apiKeyID, cost)
if err != nil {
return fmt.Errorf("increment quota used: %w", err)
}
// Check if quota is now exhausted and update status if needed
apiKey, err := s.apiKeyRepo.GetByID(ctx, apiKeyID)
if err != nil {
return nil // Don't fail the request, just log
}
// If quota is set and now exhausted, update status
if apiKey.Quota > 0 && newQuotaUsed >= apiKey.Quota {
apiKey.Status = StatusAPIKeyQuotaExhausted
// 只写 status:这条位于计费热路径,若整行回写会把刚刚原子递增的
// quota_used 与限流用量按快照覆盖掉。
if err := s.apiKeyRepo.Update(ctx, apiKey, APIKeyUpdateFields{Status: true}); err != nil {
return nil // Don't fail the request
}
// Invalidate cache so next request sees the new status
s.InvalidateAuthCacheByKey(ctx, apiKey.Key)
}
return nil
}
// GetRateLimitData returns rate limit usage and window state for an API key.
func (s *APIKeyService) GetRateLimitData(ctx context.Context, id int64) (*APIKeyRateLimitData, error) {
return s.apiKeyRepo.GetRateLimitData(ctx, id)
}
// UpdateRateLimitUsage atomically increments rate limit usage counters in the DB.
func (s *APIKeyService) UpdateRateLimitUsage(ctx context.Context, apiKeyID int64, cost float64) error {
if cost <= 0 {
return nil
}
return s.apiKeyRepo.IncrementRateLimitUsage(ctx, apiKeyID, cost)
}