backend-app-api: move jwks to JwksClient

Co-authored-by: Patrik Oldsberg <poldsberg@gmail.com>
Signed-off-by: Vincenzo Scamporlino <vincenzos@spotify.com>

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

@@ -0,0 +1,75 @@
+/*
+ * 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 { AuthenticationError } from '@backstage/errors';
+import {
+  createRemoteJWKSet,
+  decodeJwt,
+  decodeProtectedHeader,
+  FlattenedJWSInput,
+  JWSHeaderParameters,
+} from 'jose';
+import { GetKeyFunction } from 'jose/dist/types/types';
+
+const CLOCK_MARGIN_S = 10;
+
+export class JwksClient {
+  #keyStore?: GetKeyFunction<JWSHeaderParameters, FlattenedJWSInput>;
+  #keyStoreUpdated: number = 0;
+
+  constructor(private readonly getEndpoint: () => Promise<URL>) {}
+
+  get getKey() {
+    if (!this.#keyStore) {
+      throw new AuthenticationError(
+        'refreshKeyStore must be called before jwksClient.getKey',
+      );
+    }
+    return this.#keyStore;
+  }
+
+  /**
+   * If the last keystore refresh is stale, update the keystore URL to the latest
+   */
+  async refreshKeyStore(rawJwtToken: string): Promise<void> {
+    const payload = await decodeJwt(rawJwtToken);
+    const header = await decodeProtectedHeader(rawJwtToken);
+
+    // Refresh public keys if needed
+    let keyStoreHasKey;
+    try {
+      if (this.#keyStore) {
+        // Check if the key is present in the keystore
+        const [_, rawPayload, rawSignature] = rawJwtToken.split('.');
+        keyStoreHasKey = await this.#keyStore(header, {
+          payload: rawPayload,
+          signature: rawSignature,
+        });
+      }
+    } catch (error) {
+      keyStoreHasKey = false;
+    }
+    // Refresh public key URL if needed
+    // Add a small margin in case clocks are out of sync
+    const issuedAfterLastRefresh =
+      payload?.iat && payload.iat > this.#keyStoreUpdated - CLOCK_MARGIN_S;
+    if (!this.#keyStore || (!keyStoreHasKey && issuedAfterLastRefresh)) {
+      const endpoint = await this.getEndpoint();
+      this.#keyStore = createRemoteJWKSet(endpoint);
+      this.#keyStoreUpdated = Date.now() / 1000;
+    }
+  }
+}

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

@@ -28,8 +28,9 @@ import { DateTime } from 'luxon';
 import { v4 as uuid } from 'uuid';
 import { InternalKey, KeyStore } from './types';
 import { AuthenticationError } from '@backstage/errors';
-import { createRemoteJWKSet, jwtVerify } from 'jose';
+import { jwtVerify } from 'jose';
 import { tokenTypes } from '@backstage/plugin-auth-node';
+import { JwksClient } from './JwksClient';
 
 const ALLOWED_PLUGIN_ID_PATTERN = /^[a-z0-9_-]+$/i;
 
