From d8ea1edcdbead3460d1cbfbfc56d067b5b476c87 Mon Sep 17 00:00:00 2001
From: Marcus Eide <eide@spotify.com>
Date: Tue, 12 Oct 2021 16:11:29 +0200
Subject: [PATCH 01/14] Add support for auth.keyStore in application config

Signed-off-by: Marcus Eide <eide@spotify.com>
---
 app-config.yaml                  |  9 +++++++++
 plugins/auth-backend/config.d.ts | 18 ++++++++++++++++++
 2 files changed, 27 insertions(+)

diff --git a/app-config.yaml b/app-config.yaml
index c24ac1701b..60e2a2b9e0 100644
--- a/app-config.yaml
+++ b/app-config.yaml
@@ -292,6 +292,15 @@ scaffolder:
     visibility: public # or or 'private'
 
 auth:
+  ### Add auth.keyStore.provider to more granularly control how to store JWK data when running
+  # the auth-backend.
+  #
+  # keyStore:
+  #  provider: firestore
+  #  firestore:
+  #    projectId: my-project
+  #    path: my-sessions
+
   environment: development
   ### Providing an auth.session.secret will enable session support in the auth-backend
   # session:
diff --git a/plugins/auth-backend/config.d.ts b/plugins/auth-backend/config.d.ts
index 61bbe15d05..b15b0da628 100644
--- a/plugins/auth-backend/config.d.ts
+++ b/plugins/auth-backend/config.d.ts
@@ -31,6 +31,24 @@ export interface Config {
       secret?: string;
     };
 
+    /** To control how to store JWK data in auth-backend */
+    keyStore?: {
+      provider?: 'postgres' | 'memory' | 'firestore';
+      firestore?: {
+        /** The Google Cloud Project ID */
+        projectId?: string;
+        /**
+         * Local file containing the Service Account credentials.
+         * You can omit this value to automatically read from
+         * GOOGLE_APPLICATION_CREDENTIALS env which is useful for local
+         * development.
+         */
+        keyFilename?: string;
+        /** The path to use for the collection. Defaults to 'sessions' */
+        path?: string;
+      };
+    };
+
     /**
      * The available auth-provider options and attributes
      */

From 3a62f7619a9a915dec5bfb22ce3164c3a1e947c8 Mon Sep 17 00:00:00 2001
From: Marcus Eide <eide@spotify.com>
Date: Tue, 12 Oct 2021 16:14:30 +0200
Subject: [PATCH 02/14] Add class to use Firestore as a key-store

Signed-off-by: Marcus Eide <eide@spotify.com>
---
 .../src/identity/FirestoreKeyStore.test.ts    | 152 ++++++++++++++++++
 .../src/identity/FirestoreKeyStore.ts         |  90 +++++++++++
 yarn.lock                                     |  52 +++++-
 3 files changed, 293 insertions(+), 1 deletion(-)
 create mode 100644 plugins/auth-backend/src/identity/FirestoreKeyStore.test.ts
 create mode 100644 plugins/auth-backend/src/identity/FirestoreKeyStore.ts

