Wrap operations in a configurable timeout and add method to verify the database connection

Signed-off-by: Marcus Eide <eide@spotify.com>

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

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

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<FirestoreKeyStore> {
-    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<void> {
+    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<void> {
-    await this.database
-      .collection(this.path)
-      .doc(key.kid)
-      .set({
-        kid: key.kid,
-        key: JSON.stringify(key),
-      });
+    await this.withTimeout<WriteResult>(
+      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<QuerySnapshot<DocumentData>>(
+      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<void> {
     // 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<WriteResult>(
+        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<T>(operation: Promise<T>): Promise<T> {
+    const timer = new Promise<never>((_, reject) =>
+      setTimeout(() => {
+        reject(new Error(`Operation timed out after ${this.timeout}ms`));
+      }, this.timeout),
+    );
+    return Promise.race<T>([operation, timer]);
+  }
+
+  /**
+   * Used to verify that the database is reachable.
+   */
+  private async verify(): Promise<void> {
+    await this.withTimeout(this.database.collection(this.path).limit(1).get());
+  }
 }

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: {

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