From 7edb48b66d9a1a1c55bbab377de11cbbbade787f Mon Sep 17 00:00:00 2001
From: Arno V <versini@gmail.com>
Date: Thu, 4 Jul 2024 20:09:18 -0400
Subject: [PATCH] fix: refactor getAccessToken to prevent async issues (#69)

* fix: refactor getAccessToken to prevent async issues

* Update bundlesize.config.js
---
 packages/auth-provider/bundlesize.config.js   |  2 +-
 .../auth-provider/src/common/utilities.ts     | 84 +++++++++++++++++++
 .../components/AuthProvider/AuthProvider.tsx  | 28 +++----
 3 files changed, 96 insertions(+), 18 deletions(-)

diff --git a/packages/auth-provider/bundlesize.config.js b/packages/auth-provider/bundlesize.config.js
index 663cc57..8702045 100644
--- a/packages/auth-provider/bundlesize.config.js
+++ b/packages/auth-provider/bundlesize.config.js
@@ -9,7 +9,7 @@ export default {
 		 */
 		{
 			path: "dist/index.js",
-			limit: "10 kb",
+			limit: "11 kb",
 		},
 	],
 };
diff --git a/packages/auth-provider/src/common/utilities.ts b/packages/auth-provider/src/common/utilities.ts
index 07c1317..e4cf952 100644
--- a/packages/auth-provider/src/common/utilities.ts
+++ b/packages/auth-provider/src/common/utilities.ts
@@ -216,3 +216,87 @@ export const getAccessTokenSilently = async ({
 		};
 	}
 };
+
+type RefreshTokenResponse = {
+	status: "success" | "failure";
+	newAccessToken?: string;
+	newRefreshToken?: string;
+};
+type RefreshTokenProps = {
+	clientId: string;
+	userId: string;
+	nonce: string;
+};
+export class TokenManager {
+	private refreshTokenPromise: Promise<any> | null = null;
+	private accessToken: string;
+	private refreshToken: string;
+
+	constructor(
+		accessToken: string | null = null,
+		refreshToken: string | null = null,
+	) {
+		this.accessToken = accessToken || "";
+		this.refreshToken = refreshToken || "";
+	}
+
+	async refreshtoken({
+		clientId,
+		userId,
+		nonce,
+	}: RefreshTokenProps): Promise<RefreshTokenResponse> {
+		if (!this.refreshTokenPromise) {
+			// No existing refresh in progress, start a new one
+			this.refreshTokenPromise = this._refreshToken({
+				clientId,
+				userId,
+				nonce,
+			});
+		}
+
+		try {
+			// Wait for the existing refresh or start a new one
+			return await this.refreshTokenPromise;
+		} finally {
+			// Clear the promise to allow subsequent calls
+			this.refreshTokenPromise = null;
+		}
+	}
+
+	private async _refreshToken({
+		clientId,
+		userId,
+		nonce,
+	}: RefreshTokenProps): Promise<RefreshTokenResponse> {
+		const jwtRefresh = await verifyAndExtractToken(this.refreshToken);
+		if (jwtRefresh && jwtRefresh.payload[JWT.USER_ID_KEY] !== "") {
+			// Refresh token is valid, refreshing access token...
+			const response = await getAccessTokenSilently({
+				clientId,
+				userId,
+				nonce,
+				refreshToken: this.refreshToken,
+				accessToken: this.accessToken,
+			});
+			if (response.status) {
+				this.accessToken = response.accessToken;
+				this.refreshToken = response.refreshToken;
+				return {
+					status: "success",
+					newAccessToken: response.accessToken,
+					newRefreshToken: response.refreshToken,
+				};
+			} else {
+				// Access token could not be refreshed, re-authenticate the user...
+				return {
+					status: "failure",
+				};
+			}
+		} else {
+			// Refresh token is not valid, re-authenticate the user...
+			return {
+				status: "failure",
+			};
+		}
+	}
+}
diff --git a/packages/auth-provider/src/components/AuthProvider/AuthProvider.tsx b/packages/auth-provider/src/components/AuthProvider/AuthProvider.tsx
index 4fb712c..b8a72e3 100644
--- a/packages/auth-provider/src/components/AuthProvider/AuthProvider.tsx
+++ b/packages/auth-provider/src/components/AuthProvider/AuthProvider.tsx
@@ -21,8 +21,8 @@ import type {
 	LoginType,
 } from "../../common/types";
 import {
+	TokenManager,
 	authenticateUser,
-	getAccessTokenSilently,
 	getPreAuthCode,
 	logoutUser,
 } from "../../common/utilities";
@@ -47,6 +47,7 @@ export const AuthProvider = ({
 	const [nonce, setNonce, , removeNonce] = useLocalStorage({
 		key: `${LOCAL_STORAGE_PREFIX}::${clientId}::@@nonce@@`,
 	});
+	const tokenManager = new TokenManager(accessToken, refreshToken);
 
 	const [authState, setAuthState] = useState<AuthState>({
 		isLoading: true,
@@ -216,22 +217,15 @@ export const AuthProvider = ({
 				 * accessToken is not valid, so we need to try to refresh it using the
 				 * refreshToken - this is a silent refresh.
 				 */
-				const jwtRefresh = await verifyAndExtractToken(refreshToken);
-				if (jwtRefresh && jwtRefresh.payload[JWT.USER_ID_KEY] !== "") {
-					const response = await getAccessTokenSilently({
-						clientId,
-						userId: user.userId as string,
-						nonce,
-						refreshToken,
-						accessToken,
-					});
-					if (response.status) {
-						setAccessToken(response.accessToken);
-						setRefreshToken(response.refreshToken);
-						return response.accessToken;
-					}
-					removeStateAndLocalStorage(ACCESS_TOKEN_ERROR);
-					return "";
+				const res = await tokenManager.refreshtoken({
+					clientId,
+					userId: user.userId as string,
+					nonce,
+				});
+				if (res.status) {
+					setAccessToken(res.newAccessToken);
+					setRefreshToken(res.newRefreshToken);
+					return res.newAccessToken;
 				}
 				/**
 				 * refreshToken is not valid, so we need to re-authenticate the user.