diff --git a/plugins/auth-backend/src/identity/FirestoreKeyStore.test.ts b/plugins/auth-backend/src/identity/FirestoreKeyStore.test.ts
new file mode 100644
index 0000000000..03fcc15a77
--- /dev/null
+++ b/plugins/auth-backend/src/identity/FirestoreKeyStore.test.ts
@@ -0,0 +1,152 @@
+/*
+ * Copyright 2021 The Backstage Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import { Firestore } from '@google-cloud/firestore';
+
+import { FirestoreKeyStore } from './FirestoreKeyStore';
+import { AnyJWK } from './types';
+
+const data = jest.fn().mockReturnValue('data');
+const toDate = jest.fn().mockReturnValue('date');
+
+const firestoreMock = {
+  collection: jest.fn().mockReturnThis(),
+  delete: jest.fn(),
+  doc: jest.fn().mockReturnThis(),
+  get: jest.fn().mockReturnValue({
+    docs: [{ data, createTime: { toDate } }],
+  }),
+  set: jest.fn(),
+};
+
+jest.mock('@google-cloud/firestore', () => ({
+  Firestore: jest.fn().mockImplementation(() => firestoreMock),
+}));
+
+describe('FirestoreKeyStore', () => {
+  const OLD_ENV = process.env;
+  const key = {
+    kid: '123',
+    use: 'sig',
+    kty: 'plain',
+    alg: 'Base64',
+  } as AnyJWK;
+
+  beforeEach(() => {
+    jest.clearAllMocks();
+  });
+
+  afterAll(() => {
+    process.env = OLD_ENV;
+  });
+
+  it('can create an instance without settings', () => {
+    const keyStore = FirestoreKeyStore.create();
+
+    expect(keyStore).toBeInstanceOf(FirestoreKeyStore);
+  });
+
+  it('can set the project id', async () => {
+    FirestoreKeyStore.create({ projectId: 'my-project' });
+
+    expect(Firestore).toHaveBeenCalledWith(
+      expect.objectContaining({
+        projectId: 'my-project',
+      }),
+    );
+  });
+
+  it('can handle keyfile file', async () => {
+    FirestoreKeyStore.create({ keyFilename: 'keyFile.json' });
+
+    expect(Firestore).toHaveBeenCalledWith(
+      expect.objectContaining({
+        keyFilename: 'keyFile.json',
+      }),
+    );
+  });
+
+  it('can use default google credentials', () => {
+    process.env.GOOGLE_APPLICATION_CREDENTIALS = 'cred.json';
+    FirestoreKeyStore.create();
+
+    expect(Firestore).toHaveBeenCalledWith(
+      expect.objectContaining({
+        keyFilename: 'cred.json',
+      }),
+    );
+  });
+
+  it('can uses a default path', async () => {
+    const keyStore = FirestoreKeyStore.create();
+    await keyStore.addKey(key);
+
+    expect(firestoreMock.collection).toBeCalledWith('sessions');
+  });
+
+  it('can set the path', async () => {
+    const keyStore = FirestoreKeyStore.create({
+      path: 'my-path',
+    });
+    await keyStore.addKey(key);
+
+    expect(firestoreMock.collection).toBeCalledWith('my-path');
+  });
+
+  it('can add keys', async () => {
+    const keyStore = FirestoreKeyStore.create();
+    await keyStore.addKey(key);
+
+    expect(firestoreMock.collection).toBeCalledWith('sessions');
+    expect(firestoreMock.doc).toBeCalledWith(key.kid);
+    expect(firestoreMock.set).toHaveBeenCalledWith({
+      kid: key.kid,
+      key: JSON.stringify(key),
+    });
+  });
+
+  it('can delete a single key', async () => {
+    const keyStore = FirestoreKeyStore.create();
+    await keyStore.removeKeys(['123']);
+
+    expect(firestoreMock.collection).toBeCalledWith('sessions');
+    expect(firestoreMock.doc).toBeCalledWith('123');
+    expect(firestoreMock.delete).toBeCalledTimes(1);
+  });
+
+  it('can delete a multiple keys', async () => {
+    const keyStore = FirestoreKeyStore.create();
+    await keyStore.removeKeys(['123', '456']);
+
+    expect(firestoreMock.collection).toBeCalledWith('sessions');
+    expect(firestoreMock.doc).toBeCalledWith('123');
+    expect(firestoreMock.doc).toBeCalledWith('456');
+    expect(firestoreMock.delete).toBeCalledTimes(2);
+  });
+
+  it('can list keys', async () => {
+    const keyStore = FirestoreKeyStore.create();
+    const items = await keyStore.listKeys();
+
+    expect(firestoreMock.collection).toBeCalledWith('sessions');
+    expect(firestoreMock.get).toBeCalledTimes(1);
+    expect(data).toBeCalledTimes(1);
+    expect(toDate).toBeCalledTimes(1);
+    expect(items).toMatchObject({
+      items: [{ key: 'data', createdAt: 'date' }],
+    });
+  });
+});
diff --git a/plugins/auth-backend/src/identity/FirestoreKeyStore.ts b/plugins/auth-backend/src/identity/FirestoreKeyStore.ts
new file mode 100644
index 0000000000..cf8c082540
--- /dev/null
+++ b/plugins/auth-backend/src/identity/FirestoreKeyStore.ts
@@ -0,0 +1,90 @@
+/*
+ * Copyright 2021 The Backstage Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import { Firestore, Settings } from '@google-cloud/firestore';
+
+import { AnyJWK, KeyStore, StoredKey } from './types';
+
+type FirestoreSettings = Settings & {
+  path?: string;
+};
+
+export class FirestoreKeyStore implements KeyStore {
+  static create(settings?: FirestoreSettings): FirestoreKeyStore {
+    const { projectId, keyFilename, path } = settings ?? {};
+    const database = new Firestore({
+      projectId,
+      keyFilename: keyFilename ?? process.env.GOOGLE_APPLICATION_CREDENTIALS,
+    });
+
+    return new FirestoreKeyStore(database, path ?? 'sessions');
+  }
+
+  private constructor(
+    private readonly database: Firestore,
+    private readonly path: string,
+  ) {}
+
+  async addKey(key: AnyJWK): Promise<void> {
+    await 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();
+    return {
+      items: docs.docs.map(doc => ({
+        key: doc.data() as AnyJWK,
+        createdAt: doc.createTime.toDate(),
+      })),
+    };
+  }
+
+  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();
+    }
+
+    /**
+     * This could be achieved with batching but there's a couple of limitations with that:
+     *
+     * - A batched write can contain a maximum of 500 operations
+     *  https://firebase.google.com/docs/firestore/manage-data/transactions#batched-writes
+     *
+     * - The "in" operator can combine a maximum of 10 equality clauses
+     *  https://firebase.google.com/docs/firestore/query-data/queries#in_not-in_and_array-contains-any
+     *
+     * Example:
+     *
+     *  const batch = this.database.batch();
+     *  const docs = await this.database
+     *    .collection(this.path)
+     *    .where('kid', 'in', kids)
+     *    .get();
+     *  docs.forEach(doc => {
+     *    batch.delete(doc.ref);
+     *  });
+     *  await batch.commit();
+     *
+     */
+  }
+}
diff --git a/yarn.lock b/yarn.lock
index 9915fb6279..b110393471 100644
--- a/yarn.lock
+++ b/yarn.lock
@@ -3000,6 +3000,16 @@
   dependencies:
     google-gax "^2.12.0"
 
+"@google-cloud/firestore@^4.15.1":
+  version "4.15.1"
+  resolved "https://registry.npmjs.org/@google-cloud/firestore/-/firestore-4.15.1.tgz#ed764fc76823ce120e68fe8c27ef1edd0650cd93"
+  integrity sha512-2PWsCkEF1W02QbghSeRsNdYKN1qavrHBP3m72gPDMHQSYrGULOaTi7fSJquQmAtc4iPVB2/x6h80rdLHTATQtA==
+  dependencies:
+    fast-deep-equal "^3.1.1"
+    functional-red-black-tree "^1.0.1"
+    google-gax "^2.24.1"
+    protobufjs "^6.8.6"
+
 "@google-cloud/paginator@^3.0.0":
   version "3.0.5"
   resolved "https://registry.npmjs.org/@google-cloud/paginator/-/paginator-3.0.5.tgz#9d6b96c421a89bd560c1bc2c197c7611ef21db6c"
@@ -15053,6 +15063,21 @@ google-auth-library@^7.0.0, google-auth-library@^7.0.2:
     jws "^4.0.0"
     lru-cache "^6.0.0"
 
+google-auth-library@^7.6.1:
+  version "7.10.0"
+  resolved "https://registry.npmjs.org/google-auth-library/-/google-auth-library-7.10.0.tgz#6ab852f8e1abbe425aec695ac6501f80bd5eba56"
+  integrity sha512-ICsqaU+lxMHVlDUzMrfVIEqnARw2AwBiZ/2KnNM6BcTf9Nott+Af87DTIzmlnW865p3REUP2MVL0xkPC3a61aQ==
+  dependencies:
+    arrify "^2.0.0"
+    base64-js "^1.3.0"
+    ecdsa-sig-formatter "^1.0.11"
+    fast-text-encoding "^1.0.0"
+    gaxios "^4.0.0"
+    gcp-metadata "^4.2.0"
+    gtoken "^5.0.4"
+    jws "^4.0.0"
+    lru-cache "^6.0.0"
+
 google-gax@^2.12.0:
   version "2.14.1"
   resolved "https://registry.npmjs.org/google-gax/-/google-gax-2.14.1.tgz#74885c5d9f01db412917fc49bbf20c4884828d36"
@@ -15071,6 +15096,25 @@ google-gax@^2.12.0:
     protobufjs "^6.10.2"
     retry-request "^4.0.0"
 
