backend-defaults: Add support for decoration of the default PluginTokenHandler

Co-authored-by: Patrik Oldsberg <poldsberg@gmail.com>
Signed-off-by: Johan Haals <johan.haals@gmail.com>

.changeset/cuddly-chicken-wink.md

@@ -0,0 +1,5 @@
+---
+'@backstage/backend-defaults': patch
+---
+
+Export `PluginTokenHandler` and `pluginTokenHandlerDecoratorServiceRef` to allow for custom decoration of the plugin token handler without having to re-implement the entire handler.

packages/backend-defaults/report-auth.api.md

@@ -5,6 +5,7 @@
 ```ts
 import { AuthService } from '@backstage/backend-plugin-api';
 import { ServiceFactory } from '@backstage/backend-plugin-api';
+import { ServiceRef } from '@backstage/backend-plugin-api';
 
 // @public
 export const authServiceFactory: ServiceFactory<
@@ -13,5 +14,35 @@ export const authServiceFactory: ServiceFactory<
   'singleton'
 >;
 
+// @public
+export interface PluginTokenHandler {
+  // (undocumented)
+  issueToken(options: {
+    pluginId: string;
+    targetPluginId: string;
+    onBehalfOf?: {
+      limitedUserToken: string;
+      expiresAt: Date;
+    };
+  }): Promise<{
+    token: string;
+  }>;
+  // (undocumented)
+  verifyToken(token: string): Promise<
+    | {
+        subject: string;
+        limitedUserToken?: string;
+      }
+    | undefined
+  >;
+}
+
+// @public
+export const pluginTokenHandlerDecoratorServiceRef: ServiceRef<
+  (defaultImplementation: PluginTokenHandler) => PluginTokenHandler,
+  'plugin',
+  'singleton'
+>;
+
 // (No @packageDocumentation comment for this package)
 ```

packages/backend-defaults/src/entrypoints/auth/DefaultAuthService.ts

@@ -169,7 +169,10 @@ export class DefaultAuthService implements AuthService {
         return this.pluginTokenHandler.issueToken({
           pluginId: this.pluginId,
           targetPluginId,
-          onBehalfOf,
+          onBehalfOf: {
+            limitedUserToken: onBehalfOf.token,
+            expiresAt: onBehalfOf.expiresAt,
+          },
         });
       }
       default:

packages/backend-defaults/src/entrypoints/auth/authServiceFactory.test.ts

@@ -19,12 +19,17 @@ import {
   mockServices,
   registerMswTestHooks,
 } from '@backstage/backend-test-utils';
-import { authServiceFactory } from './authServiceFactory';
+import {
+  authServiceFactory,
+  pluginTokenHandlerDecoratorServiceRef,
+} from './authServiceFactory';
 import { base64url, decodeJwt } from 'jose';
 import { discoveryServiceFactory } from '../discovery';
 import { rest } from 'msw';
 import { setupServer } from 'msw/node';
 import { toInternalBackstageCredentials } from './helpers';
+import { PluginTokenHandler } from './plugin/PluginTokenHandler';
+import { createServiceFactory } from '@backstage/backend-plugin-api';
 
 const server = setupServer();
 
@@ -407,4 +412,41 @@ describe('authServiceFactory', () => {
       principal: { subject: 'unlimited-static-subject' },
     });
   });
+
+  describe('decorate PluginTokenHandler', () => {
+    it('should allow custom logic to be injected into the plugin token handler', async () => {
+      const customLogic = jest.fn();
+      const customPluginTokenHandler = createServiceFactory({
+        service: pluginTokenHandlerDecoratorServiceRef,
+        deps: {},
+        async factory() {
+          return (defaultImplementation: PluginTokenHandler) =>
+            new (class CustomHandler implements PluginTokenHandler {
+              verifyToken(
+                token: string,
+              ): Promise<
+                { subject: string; limitedUserToken?: string } | undefined
+              > {
+                customLogic(token);
+                // check if token is iam/auth or basicAuth, verify.
+                return defaultImplementation.verifyToken(token);
+              }
+              issueToken(options: {
+                pluginId: string;
+                targetPluginId: string;
+                limitedUserToken?: { token: string; expiresAt: Date };
+              }): Promise<{ token: string }> {
+                return defaultImplementation.issueToken(options);
+              }
+            })();
+        },
+      });
+      const tester = ServiceFactoryTester.from(authServiceFactory, {
+        dependencies: [...mockDeps, customPluginTokenHandler],
+      });
+      const searchAuth = await tester.getSubject('search');
+      searchAuth.authenticate('unlimited-static-token');
+      expect(customLogic).toHaveBeenCalledWith('unlimited-static-token');
+    });
+  });
 });

packages/backend-defaults/src/entrypoints/auth/authServiceFactory.ts

@@ -17,13 +17,35 @@
 import {
   coreServices,
   createServiceFactory,
+  createServiceRef,
 } from '@backstage/backend-plugin-api';
 import { DefaultAuthService } from './DefaultAuthService';
 import { ExternalTokenHandler } from './external/ExternalTokenHandler';
-import { PluginTokenHandler } from './plugin/PluginTokenHandler';
+import {
+  DefaultPluginTokenHandler,
+  PluginTokenHandler,
+} from './plugin/PluginTokenHandler';
 import { createPluginKeySource } from './plugin/keys/createPluginKeySource';
 import { UserTokenHandler } from './user/UserTokenHandler';
 
