chore: cache query misses to protect against DDoS (#6771)

## About the changes
This PR establishes a simple yet effective mechanism to avoid DDoS
against our DB while also protecting against memory leaks.

This will enable us to release the flag `queryMissingTokens` to make our
token validation consistent across different nodes

---------

Co-authored-by: Nuno Góis <github@nunogois.com>

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);
+    });
 });

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<string, Date>();
+
     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);

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);
+});

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<IApiToken> {
-        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<IApiToken[]> {