+google-gax@^2.24.1:
+  version "2.27.0"
+  resolved "https://registry.npmjs.org/google-gax/-/google-gax-2.27.0.tgz#78ba655ae7707cb92ba37c35932381b5970c1772"
+  integrity sha512-xcLCeNKCqNm/w0At7/vdZHV/zol/iRS+PSAZTu7i6xNGBra/kWI3cfn4M6ZLQXeUEGbTVLJ4zGm53TVc4lvbDA==
+  dependencies:
+    "@grpc/grpc-js" "~1.3.0"
+    "@grpc/proto-loader" "^0.6.1"
+    "@types/long" "^4.0.0"
+    abort-controller "^3.0.0"
+    duplexify "^4.0.0"
+    fast-text-encoding "^1.0.3"
+    google-auth-library "^7.6.1"
+    is-stream-ended "^0.1.4"
+    node-fetch "^2.6.1"
+    object-hash "^2.1.1"
+    proto3-json-serializer "^0.1.1"
+    protobufjs "6.11.2"
+    retry-request "^4.0.0"
+
 google-p12-pem@^3.0.3:
   version "3.0.3"
   resolved "https://registry.npmjs.org/google-p12-pem/-/google-p12-pem-3.0.3.tgz#673ac3a75d3903a87f05878f3c75e06fc151669e"
@@ -19722,6 +19766,7 @@ minipass-fetch@^1.3.0, minipass-fetch@^1.3.2:
   resolved "https://registry.npmjs.org/minipass-fetch/-/minipass-fetch-1.3.3.tgz#34c7cea038c817a8658461bf35174551dce17a0a"
   integrity sha512-akCrLDWfbdAWkMLBxJEeWTdNsjML+dt5YgOI4gJ53vuO0vrmYQkUPxa6j6V65s9CcePIr2SSWqjT2EcrNseryQ==
   dependencies:
+    encoding "^0.1.12"
     minipass "^3.1.0"
     minipass-sized "^1.0.3"
     minizlib "^2.0.0"
@@ -22556,7 +22601,12 @@ proto-list@~1.2.1:
   resolved "https://registry.npmjs.org/proto-list/-/proto-list-1.2.4.tgz#212d5bfe1318306a420f6402b8e26ff39647a849"
   integrity sha1-IS1b/hMYMGpCD2QCuOJv85ZHqEk=
 
-protobufjs@^6.10.0, protobufjs@^6.10.2:
+proto3-json-serializer@^0.1.1:
+  version "0.1.4"
+  resolved "https://registry.npmjs.org/proto3-json-serializer/-/proto3-json-serializer-0.1.4.tgz#aa2dc4c9c9b7ea05631354b2c2e52c227539a7f0"
+  integrity sha512-bFzdsKU/zaTobWrRxRniMZIzzcgKYlmBWL1gAcTXZ2M7TQTGPI0JoYYs6bN7tpWj59ZCfwg7Ii/A2e8BbQGYnQ==
+
+protobufjs@6.11.2, protobufjs@^6.10.0, protobufjs@^6.10.2, protobufjs@^6.8.6:
   version "6.11.2"
   resolved "https://registry.npmjs.org/protobufjs/-/protobufjs-6.11.2.tgz#de39fabd4ed32beaa08e9bb1e30d08544c1edf8b"
   integrity sha512-4BQJoPooKJl2G9j3XftkIXjoC9C0Av2NOrWmbLWT1vH32GcSUHjM0Arra6UfTsVyfMAuFzaLucXn1sadxJydAw==

From 16288b6580c8701ba0d3bdac26e04f597916a3ae Mon Sep 17 00:00:00 2001
From: Marcus Eide <eide@spotify.com>
Date: Tue, 12 Oct 2021 16:16:15 +0200
Subject: [PATCH 03/14] Add class that will return a key-store based on
 application config

Signed-off-by: Marcus Eide <eide@spotify.com>
---
 .../src/identity/KeyStores.test.ts            | 86 +++++++++++++++++++
 .../auth-backend/src/identity/KeyStores.ts    | 75 ++++++++++++++++
 plugins/auth-backend/src/identity/index.ts    |  3 +
 3 files changed, 164 insertions(+)
 create mode 100644 plugins/auth-backend/src/identity/KeyStores.test.ts
 create mode 100644 plugins/auth-backend/src/identity/KeyStores.ts

diff --git a/plugins/auth-backend/src/identity/KeyStores.test.ts b/plugins/auth-backend/src/identity/KeyStores.test.ts
new file mode 100644
index 0000000000..957c4c7953
--- /dev/null
+++ b/plugins/auth-backend/src/identity/KeyStores.test.ts
@@ -0,0 +1,86 @@
+/*
+ * Copyright 2021 The Backstage Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import { DatabaseManager } from '@backstage/backend-common';
+import { ConfigReader } from '@backstage/config';
+
+import { MemoryKeyStore } from './MemoryKeyStore';
+import { DatabaseKeyStore } from './DatabaseKeyStore';
+import { FirestoreKeyStore } from './FirestoreKeyStore';
+import { KeyStores } from './KeyStores';
+
+describe('KeyStores', () => {
+  const defaultConfigOptions = {
+    auth: {
+      keyStore: {
+        provider: 'memory',
+      },
+    },
+  };
+  const defaultConfig = new ConfigReader(defaultConfigOptions);
+
+  it('reads auth section from config', async () => {
+    const configSpy = jest.spyOn(defaultConfig, 'getOptionalConfig');
+    const keyStore = await KeyStores.fromConfig(defaultConfig);
+
+    expect(keyStore).toBeInstanceOf(MemoryKeyStore);
+    expect(configSpy).toHaveBeenCalledWith('auth.keyStore');
+    expect(
+      defaultConfig
+        .getOptionalConfig('auth.keyStore')
+        ?.getOptionalString('provider'),
+    ).toBe(defaultConfigOptions.auth.keyStore.provider);
+  });
+
+  it('can handle without auth config', async () => {
+    const config = new ConfigReader({
+      backend: {
+        database: {
+          client: 'sqlite3',
+          connection: ':memory:',
+        },
+      },
+    });
+    const database =
+      DatabaseManager.fromConfig(config).forPlugin('auth-backend');
+    const keyStore = await KeyStores.fromConfig(config, { database });
+
+    expect(keyStore).toBeInstanceOf(DatabaseKeyStore);
+  });
+
+  it('can handle additional provider config', async () => {
+    const configOptions = {
+      auth: {
+        keyStore: {
+          provider: 'firestore',
+          firestore: {
+            projectId: 'my-project',
+          },
+        },
+      },
+    };
+    const config = new ConfigReader(configOptions);
+    const keyStore = await KeyStores.fromConfig(config);
+
+    expect(keyStore).toBeInstanceOf(FirestoreKeyStore);
+    expect(
+      config
+        .getOptionalConfig('auth.keyStore')
+        ?.getOptionalConfig('firestore')
+        ?.getOptionalString('projectId'),
+    ).toBe(configOptions.auth.keyStore.firestore.projectId);
+  });
+});
diff --git a/plugins/auth-backend/src/identity/KeyStores.ts b/plugins/auth-backend/src/identity/KeyStores.ts
new file mode 100644
index 0000000000..0991f7ea7b
--- /dev/null
+++ b/plugins/auth-backend/src/identity/KeyStores.ts
@@ -0,0 +1,75 @@
+/*
+ * Copyright 2021 The Backstage Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import { Logger } from 'winston';
+
+import { PluginDatabaseManager } from '@backstage/backend-common';
+import { Config } from '@backstage/config';
+
+import { DatabaseKeyStore } from './DatabaseKeyStore';
+import { MemoryKeyStore } from './MemoryKeyStore';
+import { FirestoreKeyStore } from './FirestoreKeyStore';
+import { KeyStore } from './types';
+
+type Options = {
+  logger?: Logger;
+  database?: PluginDatabaseManager;
+};
+
+export class KeyStores {
+  /**
+   * Looks at the `auth.keyStore` section in the application configuration
+   * and returns a KeyStore store. Defaults to `postgres`
+   *
+   * @returns a KeyStore store
+   */
+  static async fromConfig(
+    config: Config,
+    options?: Options,
+  ): Promise<KeyStore> {
+    const { logger, database } = options ?? {};
+
+    const ks = config.getOptionalConfig('auth.keyStore');
+    const provider = ks?.getOptionalString('provider') ?? 'postgres';
+    const providerConfig = ks?.getOptionalConfig(provider);
+
+    logger?.info(`Configuring "${provider}" as KeyStore provider`);
+
+    if (provider === 'postgres') {
+      if (!database) {
+        throw new Error('This KeyStore provider requires a database');
+      }
+
+      return await DatabaseKeyStore.create({
+        database: await database.getClient(),
+      });
+    }
+
+    if (provider === 'memory') {
+      return new MemoryKeyStore();
+    }
+
+    if (provider === 'firestore') {
+      return FirestoreKeyStore.create({
+        projectId: providerConfig?.getOptionalString('projectId'),
+        keyFilename: providerConfig?.getOptionalString('keyFilename'),
+        path: providerConfig?.getOptionalString('path'),
+      });
+    }
+
+    throw new Error(`Unknown KeyStore provider: ${provider}`);
+  }
+}
diff --git a/plugins/auth-backend/src/identity/index.ts b/plugins/auth-backend/src/identity/index.ts
index d88339e8e4..73858d7e07 100644
--- a/plugins/auth-backend/src/identity/index.ts
+++ b/plugins/auth-backend/src/identity/index.ts
@@ -18,4 +18,7 @@ export { createOidcRouter } from './router';
 export { IdentityClient } from './IdentityClient';
 export { TokenFactory } from './TokenFactory';
 export { DatabaseKeyStore } from './DatabaseKeyStore';
