Merge pull request #24729 from kuangp/feat/userInfo

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

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

@@ -0,0 +1,49 @@
+/*
+ * 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');
+
+    table
+      .timestamp('exp')
+      .notNullable()
+      .comment('Expiration timestamp of the user info');
+  });
+};
+
+/**
+ * @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

@@ -21,8 +21,10 @@ import {
   decodeProtectedHeader,
   jwtVerify,
 } from 'jose';
+import { omit } from 'lodash';
 import { MemoryKeyStore } from './MemoryKeyStore';
 import { TokenFactory } from './TokenFactory';
+import { UserInfoDatabaseHandler } from './UserInfoDatabaseHandler';
 import { tokenTypes } from '@backstage/plugin-auth-node';
 import { mockServices } from '@backstage/backend-test-utils';
 
@@ -43,6 +45,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 +56,7 @@ describe('TokenFactory', () => {
       keyStore: new MemoryKeyStore(),
       keyDurationSeconds,
       logger,
+      userInfoDatabaseHandler: mockUserInfoDatabaseHandler,
     });
 
     await expect(factory.listPublicKeys()).resolves.toEqual({ keys: [] });
@@ -81,6 +88,10 @@ describe('TokenFactory', () => {
       verifyResult.payload.iat! + keyDurationSeconds,
     );
 
+    expect(mockUserInfoDatabaseHandler.addUserInfo).toHaveBeenCalledWith({
+      claims: omit(verifyResult.payload, ['aud', 'iat', 'iss', 'uip']),
+    });
+
     // 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

@@ -25,22 +25,22 @@ import {
   GeneralSign,
   KeyLike,
 } from 'jose';
+import { omit } from 'lodash';
 import { DateTime } from 'luxon';
 import { v4 as uuid } from 'uuid';
 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
 
 /**
  * The payload contents of a valid Backstage JWT token
- *
- * @internal
  */
