feat(userInfo): implement persisting user info to support limited tokens

Signed-off-by: Phil Kuang <pkuang@factset.com>

plugins/auth-backend/migrations/20240510120825_user_info.js

@@ -0,0 +1,44 @@
+/*
+ * 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.
+ */
+
+// @ts-check
+
+/**
+ * @param {import('knex').Knex} knex
+ */
+exports.up = async function up(knex) {
+  await knex.schema.createTable('user_info', table => {
+    table.comment('User information');
+
+    table
+      .string('user_entity_ref')
+      .primary()
+      .notNullable()
+      .comment('User entity reference');
+
+    table
+      .text('user_info', 'longtext')
+      .notNullable()
+      .comment('User info blob, JSON serialized');
+  });
+};
+
+/**
+ * @param {import('knex').Knex} knex
+ */
+exports.down = async function down(knex) {
+  await knex.schema.dropTable('user_info');
+};

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

@@ -24,6 +24,7 @@ import {
 } from 'jose';
 import { MemoryKeyStore } from './MemoryKeyStore';
 import { TokenFactory } from './TokenFactory';
+import { UserInfoDatabaseHandler } from './UserInfoDatabaseHandler';
 import { tokenTypes } from '@backstage/plugin-auth-node';
 
 const logger = getVoidLogger();