+export { MemoryKeyStore } from './MemoryKeyStore';
+export { FirestoreKeyStore } from './FirestoreKeyStore';
+export { KeyStores } from './KeyStores';
 export type { KeyStore, TokenIssuer, TokenParams } from './types';

From 12d4abe7b25319a085e2cfeab73f52e60004f747 Mon Sep 17 00:00:00 2001
From: Marcus Eide <eide@spotify.com>
Date: Tue, 12 Oct 2021 16:18:30 +0200
Subject: [PATCH 04/14] Add project dependency to package.json

Signed-off-by: Marcus Eide <eide@spotify.com>
---
 plugins/auth-backend/package.json | 1 +
 1 file changed, 1 insertion(+)

diff --git a/plugins/auth-backend/package.json b/plugins/auth-backend/package.json
index 342718db55..8f7f623652 100644
--- a/plugins/auth-backend/package.json
+++ b/plugins/auth-backend/package.json
@@ -36,6 +36,7 @@
     "@backstage/config": "^0.1.10",
     "@backstage/errors": "^0.1.1",
     "@backstage/test-utils": "^0.1.19",
+    "@google-cloud/firestore": "^4.15.1",
     "@types/express": "^4.17.6",
     "@types/passport": "^1.0.3",
     "compression": "^1.7.4",

From 4f4e22d176d12195eb0ee007a0f25423e01d0387 Mon Sep 17 00:00:00 2001
From: Marcus Eide <eide@spotify.com>
Date: Tue, 12 Oct 2021 16:19:45 +0200
Subject: [PATCH 05/14] Use KeyStores.fromConfig() in router

Signed-off-by: Marcus Eide <eide@spotify.com>
---
 plugins/auth-backend/src/service/router.ts | 8 +++-----
 1 file changed, 3 insertions(+), 5 deletions(-)