-interface BackstageTokenPayload {
+export interface BackstageTokenPayload {
   /**
    * The issuer of the token, currently the discovery URL of the auth backend
    */
@@ -93,11 +93,6 @@ interface BackstageUserIdentityProofPayload {
    */
   sub: string;
 
-  /**
-   * The ownership entity refs of the user
-   */
-  ent?: string[];
-
   /**
    * Standard expiry in epoch seconds
    */
@@ -124,6 +119,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 +142,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 +153,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 +188,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 +219,12 @@ 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({
+      claims: omit(claims, ['aud', 'iat', 'iss', 'uip']),
+    });
+
     return token;
   }
 
@@ -342,7 +346,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,110 @@
+/*
+ * 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',
+);
+
+jest.setTimeout(60_000);
+
+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())(
+    'should support database %p',
+    databaseId => {
+      let knex: Knex;
+      let dbHandler: UserInfoDatabaseHandler;
+
+      beforeEach(async () => {
+        ({ knex, dbHandler } = await createDatabaseHandler(databaseId));
+      });
+
+      it('addUserInfo', async () => {
+        const userInfo = {
+          claims: {
+            sub: 'user:default/foo',
+            ent: ['group:default/foo-group', 'group:default/bar'],
+            exp: 1234567890,
+          },
+        };
+
+        await dbHandler.addUserInfo(userInfo);
+
+        const savedUserInfo = await knex('user_info')
+          .where('user_entity_ref', 'user:default/foo')
+          .first();
+        expect(savedUserInfo).toEqual({
+          user_entity_ref: 'user:default/foo',
+          user_info: JSON.stringify(userInfo),
+          exp: expect.anything(),
+        });
+
+        userInfo.claims.ent = ['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: 'user:default/foo',
+          user_info: JSON.stringify(userInfo),
+          exp: expect.anything(),
+        });
+      });
+
+      it('getUserInfo', async () => {
+        const userInfo = {
+          claims: {
+            sub: 'user:default/backstage-user',
+            ent: ['group:default/group1', 'group:default/group2'],
+            exp: 1234567890,
+          },
+        };
+
+        await knex('user_info').insert({
+          user_entity_ref: 'user:default/backstage-user',
+          user_info: JSON.stringify(userInfo),
+          exp: knex.fn.now(),
+        });
+
+        const savedUserInfo = await dbHandler.getUserInfo(
+          'user:default/backstage-user',
+        );
+        expect(savedUserInfo).toEqual(userInfo);
+      });
+    },
+  );
+});

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

@@ -0,0 +1,62 @@
+/*
+ * 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 { DateTime } from 'luxon';
+import { Knex } from 'knex';
+
+import { BackstageTokenPayload } from './TokenFactory';
+
+const TABLE = 'user_info';
+
+type Row = {
+  user_entity_ref: string;
+  user_info: string;
+  exp: string;
+};
+
+type UserInfo = {
+  claims: Omit<BackstageTokenPayload, 'aud' | 'iat' | 'iss' | 'uip'>;
+};
+
+export class UserInfoDatabaseHandler {
+  constructor(private readonly client: Knex) {}
+
+  async addUserInfo(userInfo: UserInfo): Promise<void> {
+    await this.client<Row>(TABLE)
+      .insert({
+        user_entity_ref: userInfo.claims.sub as string,
+        user_info: JSON.stringify(userInfo),
+        exp: DateTime.fromSeconds(userInfo.claims.exp as number, {
+          zone: 'utc',
+        }).toSQL({ includeOffset: false }),
+      })
+      .onConflict('user_entity_ref')
+      .merge();
+  }
+
+  async getUserInfo(userEntityRef: string): Promise<UserInfo | undefined> {
+    const info = await this.client<Row>(TABLE)
+      .where({ user_entity_ref: userEntityRef })
+      .first();
+
+    if (!info) {
+      return undefined;
+    }
+
+    const userInfo = JSON.parse(info.user_info);
+    return userInfo;
+  }
+}

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,20 @@ 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({
+        claims: {
+          sub: 'k/ns:n',
+          ent: ['k/ns:a', 'k/ns:b'],
+        },
+      }),
+    } as unknown as UserInfoDatabaseHandler;
+
     const { server } = await startTestBackend({
       features: [
         createBackendPlugin({
@@ -39,6 +49,7 @@ describe('bindOidcRouter', () => {
                   baseUrl: 'http://localhost:7000',
                   auth,
                   tokenIssuer: {} as any,
+                  userInfoDatabaseHandler: mockUserInfoDatabaseHandler,
                 });
                 httpRouter.use(router);
                 httpRouter.addAuthPolicy({
@@ -64,10 +75,73 @@ describe('bindOidcRouter', () => {
         )}.s`,
       )
       .expect(200, {
-        sub: 'k/ns:n',
-        ent: ['k/ns:a', 'k/ns:b'],
+        claims: {
+          sub: 'k/ns:n',
+          ent: ['k/ns:a', 'k/ns:b'],
+        },
       });
 
-    expect('test').toBe('test');
+    expect(mockUserInfoDatabaseHandler.getUserInfo).toHaveBeenCalledWith(
+      'k/ns:n',
+    );
+  });
+
+  it('should return user info for limited tokens', async () => {
+    const auth = mockServices.auth.mock();
+    const mockUserInfoDatabaseHandler = {
+      getUserInfo: jest.fn().mockResolvedValue({
+        claims: {
+          sub: 'k/ns:n',
+          ent: ['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, {
+        claims: {
+          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,18 @@ export function bindOidcRouter(
       );
     }
 
-    const { sub: userEntityRef, ent: ownershipEntityRefs = [] } =
-      decodeJwt(token);
+    const { sub: userEntityRef } = decodeJwt(token);
 
     if (typeof userEntityRef !== 'string') {
       throw new Error('Invalid user token, user entity ref must be a string');
     }
-    if (
-      !Array.isArray(ownershipEntityRefs) ||
-      ownershipEntityRefs.some(ref => typeof ref !== 'string')
-    ) {
-      throw new Error(
-        'Invalid user token, ownership entity refs must be an array of strings',
-      );
+
+    const userInfo = await userInfoDatabaseHandler.getUserInfo(userEntityRef);
+    if (!userInfo) {
+      res.status(404).send('User info not found');
+      return;
     }
 
-    res.json({ sub: userEntityRef, ent: ownershipEntityRefs });
+    res.json(userInfo);
   });
 }

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

@@ -71,4 +71,40 @@ 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({
+        claims: {
+          ent: ['group:default/group1', 'group:default/group2'],
+        },
+      });
+
+      await knex
+        .insert({
+          user_entity_ref: 'user:default/backstage-user',
+          user_info,
+          exp: knex.fn.now(),
+        })
+        .into('user_info');
+
+      await expect(knex('user_info')).resolves.toEqual([
+        {
+          user_entity_ref: 'user:default/backstage-user',
+          user_info,
+          exp: expect.anything(),
+        },
+      ]);
+
+      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