| 12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187 |
- 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")
- 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
- 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 {
- 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
- DeleteCreateAttemptCount(ctx context.Context, userID 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
- }
- // 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)
- }
- }
- // 创建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.AuthCacheInvalidationKey())
- 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
- // 如果状态改变,清除Redis缓存
- if s.cache != nil {
- _ = s.cache.DeleteCreateAttemptCount(ctx, apiKey.UserID)
- }
- }
- // 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.AuthCacheInvalidationKey())
- 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)
- }
- // 删除成功后再清理缓存,避免"缓存已清但删除失败"的竞态。
- if s.cache != nil {
- _ = s.cache.DeleteCreateAttemptCount(ctx, userID)
- }
- 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.AuthCacheInvalidationKey())
- }
- 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)
- }
|