Clean up as per review comments

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

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

@@ -24,7 +24,10 @@ const logger = getVoidLogger();
 
 function jwtKid(jwt: string): string {
   const header = decodeProtectedHeader(jwt);
-  return header.kid ?? '';
+  if (!header.kid) {
+    throw new Error('JWT Header did not contain a key ID (kid)');
+  }
+  return header.kid;
 }
 
 const entityRef = stringifyEntityRef({

plugins/auth-backend/src/providers/aws-alb/provider.ts

@@ -44,11 +44,6 @@ type Options = {
   resolverContext: AuthResolverContext;
 };
 
-export const getJWTHeaders = (input: string): AwsAlbHeaders => {
-  const encoded = input.split('.')[0];
-  return JSON.parse(Buffer.from(encoded, 'base64').toString('utf8'));
-};
-
 export type AwsAlbHeaders = {
   alg: string;
   kid: string;

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

@@ -248,6 +248,38 @@ describe('IdentityClient', () => {
         return await client.authenticate(fakeToken);
       }).rejects.toThrow();
     });
+
+    it('should use an updated endpoint', async () => {
+      const updatedURL = 'http://backstage:9191/an-updated-base';
+      const getBaseUrl = discovery.getBaseUrl;
+      const getExternalBaseUrl = discovery.getExternalBaseUrl;
+      discovery.getBaseUrl = async () => {
+        return updatedURL;
+      };
+      discovery.getExternalBaseUrl = async () => {
+        return updatedURL;
+      };
+      server.use(
+        rest.get(`${updatedURL}/.well-known/jwks.json`, async (_, res, ctx) => {
+          const keys = await factory.listPublicKeys();
+          return res(ctx.json(keys));
+        }),
+      );
+      const token = await factory.issueToken({ claims: { sub: 'foo' } });
+      const response = await client.authenticate(token);
+      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',
+          ownershipEntityRefs: [],
+        },
+      });
+      discovery.getBaseUrl = getBaseUrl;
+      discovery.getExternalBaseUrl = getExternalBaseUrl;
+    });
   });
 
   describe('listPublicKeys', () => {

plugins/auth-node/src/IdentityClient.ts

@@ -72,6 +72,14 @@ 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
     // Note: verify throws if verification fails