auth-backend: refactor identity to remove logic from storage layer

plugins/auth-backend/src/identity/DatabaseKeyStore.ts

@@ -18,7 +18,7 @@ import Knex from 'knex';
 import path from 'path';
 import { utc } from 'moment';
 import { Logger } from 'winston';
-import { AnyJWK, KeyStore } from './types';
+import { AnyJWK, KeyStore, StoredKey } from './types';
 
 const migrationsDir = path.resolve(
   require.resolve('@backstage/plugin-auth-backend/package.json'),
@@ -36,8 +36,6 @@ type Row = {
 type Options = {
   logger: Logger;
   database: Knex;
-  /** Expiration time of signing keys in seconds */
-  keyDuration: number;
 };
 
 export class DatabaseKeyStore implements KeyStore {
@@ -53,81 +51,36 @@ export class DatabaseKeyStore implements KeyStore {
 
   private readonly logger: Logger;
   private readonly database: Knex;
-  private readonly keyDuration: number;
-
-  private removingExpiredRows: boolean = false;
 
   private constructor(options: Options) {
-    const { logger, database, keyDuration } = options;
-
-    this.database = database;
-    this.keyDuration = keyDuration;
-    this.logger = logger.child({ service: 'key-store' });
+    this.database = options.database;
+    this.logger = options.logger;
   }
 
-  async storeKey({ key }: { key: AnyJWK }): Promise<void> {
-    this.logger.info(`Storing public key ${key.kid}`);
-
+  async addKey(key: AnyJWK): Promise<void> {
     await this.database<Row>(TABLE).insert({
       kid: key.kid,
       key: JSON.stringify(key),
     });
   }
 
-  async listKeys(): Promise<{ keys: AnyJWK[] }> {
+  async listKeys(): Promise<{ items: StoredKey[] }> {
     const rows = await this.database<Row>(TABLE).select();
 
-    const [validRows, expiredRows] = this.splitExpiredRows(rows);
-    if (expiredRows.length > 0) {
-      // We don't await this, just let it run in the background
-      this.removeExpiredRows(expiredRows);
-    }
-
     return {
-      keys: validRows.map(row => JSON.parse(row.key)),
+      items: rows.map(row => ({
+        key: JSON.parse(row.key),
+        createdAt: utc(row.created_at),
+      })),
     };
   }
 
-  private splitExpiredRows(rows: Row[]) {
-    const validRows = [];
-    const expiredRows = [];
-
-    for (const row of rows) {
-      const createdAt = utc(row.created_at);
-      const expireAt = createdAt.add(3 * this.keyDuration, 'seconds');
-      const isExpired = expireAt.isBefore();
-
-      if (isExpired) {
-        expiredRows.push(row);
-      } else {
-        validRows.push(row);
-      }
-    }
-
-    return [validRows, expiredRows];
-  }
-
-  private async removeExpiredRows(rows: Row[]) {
-    if (this.removingExpiredRows) {
-      return;
-    }
-
-    try {
-      this.removingExpiredRows = true;
-
-      const kids = rows.map(row => row.kid);
-      this.logger.info(`Removing expired signing keys, '${kids.join(', ')}'`);
-
-      const result = await this.database(TABLE).delete().whereIn('kid', kids);
-      if (result !== kids.length) {
-        this.logger.warn(
-          `Wanted to remove ${kids.length} expired signing, but removed ${result} instead`,
-        );
-      }
-    } catch (error) {
-      this.logger.error(`Failed to remove expired signing keys, ${error}`);
-    } finally {
-      this.removingExpiredRows = false;
+  async removeKeys(kids: string[]): Promise<void> {
+    const result = await this.database(TABLE).delete().whereIn('kid', kids);
+    if (result !== kids.length) {
+      this.logger.warn(
+        `Wanted to remove ${kids.length} keys, but only removed ${result}`,
+      );
     }
   }
 }

plugins/auth-backend/src/identity/TokenFactory.ts

@@ -62,6 +62,37 @@ export class TokenFactory implements TokenIssuer {
     });
   }
 
+  async listPublicKeys(): Promise<{ keys: AnyJWK[] }> {
+    const { items: keys } = await this.keyStore.listKeys();
+
+    const validKeys = [];
+    const expiredKeys = [];
+
+    for (const key of keys) {
+      // Allow for a grace period of another full key duration before we remove the keys from the database
+      const expireAt = key.createdAt.add(3 * this.keyDuration, 'seconds');
+      if (expireAt.isBefore()) {
+        expiredKeys.push(key);
+      } else {
+        validKeys.push(key);
+      }
+    }
+
+    if (expiredKeys.length > 0) {
+      const kids = expiredKeys.map(({ key }) => key.kid);
+
+      this.logger.info(`Removing expired signing keys, '${kids.join("', '")}'`);
+
+      // We don't await this, just let it run in the background
+      this.keyStore.removeKeys(kids).catch(error => {
+        this.logger.error(`Failed to remove expired keys, ${error}`);
+      });
+    }
+
+    // NOTE: we're currently only storing public keys, but if we start storing private keys we'd have to convert here
+    return { keys: validKeys.map(({ key }) => key) };
+  }
+
   private async getKey(): Promise<JSONWebKey> {
     if (this.privateKeyPromise) {
       if (this.keyExpiry && Date.now() < this.keyExpiry) {
@@ -79,9 +110,8 @@ export class TokenFactory implements TokenIssuer {
         alg: 'ES256',
       });
 
-      await this.keyStore.storeKey({
-        key: (key.toJWK(false) as unknown) as AnyJWK,
-      });
+      this.logger.info(`Created new signing key ${key.kid}`);
+      await this.keyStore.addKey((key.toJWK(false) as unknown) as AnyJWK);
 
       return key as JSONWebKey;
     })();

plugins/auth-backend/src/identity/router.ts

@@ -16,16 +16,16 @@
 
 import Router from 'express-promise-router';
 import { Logger } from 'winston';
-import { KeyStore } from './types';
+import { TokenIssuer } from './types';
 
 export type Options = {
   logger: Logger;
   baseUrl: string;
-  keyStore: KeyStore;
+  tokenIssuer: TokenIssuer;
 };
 
 export function createOidcRouter(options: Options) {
-  const { logger, keyStore, baseUrl } = options;
+  const { logger, tokenIssuer, baseUrl } = options;
 
   const router = Router();
 
@@ -50,7 +50,7 @@ export function createOidcRouter(options: Options) {
 
   router.get('/.well-known/jwks.json', async (_req, res) => {
     logger.info('request certs');
-    const { keys } = await keyStore.listKeys();
+    const { keys } = await tokenIssuer.listPublicKeys();
     res.json({ keys });
   });
 

plugins/auth-backend/src/identity/types.ts

@@ -22,21 +22,6 @@ export interface AnyJWK extends Record<string, string> {
   kty: string;
 }
 
-/** A KeyStore can store and keep track of recently used JWKs */
-export type KeyStore = {
-  /**
-   * Store a new key to be used for signing. Before the key has been successfully
-   * stored it should not be used for signing.
-   */
-  storeKey(params: { key: AnyJWK }): Promise<void>;
-
-  /**
-   * List all keys that are currently being used to sign tokens, or have been used
-   * in the past within the token expiration time, including a grace period.
-   */
-  listKeys(): Promise<{ keys: AnyJWK[] }>;
-};
-
 /** Parameters used to issue new ID Tokens */
 export type TokenParams = {
   /** The claims that will be embedded within the token */
@@ -46,10 +31,46 @@ export type TokenParams = {
   };
 };
 
-/** A TokenIssuer is able to issue verifiable ID Tokens on demand */
+/**
+ * A TokenIssuer is able to issue verifiable ID Tokens on demand.
+ */
 export type TokenIssuer = {
   /**
    * Issues a new ID Token
    */
   issueToken(params: TokenParams): Promise<string>;
+
+  /**
+   * List all public keys that are currently being used to sign tokens, or have been used
+   * in the past within the token expiration time, including a grace period.
+   */
+  listPublicKeys(): Promise<{ keys: AnyJWK[] }>;
+};
+
+/**
+ * A JWK stored by a KeyStore
+ */
+export type StoredKey = {
+  key: AnyJWK;
+  createdAt: moment.Moment;
+};
+
+/**
+ * A KeyStore stores JWKs for later and shared use.
+ */
+export type KeyStore = {
+  /**
+   * Store a new key to be used for signing.
+   */
+  addKey(key: AnyJWK): Promise<void>;
+
+  /**
+   * Remove all keys with the provided kids.
+   */
+  removeKeys(kids: string[]): Promise<void>;
+
+  /**
+   * List all stored keys.
+   */
+  listKeys(): Promise<{ items: StoredKey[] }>;
 };

plugins/auth-backend/src/lib/OAuthProvider.test.ts

@@ -182,6 +182,7 @@ describe('OAuthProvider', () => {
     appOrigin: 'http://localhost:3000',
     tokenIssuer: {
       issueToken: async () => 'my-id-token',
+      listPublicKeys: async () => ({ keys: [] }),
     },
   };
 

plugins/auth-backend/src/service/router.ts

@@ -42,13 +42,12 @@ export async function createRouter(
   const keyStore = await DatabaseKeyStore.create({
     database: options.database,
     logger: logger.child({ component: 'db-key-store' }),
-    keyDuration,
   });
-  const issuer = new TokenFactory({
+  const tokenIssuer = new TokenFactory({
     issuer: baseUrl,
     keyStore,
-    logger: logger.child({ component: 'token-factory' }),
     keyDuration,
+    logger: logger.child({ component: 'token-factory' }),
   });
 
   router.use(cookieParser());
@@ -104,7 +103,7 @@ export async function createRouter(
         { baseUrl },
         providerConfig,
         logger,
-        issuer,
+        tokenIssuer,
       );
       router.use(`/${providerId}`, providerRouter);
     } catch (e) {
@@ -115,7 +114,7 @@ export async function createRouter(
   router.use(
     createOidcRouter({
       logger: logger.child({ component: 'oidc-router' }),
-      keyStore,
+      tokenIssuer,
       baseUrl,
     }),
   );