diff --git a/src/lib/services/api-token-service.test.ts b/src/lib/services/api-token-service.test.ts index 4512cdfe53c..c9fb871ba45 100644 --- a/src/lib/services/api-token-service.test.ts +++ b/src/lib/services/api-token-service.test.ts @@ -11,7 +11,7 @@ import { API_TOKEN_DELETED, API_TOKEN_UPDATED, } from '../types'; -import { addDays } from 'date-fns'; +import { addDays, minutesToMilliseconds } from 'date-fns'; import EventService from '../features/events/event-service'; import FakeFeatureTagStore from '../../test/fixtures/fake-feature-tag-store'; import { createFakeEventsService } from '../../lib/features'; @@ -243,4 +243,33 @@ describe('When token is added by another instance', () => { expect(found).toBeDefined(); expect(found?.username).toBe(token.tokenName); }); + + test('should query the db only once for invalid tokens', async () => { + jest.useFakeTimers(); + const { apiTokenService, apiTokenStore } = setup({ + experimental: { + flags: { + queryMissingTokens: true, + }, + }, + }); + const apiTokenStoreGet = jest.spyOn(apiTokenStore, 'get'); + + const invalidToken = 'invalid-token'; + for (let i = 0; i < 10; i++) { + expect( + await apiTokenService.getUserForToken(invalidToken), + ).toBeUndefined(); + } + expect(apiTokenStoreGet).toHaveBeenCalledTimes(1); + + // after more than 5 minutes we should be able to query again + jest.advanceTimersByTime(minutesToMilliseconds(6)); + for (let i = 0; i < 10; i++) { + expect( + await apiTokenService.getUserForToken(invalidToken), + ).toBeUndefined(); + } + expect(apiTokenStoreGet).toHaveBeenCalledTimes(2); + }); }); diff --git a/src/lib/services/api-token-service.ts b/src/lib/services/api-token-service.ts index 5d76e3201ab..46f03ade4cc 100644 --- a/src/lib/services/api-token-service.ts +++ b/src/lib/services/api-token-service.ts @@ -36,6 +36,7 @@ import { omitKeys, } from '../util'; import type EventService from '../features/events/event-service'; +import { addMinutes, isPast } from 'date-fns'; const resolveTokenPermissions = (tokenType: string) => { if (tokenType === ApiTokenType.ADMIN) { @@ -62,6 +63,8 @@ export class ApiTokenService { private activeTokens: IApiToken[] = []; + private queryAfter = new Map(); + private initialized = false; private eventService: EventService; @@ -185,10 +188,19 @@ export class ApiTokenService { ); } + const nextAllowedQuery = this.queryAfter.get(secret) ?? 0; if ( !token && + isPast(nextAllowedQuery) && this.flagResolver.isEnabled('queryMissingTokens', flagContext) ) { + if (this.queryAfter.size > 1000) { + // establish a max limit for queryAfter size to prevent memory leak + this.queryAfter.clear(); + } + // prevent querying the same invalid secret multiple times. Expire after 5 minutes + this.queryAfter.set(secret, addMinutes(new Date(), 5)); + token = await this.store.get(secret); if (token) { this.activeTokens.push(token); diff --git a/src/test/e2e/stores/api-token-store.e2e.test.ts b/src/test/e2e/stores/api-token-store.e2e.test.ts new file mode 100644 index 00000000000..a8b0fb28b81 --- /dev/null +++ b/src/test/e2e/stores/api-token-store.e2e.test.ts @@ -0,0 +1,37 @@ +import dbInit, { type ITestDb } from '../helpers/database-init'; +import getLogger from '../../fixtures/no-logger'; +import type { IUnleashStores } from '../../../lib/types'; +import { ApiTokenType } from '../../../lib/types/models/api-token'; + +let stores: IUnleashStores; +let db: ITestDb; + +beforeAll(async () => { + db = await dbInit('api_token_store_serial', getLogger); + stores = db.stores; +}); + +afterAll(async () => { + await db.destroy(); +}); + +test('get token is undefined when not exist', async () => { + const token = await stores.apiTokenStore.get('abcde123'); + expect(token).toBeUndefined(); +}); + +test('get token returns the token when exists', async () => { + const newToken = await stores.apiTokenStore.insert({ + secret: 'abcde321', + environment: 'default', + type: ApiTokenType.ADMIN, + projects: [], + tokenName: 'admin-test-token', + }); + const foundToken = await stores.apiTokenStore.get('abcde321'); + expect(foundToken).toBeDefined(); + expect(foundToken.secret).toBe(newToken.secret); + expect(foundToken.environment).toBe(newToken.environment); + expect(foundToken.tokenName).toBe(newToken.tokenName); + expect(foundToken.type).toBe(newToken.type); +}); diff --git a/src/test/fixtures/fake-api-token-store.ts b/src/test/fixtures/fake-api-token-store.ts index b4788cabacf..13f60c67ad7 100644 --- a/src/test/fixtures/fake-api-token-store.ts +++ b/src/test/fixtures/fake-api-token-store.ts @@ -5,7 +5,6 @@ import type { IApiTokenCreate, } from '../../lib/types/models/api-token'; -import NotFoundError from '../../lib/error/notfound-error'; import EventEmitter from 'events'; export default class FakeApiTokenStore @@ -39,11 +38,8 @@ export default class FakeApiTokenStore } async get(key: string): Promise { - const token = this.tokens.find((t) => t.secret === key); - if (token) { - return token; - } - throw new NotFoundError(`Could not find token with secret ${key}`); + // get can return undefined. See api-token-store.e2e.test.ts + return this.tokens.find((t) => t.secret === key); } async getAll(): Promise {