api_key_service_cache_test.go 21 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694
  1. //go:build unit
  2. package service
  3. import (
  4. "context"
  5. "errors"
  6. "strings"
  7. "sync"
  8. "sync/atomic"
  9. "testing"
  10. "time"
  11. "github.com/Wei-Shaw/sub2api/internal/config"
  12. "github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
  13. "github.com/redis/go-redis/v9"
  14. "github.com/stretchr/testify/require"
  15. )
  16. type authRepoStub struct {
  17. getByKeyForAuth func(ctx context.Context, key string) (*APIKey, error)
  18. listKeysByUserID func(ctx context.Context, userID int64) ([]string, error)
  19. listKeysByGroupID func(ctx context.Context, groupID int64) ([]string, error)
  20. }
  21. func (s *authRepoStub) Create(ctx context.Context, key *APIKey) error {
  22. panic("unexpected Create call")
  23. }
  24. func (s *authRepoStub) GetByID(ctx context.Context, id int64) (*APIKey, error) {
  25. panic("unexpected GetByID call")
  26. }
  27. func (s *authRepoStub) GetKeyAndOwnerID(ctx context.Context, id int64) (string, int64, error) {
  28. panic("unexpected GetKeyAndOwnerID call")
  29. }
  30. func (s *authRepoStub) GetByKey(ctx context.Context, key string) (*APIKey, error) {
  31. panic("unexpected GetByKey call")
  32. }
  33. func (s *authRepoStub) GetByKeyForAuth(ctx context.Context, key string) (*APIKey, error) {
  34. if s.getByKeyForAuth == nil {
  35. panic("unexpected GetByKeyForAuth call")
  36. }
  37. return s.getByKeyForAuth(ctx, key)
  38. }
  39. func (s *authRepoStub) Update(ctx context.Context, key *APIKey, _ APIKeyUpdateFields) error {
  40. panic("unexpected Update call")
  41. }
  42. func (s *authRepoStub) Delete(ctx context.Context, id int64) error {
  43. panic("unexpected Delete call")
  44. }
  45. func (s *authRepoStub) DeleteWithAudit(ctx context.Context, id int64) error {
  46. panic("unexpected DeleteWithAudit call")
  47. }
  48. func (s *authRepoStub) ListByUserID(ctx context.Context, userID int64, params pagination.PaginationParams, filters APIKeyListFilters) ([]APIKey, *pagination.PaginationResult, error) {
  49. panic("unexpected ListByUserID call")
  50. }
  51. func (s *authRepoStub) VerifyOwnership(ctx context.Context, userID int64, apiKeyIDs []int64) ([]int64, error) {
  52. panic("unexpected VerifyOwnership call")
  53. }
  54. func (s *authRepoStub) CountByUserID(ctx context.Context, userID int64) (int64, error) {
  55. panic("unexpected CountByUserID call")
  56. }
  57. func (s *authRepoStub) ExistsByKey(ctx context.Context, key string) (bool, error) {
  58. panic("unexpected ExistsByKey call")
  59. }
  60. func (s *authRepoStub) ListByGroupID(ctx context.Context, groupID int64, params pagination.PaginationParams) ([]APIKey, *pagination.PaginationResult, error) {
  61. panic("unexpected ListByGroupID call")
  62. }
  63. func (s *authRepoStub) SearchAPIKeys(ctx context.Context, userID int64, keyword string, limit int) ([]APIKey, error) {
  64. panic("unexpected SearchAPIKeys call")
  65. }
  66. func (s *authRepoStub) ClearGroupIDByGroupID(ctx context.Context, groupID int64) (int64, error) {
  67. panic("unexpected ClearGroupIDByGroupID call")
  68. }
  69. func (s *authRepoStub) UpdateGroupIDByUserAndGroup(ctx context.Context, userID, oldGroupID, newGroupID int64) (int64, error) {
  70. panic("unexpected UpdateGroupIDByUserAndGroup call")
  71. }
  72. func (s *authRepoStub) CountByGroupID(ctx context.Context, groupID int64) (int64, error) {
  73. panic("unexpected CountByGroupID call")
  74. }
  75. func (s *authRepoStub) ListKeysByUserID(ctx context.Context, userID int64) ([]string, error) {
  76. if s.listKeysByUserID == nil {
  77. panic("unexpected ListKeysByUserID call")
  78. }
  79. return s.listKeysByUserID(ctx, userID)
  80. }
  81. func (s *authRepoStub) ListKeysByGroupID(ctx context.Context, groupID int64) ([]string, error) {
  82. if s.listKeysByGroupID == nil {
  83. panic("unexpected ListKeysByGroupID call")
  84. }
  85. return s.listKeysByGroupID(ctx, groupID)
  86. }
  87. func (s *authRepoStub) IncrementQuotaUsed(ctx context.Context, id int64, amount float64) (float64, error) {
  88. panic("unexpected IncrementQuotaUsed call")
  89. }
  90. func (s *authRepoStub) UpdateLastUsed(ctx context.Context, id int64, usedAt time.Time) error {
  91. panic("unexpected UpdateLastUsed call")
  92. }
  93. func (s *authRepoStub) IncrementRateLimitUsage(ctx context.Context, id int64, cost float64) error {
  94. panic("unexpected IncrementRateLimitUsage call")
  95. }
  96. func (s *authRepoStub) ResetRateLimitWindows(ctx context.Context, id int64) error {
  97. panic("unexpected ResetRateLimitWindows call")
  98. }
  99. func (s *authRepoStub) GetRateLimitData(ctx context.Context, id int64) (*APIKeyRateLimitData, error) {
  100. panic("unexpected GetRateLimitData call")
  101. }
  102. type authCacheStub struct {
  103. getAuthCache func(ctx context.Context, key string) (*APIKeyAuthCacheEntry, error)
  104. setAuthKeys []string
  105. deleteAuthKeys []string
  106. }
  107. func (s *authCacheStub) GetCreateAttemptCount(ctx context.Context, userID int64) (int, error) {
  108. return 0, nil
  109. }
  110. func (s *authCacheStub) IncrementCreateAttemptCount(ctx context.Context, userID int64) error {
  111. return nil
  112. }
  113. func (s *authCacheStub) DeleteCreateAttemptCount(ctx context.Context, userID int64) error {
  114. return nil
  115. }
  116. func (s *authCacheStub) IncrementDailyUsage(ctx context.Context, apiKey string) error {
  117. return nil
  118. }
  119. func (s *authCacheStub) SetDailyUsageExpiry(ctx context.Context, apiKey string, ttl time.Duration) error {
  120. return nil
  121. }
  122. func (s *authCacheStub) GetAuthCache(ctx context.Context, key string) (*APIKeyAuthCacheEntry, error) {
  123. if s.getAuthCache == nil {
  124. return nil, redis.Nil
  125. }
  126. return s.getAuthCache(ctx, key)
  127. }
  128. func (s *authCacheStub) SetAuthCache(ctx context.Context, key string, entry *APIKeyAuthCacheEntry, ttl time.Duration) error {
  129. s.setAuthKeys = append(s.setAuthKeys, key)
  130. return nil
  131. }
  132. func (s *authCacheStub) DeleteAuthCache(ctx context.Context, key string) error {
  133. s.deleteAuthKeys = append(s.deleteAuthKeys, key)
  134. return nil
  135. }
  136. func (s *authCacheStub) PublishAuthCacheInvalidation(ctx context.Context, cacheKey string) error {
  137. return nil
  138. }
  139. func (s *authCacheStub) SubscribeAuthCacheInvalidation(ctx context.Context, handler func(cacheKey string)) error {
  140. return nil
  141. }
  142. func TestAPIKeyService_GetByKey_UsesL2Cache(t *testing.T) {
  143. cache := &authCacheStub{}
  144. repo := &authRepoStub{
  145. getByKeyForAuth: func(ctx context.Context, key string) (*APIKey, error) {
  146. return nil, errors.New("unexpected repo call")
  147. },
  148. }
  149. cfg := &config.Config{
  150. APIKeyAuth: config.APIKeyAuthCacheConfig{
  151. L2TTLSeconds: 60,
  152. NegativeTTLSeconds: 30,
  153. },
  154. }
  155. svc := NewAPIKeyService(repo, nil, nil, nil, nil, cache, cfg)
  156. groupID := int64(9)
  157. cacheEntry := &APIKeyAuthCacheEntry{
  158. Snapshot: &APIKeyAuthSnapshot{
  159. Version: apiKeyAuthSnapshotVersion,
  160. APIKeyID: 1,
  161. UserID: 2,
  162. GroupID: &groupID,
  163. Status: StatusActive,
  164. User: APIKeyAuthUserSnapshot{
  165. ID: 2,
  166. Status: StatusActive,
  167. Role: RoleUser,
  168. Balance: 10,
  169. Concurrency: 3,
  170. },
  171. Group: &APIKeyAuthGroupSnapshot{
  172. ID: groupID,
  173. Name: "g",
  174. Platform: PlatformAnthropic,
  175. Status: StatusActive,
  176. SubscriptionType: SubscriptionTypeStandard,
  177. RateMultiplier: 1,
  178. ModelRoutingEnabled: true,
  179. ModelRouting: map[string][]int64{
  180. "claude-opus-*": {1, 2},
  181. },
  182. },
  183. },
  184. }
  185. cache.getAuthCache = func(ctx context.Context, key string) (*APIKeyAuthCacheEntry, error) {
  186. return cacheEntry, nil
  187. }
  188. apiKey, err := svc.GetByKey(context.Background(), "k1")
  189. require.NoError(t, err)
  190. require.Equal(t, int64(1), apiKey.ID)
  191. require.Equal(t, int64(2), apiKey.User.ID)
  192. require.Equal(t, groupID, apiKey.Group.ID)
  193. require.True(t, apiKey.Group.ModelRoutingEnabled)
  194. require.Equal(t, map[string][]int64{"claude-opus-*": {1, 2}}, apiKey.Group.ModelRouting)
  195. }
  196. func TestAPIKeyService_SnapshotRoundTrip_PreservesMessagesDispatchModelConfig(t *testing.T) {
  197. svc := NewAPIKeyService(nil, nil, nil, nil, nil, nil, &config.Config{})
  198. groupID := int64(9)
  199. apiKey := &APIKey{
  200. ID: 1,
  201. UserID: 2,
  202. GroupID: &groupID,
  203. Key: "k-roundtrip",
  204. Name: "Audit Key",
  205. Status: StatusActive,
  206. User: &User{
  207. ID: 2,
  208. Status: StatusActive,
  209. Role: RoleUser,
  210. Balance: 10,
  211. Concurrency: 3,
  212. },
  213. Group: &Group{
  214. ID: groupID,
  215. Name: "openai",
  216. Platform: PlatformOpenAI,
  217. Status: StatusActive,
  218. SubscriptionType: SubscriptionTypeStandard,
  219. RateMultiplier: 1,
  220. AllowMessagesDispatch: true,
  221. DefaultMappedModel: "gpt-5.4",
  222. MessagesDispatchModelConfig: OpenAIMessagesDispatchModelConfig{
  223. OpusMappedModel: "gpt-5.4-nano",
  224. SonnetMappedModel: "gpt-5.3-codex",
  225. HaikuMappedModel: "gpt-5.4-mini",
  226. ExactModelMappings: map[string]string{
  227. "claude-sonnet-4.5": "gpt-5.4-nano",
  228. },
  229. },
  230. },
  231. }
  232. snapshot := svc.snapshotFromAPIKey(context.Background(), apiKey)
  233. roundTrip := svc.snapshotToAPIKey(apiKey.Key, snapshot)
  234. require.NotNil(t, roundTrip)
  235. require.Equal(t, apiKey.Name, roundTrip.Name)
  236. require.NotNil(t, roundTrip.Group)
  237. require.Equal(t, apiKey.Group.MessagesDispatchModelConfig, roundTrip.Group.MessagesDispatchModelConfig)
  238. }
  239. func TestAPIKeyService_SnapshotRoundTrip_PreservesReasoningEffortPolicy(t *testing.T) {
  240. svc := NewAPIKeyService(nil, nil, nil, nil, nil, nil, &config.Config{})
  241. groupID := int64(9)
  242. apiKey := &APIKey{
  243. ID: 1,
  244. UserID: 2,
  245. GroupID: &groupID,
  246. Key: "k-reasoning-policy",
  247. Status: StatusActive,
  248. User: &User{
  249. ID: 2,
  250. Status: StatusActive,
  251. Role: RoleUser,
  252. Balance: 10,
  253. Concurrency: 3,
  254. },
  255. Group: &Group{
  256. ID: groupID,
  257. Name: "composite",
  258. Platform: PlatformComposite,
  259. Status: StatusActive,
  260. SubscriptionType: SubscriptionTypeStandard,
  261. RateMultiplier: 1,
  262. MaxReasoningEffort: "medium",
  263. MaxReasoningEffortOverLimit: ReasoningEffortOverLimitDeny,
  264. ReasoningEffortMappings: []ReasoningEffortMapping{
  265. {From: "max", To: "xhigh"},
  266. },
  267. },
  268. }
  269. snapshot := svc.snapshotFromAPIKey(context.Background(), apiKey)
  270. roundTrip := svc.snapshotToAPIKey(apiKey.Key, snapshot)
  271. require.NotNil(t, roundTrip)
  272. require.NotNil(t, roundTrip.Group)
  273. require.Equal(t, PlatformComposite, roundTrip.Group.Platform)
  274. require.Equal(t, "medium", roundTrip.Group.MaxReasoningEffort)
  275. require.Equal(t, ReasoningEffortOverLimitDeny, roundTrip.Group.MaxReasoningEffortOverLimit)
  276. require.Equal(t, apiKey.Group.ReasoningEffortMappings, roundTrip.Group.ReasoningEffortMappings)
  277. }
  278. func TestAPIKeyService_GetByKey_IgnoresLegacyAuthCacheSnapshotWithoutMessagesDispatchConfig(t *testing.T) {
  279. cache := &authCacheStub{}
  280. var repoCalls int32
  281. repo := &authRepoStub{
  282. getByKeyForAuth: func(ctx context.Context, key string) (*APIKey, error) {
  283. atomic.AddInt32(&repoCalls, 1)
  284. groupID := int64(9)
  285. return &APIKey{
  286. ID: 1,
  287. UserID: 2,
  288. GroupID: &groupID,
  289. Status: StatusActive,
  290. User: &User{
  291. ID: 2,
  292. Status: StatusActive,
  293. Role: RoleUser,
  294. Balance: 10,
  295. Concurrency: 3,
  296. },
  297. Group: &Group{
  298. ID: groupID,
  299. Name: "openai",
  300. Platform: PlatformOpenAI,
  301. Status: StatusActive,
  302. Hydrated: true,
  303. SubscriptionType: SubscriptionTypeStandard,
  304. RateMultiplier: 1,
  305. AllowMessagesDispatch: true,
  306. DefaultMappedModel: "gpt-5.4",
  307. MessagesDispatchModelConfig: OpenAIMessagesDispatchModelConfig{
  308. OpusMappedModel: "gpt-5.4-nano",
  309. },
  310. },
  311. }, nil
  312. },
  313. }
  314. cfg := &config.Config{
  315. APIKeyAuth: config.APIKeyAuthCacheConfig{
  316. L2TTLSeconds: 60,
  317. },
  318. }
  319. svc := NewAPIKeyService(repo, nil, nil, nil, nil, cache, cfg)
  320. groupID := int64(9)
  321. cache.getAuthCache = func(ctx context.Context, key string) (*APIKeyAuthCacheEntry, error) {
  322. return &APIKeyAuthCacheEntry{
  323. Snapshot: &APIKeyAuthSnapshot{
  324. APIKeyID: 1,
  325. UserID: 2,
  326. GroupID: &groupID,
  327. Status: StatusActive,
  328. User: APIKeyAuthUserSnapshot{
  329. ID: 2,
  330. Status: StatusActive,
  331. Role: RoleUser,
  332. Balance: 10,
  333. Concurrency: 3,
  334. },
  335. Group: &APIKeyAuthGroupSnapshot{
  336. ID: groupID,
  337. Name: "openai",
  338. Platform: PlatformOpenAI,
  339. Status: StatusActive,
  340. SubscriptionType: SubscriptionTypeStandard,
  341. RateMultiplier: 1,
  342. AllowMessagesDispatch: true,
  343. DefaultMappedModel: "gpt-5.4",
  344. },
  345. },
  346. }, nil
  347. }
  348. apiKey, err := svc.GetByKey(context.Background(), "k-legacy")
  349. require.NoError(t, err)
  350. require.Equal(t, int32(1), atomic.LoadInt32(&repoCalls))
  351. require.NotNil(t, apiKey.Group)
  352. require.Equal(t, "gpt-5.4-nano", apiKey.Group.MessagesDispatchModelConfig.OpusMappedModel)
  353. }
  354. func TestAPIKeyService_GetByKey_NegativeCache(t *testing.T) {
  355. cache := &authCacheStub{}
  356. repo := &authRepoStub{
  357. getByKeyForAuth: func(ctx context.Context, key string) (*APIKey, error) {
  358. return nil, errors.New("unexpected repo call")
  359. },
  360. }
  361. cfg := &config.Config{
  362. APIKeyAuth: config.APIKeyAuthCacheConfig{
  363. L2TTLSeconds: 60,
  364. NegativeTTLSeconds: 30,
  365. },
  366. }
  367. svc := NewAPIKeyService(repo, nil, nil, nil, nil, cache, cfg)
  368. cache.getAuthCache = func(ctx context.Context, key string) (*APIKeyAuthCacheEntry, error) {
  369. return &APIKeyAuthCacheEntry{NotFound: true}, nil
  370. }
  371. _, err := svc.GetByKey(context.Background(), "missing")
  372. require.ErrorIs(t, err, ErrAPIKeyNotFound)
  373. }
  374. func TestAPIKeyService_GetByKey_CacheMissStoresL2(t *testing.T) {
  375. cache := &authCacheStub{}
  376. repo := &authRepoStub{
  377. getByKeyForAuth: func(ctx context.Context, key string) (*APIKey, error) {
  378. return &APIKey{
  379. ID: 5,
  380. UserID: 7,
  381. Status: StatusActive,
  382. User: &User{
  383. ID: 7,
  384. Status: StatusActive,
  385. Role: RoleUser,
  386. Balance: 12,
  387. Concurrency: 2,
  388. },
  389. }, nil
  390. },
  391. }
  392. cfg := &config.Config{
  393. APIKeyAuth: config.APIKeyAuthCacheConfig{
  394. L2TTLSeconds: 60,
  395. NegativeTTLSeconds: 30,
  396. },
  397. }
  398. svc := NewAPIKeyService(repo, nil, nil, nil, nil, cache, cfg)
  399. cache.getAuthCache = func(ctx context.Context, key string) (*APIKeyAuthCacheEntry, error) {
  400. return nil, redis.Nil
  401. }
  402. apiKey, err := svc.GetByKey(context.Background(), "k2")
  403. require.NoError(t, err)
  404. require.Equal(t, int64(5), apiKey.ID)
  405. require.Len(t, cache.setAuthKeys, 1)
  406. }
  407. func TestAPIKeyService_GetByKey_UsesL1Cache(t *testing.T) {
  408. var calls int32
  409. cache := &authCacheStub{}
  410. repo := &authRepoStub{
  411. getByKeyForAuth: func(ctx context.Context, key string) (*APIKey, error) {
  412. atomic.AddInt32(&calls, 1)
  413. return &APIKey{
  414. ID: 21,
  415. UserID: 3,
  416. Status: StatusActive,
  417. User: &User{
  418. ID: 3,
  419. Status: StatusActive,
  420. Role: RoleUser,
  421. Balance: 5,
  422. Concurrency: 2,
  423. },
  424. }, nil
  425. },
  426. }
  427. cfg := &config.Config{
  428. APIKeyAuth: config.APIKeyAuthCacheConfig{
  429. L1Size: 1000,
  430. L1TTLSeconds: 60,
  431. },
  432. }
  433. svc := NewAPIKeyService(repo, nil, nil, nil, nil, cache, cfg)
  434. require.NotNil(t, svc.authCacheL1)
  435. _, err := svc.GetByKey(context.Background(), "k-l1")
  436. require.NoError(t, err)
  437. svc.authCacheL1.Wait()
  438. cacheKey := svc.authCacheKey("k-l1")
  439. _, ok := svc.authCacheL1.Get(cacheKey)
  440. require.True(t, ok)
  441. _, err = svc.GetByKey(context.Background(), "k-l1")
  442. require.NoError(t, err)
  443. require.Equal(t, int32(1), atomic.LoadInt32(&calls))
  444. }
  445. func TestAPIKeyService_InvalidateAuthCacheByUserID(t *testing.T) {
  446. cache := &authCacheStub{}
  447. repo := &authRepoStub{
  448. listKeysByUserID: func(ctx context.Context, userID int64) ([]string, error) {
  449. return []string{"k1", "k2"}, nil
  450. },
  451. }
  452. cfg := &config.Config{
  453. APIKeyAuth: config.APIKeyAuthCacheConfig{
  454. L2TTLSeconds: 60,
  455. NegativeTTLSeconds: 30,
  456. },
  457. }
  458. svc := NewAPIKeyService(repo, nil, nil, nil, nil, cache, cfg)
  459. svc.InvalidateAuthCacheByUserID(context.Background(), 7)
  460. require.Len(t, cache.deleteAuthKeys, 2)
  461. }
  462. func TestAPIKeyService_InvalidateAuthCacheByGroupID(t *testing.T) {
  463. cache := &authCacheStub{}
  464. repo := &authRepoStub{
  465. listKeysByGroupID: func(ctx context.Context, groupID int64) ([]string, error) {
  466. return []string{"k1", "k2"}, nil
  467. },
  468. }
  469. cfg := &config.Config{
  470. APIKeyAuth: config.APIKeyAuthCacheConfig{
  471. L2TTLSeconds: 60,
  472. },
  473. }
  474. svc := NewAPIKeyService(repo, nil, nil, nil, nil, cache, cfg)
  475. svc.InvalidateAuthCacheByGroupID(context.Background(), 9)
  476. require.Len(t, cache.deleteAuthKeys, 2)
  477. }
  478. func TestAPIKeyService_InvalidateAuthCacheByKey(t *testing.T) {
  479. cache := &authCacheStub{}
  480. repo := &authRepoStub{
  481. listKeysByUserID: func(ctx context.Context, userID int64) ([]string, error) {
  482. return nil, nil
  483. },
  484. }
  485. cfg := &config.Config{
  486. APIKeyAuth: config.APIKeyAuthCacheConfig{
  487. L2TTLSeconds: 60,
  488. },
  489. }
  490. svc := NewAPIKeyService(repo, nil, nil, nil, nil, cache, cfg)
  491. svc.InvalidateAuthCacheByKey(context.Background(), "k1")
  492. require.Len(t, cache.deleteAuthKeys, 1)
  493. }
  494. func TestAPIKeyService_GetByKey_CachesNegativeOnRepoMiss(t *testing.T) {
  495. var repoCalls atomic.Int32
  496. cache := &authCacheStub{}
  497. repo := &authRepoStub{
  498. getByKeyForAuth: func(ctx context.Context, key string) (*APIKey, error) {
  499. repoCalls.Add(1)
  500. return nil, ErrAPIKeyNotFound
  501. },
  502. }
  503. cfg := &config.Config{
  504. APIKeyAuth: config.APIKeyAuthCacheConfig{
  505. L1Size: 100,
  506. L1TTLSeconds: 60,
  507. L2TTLSeconds: 60,
  508. NegativeTTLSeconds: 30,
  509. },
  510. }
  511. svc := NewAPIKeyService(repo, nil, nil, nil, nil, cache, cfg)
  512. cache.getAuthCache = func(ctx context.Context, key string) (*APIKeyAuthCacheEntry, error) {
  513. return nil, redis.Nil
  514. }
  515. _, err := svc.GetByKey(context.Background(), "missing")
  516. require.ErrorIs(t, err, ErrAPIKeyNotFound)
  517. require.Empty(t, cache.setAuthKeys, "attacker-controlled misses must not be written to Redis")
  518. svc.authNegativeCacheL1.Wait()
  519. _, err = svc.GetByKey(context.Background(), "missing")
  520. require.ErrorIs(t, err, ErrAPIKeyNotFound)
  521. require.Equal(t, int32(1), repoCalls.Load())
  522. }
  523. func TestAPIKeyService_GetByKeyRejectsInvalidLengthBeforeCaches(t *testing.T) {
  524. var cacheCalls atomic.Int32
  525. cache := &authCacheStub{getAuthCache: func(context.Context, string) (*APIKeyAuthCacheEntry, error) {
  526. cacheCalls.Add(1)
  527. return nil, redis.Nil
  528. }}
  529. repo := &authRepoStub{getByKeyForAuth: func(context.Context, string) (*APIKey, error) {
  530. t.Fatal("invalid credential reached repository")
  531. return nil, nil
  532. }}
  533. svc := NewAPIKeyService(repo, nil, nil, nil, nil, cache, &config.Config{APIKeyAuth: config.APIKeyAuthCacheConfig{L2TTLSeconds: 60}})
  534. for _, key := range []string{"", strings.Repeat("x", MaxAPIKeyCredentialBytes+1)} {
  535. _, err := svc.GetByKey(context.Background(), key)
  536. require.ErrorIs(t, err, ErrAPIKeyNotFound)
  537. }
  538. require.Zero(t, cacheCalls.Load())
  539. }
  540. func TestAPIKeyService_GetByKeyAllowsMaximumLength(t *testing.T) {
  541. key := strings.Repeat("x", MaxAPIKeyCredentialBytes)
  542. var repoCalls atomic.Int32
  543. repo := &authRepoStub{getByKeyForAuth: func(_ context.Context, got string) (*APIKey, error) {
  544. repoCalls.Add(1)
  545. require.Equal(t, key, got)
  546. return nil, ErrAPIKeyNotFound
  547. }}
  548. svc := NewAPIKeyService(repo, nil, nil, nil, nil, nil, &config.Config{})
  549. _, err := svc.GetByKey(context.Background(), key)
  550. require.ErrorIs(t, err, ErrAPIKeyNotFound)
  551. require.Equal(t, int32(1), repoCalls.Load())
  552. }
  553. func TestAPIKeyService_AuthLookupBulkheadRejectsExcessMisses(t *testing.T) {
  554. entered := make(chan struct{})
  555. release := make(chan struct{})
  556. repo := &authRepoStub{getByKeyForAuth: func(context.Context, string) (*APIKey, error) {
  557. close(entered)
  558. <-release
  559. return nil, ErrAPIKeyNotFound
  560. }}
  561. svc := NewAPIKeyService(repo, nil, nil, nil, nil, nil, &config.Config{APIKeyAuth: config.APIKeyAuthCacheConfig{LookupConcurrency: 1}})
  562. done := make(chan error, 1)
  563. go func() {
  564. _, err := svc.GetByKey(context.Background(), "first")
  565. done <- err
  566. }()
  567. <-entered
  568. _, err := svc.GetByKey(context.Background(), "second")
  569. require.ErrorIs(t, err, ErrAPIKeyAuthOverloaded)
  570. metrics := svc.AuthLookupMetrics()
  571. require.Equal(t, uint64(2), metrics.Total)
  572. require.Equal(t, uint64(1), metrics.Rejected)
  573. require.Equal(t, int64(1), metrics.InFlight)
  574. require.Equal(t, 1, metrics.Capacity)
  575. close(release)
  576. require.ErrorIs(t, <-done, ErrAPIKeyNotFound)
  577. }
  578. func TestAPIKeyService_GetByKey_SingleflightCollapses(t *testing.T) {
  579. var calls int32
  580. cache := &authCacheStub{}
  581. repo := &authRepoStub{
  582. getByKeyForAuth: func(ctx context.Context, key string) (*APIKey, error) {
  583. atomic.AddInt32(&calls, 1)
  584. time.Sleep(50 * time.Millisecond)
  585. return &APIKey{
  586. ID: 11,
  587. UserID: 2,
  588. Status: StatusActive,
  589. User: &User{
  590. ID: 2,
  591. Status: StatusActive,
  592. Role: RoleUser,
  593. Balance: 1,
  594. Concurrency: 1,
  595. },
  596. }, nil
  597. },
  598. }
  599. cfg := &config.Config{
  600. APIKeyAuth: config.APIKeyAuthCacheConfig{
  601. Singleflight: true,
  602. },
  603. }
  604. svc := NewAPIKeyService(repo, nil, nil, nil, nil, cache, cfg)
  605. start := make(chan struct{})
  606. wg := sync.WaitGroup{}
  607. errs := make([]error, 5)
  608. for i := 0; i < 5; i++ {
  609. wg.Add(1)
  610. go func(idx int) {
  611. defer wg.Done()
  612. <-start
  613. _, err := svc.GetByKey(context.Background(), "k1")
  614. errs[idx] = err
  615. }(i)
  616. }
  617. close(start)
  618. wg.Wait()
  619. for _, err := range errs {
  620. require.NoError(t, err)
  621. }
  622. require.Equal(t, int32(1), atomic.LoadInt32(&calls))
  623. }