@@ -52,7 +53,7 @@ type Options = {
 export class PluginTokenHandler {
   private privateKeyPromise?: Promise<JWK>;
   private keyExpiry?: Date;
-  private jwksMap: Record<string, ReturnType<typeof createRemoteJWKSet>> = {};
+  private jwksMap = new Map<string, JwksClient>();
 
   static create(options: Options) {
     return new PluginTokenHandler(
@@ -94,12 +95,18 @@ export class PluginTokenHandler {
       );
     }
 
-    const JWKS = await this.getJWKS(pluginId);
-    const { payload } = await jwtVerify<{ sub: string }>(token, JWKS, {
-      typ: tokenTypes.plugin.typParam,
-      audience: this.ownPluginId,
-      requiredClaims: ['iat', 'exp', 'sub', 'aud'],
-    }).catch(e => {
+    const jwksClient = await this.getJwksClient(pluginId);
+    await jwksClient.refreshKeyStore(token); // TODO(Rugvip): Refactor so that this isn't needed
+
+    const { payload } = await jwtVerify<{ sub: string }>(
+      token,
+      jwksClient.getKey,
+      {
+        typ: tokenTypes.plugin.typParam,
+        audience: this.ownPluginId,
+        requiredClaims: ['iat', 'exp', 'sub', 'aud'],
+      },
+    ).catch(e => {
       throw new AuthenticationError('Invalid plugin token', e);
     });
 
@@ -133,15 +140,25 @@ export class PluginTokenHandler {
     return { token };
   }
 
-  private async getJWKS(pluginId: string) {
-    if (!this.jwksMap[pluginId]) {
-      const url = `${await this.discovery.getBaseUrl(
-        pluginId,
-      )}/.backstage/auth/v1/jwks.json`;
-      this.jwksMap[pluginId] = createRemoteJWKSet(new URL(url));
+  private async getJwksClient(pluginId: string) {
+    // TODO(Rugvip): Make this more resilient to forged tokens, making sure we don't fill up the map
+    const client = this.jwksMap.get(pluginId);
+    if (client) {
+      return client;
     }
-    return this.jwksMap[pluginId];
+
+    const newClient = new JwksClient(async () => {
+      return new URL(
+        `${await this.discovery.getBaseUrl(
+          pluginId,
+        )}/.backstage/auth/v1/jwks.json`,
+      );
+    });
+
+    this.jwksMap.set(pluginId, newClient);
+    return newClient;
   }
+
   private async getKey(): Promise<JWK> {
     // Make sure that we only generate one key at a time
     if (this.privateKeyPromise) {
@@ -178,7 +195,7 @@ export class PluginTokenHandler {
       //       the new one. This also needs to be implemented cross-service though, meaning new services
       //       that boot up need to be able to grab an existing key to use for signing.
       this.logger.info(`Created new signing key ${kid}`);
-      console.log(`DEBUG: publicKey=`, publicKey);
+
       await this.publicKeyStore.addKey({
         id: kid,
         key: publicKey as InternalKey,

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

@@ -14,23 +14,17 @@
  * limitations under the License.
  */
 
-import { PluginEndpointDiscovery } from '@backstage/backend-common';
 import { DiscoveryService } from '@backstage/backend-plugin-api';
 import { AuthenticationError } from '@backstage/errors';
 import { tokenTypes } from '@backstage/plugin-auth-node';
 import {
   base64url,
-  createRemoteJWKSet,
   decodeJwt,
   decodeProtectedHeader,
-  FlattenedJWSInput,
-  JWSHeaderParameters,
   jwtVerify,
   JWTVerifyOptions,
 } from 'jose';
-import { GetKeyFunction } from 'jose/dist/types/types';
-
-const CLOCK_MARGIN_S = 10;
+import { JwksClient } from './JwksClient';
 
 /**
  * An identity client to interact with auth-backend and authenticate Backstage
@@ -39,15 +33,15 @@ const CLOCK_MARGIN_S = 10;
  * @internal
  */
 export class UserTokenHandler {
-  readonly #discovery: PluginEndpointDiscovery;
+  readonly #jwksClient: JwksClient;
   readonly #algorithms?: string[];
 
-  #keyStore?: GetKeyFunction<JWSHeaderParameters, FlattenedJWSInput>;
-  #keyStoreUpdated: number = 0;
-
   constructor(options: { discovery: DiscoveryService }) {
-    this.#discovery = options.discovery;
     this.#algorithms = ['ES256']; // TODO: configurable?
+    this.#jwksClient = new JwksClient(async () => {
+      const url = await options.discovery.getBaseUrl('auth');
+      return new URL(`${url}/.well-known/jwks.json`);
+    });
   }
 
   async verifyToken(token: string) {
@@ -56,15 +50,12 @@ export class UserTokenHandler {
       return undefined;
     }
 
-    await this.refreshKeyStore(token);
-    if (!this.#keyStore) {
-      throw new AuthenticationError('No keystore exists');
-    }
+    await this.#jwksClient.refreshKeyStore(token);
 
     // Verify a limited token, ensuring the necessarily claims are present and token type is correct
     const { payload } = await jwtVerify(
       token,
-      this.#keyStore,
+      this.#jwksClient.getKey,
       verifyOpts,
     ).catch(e => {
       throw new AuthenticationError('Invalid token', e);
@@ -160,37 +151,4 @@ export class UserTokenHandler {
 
     return { token: limitedUserToken, expiresAt: new Date(payload.exp * 1000) };
   }
-
-  /**
-   * If the last keystore refresh is stale, update the keystore URL to the latest
-   */
-  private async refreshKeyStore(rawJwtToken: string): Promise<void> {
-    const payload = await decodeJwt(rawJwtToken);
-    const header = await decodeProtectedHeader(rawJwtToken);
-
-    // Refresh public keys if needed
-    let keyStoreHasKey;
-    try {
-      if (this.#keyStore) {
-        // Check if the key is present in the keystore
-        const [_, rawPayload, rawSignature] = rawJwtToken.split('.');
-        keyStoreHasKey = await this.#keyStore(header, {
-          payload: rawPayload,
-          signature: rawSignature,
-        });
-      }
-    } catch (error) {
-      keyStoreHasKey = false;
-    }
-    // Refresh public key URL if needed
-    // Add a small margin in case clocks are out of sync
-    const issuedAfterLastRefresh =
-      payload?.iat && payload.iat > this.#keyStoreUpdated - CLOCK_MARGIN_S;
-    if (!this.#keyStore || (!keyStoreHasKey && issuedAfterLastRefresh)) {
-      const url = await this.#discovery.getBaseUrl('auth');
-      const endpoint = new URL(`${url}/.well-known/jwks.json`);
-      this.#keyStore = createRemoteJWKSet(endpoint);
-      this.#keyStoreUpdated = Date.now() / 1000;
-    }
-  }
 }