Clean up as per review comments

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

packages/backend-common/src/tokens/ServerTokenManager.ts

@@ -14,14 +14,7 @@
  * limitations under the License.
  */
 
-import {
-  base64url,
-  JWSHeaderParameters,
-  generateSecret,
-  SignJWT,
-  jwtVerify,
-  exportJWK,
-} from 'jose';
+import { base64url, generateSecret, SignJWT, jwtVerify, exportJWK } from 'jose';
 import { Config } from '@backstage/config';
 import { AuthenticationError } from '@backstage/errors';
 import { TokenManager } from './types';
@@ -76,8 +69,7 @@ export class ServerTokenManager implements TokenManager {
         'No secrets provided when constructing ServerTokenManager',
       );
     }
-    this.verificationKeys = new Array<Uint8Array>();
-    secrets.map(k => this.verificationKeys.push(base64url.decode(k)));
+    this.verificationKeys = secrets.map(s => base64url.decode(s));
     this.signingKey = this.verificationKeys[0];
   }
 
@@ -109,11 +101,8 @@ export class ServerTokenManager implements TokenManager {
   async authenticate(token: string): Promise<void> {
     let verifyError = undefined;
     for (const key of this.verificationKeys) {
-      const lookup = (_protectedHeader: JWSHeaderParameters): Uint8Array => {
-        return key;
-      };
       try {
-        await jwtVerify(token, lookup);
+        await jwtVerify(token, key);
         // If the verify succeeded, return
         return;
       } catch (e) {

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