Handle concurrent refresh requests

dust-tt · Dec 10, 2024 · 88842cc · 88842cc
1 parent 35b85d9
commit 88842cc
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 19 deletions.
diff --git a/extension/app/background.ts b/extension/app/background.ts
@@ -23,9 +23,13 @@ const log = console.error;
 
 const state: {
   refreshingToken: boolean;
+  refreshRequests: ((
+    auth: Auth0AuthorizeResponse | AuthBackgroundResponse
+  ) => void)[];
   lastHandler: (() => void) | undefined;
 } = {
   refreshingToken: false,
+  refreshRequests: [],
   lastHandler: undefined,
 };
 
@@ -385,9 +389,8 @@ const refreshToken = async (
   refreshToken: string,
   sendResponse: (auth: Auth0AuthorizeResponse | AuthBackgroundResponse) => void
 ) => {
-  if (state.refreshingToken) {
-    return false;
-  } else {
+  state.refreshRequests.push(sendResponse);
+  if (!state.refreshingToken) {
     state.refreshingToken = true;
     try {
       const tokenUrl = `https://${AUTH0_CLIENT_DOMAIN}/oauth/token`;
@@ -409,15 +412,23 @@ const refreshToken = async (
       }
 
       const data = await response.json();
-      sendResponse({
-        idToken: data.id_token,
-        accessToken: data.access_token,
-        refreshToken: data.refresh_token || refreshToken,
-        expiresIn: data.expires_in,
+      const handlers = state.refreshRequests;
+      state.refreshRequests = [];
+      handlers.forEach((sendResponse) => {
+        sendResponse({
+          idToken: data.id_token,
+          accessToken: data.access_token,
+          refreshToken: data.refresh_token || refreshToken,
+          expiresIn: data.expires_in,
+        });
       });
     } catch (error) {
       log("Token refresh failed: unknown error", error);
-      sendResponse({ success: false });
+      const handlers = state.refreshRequests;
+      state.refreshRequests = [];
+      handlers.forEach((sendResponse) => {
+        sendResponse({ success: false });
+      });
     } finally {
       state.refreshingToken = false;
     }

diff --git a/extension/app/src/hooks/useAuthErrorCheck.ts b/extension/app/src/hooks/useAuthErrorCheck.ts
@@ -1,11 +1,9 @@
 import { useAuth } from "@extension/components/auth/AuthProvider";
 import { logout, refreshToken } from "@extension/lib/auth";
 import { useEffect } from "react";
-import { useNavigate } from "react-router-dom";
 
 export const useAuthErrorCheck = (error: any, mutate: () => any) => {
-  const { setAuthError, handleLogout } = useAuth();
-  const navigate = useNavigate();
+  const { setAuthError } = useAuth();
   useEffect(() => {
     const handleError = async () => {
       if (error) {
@@ -20,8 +18,7 @@ export const useAuthErrorCheck = (error: any, mutate: () => any) => {
             if (res.isOk()) {
               mutate();
             } else {
-              handleLogout();
-              navigate("/login");
+              void logout();
             }
             break;
           case "user_not_found":

diff --git a/extension/app/src/lib/auth.ts b/extension/app/src/lib/auth.ts
@@ -79,11 +79,11 @@ export const logout = async (): Promise<boolean> => {
 // Refresh token sends a message to the background script to call the auth0 refresh token endpoint.
 // It updates the stored tokens with the new access token.
 // If the refresh token is invalid, it will call handleLogout.
-export const refreshToken = async (): Promise<
-  Result<StoredTokens, AuthError>
-> => {
+export const refreshToken = async (
+  tokens?: StoredTokens | null
+): Promise<Result<StoredTokens, AuthError>> => {
   try {
-    const tokens = await getStoredTokens();
+    tokens = tokens ?? (await getStoredTokens());
     if (!tokens) {
       return new Err(new AuthError("not_authenticated", "No tokens found."));
     }
@@ -103,7 +103,7 @@ export const refreshToken = async (): Promise<
 export const getAccessToken = async (): Promise<string | null> => {
   let tokens = await getStoredTokens();
   if (!tokens || !tokens.accessToken || tokens.expiresAt < Date.now()) {
-    const refreshRes = await refreshToken();
+    const refreshRes = await refreshToken(tokens);
     if (refreshRes.isOk()) {
       tokens = refreshRes.value;
     }