diff --git a/plugins/auth-backend/src/service/router.ts b/plugins/auth-backend/src/service/router.ts
index 423975aa75..8cce643987 100644
--- a/plugins/auth-backend/src/service/router.ts
+++ b/plugins/auth-backend/src/service/router.ts
@@ -29,7 +29,7 @@ import {
 import { NotFoundError } from '@backstage/errors';
 import { CatalogClient } from '@backstage/catalog-client';
 import { Config } from '@backstage/config';
-import { createOidcRouter, DatabaseKeyStore, TokenFactory } from '../identity';
+import { createOidcRouter, TokenFactory, KeyStores } from '../identity';
 import session from 'express-session';
 import passport from 'passport';
 import { Minimatch } from 'minimatch';
@@ -38,7 +38,7 @@ type ProviderFactories = { [s: string]: AuthProviderFactory };
 
 export interface RouterOptions {
   logger: Logger;
-  database: PluginDatabaseManager;
+  database?: PluginDatabaseManager;
   config: Config;
   discovery: PluginEndpointDiscovery;
   providerFactories?: ProviderFactories;
@@ -56,11 +56,9 @@ export async function createRouter({
   const appUrl = config.getString('app.baseUrl');
   const authUrl = await discovery.getExternalBaseUrl('auth');
 
+  const keyStore = await KeyStores.fromConfig(config, { logger, database });
   const keyDurationSeconds = 3600;
 
-  const keyStore = await DatabaseKeyStore.create({
-    database: await database.getClient(),
-  });
   const tokenIssuer = new TokenFactory({
     issuer: authUrl,
     keyStore,

From d2f755fa732b6d777f51a35630cb1110b5b9ba54 Mon Sep 17 00:00:00 2001
From: Marcus Eide <eide@spotify.com>
Date: Wed, 13 Oct 2021 09:21:30 +0200
Subject: [PATCH 06/14] Update api-reports

Signed-off-by: Marcus Eide <eide@spotify.com>
---
 plugins/auth-backend/api-report.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/plugins/auth-backend/api-report.md b/plugins/auth-backend/api-report.md
index 0bdd73f3e8..a20538e223 100644
--- a/plugins/auth-backend/api-report.md
+++ b/plugins/auth-backend/api-report.md
@@ -501,7 +501,7 @@ export interface RouterOptions {
   // (undocumented)
   config: Config;
   // (undocumented)
-  database: PluginDatabaseManager;
+  database?: PluginDatabaseManager;
   // (undocumented)
   discovery: PluginEndpointDiscovery;
   // (undocumented)

From bf259c989020210b4879aa8bec6e6024d3292814 Mon Sep 17 00:00:00 2001
From: Marcus Eide <eide@spotify.com>
Date: Wed, 13 Oct 2021 09:58:47 +0200
Subject: [PATCH 07/14] Add Firestore to vale vocab

Signed-off-by: Marcus Eide <eide@spotify.com>
---
 .github/styles/vocab.txt | 1 +
 1 file changed, 1 insertion(+)

diff --git a/.github/styles/vocab.txt b/.github/styles/vocab.txt
index a38baab3dc..a93f75192b 100644
--- a/.github/styles/vocab.txt
+++ b/.github/styles/vocab.txt
@@ -310,3 +310,4 @@ Zalando
 Zhou
 zoomable
 zsh
+Firestore

From ab9b4a6ea6a441e96e429def3bf43548ac089a98 Mon Sep 17 00:00:00 2001
From: Marcus Eide <eide@spotify.com>
Date: Wed, 13 Oct 2021 09:59:00 +0200
Subject: [PATCH 08/14] Add changeset

Signed-off-by: Marcus Eide <eide@spotify.com>
---
 .changeset/pretty-ways-search.md | 6 ++++++
 1 file changed, 6 insertions(+)
 create mode 100644 .changeset/pretty-ways-search.md

diff --git a/.changeset/pretty-ways-search.md b/.changeset/pretty-ways-search.md
new file mode 100644
index 0000000000..212c23b09a
--- /dev/null
+++ b/.changeset/pretty-ways-search.md
@@ -0,0 +1,6 @@
+---
+'@backstage/plugin-auth-backend': patch
+---
+
+Add Firestore as key-store provider.
+Add `auth.keyStore` section to application config.

From 5ad1f523d569e020548fe17713165775d53995b8 Mon Sep 17 00:00:00 2001
From: Marcus Eide <eide@spotify.com>
Date: Wed, 13 Oct 2021 14:25:06 +0200
Subject: [PATCH 09/14] Make the create function async

Signed-off-by: Marcus Eide <eide@spotify.com>
---
 .../src/identity/FirestoreKeyStore.test.ts    | 24 +++++++++----------
 .../src/identity/FirestoreKeyStore.ts         |  4 +++-
 .../auth-backend/src/identity/KeyStores.ts    |  2 +-
 3 files changed, 16 insertions(+), 14 deletions(-)

diff --git a/plugins/auth-backend/src/identity/FirestoreKeyStore.test.ts b/plugins/auth-backend/src/identity/FirestoreKeyStore.test.ts
index 03fcc15a77..a12e2ed40e 100644
--- a/plugins/auth-backend/src/identity/FirestoreKeyStore.test.ts
+++ b/plugins/auth-backend/src/identity/FirestoreKeyStore.test.ts
@@ -53,14 +53,14 @@ describe('FirestoreKeyStore', () => {
     process.env = OLD_ENV;
   });
 
-  it('can create an instance without settings', () => {
-    const keyStore = FirestoreKeyStore.create();
+  it('can create an instance without settings', async () => {
+    const keyStore = await FirestoreKeyStore.create();
 
     expect(keyStore).toBeInstanceOf(FirestoreKeyStore);
   });
 
   it('can set the project id', async () => {
-    FirestoreKeyStore.create({ projectId: 'my-project' });
+    await FirestoreKeyStore.create({ projectId: 'my-project' });
 
     expect(Firestore).toHaveBeenCalledWith(
       expect.objectContaining({
@@ -70,7 +70,7 @@ describe('FirestoreKeyStore', () => {
   });
 
   it('can handle keyfile file', async () => {
-    FirestoreKeyStore.create({ keyFilename: 'keyFile.json' });
+    await FirestoreKeyStore.create({ keyFilename: 'keyFile.json' });
 
     expect(Firestore).toHaveBeenCalledWith(
       expect.objectContaining({
@@ -79,9 +79,9 @@ describe('FirestoreKeyStore', () => {
     );
   });
 
-  it('can use default google credentials', () => {
+  it('can use default google credentials', async () => {
     process.env.GOOGLE_APPLICATION_CREDENTIALS = 'cred.json';
-    FirestoreKeyStore.create();
+    await FirestoreKeyStore.create();
 
     expect(Firestore).toHaveBeenCalledWith(
       expect.objectContaining({
@@ -91,14 +91,14 @@ describe('FirestoreKeyStore', () => {
   });
 
   it('can uses a default path', async () => {
-    const keyStore = FirestoreKeyStore.create();
+    const keyStore = await FirestoreKeyStore.create();
     await keyStore.addKey(key);
 
     expect(firestoreMock.collection).toBeCalledWith('sessions');
   });
 
   it('can set the path', async () => {
-    const keyStore = FirestoreKeyStore.create({
+    const keyStore = await FirestoreKeyStore.create({
       path: 'my-path',
     });
     await keyStore.addKey(key);
@@ -107,7 +107,7 @@ describe('FirestoreKeyStore', () => {
   });
 
   it('can add keys', async () => {
-    const keyStore = FirestoreKeyStore.create();
+    const keyStore = await FirestoreKeyStore.create();
     await keyStore.addKey(key);
 
     expect(firestoreMock.collection).toBeCalledWith('sessions');
@@ -119,7 +119,7 @@ describe('FirestoreKeyStore', () => {
   });
 
   it('can delete a single key', async () => {
-    const keyStore = FirestoreKeyStore.create();
+    const keyStore = await FirestoreKeyStore.create();
     await keyStore.removeKeys(['123']);
 
     expect(firestoreMock.collection).toBeCalledWith('sessions');
@@ -128,7 +128,7 @@ describe('FirestoreKeyStore', () => {
   });
 
   it('can delete a multiple keys', async () => {
-    const keyStore = FirestoreKeyStore.create();
+    const keyStore = await FirestoreKeyStore.create();
     await keyStore.removeKeys(['123', '456']);
 
     expect(firestoreMock.collection).toBeCalledWith('sessions');
@@ -138,7 +138,7 @@ describe('FirestoreKeyStore', () => {
   });
 
   it('can list keys', async () => {
-    const keyStore = FirestoreKeyStore.create();
+    const keyStore = await FirestoreKeyStore.create();
     const items = await keyStore.listKeys();
 
     expect(firestoreMock.collection).toBeCalledWith('sessions');
diff --git a/plugins/auth-backend/src/identity/FirestoreKeyStore.ts b/plugins/auth-backend/src/identity/FirestoreKeyStore.ts
index cf8c082540..74122f2849 100644
--- a/plugins/auth-backend/src/identity/FirestoreKeyStore.ts
+++ b/plugins/auth-backend/src/identity/FirestoreKeyStore.ts
@@ -23,7 +23,9 @@ type FirestoreSettings = Settings & {
 };
 
 export class FirestoreKeyStore implements KeyStore {
-  static create(settings?: FirestoreSettings): FirestoreKeyStore {
+  static async create(
+    settings?: FirestoreSettings,
+  ): Promise<FirestoreKeyStore> {
     const { projectId, keyFilename, path } = settings ?? {};
     const database = new Firestore({
       projectId,
diff --git a/plugins/auth-backend/src/identity/KeyStores.ts b/plugins/auth-backend/src/identity/KeyStores.ts
index 0991f7ea7b..ec236e2074 100644
--- a/plugins/auth-backend/src/identity/KeyStores.ts
+++ b/plugins/auth-backend/src/identity/KeyStores.ts
@@ -63,7 +63,7 @@ export class KeyStores {
     }
 
     if (provider === 'firestore') {
-      return FirestoreKeyStore.create({
+      return await FirestoreKeyStore.create({
         projectId: providerConfig?.getOptionalString('projectId'),
         keyFilename: providerConfig?.getOptionalString('keyFilename'),
         path: providerConfig?.getOptionalString('path'),

From 33b9694f5ce8c920daae2d9370bfb29d085ad0fc Mon Sep 17 00:00:00 2001
From: Marcus Eide <eide@spotify.com>
Date: Thu, 14 Oct 2021 14:43:56 +0200
Subject: [PATCH 10/14] Add support for more settings

Signed-off-by: Marcus Eide <eide@spotify.com>
---
 plugins/auth-backend/config.d.ts              |  6 ++
 .../src/identity/FirestoreKeyStore.test.ts    | 79 +++++++------------
 .../src/identity/FirestoreKeyStore.ts         | 11 +--
 .../auth-backend/src/identity/KeyStores.ts    | 14 ++--
 4 files changed, 44 insertions(+), 66 deletions(-)

diff --git a/plugins/auth-backend/config.d.ts b/plugins/auth-backend/config.d.ts
index b15b0da628..280d92a261 100644
--- a/plugins/auth-backend/config.d.ts
+++ b/plugins/auth-backend/config.d.ts
@@ -35,6 +35,12 @@ export interface Config {
     keyStore?: {
       provider?: 'postgres' | 'memory' | 'firestore';
       firestore?: {
+        /** The host to connect to */
+        host?: string;
+        /** The port to connect to */
+        port?: number;
+        /** Whether to use SSL when connecting. */
+        ssl?: boolean;
         /** The Google Cloud Project ID */
         projectId?: string;
         /**
diff --git a/plugins/auth-backend/src/identity/FirestoreKeyStore.test.ts b/plugins/auth-backend/src/identity/FirestoreKeyStore.test.ts
index a12e2ed40e..398e3c2b34 100644
--- a/plugins/auth-backend/src/identity/FirestoreKeyStore.test.ts
+++ b/plugins/auth-backend/src/identity/FirestoreKeyStore.test.ts
@@ -16,7 +16,10 @@
 
 import { Firestore } from '@google-cloud/firestore';
 
-import { FirestoreKeyStore } from './FirestoreKeyStore';
+import {
+  FirestoreKeyStore,
+  FirestoreKeyStoreSettings,
+} from './FirestoreKeyStore';
 import { AnyJWK } from './types';
 
 const data = jest.fn().mockReturnValue('data');
@@ -37,7 +40,6 @@ jest.mock('@google-cloud/firestore', () => ({
 }));
 
 describe('FirestoreKeyStore', () => {
-  const OLD_ENV = process.env;
   const key = {
     kid: '123',
     use: 'sig',
@@ -45,72 +47,45 @@ describe('FirestoreKeyStore', () => {
     alg: 'Base64',
   } as AnyJWK;
 
+  const settings = {
+    projectId: 'my-project',
+    host: 'my-host',
+    port: 8080,
+    ssl: false,
+    keyFilename: 'cred.json',
+  };
+  const path = 'my-path';
+  const firestoreSettings = { ...settings, path } as FirestoreKeyStoreSettings;
+
   beforeEach(() => {
     jest.clearAllMocks();
   });
 
-  afterAll(() => {
-    process.env = OLD_ENV;
-  });
-
   it('can create an instance without settings', async () => {
     const keyStore = await FirestoreKeyStore.create();
 
     expect(keyStore).toBeInstanceOf(FirestoreKeyStore);
+    expect(Firestore).toHaveBeenCalledWith({});
   });
 
-  it('can set the project id', async () => {
-    await FirestoreKeyStore.create({ projectId: 'my-project' });
+  it('can create an instance with settings', async () => {
+    await FirestoreKeyStore.create(firestoreSettings);
 
-    expect(Firestore).toHaveBeenCalledWith(
-      expect.objectContaining({
-        projectId: 'my-project',
-      }),
-    );
+    expect(Firestore).toHaveBeenCalledWith(settings);
   });
 
-  it('can handle keyfile file', async () => {
-    await FirestoreKeyStore.create({ keyFilename: 'keyFile.json' });
-
-    expect(Firestore).toHaveBeenCalledWith(
-      expect.objectContaining({
-        keyFilename: 'keyFile.json',
-      }),
-    );
-  });
-
-  it('can use default google credentials', async () => {
-    process.env.GOOGLE_APPLICATION_CREDENTIALS = 'cred.json';
-    await FirestoreKeyStore.create();
-
-    expect(Firestore).toHaveBeenCalledWith(
-      expect.objectContaining({
-        keyFilename: 'cred.json',
-      }),
-    );
-  });
-
-  it('can uses a default path', async () => {
+  it('can use a default path', async () => {
     const keyStore = await FirestoreKeyStore.create();
     await keyStore.addKey(key);
 
     expect(firestoreMock.collection).toBeCalledWith('sessions');
   });
 
-  it('can set the path', async () => {
-    const keyStore = await FirestoreKeyStore.create({
-      path: 'my-path',
-    });
-    await keyStore.addKey(key);
-
-    expect(firestoreMock.collection).toBeCalledWith('my-path');
-  });
-
   it('can add keys', async () => {
-    const keyStore = await FirestoreKeyStore.create();
+    const keyStore = await FirestoreKeyStore.create(firestoreSettings);
     await keyStore.addKey(key);
 
-    expect(firestoreMock.collection).toBeCalledWith('sessions');
+    expect(firestoreMock.collection).toBeCalledWith(path);
     expect(firestoreMock.doc).toBeCalledWith(key.kid);
     expect(firestoreMock.set).toHaveBeenCalledWith({
       kid: key.kid,
@@ -119,29 +94,29 @@ describe('FirestoreKeyStore', () => {
   });
 
   it('can delete a single key', async () => {
-    const keyStore = await FirestoreKeyStore.create();
+    const keyStore = await FirestoreKeyStore.create(firestoreSettings);
     await keyStore.removeKeys(['123']);
 
-    expect(firestoreMock.collection).toBeCalledWith('sessions');
+    expect(firestoreMock.collection).toBeCalledWith(path);
     expect(firestoreMock.doc).toBeCalledWith('123');
     expect(firestoreMock.delete).toBeCalledTimes(1);
   });
 
   it('can delete a multiple keys', async () => {
-    const keyStore = await FirestoreKeyStore.create();
+    const keyStore = await FirestoreKeyStore.create(firestoreSettings);
     await keyStore.removeKeys(['123', '456']);
 
-    expect(firestoreMock.collection).toBeCalledWith('sessions');
+    expect(firestoreMock.collection).toBeCalledWith(path);
     expect(firestoreMock.doc).toBeCalledWith('123');
     expect(firestoreMock.doc).toBeCalledWith('456');
     expect(firestoreMock.delete).toBeCalledTimes(2);
   });
 
   it('can list keys', async () => {
-    const keyStore = await FirestoreKeyStore.create();
+    const keyStore = await FirestoreKeyStore.create(firestoreSettings);
     const items = await keyStore.listKeys();
 
-    expect(firestoreMock.collection).toBeCalledWith('sessions');
+    expect(firestoreMock.collection).toBeCalledWith(path);
     expect(firestoreMock.get).toBeCalledTimes(1);
     expect(data).toBeCalledTimes(1);
     expect(toDate).toBeCalledTimes(1);
diff --git a/plugins/auth-backend/src/identity/FirestoreKeyStore.ts b/plugins/auth-backend/src/identity/FirestoreKeyStore.ts
index 74122f2849..f606bacd38 100644
--- a/plugins/auth-backend/src/identity/FirestoreKeyStore.ts
+++ b/plugins/auth-backend/src/identity/FirestoreKeyStore.ts
@@ -18,19 +18,16 @@ import { Firestore, Settings } from '@google-cloud/firestore';
 
 import { AnyJWK, KeyStore, StoredKey } from './types';
 
-type FirestoreSettings = Settings & {
+export type FirestoreKeyStoreSettings = Settings & {
   path?: string;
 };
 
 export class FirestoreKeyStore implements KeyStore {
   static async create(
-    settings?: FirestoreSettings,
+    settings?: FirestoreKeyStoreSettings,
   ): Promise<FirestoreKeyStore> {
-    const { projectId, keyFilename, path } = settings ?? {};
-    const database = new Firestore({
-      projectId,
-      keyFilename: keyFilename ?? process.env.GOOGLE_APPLICATION_CREDENTIALS,
-    });
+    const { path, ...firestoreSettings } = settings ?? {};
+    const database = new Firestore(firestoreSettings);
 
     return new FirestoreKeyStore(database, path ?? 'sessions');
   }
diff --git a/plugins/auth-backend/src/identity/KeyStores.ts b/plugins/auth-backend/src/identity/KeyStores.ts
index ec236e2074..1b19bde911 100644
--- a/plugins/auth-backend/src/identity/KeyStores.ts
+++ b/plugins/auth-backend/src/identity/KeyStores.ts
@@ -21,7 +21,10 @@ import { Config } from '@backstage/config';
 
 import { DatabaseKeyStore } from './DatabaseKeyStore';
 import { MemoryKeyStore } from './MemoryKeyStore';
-import { FirestoreKeyStore } from './FirestoreKeyStore';
+import {
+  FirestoreKeyStore,
+  FirestoreKeyStoreSettings,
+} from './FirestoreKeyStore';
 import { KeyStore } from './types';
 
 type Options = {
@@ -44,7 +47,6 @@ export class KeyStores {
 
     const ks = config.getOptionalConfig('auth.keyStore');
     const provider = ks?.getOptionalString('provider') ?? 'postgres';
-    const providerConfig = ks?.getOptionalConfig(provider);
 
     logger?.info(`Configuring "${provider}" as KeyStore provider`);
 
@@ -63,11 +65,9 @@ export class KeyStores {
     }
 
     if (provider === 'firestore') {
-      return await FirestoreKeyStore.create({
-        projectId: providerConfig?.getOptionalString('projectId'),
-        keyFilename: providerConfig?.getOptionalString('keyFilename'),
-        path: providerConfig?.getOptionalString('path'),
-      });
+      const settings = ks?.getOptional(provider) as FirestoreKeyStoreSettings;
+
+      return await FirestoreKeyStore.create(settings);
     }
 
     throw new Error(`Unknown KeyStore provider: ${provider}`);

From b294f6056f6ca835eb4f71bd22b03b2d5c6b8b28 Mon Sep 17 00:00:00 2001
From: Marcus Eide <eide@spotify.com>
Date: Tue, 19 Oct 2021 16:42:40 +0200
Subject: [PATCH 11/14] 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              |   2 +
 .../src/identity/FirestoreKeyStore.test.ts    |  71 +++++++++++--
 .../src/identity/FirestoreKeyStore.ts         | 100 +++++++++++++++---
 .../src/identity/KeyStores.test.ts            |   2 +
 .../auth-backend/src/identity/KeyStores.ts    |   4 +-
 5 files changed, 155 insertions(+), 24 deletions(-)

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

From 7e707db187c94d66570c0a36215726a145412e93 Mon Sep 17 00:00:00 2001
From: Marcus Eide <eide@spotify.com>
Date: Thu, 21 Oct 2021 10:01:12 +0200
Subject: [PATCH 12/14] Undo making database optional

Signed-off-by: Marcus Eide <eide@spotify.com>
---
 plugins/auth-backend/api-report.md         | 2 +-
 plugins/auth-backend/src/service/router.ts | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/plugins/auth-backend/api-report.md b/plugins/auth-backend/api-report.md
index a20538e223..0bdd73f3e8 100644
--- a/plugins/auth-backend/api-report.md
+++ b/plugins/auth-backend/api-report.md
@@ -501,7 +501,7 @@ export interface RouterOptions {
   // (undocumented)
   config: Config;
   // (undocumented)
-  database?: PluginDatabaseManager;
+  database: PluginDatabaseManager;
   // (undocumented)
   discovery: PluginEndpointDiscovery;
   // (undocumented)
diff --git a/plugins/auth-backend/src/service/router.ts b/plugins/auth-backend/src/service/router.ts
index 8cce643987..23311f5e56 100644
--- a/plugins/auth-backend/src/service/router.ts
+++ b/plugins/auth-backend/src/service/router.ts
@@ -38,7 +38,7 @@ type ProviderFactories = { [s: string]: AuthProviderFactory };
 
 export interface RouterOptions {
   logger: Logger;
-  database?: PluginDatabaseManager;
+  database: PluginDatabaseManager;
   config: Config;
   discovery: PluginEndpointDiscovery;
   providerFactories?: ProviderFactories;

From 6237e636eb3f4cab189b6503298f0f45a878b31b Mon Sep 17 00:00:00 2001
From: Marcus Eide <eide@spotify.com>
Date: Thu, 21 Oct 2021 10:30:29 +0200
Subject: [PATCH 13/14] Be more specific when accessing firestore provider
 config

Signed-off-by: Marcus Eide <eide@spotify.com>
---
 .../src/identity/KeyStores.test.ts            | 10 ++++++++++
 .../auth-backend/src/identity/KeyStores.ts    | 19 +++++++++++++------
 2 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/plugins/auth-backend/src/identity/KeyStores.test.ts b/plugins/auth-backend/src/identity/KeyStores.test.ts
index bb91d2cae9..f30263e27f 100644
--- a/plugins/auth-backend/src/identity/KeyStores.test.ts
+++ b/plugins/auth-backend/src/identity/KeyStores.test.ts
@@ -63,6 +63,7 @@ describe('KeyStores', () => {
 
   it('can handle additional provider config', async () => {
     jest.spyOn(FirestoreKeyStore, 'verifyConnection').mockResolvedValue();
+    const createSpy = jest.spyOn(FirestoreKeyStore, 'create');
 
     const configOptions = {
       auth: {
@@ -70,6 +71,12 @@ describe('KeyStores', () => {
           provider: 'firestore',
           firestore: {
             projectId: 'my-project',
+            keyFilename: 'cred.json',
+            path: 'my-path',
+            timeout: 100,
+            host: 'localhost',
+            port: 8088,
+            ssl: false,
           },
         },
       },
@@ -78,6 +85,9 @@ describe('KeyStores', () => {
     const keyStore = await KeyStores.fromConfig(config);
 
     expect(keyStore).toBeInstanceOf(FirestoreKeyStore);
+    expect(createSpy).toHaveBeenCalledWith(
+      configOptions.auth.keyStore.firestore,
+    );
     expect(
       config
         .getOptionalConfig('auth.keyStore')
diff --git a/plugins/auth-backend/src/identity/KeyStores.ts b/plugins/auth-backend/src/identity/KeyStores.ts
index 8ac1cf1ec6..5ec63f5439 100644
--- a/plugins/auth-backend/src/identity/KeyStores.ts
+++ b/plugins/auth-backend/src/identity/KeyStores.ts
@@ -21,10 +21,7 @@ import { Config } from '@backstage/config';
 
 import { DatabaseKeyStore } from './DatabaseKeyStore';
 import { MemoryKeyStore } from './MemoryKeyStore';
-import {
-  FirestoreKeyStore,
-  FirestoreKeyStoreSettings,
-} from './FirestoreKeyStore';
+import { FirestoreKeyStore } from './FirestoreKeyStore';
 import { KeyStore } from './types';
 
 type Options = {
@@ -65,8 +62,18 @@ export class KeyStores {
     }
 
     if (provider === 'firestore') {
-      const settings = ks?.getOptional(provider) as FirestoreKeyStoreSettings;
-      const keyStore = await FirestoreKeyStore.create(settings);
+      const settings = ks?.getConfig(provider);
+
+      const keyStore = await FirestoreKeyStore.create({
+        projectId: settings?.getOptionalString('projectId'),
+        keyFilename: settings?.getOptionalString('keyFilename'),
+        host: settings?.getOptionalString('host'),
+        port: settings?.getOptionalNumber('port'),
+        ssl: settings?.getOptionalBoolean('ssl'),
+        path: settings?.getOptionalString('path'),
+        timeout: settings?.getOptionalNumber('timeout'),
+      });
+
       await FirestoreKeyStore.verifyConnection(keyStore, logger);
 
       return keyStore;

From 34da1574a3d768e7076ee2160e882e9b1d010d6b Mon Sep 17 00:00:00 2001
From: Patrik Oldsberg <poldsberg@gmail.com>
Date: Thu, 21 Oct 2021 13:48:49 +0200
Subject: [PATCH 14/14] auth-backend: rename postgres keystore provider to
 database

Signed-off-by: Patrik Oldsberg <poldsberg@gmail.com>
---
 plugins/auth-backend/config.d.ts               | 2 +-
 plugins/auth-backend/src/identity/KeyStores.ts | 6 +++---
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/plugins/auth-backend/config.d.ts b/plugins/auth-backend/config.d.ts
index 6f31cd07df..25da7fa332 100644
--- a/plugins/auth-backend/config.d.ts
+++ b/plugins/auth-backend/config.d.ts
@@ -33,7 +33,7 @@ export interface Config {
 
     /** To control how to store JWK data in auth-backend */
     keyStore?: {
-      provider?: 'postgres' | 'memory' | 'firestore';
+      provider?: 'database' | 'memory' | 'firestore';
       firestore?: {
         /** The host to connect to */
         host?: string;
diff --git a/plugins/auth-backend/src/identity/KeyStores.ts b/plugins/auth-backend/src/identity/KeyStores.ts
index 5ec63f5439..74e60a0502 100644
--- a/plugins/auth-backend/src/identity/KeyStores.ts
+++ b/plugins/auth-backend/src/identity/KeyStores.ts
@@ -32,7 +32,7 @@ type Options = {
 export class KeyStores {
   /**
    * Looks at the `auth.keyStore` section in the application configuration
-   * and returns a KeyStore store. Defaults to `postgres`
+   * and returns a KeyStore store. Defaults to `database`
    *
    * @returns a KeyStore store
    */
@@ -43,11 +43,11 @@ export class KeyStores {
     const { logger, database } = options ?? {};
 
     const ks = config.getOptionalConfig('auth.keyStore');
-    const provider = ks?.getOptionalString('provider') ?? 'postgres';
+    const provider = ks?.getOptionalString('provider') ?? 'database';
 
     logger?.info(`Configuring "${provider}" as KeyStore provider`);
 
-    if (provider === 'postgres') {
+    if (provider === 'database') {
       if (!database) {
         throw new Error('This KeyStore provider requires a database');
       }