api_key_service.go 39 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187
  1. package service
  2. import (
  3. "context"
  4. "crypto/rand"
  5. "encoding/hex"
  6. "fmt"
  7. "html"
  8. "math"
  9. "sort"
  10. "strconv"
  11. "strings"
  12. "sync"
  13. "sync/atomic"
  14. "time"
  15. "github.com/Wei-Shaw/sub2api/internal/config"
  16. infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
  17. "github.com/Wei-Shaw/sub2api/internal/pkg/ip"
  18. "github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
  19. "github.com/Wei-Shaw/sub2api/internal/pkg/timezone"
  20. "github.com/dgraph-io/ristretto"
  21. "golang.org/x/sync/singleflight"
  22. )
  23. var (
  24. ErrAPIKeyNotFound = infraerrors.NotFound("API_KEY_NOT_FOUND", "api key not found")
  25. ErrGroupNotAllowed = infraerrors.Forbidden("GROUP_NOT_ALLOWED", "user is not allowed to bind this group")
  26. ErrAPIKeyExists = infraerrors.Conflict("API_KEY_EXISTS", "api key already exists")
  27. ErrAPIKeyTooShort = infraerrors.BadRequest("API_KEY_TOO_SHORT", "api key must be at least 16 characters")
  28. ErrAPIKeyInvalidChars = infraerrors.BadRequest("API_KEY_INVALID_CHARS", "api key can only contain letters, numbers, underscores, and hyphens")
  29. ErrAPIKeyRateLimited = infraerrors.TooManyRequests("API_KEY_RATE_LIMITED", "too many failed attempts, please try again later")
  30. ErrAPIKeyAuthOverloaded = infraerrors.ServiceUnavailable("API_KEY_AUTH_OVERLOADED", "api key authentication is temporarily overloaded")
  31. ErrInvalidIPPattern = infraerrors.BadRequest("INVALID_IP_PATTERN", "invalid IP or CIDR pattern")
  32. // ErrAPIKeyExpired = infraerrors.Forbidden("API_KEY_EXPIRED", "api key has expired")
  33. ErrAPIKeyExpired = infraerrors.Forbidden("API_KEY_EXPIRED", "api key 已过期")
  34. // ErrAPIKeyQuotaExhausted = infraerrors.TooManyRequests("API_KEY_QUOTA_EXHAUSTED", "api key quota exhausted")
  35. ErrAPIKeyQuotaExhausted = infraerrors.TooManyRequests("API_KEY_QUOTA_EXHAUSTED", "api key 额度已用完")
  36. // Rate limit errors
  37. ErrAPIKeyRateLimit5hExceeded = infraerrors.TooManyRequests("API_KEY_RATE_5H_EXCEEDED", "api key 5小时限额已用完")
  38. ErrAPIKeyRateLimit1dExceeded = infraerrors.TooManyRequests("API_KEY_RATE_1D_EXCEEDED", "api key 日限额已用完")
  39. ErrAPIKeyRateLimit7dExceeded = infraerrors.TooManyRequests("API_KEY_RATE_7D_EXCEEDED", "api key 7天限额已用完")
  40. )
  41. const (
  42. MaxAPIKeyCredentialBytes = 128
  43. defaultAuthLookupConcurrency = 64
  44. defaultNegativeAuthCacheSize = 16384
  45. apiKeyMaxErrorsPerHour = 20
  46. apiKeyLastUsedMinTouch = 30 * time.Second
  47. apiKeySortCurrentConcurrency = "current_concurrency"
  48. // DB 写失败后的短退避,避免请求路径持续同步重试造成写风暴与高延迟。
  49. apiKeyLastUsedFailBackoff = 5 * time.Second
  50. )
  51. // APIKeyUpdateFields 声明 APIKeyRepository.Update 允许写回的列。
  52. //
  53. // 与 UserUpdateFields 同理:api_keys 的用量列由计费热路径原子递增
  54. // (IncrementQuotaUsed / IncrementRateLimitUsage 的 quota_used、usage_5h/1d/7d),
  55. // 若编辑 Key 时无条件整行回写,并发累计的配额与限流计数就会被旧快照覆盖。
  56. // 因此调用方必须显式声明要改的列。
  57. type APIKeyUpdateFields struct {
  58. Name bool
  59. Status bool
  60. Quota bool
  61. GroupID bool
  62. ExpiresAt bool
  63. // QuotaUsed 仅供"重置配额用量"路径声明;常规计费走 IncrementQuotaUsed。
  64. QuotaUsed bool
  65. // RateLimits 覆盖 rate_limit_5h / _1d / _7d 三个阈值。
  66. RateLimits bool
  67. // RateLimitUsage 覆盖 usage_5h/_1d/_7d 与三个窗口起点,
  68. // 仅供"重置限流用量"路径声明;常规计费走 IncrementRateLimitUsage。
  69. RateLimitUsage bool
  70. // IPRules 覆盖 ip_whitelist 与 ip_blacklist。
  71. IPRules bool
  72. }
  73. // IsEmpty 报告该次 Update 是否不写任何列。
  74. func (f APIKeyUpdateFields) IsEmpty() bool {
  75. return f == APIKeyUpdateFields{}
  76. }
  77. type APIKeyRepository interface {
  78. Create(ctx context.Context, key *APIKey) error
  79. GetByID(ctx context.Context, id int64) (*APIKey, error)
  80. // GetKeyAndOwnerID 仅获取 API Key 的 key 与所有者 ID,用于删除等轻量场景
  81. GetKeyAndOwnerID(ctx context.Context, id int64) (string, int64, error)
  82. GetByKey(ctx context.Context, key string) (*APIKey, error)
  83. // GetByKeyForAuth 认证专用查询,返回最小字段集
  84. GetByKeyForAuth(ctx context.Context, key string) (*APIKey, error)
  85. // Update 只写 fields 中显式声明的列,其余列保持库中当前值。
  86. Update(ctx context.Context, key *APIKey, fields APIKeyUpdateFields) error
  87. Delete(ctx context.Context, id int64) error
  88. // DeleteWithAudit keeps the legacy interface name for rolling-upgrade compatibility.
  89. // Implementations must tombstone the key and soft-delete it atomically without
  90. // retaining the deleted credential material.
  91. DeleteWithAudit(ctx context.Context, id int64) error
  92. ListByUserID(ctx context.Context, userID int64, params pagination.PaginationParams, filters APIKeyListFilters) ([]APIKey, *pagination.PaginationResult, error)
  93. VerifyOwnership(ctx context.Context, userID int64, apiKeyIDs []int64) ([]int64, error)
  94. CountByUserID(ctx context.Context, userID int64) (int64, error)
  95. ExistsByKey(ctx context.Context, key string) (bool, error)
  96. ListByGroupID(ctx context.Context, groupID int64, params pagination.PaginationParams) ([]APIKey, *pagination.PaginationResult, error)
  97. SearchAPIKeys(ctx context.Context, userID int64, keyword string, limit int) ([]APIKey, error)
  98. ClearGroupIDByGroupID(ctx context.Context, groupID int64) (int64, error)
  99. // UpdateGroupIDByUserAndGroup 将用户下绑定 oldGroupID 的所有 Key 迁移到 newGroupID
  100. UpdateGroupIDByUserAndGroup(ctx context.Context, userID, oldGroupID, newGroupID int64) (int64, error)
  101. CountByGroupID(ctx context.Context, groupID int64) (int64, error)
  102. ListKeysByUserID(ctx context.Context, userID int64) ([]string, error)
  103. ListKeysByGroupID(ctx context.Context, groupID int64) ([]string, error)
  104. // Quota methods
  105. IncrementQuotaUsed(ctx context.Context, id int64, amount float64) (float64, error)
  106. UpdateLastUsed(ctx context.Context, id int64, usedAt time.Time) error
  107. // Rate limit methods
  108. IncrementRateLimitUsage(ctx context.Context, id int64, cost float64) error
  109. ResetRateLimitWindows(ctx context.Context, id int64) error
  110. GetRateLimitData(ctx context.Context, id int64) (*APIKeyRateLimitData, error)
  111. }
  112. type apiKeyAllByUserIDLister interface {
  113. ListAllByUserID(ctx context.Context, userID int64, filters APIKeyListFilters) ([]APIKey, error)
  114. }
  115. // APIKeyRateLimitData holds rate limit usage and window state for an API key.
  116. type APIKeyRateLimitData struct {
  117. Usage5h float64
  118. Usage1d float64
  119. Usage7d float64
  120. Window5hStart *time.Time
  121. Window1dStart *time.Time
  122. Window7dStart *time.Time
  123. }
  124. // EffectiveUsage5h returns the 5h window usage, or 0 if the window has expired.
  125. func (d *APIKeyRateLimitData) EffectiveUsage5h() float64 {
  126. if IsWindowExpired(d.Window5hStart, RateLimitWindow5h) {
  127. return 0
  128. }
  129. return d.Usage5h
  130. }
  131. // EffectiveUsage1d returns the 1d window usage, or 0 if the window has expired.
  132. func (d *APIKeyRateLimitData) EffectiveUsage1d() float64 {
  133. if IsWindowExpired(d.Window1dStart, RateLimitWindow1d) {
  134. return 0
  135. }
  136. return d.Usage1d
  137. }
  138. // EffectiveUsage7d returns the 7d window usage, or 0 if the window has expired.
  139. func (d *APIKeyRateLimitData) EffectiveUsage7d() float64 {
  140. if IsWindowExpired(d.Window7dStart, RateLimitWindow7d) {
  141. return 0
  142. }
  143. return d.Usage7d
  144. }
  145. // APIKeyQuotaUsageState captures the latest quota fields after an atomic quota update.
  146. // It is intentionally small so repositories can return it from a single SQL statement.
  147. type APIKeyQuotaUsageState struct {
  148. QuotaUsed float64
  149. Quota float64
  150. Key string
  151. Status string
  152. }
  153. // APIKeyCache defines cache operations for API key service
  154. type APIKeyCache interface {
  155. GetCreateAttemptCount(ctx context.Context, userID int64) (int, error)
  156. IncrementCreateAttemptCount(ctx context.Context, userID int64) error
  157. DeleteCreateAttemptCount(ctx context.Context, userID int64) error
  158. IncrementDailyUsage(ctx context.Context, apiKey string) error
  159. SetDailyUsageExpiry(ctx context.Context, apiKey string, ttl time.Duration) error
  160. GetAuthCache(ctx context.Context, key string) (*APIKeyAuthCacheEntry, error)
  161. SetAuthCache(ctx context.Context, key string, entry *APIKeyAuthCacheEntry, ttl time.Duration) error
  162. DeleteAuthCache(ctx context.Context, key string) error
  163. // Pub/Sub for L1 cache invalidation across instances
  164. PublishAuthCacheInvalidation(ctx context.Context, cacheKey string) error
  165. SubscribeAuthCacheInvalidation(ctx context.Context, handler func(cacheKey string)) error
  166. }
  167. type authCacheSubscriptionReadyKey struct{}
  168. func withAuthCacheSubscriptionReady(ctx context.Context, ready func()) context.Context {
  169. return context.WithValue(ctx, authCacheSubscriptionReadyKey{}, ready)
  170. }
  171. // NotifyAuthCacheSubscriptionReady lets cache implementations report that the
  172. // server acknowledged the subscription without widening the public cache API.
  173. func NotifyAuthCacheSubscriptionReady(ctx context.Context) {
  174. if ready, ok := ctx.Value(authCacheSubscriptionReadyKey{}).(func()); ok && ready != nil {
  175. ready()
  176. }
  177. }
  178. // APIKeyAuthCacheInvalidator 提供认证缓存失效能力
  179. type APIKeyAuthCacheInvalidator interface {
  180. InvalidateAuthCacheByKey(ctx context.Context, key string)
  181. InvalidateAuthCacheByUserID(ctx context.Context, userID int64)
  182. InvalidateAuthCacheByGroupID(ctx context.Context, groupID int64)
  183. }
  184. // CreateAPIKeyRequest 创建API Key请求
  185. type CreateAPIKeyRequest struct {
  186. Name string `json:"name"`
  187. GroupID *int64 `json:"group_id"`
  188. CustomKey *string `json:"custom_key"` // 可选的自定义key
  189. IPWhitelist []string `json:"ip_whitelist"` // IP 白名单
  190. IPBlacklist []string `json:"ip_blacklist"` // IP 黑名单
  191. // Quota fields
  192. Quota float64 `json:"quota"` // Quota limit in USD (0 = unlimited)
  193. ExpiresInDays *int `json:"expires_in_days"` // Days until expiry (nil = never expires)
  194. // Rate limit fields (0 = unlimited)
  195. RateLimit5h float64 `json:"rate_limit_5h"`
  196. RateLimit1d float64 `json:"rate_limit_1d"`
  197. RateLimit7d float64 `json:"rate_limit_7d"`
  198. }
  199. // UpdateAPIKeyRequest 更新API Key请求
  200. type UpdateAPIKeyRequest struct {
  201. Name *string `json:"name"`
  202. GroupID *int64 `json:"group_id"`
  203. Status *string `json:"status"`
  204. IPWhitelist *[]string `json:"ip_whitelist"` // IP 白名单(nil 不修改,空数组清空)
  205. IPBlacklist *[]string `json:"ip_blacklist"` // IP 黑名单(nil 不修改,空数组清空)
  206. // Quota fields
  207. Quota *float64 `json:"quota"` // Quota limit in USD (nil = no change, 0 = unlimited)
  208. ExpiresAt *time.Time `json:"expires_at"` // Expiration time (nil = no change)
  209. ClearExpiration bool `json:"-"` // Clear expiration (internal use)
  210. ResetQuota *bool `json:"reset_quota"` // Reset quota_used to 0
  211. // Rate limit fields (nil = no change, 0 = unlimited)
  212. RateLimit5h *float64 `json:"rate_limit_5h"`
  213. RateLimit1d *float64 `json:"rate_limit_1d"`
  214. RateLimit7d *float64 `json:"rate_limit_7d"`
  215. ResetRateLimitUsage *bool `json:"reset_rate_limit_usage"` // Reset all usage counters to 0
  216. }
  217. func validateAPIKeyLimit(v float64) error {
  218. if math.IsNaN(v) || math.IsInf(v, 0) || v < 0 {
  219. return infraerrors.BadRequest("API_KEY_LIMIT_INVALID", "API key limits must be finite and non-negative")
  220. }
  221. return nil
  222. }
  223. func validateCreateAPIKeyRequest(req CreateAPIKeyRequest) error {
  224. for _, v := range []float64{req.Quota, req.RateLimit5h, req.RateLimit1d, req.RateLimit7d} {
  225. if err := validateAPIKeyLimit(v); err != nil {
  226. return err
  227. }
  228. }
  229. if req.ExpiresInDays != nil && *req.ExpiresInDays <= 0 {
  230. return infraerrors.BadRequest("API_KEY_EXPIRY_INVALID", "expires_in_days must be greater than zero")
  231. }
  232. return nil
  233. }
  234. func validateUpdateAPIKeyRequest(req UpdateAPIKeyRequest) error {
  235. for _, v := range []*float64{req.Quota, req.RateLimit5h, req.RateLimit1d, req.RateLimit7d} {
  236. if v != nil {
  237. if err := validateAPIKeyLimit(*v); err != nil {
  238. return err
  239. }
  240. }
  241. }
  242. return nil
  243. }
  244. // APIKeyService API Key服务
  245. // RateLimitCacheInvalidator invalidates rate limit cache entries on manual reset.
  246. type RateLimitCacheInvalidator interface {
  247. InvalidateAPIKeyRateLimit(ctx context.Context, keyID int64) error
  248. }
  249. type APIKeyService struct {
  250. apiKeyRepo APIKeyRepository
  251. userRepo UserRepository
  252. groupRepo GroupRepository
  253. userSubRepo UserSubscriptionRepository
  254. userGroupRateRepo UserGroupRateRepository
  255. cache APIKeyCache
  256. rateLimitCacheInvalid RateLimitCacheInvalidator // optional: invalidate Redis rate limit cache
  257. concurrencyService *ConcurrencyService
  258. cfg *config.Config
  259. authCacheL1 *ristretto.Cache
  260. authNegativeCacheL1 *ristretto.Cache
  261. authCfg apiKeyAuthCacheConfig
  262. authGroup singleflight.Group
  263. authLookupSlots chan struct{}
  264. authLookupTotal atomic.Uint64
  265. authLookupRejected atomic.Uint64
  266. authLookupInFlight atomic.Int64
  267. invalidAuthAbuse *invalidAuthAbuseLimiter
  268. authInvalidationStart sync.Once
  269. authInvalidationStop sync.Once
  270. authInvalidationCancel context.CancelFunc
  271. authInvalidationWG sync.WaitGroup
  272. authInvalidationConnected atomic.Bool
  273. authInvalidationFailures atomic.Uint64
  274. lastUsedTouchL1 sync.Map // keyID -> nextAllowedAt(time.Time)
  275. lastUsedTouchSF singleflight.Group
  276. }
  277. type APIKeyAuthLookupMetrics struct {
  278. Total uint64 `json:"total"`
  279. Rejected uint64 `json:"rejected"`
  280. InFlight int64 `json:"in_flight"`
  281. Capacity int `json:"capacity"`
  282. }
  283. func (s *APIKeyService) AuthLookupMetrics() APIKeyAuthLookupMetrics {
  284. if s == nil {
  285. return APIKeyAuthLookupMetrics{}
  286. }
  287. return APIKeyAuthLookupMetrics{
  288. Total: s.authLookupTotal.Load(),
  289. Rejected: s.authLookupRejected.Load(),
  290. InFlight: s.authLookupInFlight.Load(),
  291. Capacity: cap(s.authLookupSlots),
  292. }
  293. }
  294. // NewAPIKeyService 创建API Key服务实例
  295. func NewAPIKeyService(
  296. apiKeyRepo APIKeyRepository,
  297. userRepo UserRepository,
  298. groupRepo GroupRepository,
  299. userSubRepo UserSubscriptionRepository,
  300. userGroupRateRepo UserGroupRateRepository,
  301. cache APIKeyCache,
  302. cfg *config.Config,
  303. ) *APIKeyService {
  304. svc := &APIKeyService{
  305. apiKeyRepo: apiKeyRepo,
  306. userRepo: userRepo,
  307. groupRepo: groupRepo,
  308. userSubRepo: userSubRepo,
  309. userGroupRateRepo: userGroupRateRepo,
  310. cache: cache,
  311. cfg: cfg,
  312. }
  313. svc.initAuthCache(cfg)
  314. lookupConcurrency := defaultAuthLookupConcurrency
  315. if cfg != nil && cfg.APIKeyAuth.LookupConcurrency > 0 {
  316. lookupConcurrency = cfg.APIKeyAuth.LookupConcurrency
  317. }
  318. svc.authLookupSlots = make(chan struct{}, lookupConcurrency)
  319. svc.invalidAuthAbuse = newInvalidAuthAbuseLimiter(cfg)
  320. return svc
  321. }
  322. // SetRateLimitCacheInvalidator sets the optional rate limit cache invalidator.
  323. // Called after construction (e.g. in wire) to avoid circular dependencies.
  324. func (s *APIKeyService) SetRateLimitCacheInvalidator(inv RateLimitCacheInvalidator) {
  325. s.rateLimitCacheInvalid = inv
  326. }
  327. func (s *APIKeyService) SetConcurrencyService(concurrencyService *ConcurrencyService) {
  328. s.concurrencyService = concurrencyService
  329. }
  330. func (s *APIKeyService) compileAPIKeyIPRules(apiKey *APIKey) {
  331. if apiKey == nil {
  332. return
  333. }
  334. apiKey.CompiledIPWhitelist = ip.CompileIPRules(apiKey.IPWhitelist)
  335. apiKey.CompiledIPBlacklist = ip.CompileIPRules(apiKey.IPBlacklist)
  336. }
  337. // GenerateKey 生成随机API Key
  338. func (s *APIKeyService) GenerateKey() (string, error) {
  339. // 生成32字节随机数据
  340. bytes := make([]byte, 32)
  341. if _, err := rand.Read(bytes); err != nil {
  342. return "", fmt.Errorf("generate random bytes: %w", err)
  343. }
  344. // 转换为十六进制字符串并添加前缀
  345. prefix := s.cfg.Default.APIKeyPrefix
  346. if prefix == "" {
  347. prefix = "sk-"
  348. }
  349. key := prefix + hex.EncodeToString(bytes)
  350. return key, nil
  351. }
  352. // ValidateCustomKey 验证自定义API Key格式
  353. func (s *APIKeyService) ValidateCustomKey(key string) error {
  354. // 检查长度
  355. if len(key) < 16 {
  356. return ErrAPIKeyTooShort
  357. }
  358. // 检查字符:只允许字母、数字、下划线、连字符
  359. for _, c := range key {
  360. if (c >= 'a' && c <= 'z') ||
  361. (c >= 'A' && c <= 'Z') ||
  362. (c >= '0' && c <= '9') ||
  363. c == '_' || c == '-' {
  364. continue
  365. }
  366. return ErrAPIKeyInvalidChars
  367. }
  368. return nil
  369. }
  370. // checkAPIKeyRateLimit 检查用户创建自定义Key的错误次数是否超限
  371. func (s *APIKeyService) checkAPIKeyRateLimit(ctx context.Context, userID int64) error {
  372. if s.cache == nil {
  373. return nil
  374. }
  375. count, err := s.cache.GetCreateAttemptCount(ctx, userID)
  376. if err != nil {
  377. // Redis 出错时不阻止用户操作
  378. return nil
  379. }
  380. if count >= apiKeyMaxErrorsPerHour {
  381. return ErrAPIKeyRateLimited
  382. }
  383. return nil
  384. }
  385. // incrementAPIKeyErrorCount 增加用户创建自定义Key的错误计数
  386. func (s *APIKeyService) incrementAPIKeyErrorCount(ctx context.Context, userID int64) {
  387. if s.cache == nil {
  388. return
  389. }
  390. _ = s.cache.IncrementCreateAttemptCount(ctx, userID)
  391. }
  392. // canUserBindGroup 检查用户是否可以绑定指定分组
  393. // 对于订阅类型分组:检查用户是否有有效订阅
  394. // 对于标准类型分组:使用原有的 AllowedGroups 和 IsExclusive 逻辑
  395. func (s *APIKeyService) canUserBindGroup(ctx context.Context, user *User, group *Group) bool {
  396. // 订阅类型分组:需要有效订阅
  397. if group.IsSubscriptionType() {
  398. _, err := s.userSubRepo.GetActiveByUserIDAndGroupID(ctx, user.ID, group.ID)
  399. return err == nil // 有有效订阅则允许
  400. }
  401. // 标准类型分组:使用原有逻辑
  402. return user.CanBindGroup(group.ID, group.IsExclusive)
  403. }
  404. // Create 创建API Key
  405. func (s *APIKeyService) Create(ctx context.Context, userID int64, req CreateAPIKeyRequest) (*APIKey, error) {
  406. if err := validateCreateAPIKeyRequest(req); err != nil {
  407. return nil, err
  408. }
  409. // 验证用户存在
  410. user, err := s.userRepo.GetByID(ctx, userID)
  411. if err != nil {
  412. return nil, fmt.Errorf("get user: %w", err)
  413. }
  414. // 验证 IP 白名单格式
  415. if len(req.IPWhitelist) > 0 {
  416. if invalid := ip.ValidateIPPatterns(req.IPWhitelist); len(invalid) > 0 {
  417. return nil, fmt.Errorf("%w: %v", ErrInvalidIPPattern, invalid)
  418. }
  419. }
  420. // 验证 IP 黑名单格式
  421. if len(req.IPBlacklist) > 0 {
  422. if invalid := ip.ValidateIPPatterns(req.IPBlacklist); len(invalid) > 0 {
  423. return nil, fmt.Errorf("%w: %v", ErrInvalidIPPattern, invalid)
  424. }
  425. }
  426. // 验证分组权限(如果指定了分组)
  427. if req.GroupID != nil {
  428. group, err := s.groupRepo.GetByID(ctx, *req.GroupID)
  429. if err != nil {
  430. return nil, fmt.Errorf("get group: %w", err)
  431. }
  432. // 检查用户是否可以绑定该分组
  433. if !s.canUserBindGroup(ctx, user, group) {
  434. return nil, ErrGroupNotAllowed
  435. }
  436. }
  437. var key string
  438. // 判断是否使用自定义Key
  439. if req.CustomKey != nil && *req.CustomKey != "" {
  440. // 检查限流(仅对自定义key进行限流)
  441. if err := s.checkAPIKeyRateLimit(ctx, userID); err != nil {
  442. return nil, err
  443. }
  444. // 验证自定义Key格式
  445. if err := s.ValidateCustomKey(*req.CustomKey); err != nil {
  446. return nil, err
  447. }
  448. // 检查Key是否已存在
  449. exists, err := s.apiKeyRepo.ExistsByKey(ctx, *req.CustomKey)
  450. if err != nil {
  451. return nil, fmt.Errorf("check key exists: %w", err)
  452. }
  453. if exists {
  454. // Key已存在,增加错误计数
  455. s.incrementAPIKeyErrorCount(ctx, userID)
  456. return nil, ErrAPIKeyExists
  457. }
  458. key = *req.CustomKey
  459. } else {
  460. // 生成随机API Key
  461. var err error
  462. key, err = s.GenerateKey()
  463. if err != nil {
  464. return nil, fmt.Errorf("generate key: %w", err)
  465. }
  466. }
  467. // 创建API Key记录
  468. apiKey := &APIKey{
  469. UserID: userID,
  470. Key: key,
  471. Name: html.EscapeString(req.Name),
  472. GroupID: req.GroupID,
  473. Status: StatusActive,
  474. IPWhitelist: req.IPWhitelist,
  475. IPBlacklist: req.IPBlacklist,
  476. Quota: req.Quota,
  477. QuotaUsed: 0,
  478. RateLimit5h: req.RateLimit5h,
  479. RateLimit1d: req.RateLimit1d,
  480. RateLimit7d: req.RateLimit7d,
  481. }
  482. // Set expiration time if specified
  483. if req.ExpiresInDays != nil && *req.ExpiresInDays > 0 {
  484. expiresAt := time.Now().AddDate(0, 0, *req.ExpiresInDays)
  485. apiKey.ExpiresAt = &expiresAt
  486. }
  487. if err := s.apiKeyRepo.Create(ctx, apiKey); err != nil {
  488. return nil, fmt.Errorf("create api key: %w", err)
  489. }
  490. s.InvalidateAuthCacheByKey(ctx, apiKey.AuthCacheInvalidationKey())
  491. s.compileAPIKeyIPRules(apiKey)
  492. return apiKey, nil
  493. }
  494. // List 获取用户的API Key列表
  495. func (s *APIKeyService) List(ctx context.Context, userID int64, params pagination.PaginationParams, filters APIKeyListFilters) ([]APIKey, *pagination.PaginationResult, error) {
  496. if normalizedAPIKeySortBy(params.SortBy) == apiKeySortCurrentConcurrency {
  497. return s.listByCurrentConcurrency(ctx, userID, params, filters)
  498. }
  499. keys, pagination, err := s.apiKeyRepo.ListByUserID(ctx, userID, params, filters)
  500. if err != nil {
  501. return nil, nil, fmt.Errorf("list api keys: %w", err)
  502. }
  503. s.fillCurrentConcurrency(ctx, keys)
  504. return keys, pagination, nil
  505. }
  506. func (s *APIKeyService) listByCurrentConcurrency(ctx context.Context, userID int64, params pagination.PaginationParams, filters APIKeyListFilters) ([]APIKey, *pagination.PaginationResult, error) {
  507. repo, ok := s.apiKeyRepo.(apiKeyAllByUserIDLister)
  508. if !ok {
  509. return nil, nil, fmt.Errorf("list api keys by current concurrency: repository does not support unpaginated API key listing")
  510. }
  511. keys, err := repo.ListAllByUserID(ctx, userID, filters)
  512. if err != nil {
  513. return nil, nil, fmt.Errorf("list api keys: %w", err)
  514. }
  515. s.fillCurrentConcurrency(ctx, keys)
  516. sortAPIKeysByCurrentConcurrency(keys, params.NormalizedSortOrder(pagination.SortOrderDesc))
  517. return paginateAPIKeys(keys, params), apiKeyPaginationResult(int64(len(keys)), params), nil
  518. }
  519. func normalizedAPIKeySortBy(sortBy string) string {
  520. return strings.ToLower(strings.TrimSpace(sortBy))
  521. }
  522. func sortAPIKeysByCurrentConcurrency(keys []APIKey, sortOrder string) {
  523. desc := sortOrder != pagination.SortOrderAsc
  524. sort.SliceStable(keys, func(i, j int) bool {
  525. if keys[i].CurrentConcurrency == keys[j].CurrentConcurrency {
  526. if desc {
  527. return keys[i].ID > keys[j].ID
  528. }
  529. return keys[i].ID < keys[j].ID
  530. }
  531. if desc {
  532. return keys[i].CurrentConcurrency > keys[j].CurrentConcurrency
  533. }
  534. return keys[i].CurrentConcurrency < keys[j].CurrentConcurrency
  535. })
  536. }
  537. func paginateAPIKeys(keys []APIKey, params pagination.PaginationParams) []APIKey {
  538. if len(keys) == 0 {
  539. return []APIKey{}
  540. }
  541. limit := params.Limit()
  542. page := params.Page
  543. if page < 1 {
  544. page = 1
  545. }
  546. offset := (page - 1) * limit
  547. if offset >= len(keys) {
  548. return []APIKey{}
  549. }
  550. end := offset + limit
  551. if end > len(keys) {
  552. end = len(keys)
  553. }
  554. return keys[offset:end]
  555. }
  556. func apiKeyPaginationResult(total int64, params pagination.PaginationParams) *pagination.PaginationResult {
  557. limit := params.Limit()
  558. pages := int(total) / limit
  559. if int(total)%limit > 0 {
  560. pages++
  561. }
  562. return &pagination.PaginationResult{
  563. Total: total,
  564. Page: params.Page,
  565. PageSize: limit,
  566. Pages: pages,
  567. }
  568. }
  569. func (s *APIKeyService) fillCurrentConcurrency(ctx context.Context, keys []APIKey) {
  570. if s == nil || s.concurrencyService == nil || len(keys) == 0 {
  571. return
  572. }
  573. ids := make([]int64, 0, len(keys))
  574. for i := range keys {
  575. if keys[i].ID > 0 {
  576. ids = append(ids, keys[i].ID)
  577. }
  578. }
  579. counts, err := s.concurrencyService.GetAPIKeyConcurrencyBatch(ctx, ids)
  580. if err != nil {
  581. return
  582. }
  583. for i := range keys {
  584. keys[i].CurrentConcurrency = counts[keys[i].ID]
  585. }
  586. }
  587. func (s *APIKeyService) currentConcurrencyForAPIKey(ctx context.Context, apiKeyID int64) int {
  588. if s == nil || s.concurrencyService == nil || apiKeyID <= 0 {
  589. return 0
  590. }
  591. counts, err := s.concurrencyService.GetAPIKeyConcurrencyBatch(ctx, []int64{apiKeyID})
  592. if err != nil {
  593. return 0
  594. }
  595. return counts[apiKeyID]
  596. }
  597. func (s *APIKeyService) VerifyOwnership(ctx context.Context, userID int64, apiKeyIDs []int64) ([]int64, error) {
  598. if len(apiKeyIDs) == 0 {
  599. return []int64{}, nil
  600. }
  601. validIDs, err := s.apiKeyRepo.VerifyOwnership(ctx, userID, apiKeyIDs)
  602. if err != nil {
  603. return nil, fmt.Errorf("verify api key ownership: %w", err)
  604. }
  605. return validIDs, nil
  606. }
  607. // GetByID 根据ID获取API Key
  608. func (s *APIKeyService) GetByID(ctx context.Context, id int64) (*APIKey, error) {
  609. apiKey, err := s.apiKeyRepo.GetByID(ctx, id)
  610. if err != nil {
  611. return nil, fmt.Errorf("get api key: %w", err)
  612. }
  613. s.compileAPIKeyIPRules(apiKey)
  614. if apiKey != nil {
  615. apiKey.CurrentConcurrency = s.currentConcurrencyForAPIKey(ctx, apiKey.ID)
  616. }
  617. return apiKey, nil
  618. }
  619. // GetByKey 根据Key字符串获取API Key(用于认证)
  620. func (s *APIKeyService) GetByKey(ctx context.Context, key string) (*APIKey, error) {
  621. if len(key) == 0 || len(key) > MaxAPIKeyCredentialBytes {
  622. return nil, ErrAPIKeyNotFound
  623. }
  624. cacheKey := s.authCacheKey(key)
  625. if entry, ok := s.getAuthCacheEntry(ctx, cacheKey); ok {
  626. if apiKey, used, err := s.applyAuthCacheEntry(key, entry); used {
  627. if err != nil {
  628. return nil, fmt.Errorf("get api key: %w", err)
  629. }
  630. s.compileAPIKeyIPRules(apiKey)
  631. return apiKey, nil
  632. }
  633. }
  634. if s.authCfg.singleflight {
  635. value, err, _ := s.authGroup.Do(cacheKey, func() (any, error) {
  636. return s.loadAuthCacheEntry(ctx, key, cacheKey)
  637. })
  638. if err != nil {
  639. return nil, err
  640. }
  641. entry, _ := value.(*APIKeyAuthCacheEntry)
  642. if apiKey, used, err := s.applyAuthCacheEntry(key, entry); used {
  643. if err != nil {
  644. return nil, fmt.Errorf("get api key: %w", err)
  645. }
  646. s.compileAPIKeyIPRules(apiKey)
  647. return apiKey, nil
  648. }
  649. } else {
  650. entry, err := s.loadAuthCacheEntry(ctx, key, cacheKey)
  651. if err != nil {
  652. return nil, err
  653. }
  654. if apiKey, used, err := s.applyAuthCacheEntry(key, entry); used {
  655. if err != nil {
  656. return nil, fmt.Errorf("get api key: %w", err)
  657. }
  658. s.compileAPIKeyIPRules(apiKey)
  659. return apiKey, nil
  660. }
  661. }
  662. apiKey, err := s.lookupAPIKeyForAuth(ctx, key)
  663. if err != nil {
  664. return nil, fmt.Errorf("get api key: %w", err)
  665. }
  666. apiKey.Key = key
  667. s.compileAPIKeyIPRules(apiKey)
  668. return apiKey, nil
  669. }
  670. // Update 更新API Key
  671. func (s *APIKeyService) Update(ctx context.Context, id int64, userID int64, req UpdateAPIKeyRequest) (*APIKey, error) {
  672. if err := validateUpdateAPIKeyRequest(req); err != nil {
  673. return nil, err
  674. }
  675. apiKey, err := s.apiKeyRepo.GetByID(ctx, id)
  676. if err != nil {
  677. return nil, fmt.Errorf("get api key: %w", err)
  678. }
  679. // 验证所有权
  680. if apiKey.UserID != userID {
  681. return nil, ErrInsufficientPerms
  682. }
  683. // 验证 IP 白名单格式
  684. if req.IPWhitelist != nil && len(*req.IPWhitelist) > 0 {
  685. if invalid := ip.ValidateIPPatterns(*req.IPWhitelist); len(invalid) > 0 {
  686. return nil, fmt.Errorf("%w: %v", ErrInvalidIPPattern, invalid)
  687. }
  688. }
  689. // 验证 IP 黑名单格式
  690. if req.IPBlacklist != nil && len(*req.IPBlacklist) > 0 {
  691. if invalid := ip.ValidateIPPatterns(*req.IPBlacklist); len(invalid) > 0 {
  692. return nil, fmt.Errorf("%w: %v", ErrInvalidIPPattern, invalid)
  693. }
  694. }
  695. // fields 只登记本次请求真正要改的列。quota_used 与 usage_5h/1d/7d 由计费热路径
  696. // 原子递增,除非用户显式点了"重置",否则这里不用快照把它们写回去。
  697. var fields APIKeyUpdateFields
  698. // 下面若干分支会顺带把 Status 改回 active(配额扩容、清除过期等),
  699. // 所以用原始值比对来决定是否写 status,而不是只看 req.Status。
  700. originalStatus := apiKey.Status
  701. // 更新字段
  702. if req.Name != nil {
  703. apiKey.Name = html.EscapeString(*req.Name)
  704. fields.Name = true
  705. }
  706. if req.GroupID != nil {
  707. // 验证分组权限
  708. user, err := s.userRepo.GetByID(ctx, userID)
  709. if err != nil {
  710. return nil, fmt.Errorf("get user: %w", err)
  711. }
  712. group, err := s.groupRepo.GetByID(ctx, *req.GroupID)
  713. if err != nil {
  714. return nil, fmt.Errorf("get group: %w", err)
  715. }
  716. if !s.canUserBindGroup(ctx, user, group) {
  717. return nil, ErrGroupNotAllowed
  718. }
  719. apiKey.GroupID = req.GroupID
  720. fields.GroupID = true
  721. }
  722. if req.Status != nil {
  723. apiKey.Status = *req.Status
  724. fields.Status = true
  725. // 如果状态改变,清除Redis缓存
  726. if s.cache != nil {
  727. _ = s.cache.DeleteCreateAttemptCount(ctx, apiKey.UserID)
  728. }
  729. }
  730. // Update quota fields
  731. if req.Quota != nil {
  732. apiKey.Quota = *req.Quota
  733. fields.Quota = true
  734. // If quota now has room, or is changed to unlimited, reactivate exhausted keys.
  735. if apiKey.Status == StatusAPIKeyQuotaExhausted && (*req.Quota <= 0 || *req.Quota > apiKey.QuotaUsed) {
  736. apiKey.Status = StatusActive
  737. }
  738. }
  739. if req.ResetQuota != nil && *req.ResetQuota {
  740. apiKey.QuotaUsed = 0
  741. fields.QuotaUsed = true
  742. // If resetting quota and status was quota_exhausted, reactivate
  743. if apiKey.Status == StatusAPIKeyQuotaExhausted {
  744. apiKey.Status = StatusActive
  745. }
  746. }
  747. if req.ClearExpiration {
  748. apiKey.ExpiresAt = nil
  749. fields.ExpiresAt = true
  750. // If clearing expiry and status was expired, reactivate
  751. if apiKey.Status == StatusAPIKeyExpired {
  752. apiKey.Status = StatusActive
  753. }
  754. } else if req.ExpiresAt != nil {
  755. apiKey.ExpiresAt = req.ExpiresAt
  756. fields.ExpiresAt = true
  757. // If extending expiry and status was expired, reactivate
  758. if apiKey.Status == StatusAPIKeyExpired && time.Now().Before(*req.ExpiresAt) {
  759. apiKey.Status = StatusActive
  760. }
  761. }
  762. // 更新 IP 限制(nil 不修改,空数组清空设置)
  763. if req.IPWhitelist != nil {
  764. apiKey.IPWhitelist = *req.IPWhitelist
  765. fields.IPRules = true
  766. }
  767. if req.IPBlacklist != nil {
  768. apiKey.IPBlacklist = *req.IPBlacklist
  769. fields.IPRules = true
  770. }
  771. // Update rate limit configuration
  772. if req.RateLimit5h != nil {
  773. apiKey.RateLimit5h = *req.RateLimit5h
  774. fields.RateLimits = true
  775. }
  776. if req.RateLimit1d != nil {
  777. apiKey.RateLimit1d = *req.RateLimit1d
  778. fields.RateLimits = true
  779. }
  780. if req.RateLimit7d != nil {
  781. apiKey.RateLimit7d = *req.RateLimit7d
  782. fields.RateLimits = true
  783. }
  784. resetRateLimit := req.ResetRateLimitUsage != nil && *req.ResetRateLimitUsage
  785. if resetRateLimit {
  786. apiKey.Usage5h = 0
  787. apiKey.Usage1d = 0
  788. apiKey.Usage7d = 0
  789. apiKey.Window5hStart = nil
  790. apiKey.Window1dStart = nil
  791. apiKey.Window7dStart = nil
  792. fields.RateLimitUsage = true
  793. }
  794. // 上面的自动复活分支可能改了 status,这里统一登记。
  795. if apiKey.Status != originalStatus {
  796. fields.Status = true
  797. }
  798. if err := s.apiKeyRepo.Update(ctx, apiKey, fields); err != nil {
  799. return nil, fmt.Errorf("update api key: %w", err)
  800. }
  801. s.InvalidateAuthCacheByKey(ctx, apiKey.AuthCacheInvalidationKey())
  802. s.compileAPIKeyIPRules(apiKey)
  803. // Invalidate Redis rate limit cache so reset takes effect immediately
  804. if resetRateLimit && s.rateLimitCacheInvalid != nil {
  805. _ = s.rateLimitCacheInvalid.InvalidateAPIKeyRateLimit(ctx, apiKey.ID)
  806. }
  807. return apiKey, nil
  808. }
  809. // Delete 删除API Key
  810. func (s *APIKeyService) Delete(ctx context.Context, id int64, userID int64) error {
  811. key, ownerID, err := s.apiKeyRepo.GetKeyAndOwnerID(ctx, id)
  812. if err != nil {
  813. return fmt.Errorf("get api key: %w", err)
  814. }
  815. // 验证当前用户是否为该 API Key 的所有者
  816. if ownerID != userID {
  817. return ErrInsufficientPerms
  818. }
  819. // 事务内:写审计 + 软删除(tombstone)。
  820. if err := s.apiKeyRepo.DeleteWithAudit(ctx, id); err != nil {
  821. return fmt.Errorf("delete api key: %w", err)
  822. }
  823. // 删除成功后再清理缓存,避免"缓存已清但删除失败"的竞态。
  824. if s.cache != nil {
  825. _ = s.cache.DeleteCreateAttemptCount(ctx, userID)
  826. }
  827. s.InvalidateAuthCacheByKey(ctx, key)
  828. s.lastUsedTouchL1.Delete(id)
  829. return nil
  830. }
  831. // ValidateKey 验证API Key是否有效(用于认证中间件)
  832. func (s *APIKeyService) ValidateKey(ctx context.Context, key string) (*APIKey, *User, error) {
  833. // 获取API Key
  834. apiKey, err := s.GetByKey(ctx, key)
  835. if err != nil {
  836. return nil, nil, err
  837. }
  838. // 检查API Key状态
  839. if !apiKey.IsActive() {
  840. return nil, nil, infraerrors.Unauthorized("API_KEY_INACTIVE", "api key is not active")
  841. }
  842. // 获取用户信息
  843. user, err := s.userRepo.GetByID(ctx, apiKey.UserID)
  844. if err != nil {
  845. return nil, nil, fmt.Errorf("get user: %w", err)
  846. }
  847. // 检查用户状态
  848. if !user.IsActive() {
  849. return nil, nil, ErrUserNotActive
  850. }
  851. return apiKey, user, nil
  852. }
  853. // TouchLastUsed 通过防抖更新 api_keys.last_used_at,减少高频写放大。
  854. // 该操作为尽力而为,不应阻塞主请求链路。
  855. func (s *APIKeyService) TouchLastUsed(ctx context.Context, keyID int64) error {
  856. if keyID <= 0 {
  857. return nil
  858. }
  859. now := time.Now()
  860. if v, ok := s.lastUsedTouchL1.Load(keyID); ok {
  861. if nextAllowedAt, ok := v.(time.Time); ok && now.Before(nextAllowedAt) {
  862. return nil
  863. }
  864. }
  865. _, err, _ := s.lastUsedTouchSF.Do(strconv.FormatInt(keyID, 10), func() (any, error) {
  866. latest := time.Now()
  867. if v, ok := s.lastUsedTouchL1.Load(keyID); ok {
  868. if nextAllowedAt, ok := v.(time.Time); ok && latest.Before(nextAllowedAt) {
  869. return nil, nil
  870. }
  871. }
  872. if err := s.apiKeyRepo.UpdateLastUsed(ctx, keyID, latest); err != nil {
  873. s.lastUsedTouchL1.Store(keyID, latest.Add(apiKeyLastUsedFailBackoff))
  874. return nil, fmt.Errorf("touch api key last used: %w", err)
  875. }
  876. s.lastUsedTouchL1.Store(keyID, latest.Add(apiKeyLastUsedMinTouch))
  877. return nil, nil
  878. })
  879. return err
  880. }
  881. // IncrementUsage 增加API Key使用次数(可选:用于统计)
  882. func (s *APIKeyService) IncrementUsage(ctx context.Context, keyID int64) error {
  883. // 使用Redis计数器
  884. if s.cache != nil {
  885. cacheKey := fmt.Sprintf("apikey:usage:%d:%s", keyID, timezone.Now().Format("2006-01-02"))
  886. if err := s.cache.IncrementDailyUsage(ctx, cacheKey); err != nil {
  887. return fmt.Errorf("increment usage: %w", err)
  888. }
  889. // 设置24小时过期
  890. _ = s.cache.SetDailyUsageExpiry(ctx, cacheKey, 24*time.Hour)
  891. }
  892. return nil
  893. }
  894. // GetAvailableGroups 获取用户有权限绑定的分组列表
  895. // 返回用户可以选择的分组:
  896. // - 标准类型分组:公开的(非专属)或用户被明确允许的
  897. // - 订阅类型分组:用户有有效订阅的
  898. func (s *APIKeyService) GetAvailableGroups(ctx context.Context, userID int64) ([]Group, error) {
  899. // 获取用户信息
  900. user, err := s.userRepo.GetByID(ctx, userID)
  901. if err != nil {
  902. return nil, fmt.Errorf("get user: %w", err)
  903. }
  904. // 获取所有活跃分组
  905. allGroups, err := s.groupRepo.ListActive(ctx)
  906. if err != nil {
  907. return nil, fmt.Errorf("list active groups: %w", err)
  908. }
  909. // 获取用户的所有有效订阅
  910. activeSubscriptions, err := s.userSubRepo.ListActiveByUserID(ctx, userID)
  911. if err != nil {
  912. return nil, fmt.Errorf("list active subscriptions: %w", err)
  913. }
  914. // 构建订阅分组 ID 集合
  915. subscribedGroupIDs := make(map[int64]bool)
  916. for _, sub := range activeSubscriptions {
  917. subscribedGroupIDs[sub.GroupID] = true
  918. }
  919. // 过滤出用户有权限的分组
  920. availableGroups := make([]Group, 0)
  921. for _, group := range allGroups {
  922. if s.canUserBindGroupInternal(user, &group, subscribedGroupIDs) {
  923. availableGroups = append(availableGroups, group)
  924. }
  925. }
  926. return availableGroups, nil
  927. }
  928. // canUserBindGroupInternal 内部方法,检查用户是否可以绑定分组(使用预加载的订阅数据)
  929. func (s *APIKeyService) canUserBindGroupInternal(user *User, group *Group, subscribedGroupIDs map[int64]bool) bool {
  930. // 订阅类型分组:需要有效订阅
  931. if group.IsSubscriptionType() {
  932. return subscribedGroupIDs[group.ID]
  933. }
  934. // 标准类型分组:使用原有逻辑
  935. return user.CanBindGroup(group.ID, group.IsExclusive)
  936. }
  937. func (s *APIKeyService) SearchAPIKeys(ctx context.Context, userID int64, keyword string, limit int) ([]APIKey, error) {
  938. keys, err := s.apiKeyRepo.SearchAPIKeys(ctx, userID, keyword, limit)
  939. if err != nil {
  940. return nil, fmt.Errorf("search api keys: %w", err)
  941. }
  942. return keys, nil
  943. }
  944. // GetUserGroupVisibility 返回 user_allowed_groups 授权及有效订阅的分组 ID 集合,
  945. // 以及该用户是否开启了公开分组限制。开启时公开分组的可见性也要落在该集合内。
  946. //
  947. // 与 GetAvailableGroups 的区别:这里保留普通授权分组的「橱窗」语义,不检查
  948. // 分组是否活跃;有效订阅也授予对应分组的可见性。返回值恒非 nil。
  949. func (s *APIKeyService) GetUserGroupVisibility(ctx context.Context, userID int64) (map[int64]struct{}, bool, error) {
  950. user, err := s.userRepo.GetByID(ctx, userID)
  951. if err != nil {
  952. return nil, false, fmt.Errorf("get user: %w", err)
  953. }
  954. allowed := make(map[int64]struct{}, len(user.AllowedGroups))
  955. for _, id := range user.AllowedGroups {
  956. allowed[id] = struct{}{}
  957. }
  958. subscriptions, err := s.userSubRepo.ListActiveByUserID(ctx, userID)
  959. if err != nil {
  960. return nil, false, fmt.Errorf("list active subscriptions: %w", err)
  961. }
  962. for _, sub := range subscriptions {
  963. allowed[sub.GroupID] = struct{}{}
  964. }
  965. return allowed, user.RestrictPublicGroups, nil
  966. }
  967. // GetUserGroupRates 获取用户的专属分组倍率配置
  968. // 返回 map[groupID]rateMultiplier
  969. func (s *APIKeyService) GetUserGroupRates(ctx context.Context, userID int64) (map[int64]float64, error) {
  970. if s.userGroupRateRepo == nil {
  971. return nil, nil
  972. }
  973. rates, err := s.userGroupRateRepo.GetByUserID(ctx, userID)
  974. if err != nil {
  975. return nil, fmt.Errorf("get user group rates: %w", err)
  976. }
  977. return rates, nil
  978. }
  979. // CheckAPIKeyQuotaAndExpiry checks if the API key is valid for use (not expired, quota not exhausted)
  980. // Returns nil if valid, error if invalid
  981. func (s *APIKeyService) CheckAPIKeyQuotaAndExpiry(apiKey *APIKey) error {
  982. // Check expiration
  983. if apiKey.IsExpired() {
  984. return ErrAPIKeyExpired
  985. }
  986. // Check quota
  987. if apiKey.IsQuotaExhausted() {
  988. return ErrAPIKeyQuotaExhausted
  989. }
  990. return nil
  991. }
  992. // UpdateQuotaUsed updates the quota_used field after a request
  993. // Also checks if quota is exhausted and updates status accordingly
  994. func (s *APIKeyService) UpdateQuotaUsed(ctx context.Context, apiKeyID int64, cost float64) error {
  995. if cost <= 0 {
  996. return nil
  997. }
  998. type quotaStateReader interface {
  999. IncrementQuotaUsedAndGetState(ctx context.Context, id int64, amount float64) (*APIKeyQuotaUsageState, error)
  1000. }
  1001. if repo, ok := s.apiKeyRepo.(quotaStateReader); ok {
  1002. state, err := repo.IncrementQuotaUsedAndGetState(ctx, apiKeyID, cost)
  1003. if err != nil {
  1004. return fmt.Errorf("increment quota used: %w", err)
  1005. }
  1006. if state != nil && state.Status == StatusAPIKeyQuotaExhausted && strings.TrimSpace(state.Key) != "" {
  1007. s.InvalidateAuthCacheByKey(ctx, state.Key)
  1008. }
  1009. return nil
  1010. }
  1011. // Use repository to atomically increment quota_used
  1012. newQuotaUsed, err := s.apiKeyRepo.IncrementQuotaUsed(ctx, apiKeyID, cost)
  1013. if err != nil {
  1014. return fmt.Errorf("increment quota used: %w", err)
  1015. }
  1016. // Check if quota is now exhausted and update status if needed
  1017. apiKey, err := s.apiKeyRepo.GetByID(ctx, apiKeyID)
  1018. if err != nil {
  1019. return nil // Don't fail the request, just log
  1020. }
  1021. // If quota is set and now exhausted, update status
  1022. if apiKey.Quota > 0 && newQuotaUsed >= apiKey.Quota {
  1023. apiKey.Status = StatusAPIKeyQuotaExhausted
  1024. // 只写 status:这条位于计费热路径,若整行回写会把刚刚原子递增的
  1025. // quota_used 与限流用量按快照覆盖掉。
  1026. if err := s.apiKeyRepo.Update(ctx, apiKey, APIKeyUpdateFields{Status: true}); err != nil {
  1027. return nil // Don't fail the request
  1028. }
  1029. // Invalidate cache so next request sees the new status
  1030. s.InvalidateAuthCacheByKey(ctx, apiKey.AuthCacheInvalidationKey())
  1031. }
  1032. return nil
  1033. }
  1034. // GetRateLimitData returns rate limit usage and window state for an API key.
  1035. func (s *APIKeyService) GetRateLimitData(ctx context.Context, id int64) (*APIKeyRateLimitData, error) {
  1036. return s.apiKeyRepo.GetRateLimitData(ctx, id)
  1037. }
  1038. // UpdateRateLimitUsage atomically increments rate limit usage counters in the DB.
  1039. func (s *APIKeyService) UpdateRateLimitUsage(ctx context.Context, apiKeyID int64, cost float64) error {
  1040. if cost <= 0 {
  1041. return nil
  1042. }
  1043. return s.apiKeyRepo.IncrementRateLimitUsage(ctx, apiKeyID, cost)
  1044. }