feat: experimental option to refresh token on expired kid header

awinogrodzki · Sep 6, 2024 · 2869531 · 2869531
1 parent 4d9c59b
commit 2869531
Show file tree

Hide file tree

Showing 7 changed files with 146 additions and 14 deletions.
diff --git a/src/auth/error.ts b/src/auth/error.ts
@@ -39,8 +39,8 @@ function mergeStackTraceAndCause(target: Error, original: unknown) {
   const originalCause =
     typeof originalError?.cause === 'string' ? originalError.cause : '';
 
-  target.stack = originalErrorStack + (target?.stack ?? '');
-  target.cause = originalCause + (target?.cause ?? '');
+  target.stack = (target?.stack ?? '') + originalErrorStack;
+  target.cause = (target?.cause ?? '') + originalCause;
 }
 
 export class AuthError extends Error {

diff --git a/src/auth/index.ts b/src/auth/index.ts
@@ -181,21 +181,36 @@ export function isInvalidCredentialError(error: unknown): error is AuthError {
   return (error as AuthError)?.code === AuthErrorCode.INVALID_CREDENTIAL;
 }
 
+async function handleVerifyTokenError<T>(
+  e: unknown,
+  onExpired: (e: AuthError) => Promise<T>,
+  onError: (e: unknown) => Promise<T>
+) {
+  try {
+    return await onExpired(e as AuthError);
+  } catch (e) {
+    return onError(e);
+  }
+}
+
 export async function handleExpiredToken<T>(
   verifyIdToken: () => Promise<T>,
   onExpired: (e: AuthError) => Promise<T>,
-  onError: (e: unknown) => Promise<T>
+  onError: (e: unknown) => Promise<T>,
+  shouldExpireOnNoMatchingKidError: boolean
 ): Promise<T> {
   try {
     return await verifyIdToken();
   } catch (e: unknown) {
-    switch ((e as AuthError).code) {
-      case AuthErrorCode.TOKEN_EXPIRED:
-        try {
-          return await onExpired(e as AuthError);
-        } catch (e) {
-          return onError(e);
+    switch ((e as AuthError)?.code) {
+      case AuthErrorCode.NO_MATCHING_KID:
+        if (shouldExpireOnNoMatchingKidError) {
+          return handleVerifyTokenError(e, onExpired, onError);
         }
+
+        return onError(e);
+      case AuthErrorCode.TOKEN_EXPIRED:
+        return handleVerifyTokenError(e, onExpired, onError);
       default:
         return onError(e);
     }
@@ -378,9 +393,20 @@ function getAuth(options: AuthOptions) {
 
         throw new InvalidTokenError(InvalidTokenReason.MISSING_REFRESH_TOKEN);
       },
-      async () => {
-        throw new InvalidTokenError(InvalidTokenReason.INVALID_CREDENTIALS);
-      }
+      async (e) => {
+        if (
+          e instanceof AuthError &&
+          e.code === AuthErrorCode.NO_MATCHING_KID
+        ) {
+          throw InvalidTokenError.fromError(e, InvalidTokenReason.INVALID_KID);
+        }
+
+        throw InvalidTokenError.fromError(
+          e,
+          InvalidTokenReason.INVALID_CREDENTIALS
+        );
+      },
+      verifyOptions.experimental_enableTokenRefreshOnExpiredKidHeader ?? false
     );
   }
 

diff --git a/src/auth/jwt/verify.ts b/src/auth/jwt/verify.ts
@@ -14,6 +14,7 @@ export interface VerifyOptions {
   currentDate?: Date;
   checkRevoked?: boolean;
   referer?: string;
+  experimental_enableTokenRefreshOnExpiredKidHeader?: boolean;
 }
 
 const keyMap: Map<string, KeyLike> = new Map();

diff --git a/src/auth/test/no-matching-kid.integration.test.ts b/src/auth/test/no-matching-kid.integration.test.ts
@@ -0,0 +1,96 @@
+import {v4} from 'uuid';
+import {CLIENT_CERT_URL} from '../firebase';
+import {customTokenToIdAndRefreshTokens, getFirebaseAuth} from '../index';
+import {InvalidTokenError, InvalidTokenReason} from '../error';
+
+const {
+  FIREBASE_API_KEY,
+  FIREBASE_PROJECT_ID,
+  FIREBASE_ADMIN_CLIENT_EMAIL,
+  FIREBASE_ADMIN_PRIVATE_KEY
+} = process.env;
+
+const TEST_SERVICE_ACCOUNT = {
+  clientEmail: FIREBASE_ADMIN_CLIENT_EMAIL!,
+  privateKey: FIREBASE_ADMIN_PRIVATE_KEY!.replace(/\\n/g, '\n'),
+  projectId: FIREBASE_PROJECT_ID!
+};
+
+const REFERER = 'http://localhost:3000';
+
+describe('no matching kid integration test', () => {
+  const {createCustomToken, verifyAndRefreshExpiredIdToken} = getFirebaseAuth(
+    TEST_SERVICE_ACCOUNT,
+    FIREBASE_API_KEY!
+  );
+
+  beforeEach(() => {
+    let numberOfCalls = 0;
+
+    const actualFetch = global.fetch;
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    global.fetch = jest.fn((url: URL, ...args: any[]) => {
+      if (url?.href === CLIENT_CERT_URL && !numberOfCalls) {
+        numberOfCalls++;
+        return {
+          ok: true,
+          headers: {
+            forEach: () => {},
+            has: () => false
+          },
+          json: () => Promise.resolve({})
+        };
+      }
+
+      return actualFetch(url, ...args);
+    }) as jest.Mock;
+  });
+
+  it('should throw invalid token error if kid header does not match public keys', async () => {
+    const userId = v4();
+    const customToken = await createCustomToken(userId, {
+      customClaim: 'customClaimValue'
+    });
+
+    const {idToken, refreshToken} = await customTokenToIdAndRefreshTokens(
+      customToken,
+      FIREBASE_API_KEY!,
+      {referer: REFERER}
+    );
+
+    return expect(() =>
+      verifyAndRefreshExpiredIdToken(
+        {idToken, refreshToken, customToken},
+        {
+          referer: REFERER
+        }
+      )
+    ).rejects.toEqual(new InvalidTokenError(InvalidTokenReason.INVALID_KID));
+  });
+
+  it('should refresh the token if kid header does not match public keys and experimental flag is provided', async () => {
+    const userId = v4();
+    const customToken = await createCustomToken(userId, {
+      customClaim: 'customClaimValue'
+    });
+
+    const {idToken, refreshToken} = await customTokenToIdAndRefreshTokens(
+      customToken,
+      FIREBASE_API_KEY!,
+      {referer: REFERER}
+    );
+
+    const onTokenRefresh = jest.fn();
+
+    const result = await verifyAndRefreshExpiredIdToken(
+      {idToken, refreshToken, customToken},
+      {
+        referer: REFERER,
+        experimental_enableTokenRefreshOnExpiredKidHeader: true,
+        onTokenRefresh
+      }
+    );
+
+    expect(onTokenRefresh).toHaveBeenCalledWith(result);
+  });
+});
diff --git a/src/auth/test/verify-token.integration.test.ts b/src/auth/test/verify-token.integration.test.ts
@@ -104,15 +104,19 @@ describe('verify token integration test', () => {
           {tenantId, appCheckToken: appCheckToken.token, referer: REFERER}
         );
 
+        const onTokenRefresh = jest.fn();
+
         const result = await verifyAndRefreshExpiredIdToken(
           {idToken, refreshToken, customToken},
           {
             currentDate: new Date(Date.now() + 7200 * 1000),
-            referer: REFERER
+            referer: REFERER,
+            onTokenRefresh
           }
         );
 
         expect(result?.decodedIdToken?.customClaim).toEqual('customClaimValue');
+        expect(onTokenRefresh).toHaveBeenCalledWith(result);
       });
 
       it('should verify token', async () => {

diff --git a/src/next/middleware.ts b/src/next/middleware.ts
@@ -128,6 +128,7 @@ export interface AuthMiddlewareOptions extends CreateAuthMiddlewareOptions {
   handleInvalidToken?: HandleInvalidToken;
   handleValidToken?: HandleValidToken;
   handleError?: HandleError;
+  experimental_enableTokenRefreshOnExpiredKidHeader?: boolean;
 }
 
 const defaultInvalidTokenHandler = async () => NextResponse.next();
@@ -269,7 +270,8 @@ export async function authMiddleware(
         debug('Authentication failed with error', {error: e});
 
         return handleError(e);
-      }
+      },
+      options.experimental_enableTokenRefreshOnExpiredKidHeader ?? false
     );
   } catch (error: unknown) {
     if (error instanceof InvalidTokenError) {

diff --git a/src/next/tokens.ts b/src/next/tokens.ts
@@ -24,6 +24,7 @@ export interface GetTokensOptions extends GetCookiesTokensOptions {
   apiKey: string;
   debug?: boolean;
   headers?: Headers;
+  experimental_enableTokenRefreshOnExpiredKidHeader?: boolean;
 }
 
 export function validateOptions(options: GetTokensOptions) {
@@ -125,6 +126,8 @@ export async function getTokens(
 
     const result = await verifyAndRefreshExpiredIdToken(tokens, {
       referer,
+      experimental_enableTokenRefreshOnExpiredKidHeader:
+        options.experimental_enableTokenRefreshOnExpiredKidHeader,
       async onTokenRefresh({idToken, refreshToken, customToken}) {
         const cookieSerializeOptions = options.cookieSerializeOptions;