diff --git a/plugins/auth-backend/config.d.ts b/plugins/auth-backend/config.d.ts index 280d92a261..6f31cd07df 100644 --- a/plugins/auth-backend/config.d.ts +++ b/plugins/auth-backend/config.d.ts @@ -52,6 +52,8 @@ export interface Config { keyFilename?: string; /** The path to use for the collection. Defaults to 'sessions' */ path?: string; + /** Timeout used for database operations. Defaults to 10000ms */ + timeout?: number; }; }; diff --git a/plugins/auth-backend/src/identity/FirestoreKeyStore.test.ts b/plugins/auth-backend/src/identity/FirestoreKeyStore.test.ts index 398e3c2b34..a7192ed4ab 100644 --- a/plugins/auth-backend/src/identity/FirestoreKeyStore.test.ts +++ b/plugins/auth-backend/src/identity/FirestoreKeyStore.test.ts @@ -17,6 +17,8 @@ import { Firestore } from '@google-cloud/firestore'; import { + DEFAULT_DOCUMENT_PATH, + DEFAULT_TIMEOUT_MS, FirestoreKeyStore, FirestoreKeyStoreSettings, } from './FirestoreKeyStore'; @@ -24,15 +26,18 @@ import { AnyJWK } from './types'; const data = jest.fn().mockReturnValue('data'); const toDate = jest.fn().mockReturnValue('date'); +const get = jest.fn().mockReturnValue({ + docs: [{ data, createTime: { toDate } }], +}); +const set = jest.fn(); const firestoreMock = { + limit: jest.fn().mockReturnThis(), collection: jest.fn().mockReturnThis(), delete: jest.fn(), doc: jest.fn().mockReturnThis(), - get: jest.fn().mockReturnValue({ - docs: [{ data, createTime: { toDate } }], - }), - set: jest.fn(), + set, + get, }; jest.mock('@google-cloud/firestore', () => ({ @@ -55,10 +60,20 @@ describe('FirestoreKeyStore', () => { keyFilename: 'cred.json', }; const path = 'my-path'; - const firestoreSettings = { ...settings, path } as FirestoreKeyStoreSettings; + const timeout = 10; + const firestoreSettings = { + ...settings, + path, + timeout, + } as FirestoreKeyStoreSettings; beforeEach(() => { jest.clearAllMocks(); + jest.useFakeTimers(); + }); + + afterEach(() => { + jest.useRealTimers(); }); it('can create an instance without settings', async () => { @@ -74,17 +89,56 @@ describe('FirestoreKeyStore', () => { expect(Firestore).toHaveBeenCalledWith(settings); }); - it('can use a default path', async () => { + it('can verify that is has a connection to the database', async () => { + const keyStore = await FirestoreKeyStore.create(); + + await expect( + FirestoreKeyStore.verifyConnection(keyStore), + ).resolves.not.toThrow(); + }); + + it('can verify that it can not connect to the database', async () => { + const keyStore = await FirestoreKeyStore.create(); + firestoreMock.get = jest.fn().mockRejectedValue(new Error()); + + await expect( + FirestoreKeyStore.verifyConnection(keyStore), + ).rejects.toThrow(); + + firestoreMock.get = get; + }); + + it('can use a default timeout and path', async () => { const keyStore = await FirestoreKeyStore.create(); await keyStore.addKey(key); - expect(firestoreMock.collection).toBeCalledWith('sessions'); + expect(setTimeout).toBeCalledWith(expect.any(Function), DEFAULT_TIMEOUT_MS); + expect(firestoreMock.collection).toBeCalledWith(DEFAULT_DOCUMENT_PATH); + }); + + it('can handle a timeout', async () => { + firestoreMock.set = jest + .fn() + .mockImplementation( + () => new Promise(resolve => setTimeout(resolve, 20)), + ); + const keyStore = await FirestoreKeyStore.create(firestoreSettings); + const add = keyStore.addKey(key); + + jest.advanceTimersByTime(50); + + await expect(add).rejects.toEqual( + new Error(`Operation timed out after ${timeout}ms`), + ); + + firestoreMock.set = set; }); it('can add keys', async () => { const keyStore = await FirestoreKeyStore.create(firestoreSettings); await keyStore.addKey(key); + expect(setTimeout).toBeCalledTimes(1); expect(firestoreMock.collection).toBeCalledWith(path); expect(firestoreMock.doc).toBeCalledWith(key.kid); expect(firestoreMock.set).toHaveBeenCalledWith({ @@ -97,6 +151,7 @@ describe('FirestoreKeyStore', () => { const keyStore = await FirestoreKeyStore.create(firestoreSettings); await keyStore.removeKeys(['123']); + expect(setTimeout).toBeCalledTimes(1); expect(firestoreMock.collection).toBeCalledWith(path); expect(firestoreMock.doc).toBeCalledWith('123'); expect(firestoreMock.delete).toBeCalledTimes(1); @@ -106,6 +161,7 @@ describe('FirestoreKeyStore', () => { const keyStore = await FirestoreKeyStore.create(firestoreSettings); await keyStore.removeKeys(['123', '456']); + expect(setTimeout).toBeCalledTimes(2); expect(firestoreMock.collection).toBeCalledWith(path); expect(firestoreMock.doc).toBeCalledWith('123'); expect(firestoreMock.doc).toBeCalledWith('456'); @@ -116,6 +172,7 @@ describe('FirestoreKeyStore', () => { const keyStore = await FirestoreKeyStore.create(firestoreSettings); const items = await keyStore.listKeys(); + expect(setTimeout).toBeCalledTimes(1); expect(firestoreMock.collection).toBeCalledWith(path); expect(firestoreMock.get).toBeCalledTimes(1); expect(data).toBeCalledTimes(1); diff --git a/plugins/auth-backend/src/identity/FirestoreKeyStore.ts b/plugins/auth-backend/src/identity/FirestoreKeyStore.ts index f606bacd38..a5479a6470 100644 --- a/plugins/auth-backend/src/identity/FirestoreKeyStore.ts +++ b/plugins/auth-backend/src/identity/FirestoreKeyStore.ts @@ -14,45 +14,86 @@ * limitations under the License. */ -import { Firestore, Settings } from '@google-cloud/firestore'; +import { Logger } from 'winston'; +import { + DocumentData, + Firestore, + QuerySnapshot, + Settings, + WriteResult, +} from '@google-cloud/firestore'; import { AnyJWK, KeyStore, StoredKey } from './types'; -export type FirestoreKeyStoreSettings = Settings & { +export type FirestoreKeyStoreSettings = Settings & Options; + +type Options = { path?: string; + timeout?: number; }; +export const DEFAULT_TIMEOUT_MS = 10000; +export const DEFAULT_DOCUMENT_PATH = 'sessions'; + export class FirestoreKeyStore implements KeyStore { static async create( settings?: FirestoreKeyStoreSettings, ): Promise { - const { path, ...firestoreSettings } = settings ?? {}; + const { path, timeout, ...firestoreSettings } = settings ?? {}; const database = new Firestore(firestoreSettings); - return new FirestoreKeyStore(database, path ?? 'sessions'); + return new FirestoreKeyStore( + database, + path ?? DEFAULT_DOCUMENT_PATH, + timeout ?? DEFAULT_TIMEOUT_MS, + ); } private constructor( private readonly database: Firestore, private readonly path: string, + private readonly timeout: number, ) {} + static async verifyConnection( + keyStore: FirestoreKeyStore, + logger?: Logger, + ): Promise { + try { + await keyStore.verify(); + } catch (error) { + if (process.env.NODE_ENV !== 'development') { + throw new Error( + `Failed to connect to database: ${(error as Error).message}`, + ); + } + logger?.warn( + `Failed to connect to database: ${(error as Error).message}`, + ); + } + } + async addKey(key: AnyJWK): Promise { - await this.database - .collection(this.path) - .doc(key.kid) - .set({ - kid: key.kid, - key: JSON.stringify(key), - }); + await this.withTimeout( + this.database + .collection(this.path) + .doc(key.kid) + .set({ + kid: key.kid, + key: JSON.stringify(key), + }), + ); } async listKeys(): Promise<{ items: StoredKey[] }> { - const docs = await this.database.collection(this.path).get(); + const keys = await this.withTimeout>( + this.database.collection(this.path).get(), + ); + return { - items: docs.docs.map(doc => ({ - key: doc.data() as AnyJWK, - createdAt: doc.createTime.toDate(), + items: keys.docs.map(key => ({ + key: key.data() as AnyJWK, + createdAt: key.createTime.toDate(), })), }; } @@ -60,7 +101,9 @@ export class FirestoreKeyStore implements KeyStore { async removeKeys(kids: string[]): Promise { // This is probably really slow, but it's done async in the background for (const kid of kids) { - await this.database.collection(this.path).doc(kid).delete(); + await this.withTimeout( + this.database.collection(this.path).doc(kid).delete(), + ); } /** @@ -86,4 +129,29 @@ export class FirestoreKeyStore implements KeyStore { * */ } + + /** + * Helper function to allow us to modify the timeout used when + * performing Firestore database operations. + * + * The reason for this is that it seems that there's no other + * practical solution to change the default timeout of 10mins + * that Firestore has. + * + */ + private async withTimeout(operation: Promise): Promise { + const timer = new Promise((_, reject) => + setTimeout(() => { + reject(new Error(`Operation timed out after ${this.timeout}ms`)); + }, this.timeout), + ); + return Promise.race([operation, timer]); + } + + /** + * Used to verify that the database is reachable. + */ + private async verify(): Promise { + await this.withTimeout(this.database.collection(this.path).limit(1).get()); + } } diff --git a/plugins/auth-backend/src/identity/KeyStores.test.ts b/plugins/auth-backend/src/identity/KeyStores.test.ts index 957c4c7953..bb91d2cae9 100644 --- a/plugins/auth-backend/src/identity/KeyStores.test.ts +++ b/plugins/auth-backend/src/identity/KeyStores.test.ts @@ -62,6 +62,8 @@ describe('KeyStores', () => { }); it('can handle additional provider config', async () => { + jest.spyOn(FirestoreKeyStore, 'verifyConnection').mockResolvedValue(); + const configOptions = { auth: { keyStore: { diff --git a/plugins/auth-backend/src/identity/KeyStores.ts b/plugins/auth-backend/src/identity/KeyStores.ts index 1b19bde911..8ac1cf1ec6 100644 --- a/plugins/auth-backend/src/identity/KeyStores.ts +++ b/plugins/auth-backend/src/identity/KeyStores.ts @@ -66,8 +66,10 @@ export class KeyStores { if (provider === 'firestore') { const settings = ks?.getOptional(provider) as FirestoreKeyStoreSettings; + const keyStore = await FirestoreKeyStore.create(settings); + await FirestoreKeyStore.verifyConnection(keyStore, logger); - return await FirestoreKeyStore.create(settings); + return keyStore; } throw new Error(`Unknown KeyStore provider: ${provider}`);