backend-app-api: update DatabaseKeyStore to trim expired keys

Signed-off-by: Patrik Oldsberg <poldsberg@gmail.com>

packages/backend-app-api/migrations/20240327104803_public_keys.js

@@ -32,7 +32,10 @@ exports.up = async function up(knex) {
 
       table.text('key').notNullable().comment('JSON serialized public key');
 
-      table.timestamp('expires_at').notNullable();
+      table
+        .timestamp('expires_at')
+        .notNullable()
+        .comment('The time that the key expires');
     },
   );
 };

packages/backend-app-api/src/services/implementations/auth/DatabaseKeyStore.test.ts

@@ -0,0 +1,124 @@
+/*
+ * Copyright 2024 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 {
+  TestDatabaseId,
+  TestDatabases,
+  mockServices,
+} from '@backstage/backend-test-utils';
+import { DatabaseKeyStore, TABLE } from './DatabaseKeyStore';
+
+const testKey = {
+  kid: 'test-key',
+  kty: 'RSA',
+  e: 'ABC',
+  n: 'test',
+};
+const testKey2 = {
+  kid: 'test-key-2',
+  kty: 'RSA',
+  e: 'XYZ',
+  n: 'test',
+};
+
+describe('DatabaseKeyStore', () => {
+  const databases = TestDatabases.create();
+
+  async function createDatabaseKeyStore(databaseId: TestDatabaseId) {
+    const knex = await databases.init(databaseId);
+    const logger = mockServices.logger.mock();
+    return {
+      knex,
+      logger,
+      keyStore: await DatabaseKeyStore.create({
+        database: { getClient: async () => knex },
+        logger,
+      }),
+    };
+  }
+
+  describe.each(databases.eachSupportedId())('%p', databaseId => {
+    it('should insert a key into the database and list it', async () => {
+      const { keyStore } = await createDatabaseKeyStore(databaseId);
+      const expiresAt = new Date(Date.now() + 3600_000);
+
+      await expect(keyStore.listKeys()).resolves.toEqual({ keys: [] });
+
+      await keyStore.addKey({
+        id: testKey.kid,
+        key: testKey,
+        expiresAt,
+      });
+
+      await expect(keyStore.listKeys()).resolves.toEqual({
+        keys: [{ id: testKey.kid, key: testKey, expiresAt }],
+      });
+    });
+
+    it('should automatically exclude and remove expired keys when listing', async () => {
+      const { keyStore, knex, logger } = await createDatabaseKeyStore(
+        databaseId,
+      );
+      const expiresAt = new Date(Date.now() + 3600_000);
+
+      await expect(keyStore.listKeys()).resolves.toEqual({ keys: [] });
+
+      await keyStore.addKey({
+        id: testKey.kid,
+        key: testKey,
+        expiresAt,
+      });
+      await keyStore.addKey({
+        id: testKey2.kid,
+        key: testKey2,
+        expiresAt: new Date(0),
+      });
+
+      await expect(knex(TABLE).select('id')).resolves.toEqual([
+        { id: testKey.kid },
+        { id: testKey2.kid },
+      ]);
+
+      expect(logger.info).not.toHaveBeenCalled();
+
+      await expect(keyStore.listKeys()).resolves.toEqual({
+        keys: [{ id: testKey.kid, key: testKey, expiresAt }],
+      });
+
+      expect(logger.info).toHaveBeenCalledWith(
+        "Removing expired plugin service keys, 'test-key-2'",
+      );
+
+      await expect(knex(TABLE).select('id')).resolves.toEqual([
+        { id: testKey.kid },
+      ]);
+    });
+
+    it('should fail to insert with invalid date', async () => {
+      const { keyStore } = await createDatabaseKeyStore(databaseId);
+
+      await expect(keyStore.listKeys()).resolves.toEqual({ keys: [] });
+
+      await expect(
+        keyStore.addKey({
+          id: testKey.kid,
+          key: testKey,
+          expiresAt: new Date(NaN),
+        }),
+      ).rejects.toThrow('Failed to format public key expiration date');
+    });
+  });
+});

packages/backend-app-api/src/services/implementations/auth/DatabaseKeyStore.ts

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-import { DatabaseService } from '@backstage/backend-plugin-api';
+import { DatabaseService, LoggerService } from '@backstage/backend-plugin-api';
 import { DateTime } from 'luxon';
 import { Knex } from 'knex';
 import { JsonObject } from '@backstage/types';
@@ -22,7 +22,8 @@ import { resolvePackagePath } from '@backstage/backend-common';
 import { KeyStore } from './types';
 
 const MIGRATIONS_TABLE = 'backstage_backend_public_keys__knex_migrations';
-const TABLE = 'backstage_backend_public_keys__keys';
+/** @internal */
+export const TABLE = 'backstage_backend_public_keys__keys';
 
 type Row = {
   id: string;
@@ -30,70 +31,6 @@ type Row = {
   expires_at: string | Date; // Needs parsing to handle different DB implementations
 };
 
-/** @internal */
-export class DatabaseKeyStore implements KeyStore {
-  private constructor(private readonly client: Knex) {}
-
-  static async create(options: { database: DatabaseService }) {
-    const { database } = options;
-
-    console.log(`DEBUG: ###### CREATING STORE`);
-
-    const client = await database.getClient();
-    if (!database.migrations?.skip) {
-      await applyDatabaseMigrations(client);
-    }
-    return new DatabaseKeyStore(client);
-  }
-
-  async addKey(options: {
-    id: string;
-    key: JsonObject & { kid: string };
-    expiresAt: Date;
-  }) {
-    console.log(`DEBUG: STORING KEY`, options);
-    await this.client<Row>(TABLE).insert({
-      id: options.key.kid,
-      key: JSON.stringify(options.key),
-      // TODO: figure out the best way to format this for the DB
-      expires_at: DateTime.fromJSDate(options.expiresAt).toSQL()!,
-    });
-  }
-
-  async listKeys() {
-    const rows = await this.client<Row>(TABLE).select();
-
-    // TODO: move over filter/delete the logic from listPublicKeys() in plugins/auth-backend/src/identity/TokenFactory.ts
-
-    return {
-      keys: rows.map(row => ({
-        id: row.id,
-        key: JSON.parse(row.key),
-        expiresAt: parseDate(row.expires_at),
-      })),
-    };
-  }
-
-  // async removeKeys(kids: string[]): Promise<void> {
-  //   await this.client(TABLE).delete().whereIn('kid', kids);
-  // }
-}
-
-function parseDate(date: string | Date) {
-  const parsedDate =
-    typeof date === 'string'
-      ? DateTime.fromSQL(date, { zone: 'UTC' })
-      : DateTime.fromJSDate(date);
-
-  if (!parsedDate.isValid) {
-    throw new Error(
-      `Failed to parse date, reason: ${parsedDate.invalidReason}, explanation: ${parsedDate.invalidExplanation}`,
-    );
-  }
-
-  return parsedDate.toJSDate();
-}
-
 export function applyDatabaseMigrations(knex: Knex): Promise<void> {
   const migrationsDir = resolvePackagePath(
     '@backstage/backend-app-api',
@@ -105,3 +42,97 @@ export function applyDatabaseMigrations(knex: Knex): Promise<void> {
     tableName: MIGRATIONS_TABLE,
   });
 }
+
+/** @internal */
+export class DatabaseKeyStore implements KeyStore {
+  static async create(options: {
+    database: DatabaseService;
+    logger: LoggerService;
+  }) {
+    const { database, logger } = options;
+
+    const client = await database.getClient();
+    if (!database.migrations?.skip) {
+      await applyDatabaseMigrations(client);
+    }
+    return new DatabaseKeyStore(client, logger);
+  }
+
+  private constructor(
+    private readonly client: Knex,
+    private readonly logger: LoggerService,
+  ) {}
+
+  async addKey(options: {
+    id: string;
+    key: JsonObject & { kid: string };
+    expiresAt: Date;
+  }) {
+    const expiresAt = DateTime.fromJSDate(options.expiresAt).toSQL();
+    if (!expiresAt) {
+      throw new Error('Failed to format public key expiration date');
+    }
+    await this.client<Row>(TABLE).insert({
+      id: options.key.kid,
+      key: JSON.stringify(options.key),
+      expires_at: expiresAt,
+    });
+  }
+
+  async listKeys() {
+    const rows = await this.client<Row>(TABLE).select();
+    const keys = rows.map(row => ({
+      id: row.id,
+      key: JSON.parse(row.key),
+      expiresAt: this.#parseDate(row.expires_at),
+    }));
+
+    const validKeys = [];
+    const expiredKeys = [];
+
+    for (const key of keys) {
+      if (DateTime.fromJSDate(key.expiresAt) < DateTime.local()) {
+        expiredKeys.push(key);
+      } else {
+        validKeys.push(key);
+      }
+    }
+
+    // Lazily prune expired keys. This may cause duplicate removals if we have concurrent callers, but w/e
+    if (expiredKeys.length > 0) {
+      const kids = expiredKeys.map(({ key }) => key.kid);
+
+      this.logger.info(
+        `Removing expired plugin service keys, '${kids.join("', '")}'`,
+      );
+
+      // We don't await this, just let it run in the background
+      this.client<Row>(TABLE)
+        .delete()
+        .whereIn('id', kids)
+        .catch(error => {
+          this.logger.error(
+            'Failed to remove expired plugin service keys',
+            error,
+          );
+        });
+    }
+
+    return { keys: validKeys };
+  }
+
+  #parseDate(date: string | Date) {
+    const parsedDate =
+      typeof date === 'string'
+        ? DateTime.fromSQL(date, { zone: 'UTC' })
+        : DateTime.fromJSDate(date);
+
+    if (!parsedDate.isValid) {
+      throw new Error(
+        `Failed to parse date, reason: ${parsedDate.invalidReason}, explanation: ${parsedDate.invalidExplanation}`,
+      );
+    }
+
+    return parsedDate.toJSDate();
+  }
+}

packages/backend-app-api/src/services/implementations/auth/authServiceFactory.ts

@@ -267,7 +267,7 @@ export const authServiceFactory = createServiceFactory({
       ),
     );
 
-    const publicKeyStore = await DatabaseKeyStore.create({ database });
+    const publicKeyStore = await DatabaseKeyStore.create({ database, logger });
 
     return new DefaultAuthService(
       tokenManager,