@@ -43,6 +44,10 @@ const entityRef = stringifyEntityRef({
 });
 
 describe('TokenFactory', () => {
+  const mockUserInfoDatabaseHandler = {
+    addUserInfo: jest.fn().mockResolvedValue(undefined),
+  } as unknown as UserInfoDatabaseHandler;
+
   it('should issue valid tokens signed by a listed key', async () => {
     const keyDurationSeconds = 5;
     const factory = new TokenFactory({
@@ -50,6 +55,7 @@ describe('TokenFactory', () => {
       keyStore: new MemoryKeyStore(),
       keyDurationSeconds,
       logger,
+      userInfoDatabaseHandler: mockUserInfoDatabaseHandler,
     });
 
     await expect(factory.listPublicKeys()).resolves.toEqual({ keys: [] });
@@ -81,6 +87,11 @@ describe('TokenFactory', () => {
       verifyResult.payload.iat! + keyDurationSeconds,
     );
 
+    expect(mockUserInfoDatabaseHandler.addUserInfo).toHaveBeenCalledWith({
+      userEntityRef: entityRef,
+      ownershipEntityRefs: [entityRef],
+    });
+
     // Emulate the reconstruction of a limited user token
     const limitedUserToken = [
       base64url.encode(
@@ -93,7 +104,6 @@ describe('TokenFactory', () => {
       base64url.encode(
         JSON.stringify({
           sub: verifyResult.payload.sub,
-          ent: verifyResult.payload.ent,
           iat: verifyResult.payload.iat,
           exp: verifyResult.payload.exp,
         }),
@@ -107,7 +117,6 @@ describe('TokenFactory', () => {
     );
     expect(verifyProofResult.payload).toEqual({
       sub: entityRef,
-      ent: [entityRef],
       iat: expect.any(Number),
       exp: expect.any(Number),
     });
@@ -125,6 +134,7 @@ describe('TokenFactory', () => {
       keyStore: new MemoryKeyStore(),
       keyDurationSeconds: 5,
       logger,
+      userInfoDatabaseHandler: mockUserInfoDatabaseHandler,
     });
 
     const token1 = await factory.issueToken({
@@ -170,6 +180,7 @@ describe('TokenFactory', () => {
       keyStore: new MemoryKeyStore(),
       keyDurationSeconds,
       logger,
+      userInfoDatabaseHandler: mockUserInfoDatabaseHandler,
     });
 
     await expect(() => {
@@ -187,6 +198,7 @@ describe('TokenFactory', () => {
       keyDurationSeconds,
       logger,
       algorithm: '',
+      userInfoDatabaseHandler: mockUserInfoDatabaseHandler,
     });
 
     await expect(() => {
@@ -202,6 +214,7 @@ describe('TokenFactory', () => {
       keyStore: new MemoryKeyStore(),
       keyDurationSeconds: 5,
       logger,
+      userInfoDatabaseHandler: mockUserInfoDatabaseHandler,
     });
 
     await expect(() => {
@@ -220,6 +233,7 @@ describe('TokenFactory', () => {
       keyStore: new MemoryKeyStore(),
       keyDurationSeconds,
       logger,
+      userInfoDatabaseHandler: mockUserInfoDatabaseHandler,
     });
 
     const token = await factory.issueToken({

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

@@ -31,6 +31,7 @@ import { LoggerService } from '@backstage/backend-plugin-api';
 import { TokenParams, tokenTypes } from '@backstage/plugin-auth-node';
 import { AnyJWK, KeyStore, TokenIssuer } from './types';
 import { JsonValue } from '@backstage/types';
+import { UserInfoDatabaseHandler } from './UserInfoDatabaseHandler';
 
 const MS_IN_S = 1000;
 const MAX_TOKEN_LENGTH = 32768; // At 64 bytes per entity ref this still leaves room for about 500 entities
@@ -93,11 +94,6 @@ interface BackstageUserIdentityProofPayload {
    */
   sub: string;
 
-  /**
-   * The ownership entity refs of the user
-   */
-  ent?: string[];
-
   /**
    * Standard expiry in epoch seconds
    */
@@ -124,6 +120,7 @@ type Options = {
    * If not, add a knex migration file in the migrations folder.
    * More info on supported algorithms: https://github.com/panva/jose */
   algorithm?: string;
+  userInfoDatabaseHandler: UserInfoDatabaseHandler;
 };
 
 /**
@@ -146,6 +143,7 @@ export class TokenFactory implements TokenIssuer {
   private readonly keyStore: KeyStore;
   private readonly keyDurationSeconds: number;
   private readonly algorithm: string;
+  private readonly userInfoDatabaseHandler: UserInfoDatabaseHandler;
 
   private keyExpiry?: Date;
   private privateKeyPromise?: Promise<JWK>;
@@ -156,6 +154,7 @@ export class TokenFactory implements TokenIssuer {
     this.keyStore = options.keyStore;
     this.keyDurationSeconds = options.keyDurationSeconds;
     this.algorithm = options.algorithm ?? 'ES256';
+    this.userInfoDatabaseHandler = options.userInfoDatabaseHandler;
   }
 
   async issueToken(params: TokenParams): Promise<string> {
@@ -190,7 +189,7 @@ export class TokenFactory implements TokenIssuer {
         alg: key.alg,
         kid: key.kid,
       },
-      payload: { sub, ent, iat, exp },
+      payload: { sub, iat, exp },
       key: signingKey,
     });
 
@@ -221,6 +220,13 @@ export class TokenFactory implements TokenIssuer {
       );
     }
 
+    // Store the user info in the database upon successful token
+    // issuance so that it can be retrieved later by limited user tokens
+    await this.userInfoDatabaseHandler.addUserInfo({
+      userEntityRef: sub,
+      ownershipEntityRefs: ent,
+    });
+
     return token;
   }
 
@@ -342,7 +348,6 @@ export class TokenFactory implements TokenIssuer {
 
     const payload = {
       sub: options.payload.sub,
-      ent: options.payload.ent,
       iat: options.payload.iat,
       exp: options.payload.exp,
     };

plugins/auth-backend/src/identity/UserInfoDatabaseHandler.test.ts

@@ -0,0 +1,109 @@
+/*
+ * 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 { resolvePackagePath } from '@backstage/backend-common';
+import { TestDatabaseId, TestDatabases } from '@backstage/backend-test-utils';
+import { Knex } from 'knex';
+import { UserInfoDatabaseHandler } from './UserInfoDatabaseHandler';
+
+const migrationsDir = resolvePackagePath(
+  '@backstage/plugin-auth-backend',
+  'migrations',
+);
+
+describe('UserInfoDatabaseHandler', () => {
+  const databases = TestDatabases.create();
+
+  async function createDatabaseHandler(databaseId: TestDatabaseId) {
+    const knex = await databases.init(databaseId);
+
+    await knex.migrate.latest({
+      directory: migrationsDir,
+    });
+
+    return {
+      knex,
+      dbHandler: new UserInfoDatabaseHandler(knex),
+    };
+  }
+
+  describe.each(databases.eachSupportedId())(
+    '%p',
+    databaseId => {
+      let knex: Knex;
+      let dbHandler: UserInfoDatabaseHandler;
+
+      beforeEach(async () => {
+        ({ knex, dbHandler } = await createDatabaseHandler(databaseId));
+      }, 30000);
+
+      it('addUserInfo', async () => {
+        const userInfo = {
+          userEntityRef: 'user:default/foo',
+          ownershipEntityRefs: ['group:default/foo-group', 'group:default/bar'],
+        };
+
+        await dbHandler.addUserInfo(userInfo);
+
+        const savedUserInfo = await knex('user_info')
+          .where('user_entity_ref', 'user:default/foo')
+          .first();
+        expect(savedUserInfo).toEqual({
+          user_entity_ref: userInfo.userEntityRef,
+          user_info: JSON.stringify({
+            ownershipEntityRefs: userInfo.ownershipEntityRefs,
+          }),
+        });
+
+        userInfo.ownershipEntityRefs = [
+          'group:default/group1',
+          'group:default/group2',
+        ];
+        await dbHandler.addUserInfo(userInfo);
+
+        const updatedUserInfo = await knex('user_info')
+          .where('user_entity_ref', 'user:default/foo')
+          .first();
+        expect(updatedUserInfo).toEqual({
+          user_entity_ref: userInfo.userEntityRef,
+          user_info: JSON.stringify({
+            ownershipEntityRefs: userInfo.ownershipEntityRefs,
+          }),
+        });
+      });
+
+      it('getUserInfo', async () => {
+        const userInfo = {
+          userEntityRef: 'user:default/backstage-user',
+          ownershipEntityRefs: ['group:default/group1', 'group:default/group2'],
+        };
+
+        await knex('user_info').insert({
+          user_entity_ref: userInfo.userEntityRef,
+          user_info: JSON.stringify({
+            ownershipEntityRefs: userInfo.ownershipEntityRefs,
+          }),
+        });
+
+        const savedUserInfo = await dbHandler.getUserInfo(
+          userInfo.userEntityRef,
+        );
+        expect(savedUserInfo).toEqual(userInfo);
+      });
+    },
+    60000,
+  );
+});

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

@@ -0,0 +1,60 @@
+/*
+ * 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 { BackstageUserInfo } from '@backstage/backend-plugin-api';
+import { Knex } from 'knex';
+
+const TABLE = 'user_info';
+
+type Row = {
+  user_entity_ref: string;
+  user_info: string;
+};
+
+// TODO: How do we prune stale users?
+export class UserInfoDatabaseHandler {
+  constructor(private readonly client: Knex) {}
+
+  async addUserInfo(userInfo: BackstageUserInfo): Promise<void> {
+    await this.client<Row>(TABLE)
+      .insert({
+        user_entity_ref: userInfo.userEntityRef,
+        user_info: JSON.stringify({
+          ownershipEntityRefs: userInfo.ownershipEntityRefs,
+        }),
+      })
+      .onConflict('user_entity_ref')
+      .merge();
+  }
+
+  async getUserInfo(
+    userEntityRef: string,
+  ): Promise<BackstageUserInfo | undefined> {
+    const info = await this.client<Row>(TABLE)
+      .where({ user_entity_ref: userEntityRef })
+      .first();
+
+    if (!info) {
+      return undefined;
+    }
+
+    const { ownershipEntityRefs } = JSON.parse(info.user_info);
+    return {
+      userEntityRef: info.user_entity_ref,
+      ownershipEntityRefs,
+    };
+  }
+}

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

@@ -21,3 +21,4 @@ export { MemoryKeyStore } from './MemoryKeyStore';
 export { FirestoreKeyStore } from './FirestoreKeyStore';
 export { KeyStores } from './KeyStores';
 export type { KeyStore, TokenParams } from './types';
+export { UserInfoDatabaseHandler } from './UserInfoDatabaseHandler';

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

@@ -22,10 +22,15 @@ import { mockServices, startTestBackend } from '@backstage/backend-test-utils';
 import Router from 'express-promise-router';
 import request from 'supertest';
 import { bindOidcRouter } from './router';
+import { UserInfoDatabaseHandler } from './UserInfoDatabaseHandler';
 
 describe('bindOidcRouter', () => {
-  it('should return user info', async () => {
+  it('should return user info for full tokens', async () => {
     const auth = mockServices.auth.mock();
+    const mockUserInfoDatabaseHandler = {
+      getUserInfo: jest.fn().mockResolvedValue(undefined),
+    } as unknown as UserInfoDatabaseHandler;
+
     const { server } = await startTestBackend({
       features: [
         createBackendPlugin({
@@ -39,6 +44,7 @@ describe('bindOidcRouter', () => {
                   baseUrl: 'http://localhost:7000',
                   auth,
                   tokenIssuer: {} as any,
+                  userInfoDatabaseHandler: mockUserInfoDatabaseHandler,
                 });
                 httpRouter.use(router);
                 httpRouter.addAuthPolicy({
@@ -68,6 +74,61 @@ describe('bindOidcRouter', () => {
         ent: ['k/ns:a', 'k/ns:b'],
       });
 
-    expect('test').toBe('test');
+    expect(mockUserInfoDatabaseHandler.getUserInfo).not.toHaveBeenCalled();
+  });
+
+  it('should return user info for limited tokens', async () => {
+    const auth = mockServices.auth.mock();
+    const mockUserInfoDatabaseHandler = {
+      getUserInfo: jest.fn().mockResolvedValue({
+        userEntityRef: 'k/ns:n',
+        ownershipEntityRefs: ['k/ns:a', 'k/ns:b'],
+      }),
+    } as unknown as UserInfoDatabaseHandler;
+
+    const { server } = await startTestBackend({
+      features: [
+        createBackendPlugin({
+          pluginId: 'auth',
+          register(reg) {
+            reg.registerInit({
+              deps: { httpRouter: coreServices.httpRouter },
+              async init({ httpRouter }) {
+                const router = Router();
+                bindOidcRouter(router, {
+                  baseUrl: 'http://localhost:7000',
+                  auth,
+                  tokenIssuer: {} as any,
+                  userInfoDatabaseHandler: mockUserInfoDatabaseHandler,
+                });
+                httpRouter.use(router);
+                httpRouter.addAuthPolicy({
+                  path: '/',
+                  allow: 'unauthenticated',
+                });
+              },
+            });
+          },
+        }),
+      ],
+    });
+
+    auth.authenticate.mockResolvedValueOnce({} as any);
+    auth.isPrincipal.mockReturnValueOnce(true);
+
+    await request(server)
+      .get('/api/auth/v1/userinfo')
+      .set(
+        'Authorization',
+        `Bearer h.${btoa(JSON.stringify({ sub: 'k/ns:n' }))}.s`,
+      )
+      .expect(200, {
+        sub: 'k/ns:n',
+        ent: ['k/ns:a', 'k/ns:b'],
+      });
+
+    expect(mockUserInfoDatabaseHandler.getUserInfo).toHaveBeenCalledWith(
+      'k/ns:n',
+    );
   });
 });

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

@@ -20,6 +20,7 @@ import { TokenIssuer } from './types';
 import { AuthService } from '@backstage/backend-plugin-api';
 import { decodeJwt } from 'jose';
 import { AuthenticationError, InputError } from '@backstage/errors';
+import { UserInfoDatabaseHandler } from './UserInfoDatabaseHandler';
 
 export function bindOidcRouter(
   targetRouter: express.Router,
@@ -27,9 +28,10 @@ export function bindOidcRouter(
     baseUrl: string;
     auth: AuthService;
     tokenIssuer: TokenIssuer;
+    userInfoDatabaseHandler: UserInfoDatabaseHandler;
   },
 ) {
-  const { baseUrl, auth, tokenIssuer } = options;
+  const { baseUrl, auth, tokenIssuer, userInfoDatabaseHandler } = options;
 
   const router = Router();
   targetRouter.use(router);
@@ -91,21 +93,30 @@ export function bindOidcRouter(
       );
     }
 
-    const { sub: userEntityRef, ent: ownershipEntityRefs = [] } =
-      decodeJwt(token);
+    const { sub: userEntityRef, ent: ownershipEntityRefs } = decodeJwt(token);
 
     if (typeof userEntityRef !== 'string') {
       throw new Error('Invalid user token, user entity ref must be a string');
     }
+
+    // Return user info if it's already available in the token (ie. it is a full token)
     if (
-      !Array.isArray(ownershipEntityRefs) ||
-      ownershipEntityRefs.some(ref => typeof ref !== 'string')
+      Array.isArray(ownershipEntityRefs) &&
+      ownershipEntityRefs.every(ref => typeof ref === 'string')
     ) {
-      throw new Error(
-        'Invalid user token, ownership entity refs must be an array of strings',
-      );
+      res.json({ sub: userEntityRef, ent: ownershipEntityRefs });
+      return;
     }
 
-    res.json({ sub: userEntityRef, ent: ownershipEntityRefs });
+    const userInfo = await userInfoDatabaseHandler.getUserInfo(userEntityRef);
+    if (!userInfo) {
+      res.status(404).send('User info not found');
+      return;
+    }
+
+    res.json({
+      sub: userInfo.userEntityRef,
+      ent: userInfo.ownershipEntityRefs,
+    });
   });
 }

plugins/auth-backend/src/migrations.test.ts

@@ -71,4 +71,30 @@ describe('migrations', () => {
       await knex.destroy();
     },
   );
+
+  it.each(databases.eachSupportedId())(
+    '20240510120825_user_info.js, %p',
+    async databaseId => {
+      const knex = await databases.init(databaseId);
+
+      await migrateUntilBefore(knex, '20240510120825_user_info.js');
+      await migrateUpOnce(knex);
+
+      const user_info = JSON.stringify({
+        ownershipEntityRefs: ['group:default/group1', 'group:default/group2'],
+      });
+
+      await knex
+        .insert({ user_entity_ref: 'user:default/backstage-user', user_info })
+        .into('user_info');
+
+      await expect(knex('user_info')).resolves.toEqual([
+        { user_entity_ref: 'user:default/backstage-user', user_info },
+      ]);
+
+      await migrateDownOnce(knex);
+
+      await knex.destroy();
+    },
+  );
 });

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

@@ -32,7 +32,12 @@ import {
 } from '@backstage/backend-common';
 import { NotFoundError } from '@backstage/errors';
 import { CatalogApi } from '@backstage/catalog-client';
-import { bindOidcRouter, KeyStores, TokenFactory } from '../identity';
+import {
+  bindOidcRouter,
+  KeyStores,
+  TokenFactory,
+  UserInfoDatabaseHandler,
+} from '../identity';
 import session from 'express-session';
 import connectSessionKnex from 'connect-session-knex';
 import passport from 'passport';
@@ -87,6 +92,10 @@ export async function createRouter(
     database: authDb,
   });
 
+  const userInfoDatabaseHandler = new UserInfoDatabaseHandler(
+    await authDb.get(),
+  );
+
   let tokenIssuer: TokenIssuer;
   if (keyStore instanceof StaticKeyStore) {
     tokenIssuer = new StaticTokenIssuer(
@@ -106,6 +115,7 @@ export async function createRouter(
       algorithm:
         tokenFactoryAlgorithm ??
         config.getOptionalString('auth.identityTokenAlgorithm'),
+      userInfoDatabaseHandler,
     });
   }
 
@@ -156,6 +166,7 @@ export async function createRouter(
     auth,
     tokenIssuer,
     baseUrl: authUrl,
+    userInfoDatabaseHandler,
   });
 
   // Gives a more helpful error message than a plain 404