feat: Add init api tokens option (#1181)

Adds support for initializing a fresh Unleash instance with predefined API tokens. Co-authored-by: sighphyre <liquidwicked64@gmail.com> Co-authored-by: Juraj Malenica <juraj.malenica@mindsmiths.com> Co-authored-by: Ivar Conradi Østhus <ivarconr@gmail.com>
Unleash · Jan 5, 2022 · e757c00 · e757c00 · vercel · Jan 5, 2022
1 parent e814a5f
commit e757c00
Show file tree

Hide file tree

Showing 12 changed files with 265 additions and 31 deletions.
diff --git a/src/lib/__snapshots__/create-config.test.ts.snap b/src/lib/__snapshots__/create-config.test.ts.snap
@@ -6,6 +6,7 @@ Object {
     "createAdminUser": true,
     "customAuthHandler": [Function],
     "enableApiToken": true,
+    "initApiTokens": Array [],
     "type": "open-source",
   },
   "db": Object {

diff --git a/src/lib/create-config.test.ts b/src/lib/create-config.test.ts
@@ -1,5 +1,5 @@
-// @ts-nocheck
 import { createConfig } from './create-config';
+import { ApiTokenType } from './types/models/api-token';
 
 test('should create default config', async () => {
     const config = createConfig({
@@ -17,3 +17,105 @@ test('should create default config', async () => {
 
     expect(config).toMatchSnapshot();
 });
+
+test('should add initApiToken from options', async () => {
+    const token = {
+        environment: '*',
+        project: '*',
+        secret: '*:*:some-random-string',
+        type: ApiTokenType.ADMIN,
+        username: 'admin',
+    };
+    const config = createConfig({
+        db: {
+            host: 'localhost',
+            port: 4242,
+            user: 'unleash',
+            password: 'password',
+            database: 'unleash_db',
+        },
+        server: {
+            port: 4242,
+        },
+        authentication: {
+            initApiTokens: [token],
+        },
+    });
+
+    expect(config.authentication.initApiTokens).toHaveLength(1);
+    expect(config.authentication.initApiTokens[0].environment).toBe(
+        token.environment,
+    );
+    expect(config.authentication.initApiTokens[0].project).toBe(token.project);
+    expect(config.authentication.initApiTokens[0].type).toBe(
+        ApiTokenType.ADMIN,
+    );
+});
+
+test('should add initApiToken from env var', async () => {
+    process.env.INIT_ADMIN_API_TOKENS = '*:*:some-token1, *:*:some-token2';
+
+    const config = createConfig({
+        db: {
+            host: 'localhost',
+            port: 4242,
+            user: 'unleash',
+            password: 'password',
+            database: 'unleash_db',
+        },
+        server: {
+            port: 4242,
+        },
+    });
+
+    expect(config.authentication.initApiTokens).toHaveLength(2);
+    expect(config.authentication.initApiTokens[0].environment).toBe('*');
+    expect(config.authentication.initApiTokens[0].project).toBe('*');
+    expect(config.authentication.initApiTokens[0].type).toBe(
+        ApiTokenType.ADMIN,
+    );
+    expect(config.authentication.initApiTokens[1].secret).toBe(
+        '*:*:some-token2',
+    );
+
+    delete process.env.INIT_ADMIN_API_TOKENS;
+});
+
+test('should validate initApiToken from env var', async () => {
+    process.env.INIT_ADMIN_API_TOKENS = 'invalidProject:*:some-token1';
+
+    expect(() => createConfig({})).toThrow(
+        'Admin token cannot be scoped to single project',
+    );
+
+    delete process.env.INIT_ADMIN_API_TOKENS;
+});
+
+test('should merge initApiToken from options and env vars', async () => {
+    process.env.INIT_ADMIN_API_TOKENS = '*:*:some-token1, *:*:some-token2';
+    const token = {
+        environment: '*',
+        project: '*',
+        secret: '*:*:some-random-string',
+        type: ApiTokenType.ADMIN,
+        username: 'admin',
+    };
+    const config = createConfig({
+        db: {
+            host: 'localhost',
+            port: 4242,
+            user: 'unleash',
+            password: 'password',
+            database: 'unleash_db',
+        },
+        server: {
+            port: 4242,
+        },
+        authentication: {
+            initApiTokens: [token],
+        },
+    });
+
+    expect(config.authentication.initApiTokens).toHaveLength(3);
+    delete process.env.INIT_ADMIN_API_TOKENS;
+});
diff --git a/src/lib/create-config.ts b/src/lib/create-config.ts
@@ -21,6 +21,7 @@ import { defaultCustomAuthDenyAll } from './default-custom-auth-deny-all';
 import { formatBaseUri } from './util/format-base-uri';
 import { minutesToMilliseconds, secondsToMilliseconds } from 'date-fns';
 import EventEmitter from 'events';
+import { ApiTokenType, validateApiToken } from './types/models/api-token';
 
 const safeToUpper = (s: string) => (s ? s.toUpperCase() : s);
 
@@ -132,6 +133,7 @@ const defaultAuthentication: IAuthOption = {
     type: authTypeFromString(process.env.AUTH_TYPE),
     customAuthHandler: defaultCustomAuthDenyAll,
     createAdminUser: true,
+    initApiTokens: [],
 };
 
 const defaultImport: IImportOption = {
@@ -179,6 +181,28 @@ const formatServerOptions = (
     };
 };
 
+const loadInitApiTokens = () => {
+    if (process.env.INIT_ADMIN_API_TOKENS) {
+        const initApiTokens = process.env.INIT_ADMIN_API_TOKENS.split(/,\s?/);
+        const tokens = initApiTokens.map((secret) => {
+            const [project = '*', environment = '*'] = secret.split(':');
+            const token = {
+                createdAt: undefined,
+                project,
+                environment,
+                secret,
+                type: ApiTokenType.ADMIN,
+                username: 'admin',
+            };
+            validateApiToken(token);
+            return token;
+        });
+        return tokens;
+    } else {
+        return [];
+    }
+};
+
 export function createConfig(options: IUnleashOptions): IUnleashConfig {
     let extraDbOptions = {};
 
@@ -227,11 +251,14 @@ export function createConfig(options: IUnleashOptions): IUnleashConfig {
         options.versionCheck,
     ]);
 
+    const initApiTokens = loadInitApiTokens();
+
     const authentication: IAuthOption = mergeAll([
         defaultAuthentication,
         options.authentication
             ? removeUndefinedKeys(options.authentication)
             : options.authentication,
+        { initApiTokens: initApiTokens },
     ]);
 
     const importSetting: IImportOption = mergeAll([

diff --git a/src/lib/db/api-token-store.ts b/src/lib/db/api-token-store.ts
@@ -64,6 +64,12 @@ export class ApiTokenStore implements IApiTokenStore {
             });
     }
 
+    count(): Promise<number> {
+        return this.db(TABLE)
+            .count('*')
+            .then((res) => Number(res[0].count));
+    }
+
     async getAll(): Promise<IApiToken[]> {
         const stopTimer = this.timer('getAll');
         const rows = await this.db<ITokenTable>(TABLE);

diff --git a/src/lib/services/api-token-service.test.ts b/src/lib/services/api-token-service.test.ts
@@ -0,0 +1,33 @@
+import { ApiTokenService } from './api-token-service';
+import { createTestConfig } from '../../test/config/test-config';
+import { IUnleashConfig } from '../server-impl';
+import { ApiTokenType } from '../types/models/api-token';
+import FakeApiTokenStore from '../../test/fixtures/fake-api-token-store';
+
+test('Should init api token', async () => {
+    const token = {
+        environment: '*',
+        project: '*',
+        secret: '*:*:some-random-string',
+        type: ApiTokenType.ADMIN,
+        username: 'admin',
+    };
+
+    const config: IUnleashConfig = createTestConfig({
+        authentication: {
+            initApiTokens: [token],
+        },
+    });
+    const apiTokenStore = new FakeApiTokenStore();
+    const insertCalled = new Promise((resolve) => {
+        apiTokenStore.on('insert', resolve);
+    });
+
+    new ApiTokenService({ apiTokenStore }, config);
+
+    await insertCalled;
+
+    const tokens = await apiTokenStore.getAll();
+
+    expect(tokens).toHaveLength(1);
+});
diff --git a/src/lib/services/api-token-service.ts b/src/lib/services/api-token-service.ts
@@ -5,10 +5,10 @@ import { IUnleashStores } from '../types/stores';
 import { IUnleashConfig } from '../types/option';
 import ApiUser from '../types/api-user';
 import {
-    ALL,
     ApiTokenType,
     IApiToken,
     IApiTokenCreate,
+    validateApiToken,
 } from '../types/models/api-token';
 import { IApiTokenStore } from '../types/stores/api-token-store';
 import { FOREIGN_KEY_VIOLATION } from '../error/db-error';
@@ -26,7 +26,7 @@ export class ApiTokenService {
 
     constructor(
         { apiTokenStore }: Pick<IUnleashStores, 'apiTokenStore'>,
-        config: Pick<IUnleashConfig, 'getLogger'>,
+        config: Pick<IUnleashConfig, 'getLogger' | 'authentication'>,
     ) {
         this.store = apiTokenStore;
         this.logger = config.getLogger('/services/api-token-service.ts');
@@ -35,6 +35,11 @@ export class ApiTokenService {
             () => this.fetchActiveTokens(),
             minutesToMilliseconds(1),
         ).unref();
+        if (config.authentication.initApiTokens.length > 0) {
+            process.nextTick(async () =>
+                this.initApiTokens(config.authentication.initApiTokens),
+            );
+        }
     }
 
     private async fetchActiveTokens(): Promise<void> {
@@ -54,6 +59,19 @@ export class ApiTokenService {
         return this.store.getAllActive();
     }
 
+    private async initApiTokens(tokens: IApiTokenCreate[]) {
+        const tokenCount = await this.store.count();
+        if (tokenCount > 0) {
+            return;
+        }
+        try {
+            const createAll = tokens.map((t) => this.insertNewApiToken(t));
+            await Promise.all(createAll);
+        } catch (e) {
+            this.logger.error('Unable to create initial Admin API tokens');
+        }
+    }
+
     public getUserForToken(secret: string): ApiUser | undefined {
         const token = this.activeTokens.find((t) => t.secret === secret);
         if (token) {
@@ -82,45 +100,30 @@ export class ApiTokenService {
         return this.store.delete(secret);
     }
 
-    private validateNewApiToken({ type, project, environment }) {
-        if (type === ApiTokenType.ADMIN && project !== ALL) {
-            throw new BadDataError(
-                'Admin token cannot be scoped to single project',
-            );
-        }
-
-        if (type === ApiTokenType.ADMIN && environment !== ALL) {
-            throw new BadDataError(
-                'Admin token cannot be scoped to single environment',
-            );
-        }
-
-        if (type === ApiTokenType.CLIENT && environment === ALL) {
-            throw new BadDataError(
-                'Client token cannot be scoped to all environments',
-            );
-        }
-    }
-
     public async createApiToken(
         newToken: Omit<IApiTokenCreate, 'secret'>,
     ): Promise<IApiToken> {
-        this.validateNewApiToken(newToken);
+        validateApiToken(newToken);
 
         const secret = this.generateSecretKey(newToken);
         const createNewToken = { ...newToken, secret };
+        return this.insertNewApiToken(createNewToken);
+    }
 
+    private async insertNewApiToken(
+        newApiToken: IApiTokenCreate,
+    ): Promise<IApiToken> {
         try {
-            const token = await this.store.insert(createNewToken);
+            const token = await this.store.insert(newApiToken);
             this.activeTokens.push(token);
             return token;
         } catch (error) {
             if (error.code === FOREIGN_KEY_VIOLATION) {
                 let { message } = error;
                 if (error.constraint === 'api_tokens_project_fkey') {
-                    message = `Project=${newToken.project} does not exist`;
+                    message = `Project=${newApiToken.project} does not exist`;
                 } else if (error.constraint === 'api_tokens_environment_fkey') {
-                    message = `Environment=${newToken.environment} does not exist`;
+                    message = `Environment=${newApiToken.environment} does not exist`;
                 }
                 throw new BadDataError(message);
             }

diff --git a/src/lib/types/models/api-token.ts b/src/lib/types/models/api-token.ts
@@ -1,3 +1,5 @@
+import BadDataError from '../../error/bad-data-error';
+
 export const ALL = '*';
 
 export enum ApiTokenType {
@@ -20,3 +22,27 @@ export interface IApiToken extends IApiTokenCreate {
     environment: string;
     project: string;
 }
+
+export const validateApiToken = ({
+    type,
+    project,
+    environment,
+}: Omit<IApiTokenCreate, 'secret'>): void => {
+    if (type === ApiTokenType.ADMIN && project !== ALL) {
+        throw new BadDataError(
+            'Admin token cannot be scoped to single project',
+        );
+    }
+
+    if (type === ApiTokenType.ADMIN && environment !== ALL) {
+        throw new BadDataError(
+            'Admin token cannot be scoped to single environment',
+        );
+    }
+
+    if (type === ApiTokenType.CLIENT && environment === ALL) {
+        throw new BadDataError(
+            'Client token cannot be scoped to all environments',
+        );
+    }
+};
diff --git a/src/lib/types/option.ts b/src/lib/types/option.ts
@@ -1,5 +1,6 @@
 import EventEmitter from 'events';
 import { LogLevel, LogProvider } from '../logger';
+import { IApiTokenCreate } from './models/api-token';
 
 export type EventHook = (eventName: string, data: object) => void;
 
@@ -53,6 +54,7 @@ export interface IAuthOption {
     type: IAuthType;
     customAuthHandler?: Function;
     createAdminUser: boolean;
+    initApiTokens: IApiTokenCreate[];
 }
 
 export interface IImportOption {

diff --git a/src/lib/types/stores/api-token-store.ts b/src/lib/types/stores/api-token-store.ts
@@ -6,4 +6,5 @@ export interface IApiTokenStore extends Store<IApiToken, string> {
     insert(newToken: IApiTokenCreate): Promise<IApiToken>;
     setExpiry(secret: string, expiresAt: Date): Promise<IApiToken>;
     markSeenAt(secrets: string[]): Promise<void>;
+    count(): Promise<number>;
 }