Skip to content

Commit

Permalink
chore: cache query misses to protect against DDoS (#6771)
Browse files Browse the repository at this point in the history
## 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>
  • Loading branch information
gastonfournier and nunogois committed Apr 3, 2024
1 parent d466f60 commit d7ab886
Show file tree
Hide file tree
Showing 4 changed files with 81 additions and 7 deletions.
31 changes: 30 additions & 1 deletion src/lib/services/api-token-service.test.ts
Expand Up @@ -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';
Expand Down Expand Up @@ -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);
});
});
12 changes: 12 additions & 0 deletions src/lib/services/api-token-service.ts
Expand Up @@ -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) {
Expand All @@ -62,6 +63,8 @@ export class ApiTokenService {

private activeTokens: IApiToken[] = [];

private queryAfter = new Map<string, Date>();

private initialized = false;

private eventService: EventService;
Expand Down Expand Up @@ -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);
Expand Down
37 changes: 37 additions & 0 deletions 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);
});
8 changes: 2 additions & 6 deletions src/test/fixtures/fake-api-token-store.ts
Expand Up @@ -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
Expand Down Expand Up @@ -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[]> {
Expand Down

0 comments on commit d7ab886

Please sign in to comment.