Update public key URL when key lookup fails and mutex dev key

Signed-off-by: Andy Caruso <macaruso@gmail.com>

plugins/auth-node/src/IdentityClient.test.ts

@@ -21,6 +21,7 @@ import {
   decodeProtectedHeader,
   exportJWK,
 } from 'jose';
+import { cloneDeep } from 'lodash';
 import { rest } from 'msw';
 import { setupServer } from 'msw/node';
 import { IdentityClient } from './IdentityClient';
@@ -249,10 +250,23 @@ describe('IdentityClient', () => {
       }).rejects.toThrow();
     });
 
-    it('should use an updated endpoint', async () => {
+    it('should use an updated endpoint when the key is not found', async () => {
       const updatedURL = 'http://backstage:9191/an-updated-base';
       const getBaseUrl = discovery.getBaseUrl;
       const getExternalBaseUrl = discovery.getExternalBaseUrl;
+      // Generate a key and sign a token with it
+      await factory.issueToken({ claims: { sub: 'foo' } });
+      // Only return the key from a single token
+      const singleKey = cloneDeep(await factory.listPublicKeys());
+      server.use(
+        rest.get(
+          `${mockBaseUrl}/.well-known/jwks.json`,
+          async (_, res, ctx) => {
+            return res(ctx.json(singleKey));
+          },
+        ),
+      );
+      // Update the discovery endpoint to point to a new URL
       discovery.getBaseUrl = async () => {
         return updatedURL;
       };
@@ -265,20 +279,29 @@ describe('IdentityClient', () => {
           return res(ctx.json(keys));
         }),
       );
-      const token = await factory.issueToken({ claims: { sub: 'foo' } });
+      // Advance time
+      const future_11s = Date.now() + 11 * 1000;
+      const dateSpy = jest
+        .spyOn(Date, 'now')
+        .mockImplementation(() => future_11s);
+      // Issue a new token
+      const token = await factory.issueToken({ claims: { sub: 'foo2' } });
       const response = await client.authenticate(token);
+      // Verify that the endpoint was updated.
       const url = (client as any).endpoint as URL;
       expect(url.toString()).toMatch(`${updatedURL}/.well-known/jwks.json`);
       expect(response).toEqual({
         token: token,
         identity: {
           type: 'user',
-          userEntityRef: 'foo',
+          userEntityRef: 'foo2',
           ownershipEntityRefs: [],
         },
       });
+      // Restore the discovery endpoint and time
       discovery.getBaseUrl = getBaseUrl;
       discovery.getExternalBaseUrl = getExternalBaseUrl;
+      dateSpy.mockClear();
     });
   });
 

plugins/auth-node/src/IdentityClient.ts

@@ -18,13 +18,17 @@ import { PluginEndpointDiscovery } from '@backstage/backend-common';
 import { AuthenticationError } from '@backstage/errors';
 import {
   createRemoteJWKSet,
+  decodeJwt,
   jwtVerify,
   FlattenedJWSInput,
   JWSHeaderParameters,
+  decodeProtectedHeader,
 } from 'jose';
 import { GetKeyFunction } from 'jose/dist/types/types';
 import { BackstageIdentityResponse } from './types';
 
+const CLOCK_MARGIN_S = 10;
+
 /**
  * An identity client to interact with auth-backend and authenticate Backstage
  * tokens
@@ -37,6 +41,7 @@ export class IdentityClient {
   private readonly issuer: string;
   private keyStore?: GetKeyFunction<JWSHeaderParameters, FlattenedJWSInput>;
   private endpoint?: URL;
+  private keyStoreUpdated: number = 0;
 
   /**
    * Create a new {@link IdentityClient} instance.
@@ -57,6 +62,7 @@ export class IdentityClient {
     this.discovery.getBaseUrl('auth').then(url => {
       this.endpoint = new URL(`${url}/.well-known/jwks.json`);
       this.keyStore = createRemoteJWKSet(this.endpoint);
+      this.keyStoreUpdated = Date.now() / 1000;
     });
   }
 
@@ -72,13 +78,6 @@ export class IdentityClient {
     if (!token) {
       throw new AuthenticationError('No token specified');
     }
-    // Check if the keystore needs to be updated
-    const url = await this.discovery.getBaseUrl('auth');
-    const endpoint = new URL(`${url}/.well-known/jwks.json`);
-    if (endpoint !== this.endpoint) {
-      this.endpoint = endpoint;
-      this.keyStore = createRemoteJWKSet(this.endpoint);
-    }
 
     // Verify token claims and signature
     // Note: Claims must match those set by TokenFactory when issuing tokens
@@ -86,6 +85,8 @@ export class IdentityClient {
     if (!this.keyStore) {
       throw new AuthenticationError('No keystore exists');
     }
+    // Check if the keystore needs to be updated
+    await this.refreshKeyStore(token);
     const decoded = await jwtVerify(token, this.keyStore, {
       algorithms: ['ES256'],
       audience: 'backstage',
@@ -109,4 +110,38 @@ export class IdentityClient {
     };
     return user;
   }
+
+  /**
+   * 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 (!keyStoreHasKey && issuedAfterLastRefresh) {
+      const url = await this.discovery.getBaseUrl('auth');
+      const endpoint = new URL(`${url}/.well-known/jwks.json`);
+      this.endpoint = endpoint;
+      this.keyStore = createRemoteJWKSet(this.endpoint);
+      this.keyStoreUpdated = Date.now() / 1000;
+    }
+  }
 }