+/**
+ * @public
+ * This service is used to decorate the default plugin token handler with custom logic.
+ */
+export const pluginTokenHandlerDecoratorServiceRef = createServiceRef<
+  (defaultImplementation: PluginTokenHandler) => PluginTokenHandler
+>({
+  id: 'core.auth.pluginTokenHandlerDecorator',
+  defaultFactory: async service =>
+    createServiceFactory({
+      service,
+      deps: {},
+      factory: async () => {
+        return impl => impl;
+      },
+    }),
+});
+
 /**
  * Handles token authentication and credentials management.
  *
@@ -41,8 +63,16 @@ export const authServiceFactory = createServiceFactory({
     discovery: coreServices.discovery,
     plugin: coreServices.pluginMetadata,
     database: coreServices.database,
+    pluginTokenHandlerDecorator: pluginTokenHandlerDecoratorServiceRef,
   },
-  async factory({ config, discovery, plugin, logger, database }) {
+  async factory({
+    config,
+    discovery,
+    plugin,
+    logger,
+    database,
+    pluginTokenHandlerDecorator,
+  }) {
     const disableDefaultAuthPolicy =
       config.getOptionalBoolean(
         'backend.auth.dangerouslyDisableDefaultAuthPolicy',
@@ -61,13 +91,15 @@ export const authServiceFactory = createServiceFactory({
       discovery,
     });
 
-    const pluginTokens = PluginTokenHandler.create({
-      ownPluginId: plugin.getId(),
-      logger,
-      keySource,
-      keyDuration,
-      discovery,
-    });
+    const pluginTokens = pluginTokenHandlerDecorator(
+      DefaultPluginTokenHandler.create({
+        ownPluginId: plugin.getId(),
+        logger,
+        keySource,
+        keyDuration,
+        discovery,
+      }),
+    );
 
     const externalTokens = ExternalTokenHandler.create({
       ownPluginId: plugin.getId(),

packages/backend-defaults/src/entrypoints/auth/index.ts

@@ -14,4 +14,9 @@
  * limitations under the License.
  */
 
-export { authServiceFactory } from './authServiceFactory';
+export {
+  authServiceFactory,
+  pluginTokenHandlerDecoratorServiceRef,
+} from './authServiceFactory';
+
+export type { PluginTokenHandler } from './plugin/PluginTokenHandler';

packages/backend-defaults/src/entrypoints/auth/plugin/PluginTokenHandler.test.ts

@@ -15,7 +15,7 @@
  */
 
 import { mockServices } from '@backstage/backend-test-utils';
-import { PluginTokenHandler } from './PluginTokenHandler';
+import { DefaultPluginTokenHandler } from './PluginTokenHandler';
 import { decodeJwt } from 'jose';
 
 describe('PluginTokenHandler', () => {
@@ -42,7 +42,7 @@ describe('PluginTokenHandler', () => {
     });
 
     const getKeyMock = jest.fn(async () => mockPrivateKey);
-    const handler = PluginTokenHandler.create({
+    const handler = DefaultPluginTokenHandler.create({
       discovery: mockServices.discovery(),
       keyDuration: { seconds: 10 },
       logger: mockServices.logger.mock(),

packages/backend-defaults/src/entrypoints/auth/plugin/PluginTokenHandler.ts

@@ -22,7 +22,6 @@ import { tokenTypes } from '@backstage/plugin-auth-node';
 import { JwksClient } from '../JwksClient';
 import { HumanDuration, durationToMilliseconds } from '@backstage/types';
 import { PluginKeySource } from './keys/types';
-import fetch from 'node-fetch';
 
 const SECONDS_IN_MS = 1000;
 
@@ -45,7 +44,22 @@ type Options = {
   algorithm?: string;
 };
 
-export class PluginTokenHandler {
+/**
+ * @public
+ * PluginTokenHandler is responsible for issuing and verifying tokens
+ */
+export interface PluginTokenHandler {
+  verifyToken(
+    token: string,
+  ): Promise<{ subject: string; limitedUserToken?: string } | undefined>;
+  issueToken(options: {
+    pluginId: string;
+    targetPluginId: string;
+    onBehalfOf?: { limitedUserToken: string; expiresAt: Date };
+  }): Promise<{ token: string }>;
+}
+
+export class DefaultPluginTokenHandler implements PluginTokenHandler {
   private jwksMap = new Map<string, JwksClient>();
 
   // Tracking state for isTargetPluginSupported
@@ -53,7 +67,7 @@ export class PluginTokenHandler {
   private targetPluginInflightChecks = new Map<string, Promise<boolean>>();
 
   static create(options: Options) {
-    return new PluginTokenHandler(
+    return new DefaultPluginTokenHandler(
       options.logger,
       options.ownPluginId,
       options.keySource,
@@ -115,7 +129,7 @@ export class PluginTokenHandler {
   async issueToken(options: {
     pluginId: string;
     targetPluginId: string;
-    onBehalfOf?: { token: string; expiresAt: Date };
+    onBehalfOf?: { limitedUserToken: string; expiresAt: Date };
   }): Promise<{ token: string }> {
     const { pluginId, targetPluginId, onBehalfOf } = options;
     const key = await this.keySource.getPrivateSigningKey();
@@ -131,7 +145,7 @@ export class PluginTokenHandler {
         )
       : ourExp;
 
-    const claims = { sub, aud, iat, exp, obo: onBehalfOf?.token };
+    const claims = { sub, aud, iat, exp, obo: onBehalfOf?.limitedUserToken };
     const token = await new SignJWT(claims)
       .setProtectedHeader({
         typ: tokenTypes.plugin.typParam,
@@ -180,7 +194,7 @@ export class PluginTokenHandler {
 
         this.supportedTargetPlugins.add(targetPluginId);
         return true;
-      } catch (error) {
+      } catch (error: any) {
         this.logger.error('Unexpected failure for target JWKS check', error);
         return false;
       } finally {