From 3938d8863ebb59ee0aaf2da72661adab1747a7b1 Mon Sep 17 00:00:00 2001
From: Wu Tianwei <30284043+WTW0313@users.noreply.github.com>
Date: Mon, 14 Oct 2024 16:54:23 +0800
Subject: [PATCH] fix: token refreshing logic issue (#9308)

---
 web/app/components/swr-initor.tsx |  6 +++---
 web/app/signin/normalForm.tsx     |  2 +-
 web/app/signin/userSSOForm.tsx    |  2 +-
 web/hooks/use-refresh-token.ts    | 33 ++++++++++++++-----------------
 4 files changed, 20 insertions(+), 23 deletions(-)

diff --git a/web/app/components/swr-initor.tsx b/web/app/components/swr-initor.tsx
index 85e05499e6..ce126512fa 100644
--- a/web/app/components/swr-initor.tsx
+++ b/web/app/components/swr-initor.tsx
@@ -14,12 +14,12 @@ const SwrInitor = ({
 }: SwrInitorProps) => {
   const router = useRouter()
   const searchParams = useSearchParams()
+  const { getNewAccessToken } = useRefreshToken()
   const consoleToken = searchParams.get('access_token')
   const refreshToken = searchParams.get('refresh_token')
   const consoleTokenFromLocalStorage = localStorage?.getItem('console_token')
   const refreshTokenFromLocalStorage = localStorage?.getItem('refresh_token')
   const [init, setInit] = useState(false)
-  const { getNewAccessToken } = useRefreshToken()
 
   useEffect(() => {
     if (!(consoleToken || refreshToken || consoleTokenFromLocalStorage || refreshTokenFromLocalStorage)) {
@@ -27,12 +27,12 @@ const SwrInitor = ({
       return
     }
     if (consoleTokenFromLocalStorage && refreshTokenFromLocalStorage)
-      getNewAccessToken(consoleTokenFromLocalStorage, refreshTokenFromLocalStorage)
+      getNewAccessToken()
 
     if (consoleToken && refreshToken) {
       localStorage.setItem('console_token', consoleToken)
       localStorage.setItem('refresh_token', refreshToken)
-      getNewAccessToken(consoleToken, refreshToken).then(() => {
+      getNewAccessToken().then(() => {
         router.replace('/apps', { forceOptimisticNavigation: false } as any)
       }).catch(() => {
         router.replace('/signin')
diff --git a/web/app/signin/normalForm.tsx b/web/app/signin/normalForm.tsx
index 0ae4eb1f43..113ed64b57 100644
--- a/web/app/signin/normalForm.tsx
+++ b/web/app/signin/normalForm.tsx
@@ -99,7 +99,7 @@ const NormalForm = () => {
       if (res.result === 'success') {
         localStorage.setItem('console_token', res.data.access_token)
         localStorage.setItem('refresh_token', res.data.refresh_token)
-        getNewAccessToken(res.data.access_token, res.data.refresh_token)
+        getNewAccessToken()
         router.replace('/apps')
       }
       else {
diff --git a/web/app/signin/userSSOForm.tsx b/web/app/signin/userSSOForm.tsx
index e4b61413bc..f01afa9eaf 100644
--- a/web/app/signin/userSSOForm.tsx
+++ b/web/app/signin/userSSOForm.tsx
@@ -31,7 +31,7 @@ const UserSSOForm: FC<UserSSOFormProps> = ({
     if (refreshToken && consoleToken) {
       localStorage.setItem('console_token', consoleToken)
       localStorage.setItem('refresh_token', refreshToken)
-      getNewAccessToken(consoleToken, refreshToken)
+      getNewAccessToken()
       router.replace('/apps')
     }
 
diff --git a/web/hooks/use-refresh-token.ts b/web/hooks/use-refresh-token.ts
index 3d8779636f..293f3159de 100644
--- a/web/hooks/use-refresh-token.ts
+++ b/web/hooks/use-refresh-token.ts
@@ -14,7 +14,6 @@ const useRefreshToken = () => {
   const router = useRouter()
   const timer = useRef<NodeJS.Timeout>()
   const advanceTime = useRef<number>(5 * 60 * 1000)
-  const interval = useRef<number>(55 * 60 * 1000)
 
   const getExpireTime = useCallback((token: string) => {
     if (!token)
@@ -31,18 +30,24 @@ const useRefreshToken = () => {
     localStorage?.removeItem('is_refreshing')
     localStorage?.removeItem('console_token')
     localStorage?.removeItem('refresh_token')
-    localStorage?.removeItem('last_refresh_time')
     router.replace('/signin')
   }, [])
 
-  const getNewAccessToken = useCallback(async (currentAccessToken: string, currentRefreshToken: string) => {
-    if (localStorage?.getItem('is_refreshing') === '1')
+  const getNewAccessToken = useCallback(async () => {
+    const currentAccessToken = localStorage?.getItem('console_token')
+    const currentRefreshToken = localStorage?.getItem('refresh_token')
+    if (!currentAccessToken || !currentRefreshToken) {
+      handleError()
+      return new Error('No access token or refresh token found')
+    }
+    if (localStorage?.getItem('is_refreshing') === '1') {
+      timer.current = setTimeout(() => {
+        getNewAccessToken()
+      }, 1000)
       return null
+    }
     const currentTokenExpireTime = getExpireTime(currentAccessToken)
-    let lastRefreshTime = parseInt(localStorage?.getItem('last_refresh_time') || '0')
-    lastRefreshTime = isNaN(lastRefreshTime) ? 0 : lastRefreshTime
-    if (getCurrentTimeStamp() + advanceTime.current > currentTokenExpireTime
-      && lastRefreshTime + interval.current < getCurrentTimeStamp()) {
+    if (getCurrentTimeStamp() + advanceTime.current > currentTokenExpireTime) {
       localStorage?.setItem('is_refreshing', '1')
       const [e, res] = await fetchWithRetry(fetchNewToken({
         body: { refresh_token: currentRefreshToken },
@@ -53,24 +58,17 @@ const useRefreshToken = () => {
       }
       const { access_token, refresh_token } = res.data
       localStorage?.setItem('is_refreshing', '0')
-      localStorage?.setItem('last_refresh_time', getCurrentTimeStamp().toString())
       localStorage?.setItem('console_token', access_token)
       localStorage?.setItem('refresh_token', refresh_token)
       const newTokenExpireTime = getExpireTime(access_token)
       timer.current = setTimeout(() => {
-        const consoleTokenFromLocalStorage = localStorage?.getItem('console_token')
-        const refreshTokenFromLocalStorage = localStorage?.getItem('refresh_token')
-        if (consoleTokenFromLocalStorage && refreshTokenFromLocalStorage)
-          getNewAccessToken(consoleTokenFromLocalStorage, refreshTokenFromLocalStorage)
+        getNewAccessToken()
       }, newTokenExpireTime - advanceTime.current - getCurrentTimeStamp())
     }
     else {
       const newTokenExpireTime = getExpireTime(currentAccessToken)
       timer.current = setTimeout(() => {
-        const consoleTokenFromLocalStorage = localStorage?.getItem('console_token')
-        const refreshTokenFromLocalStorage = localStorage?.getItem('refresh_token')
-        if (consoleTokenFromLocalStorage && refreshTokenFromLocalStorage)
-          getNewAccessToken(consoleTokenFromLocalStorage, refreshTokenFromLocalStorage)
+        getNewAccessToken()
       }, newTokenExpireTime - advanceTime.current - getCurrentTimeStamp())
     }
     return null
@@ -80,7 +78,6 @@ const useRefreshToken = () => {
     return () => {
       clearTimeout(timer.current)
       localStorage?.removeItem('is_refreshing')
-      localStorage?.removeItem('last_refresh_time')
     }
   }, [])