From 302f4407f6ec88c1cd68cc4ab8d809a8f301c472 Mon Sep 17 00:00:00 2001 From: NFish Date: Tue, 5 Nov 2024 12:38:31 +0800 Subject: [PATCH] refactor the logic of refreshing access_token (#10068) --- web/app/account/avatar.tsx | 5 +- .../header/account-dropdown/index.tsx | 5 +- web/app/components/swr-initor.tsx | 39 ++---- web/app/signin/normalForm.tsx | 5 +- web/hooks/use-refresh-token.ts | 99 -------------- web/service/base.ts | 128 +++++++++++------- web/service/refresh-token.ts | 75 ++++++++++ 7 files changed, 171 insertions(+), 185 deletions(-) delete mode 100644 web/hooks/use-refresh-token.ts create mode 100644 web/service/refresh-token.ts diff --git a/web/app/account/avatar.tsx b/web/app/account/avatar.tsx index 2b9aeba5da..544e43ab27 100644 --- a/web/app/account/avatar.tsx +++ b/web/app/account/avatar.tsx @@ -23,8 +23,9 @@ export default function AppSelector() { params: {}, }) - if (localStorage?.getItem('console_token')) - localStorage.removeItem('console_token') + localStorage.removeItem('setup_status') + localStorage.removeItem('console_token') + localStorage.removeItem('refresh_token') router.push('/signin') } diff --git a/web/app/components/header/account-dropdown/index.tsx b/web/app/components/header/account-dropdown/index.tsx index 712906ebae..14f079c0f2 100644 --- a/web/app/components/header/account-dropdown/index.tsx +++ b/web/app/components/header/account-dropdown/index.tsx @@ -47,8 +47,9 @@ export default function AppSelector({ isMobile }: IAppSelector) { params: {}, }) - if (localStorage?.getItem('console_token')) - localStorage.removeItem('console_token') + localStorage.removeItem('setup_status') + localStorage.removeItem('console_token') + localStorage.removeItem('refresh_token') router.push('/signin') } diff --git a/web/app/components/swr-initor.tsx b/web/app/components/swr-initor.tsx index ff9a7b832f..2a119df996 100644 --- a/web/app/components/swr-initor.tsx +++ b/web/app/components/swr-initor.tsx @@ -4,7 +4,6 @@ import { SWRConfig } from 'swr' import { useCallback, useEffect, useState } from 'react' import type { ReactNode } from 'react' import { usePathname, useRouter, useSearchParams } from 'next/navigation' -import useRefreshToken from '@/hooks/use-refresh-token' import { fetchSetupStatus } from '@/service/common' type SwrInitorProps = { @@ -15,12 +14,11 @@ const SwrInitor = ({ }: SwrInitorProps) => { const router = useRouter() const searchParams = useSearchParams() - const pathname = usePathname() - const { getNewAccessToken } = useRefreshToken() - const consoleToken = searchParams.get('access_token') - const refreshToken = searchParams.get('refresh_token') + const consoleToken = decodeURIComponent(searchParams.get('access_token') || '') + const refreshToken = decodeURIComponent(searchParams.get('refresh_token') || '') const consoleTokenFromLocalStorage = localStorage?.getItem('console_token') const refreshTokenFromLocalStorage = localStorage?.getItem('refresh_token') + const pathname = usePathname() const [init, setInit] = useState(false) const isSetupFinished = useCallback(async () => { @@ -41,25 +39,6 @@ const SwrInitor = ({ } }, []) - const setRefreshToken = useCallback(async () => { - try { - if (!(consoleToken || refreshToken || consoleTokenFromLocalStorage || refreshTokenFromLocalStorage)) - return Promise.reject(new Error('No token found')) - - if (consoleTokenFromLocalStorage && refreshTokenFromLocalStorage) - await getNewAccessToken() - - if (consoleToken && refreshToken) { - localStorage.setItem('console_token', consoleToken) - localStorage.setItem('refresh_token', refreshToken) - await getNewAccessToken() - } - } - catch (error) { - return Promise.reject(error) - } - }, [consoleToken, refreshToken, consoleTokenFromLocalStorage, refreshTokenFromLocalStorage, getNewAccessToken]) - useEffect(() => { (async () => { try { @@ -68,9 +47,15 @@ const SwrInitor = ({ router.replace('/install') return } - await setRefreshToken() - if (searchParams.has('access_token') || searchParams.has('refresh_token')) + if (!((consoleToken && refreshToken) || (consoleTokenFromLocalStorage && refreshTokenFromLocalStorage))) { + router.replace('/signin') + return + } + if (searchParams.has('access_token') || searchParams.has('refresh_token')) { + consoleToken && localStorage.setItem('console_token', consoleToken) + refreshToken && localStorage.setItem('refresh_token', refreshToken) router.replace(pathname) + } setInit(true) } @@ -78,7 +63,7 @@ const SwrInitor = ({ router.replace('/signin') } })() - }, [isSetupFinished, setRefreshToken, router, pathname, searchParams]) + }, [isSetupFinished, router, pathname, searchParams, consoleToken, refreshToken, consoleTokenFromLocalStorage, refreshTokenFromLocalStorage]) return init ? ( diff --git a/web/app/signin/normalForm.tsx b/web/app/signin/normalForm.tsx index c0f2d89b37..f4f46c68ba 100644 --- a/web/app/signin/normalForm.tsx +++ b/web/app/signin/normalForm.tsx @@ -12,11 +12,9 @@ import cn from '@/utils/classnames' import { getSystemFeatures, invitationCheck } from '@/service/common' import { defaultSystemFeatures } from '@/types/feature' import Toast from '@/app/components/base/toast' -import useRefreshToken from '@/hooks/use-refresh-token' import { IS_CE_EDITION } from '@/config' const NormalForm = () => { - const { getNewAccessToken } = useRefreshToken() const { t } = useTranslation() const router = useRouter() const searchParams = useSearchParams() @@ -38,7 +36,6 @@ const NormalForm = () => { if (consoleToken && refreshToken) { localStorage.setItem('console_token', consoleToken) localStorage.setItem('refresh_token', refreshToken) - getNewAccessToken() router.replace('/apps') return } @@ -71,7 +68,7 @@ const NormalForm = () => { setSystemFeatures(defaultSystemFeatures) } finally { setIsLoading(false) } - }, [consoleToken, refreshToken, message, router, invite_token, isInviteLink, getNewAccessToken]) + }, [consoleToken, refreshToken, message, router, invite_token, isInviteLink]) useEffect(() => { init() }, [init]) diff --git a/web/hooks/use-refresh-token.ts b/web/hooks/use-refresh-token.ts deleted file mode 100644 index 53dc4faf00..0000000000 --- a/web/hooks/use-refresh-token.ts +++ /dev/null @@ -1,99 +0,0 @@ -'use client' -import { useCallback, useEffect, useRef } from 'react' -import { jwtDecode } from 'jwt-decode' -import dayjs from 'dayjs' -import utc from 'dayjs/plugin/utc' -import { useRouter } from 'next/navigation' -import type { CommonResponse } from '@/models/common' -import { fetchNewToken } from '@/service/common' -import { fetchWithRetry } from '@/utils' - -dayjs.extend(utc) - -const useRefreshToken = () => { - const router = useRouter() - const timer = useRef() - const advanceTime = useRef(5 * 60 * 1000) - - const getExpireTime = useCallback((token: string) => { - if (!token) - return 0 - const decoded = jwtDecode(token) - return (decoded.exp || 0) * 1000 - }, []) - - const getCurrentTimeStamp = useCallback(() => { - return dayjs.utc().valueOf() - }, []) - - const handleError = useCallback(() => { - localStorage?.removeItem('is_refreshing') - localStorage?.removeItem('console_token') - localStorage?.removeItem('refresh_token') - router.replace('/signin') - }, []) - - 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') { - clearTimeout(timer.current) - timer.current = setTimeout(() => { - getNewAccessToken() - }, 1000) - return null - } - const currentTokenExpireTime = getExpireTime(currentAccessToken) - if (getCurrentTimeStamp() + advanceTime.current > currentTokenExpireTime) { - localStorage?.setItem('is_refreshing', '1') - const [e, res] = await fetchWithRetry(fetchNewToken({ - body: { refresh_token: currentRefreshToken }, - }) as Promise) - if (e) { - handleError() - return e - } - const { access_token, refresh_token } = res.data - localStorage?.setItem('is_refreshing', '0') - localStorage?.setItem('console_token', access_token) - localStorage?.setItem('refresh_token', refresh_token) - const newTokenExpireTime = getExpireTime(access_token) - clearTimeout(timer.current) - timer.current = setTimeout(() => { - getNewAccessToken() - }, newTokenExpireTime - advanceTime.current - getCurrentTimeStamp()) - } - else { - const newTokenExpireTime = getExpireTime(currentAccessToken) - clearTimeout(timer.current) - timer.current = setTimeout(() => { - getNewAccessToken() - }, newTokenExpireTime - advanceTime.current - getCurrentTimeStamp()) - } - return null - }, [getExpireTime, getCurrentTimeStamp, handleError]) - - const handleVisibilityChange = useCallback(() => { - if (document.visibilityState === 'visible') - getNewAccessToken() - }, []) - - useEffect(() => { - window.addEventListener('visibilitychange', handleVisibilityChange) - return () => { - window.removeEventListener('visibilitychange', handleVisibilityChange) - clearTimeout(timer.current) - localStorage?.removeItem('is_refreshing') - } - }, []) - - return { - getNewAccessToken, - } -} - -export default useRefreshToken diff --git a/web/service/base.ts b/web/service/base.ts index fbdd5c1fd3..fcf8d8bd7d 100644 --- a/web/service/base.ts +++ b/web/service/base.ts @@ -1,3 +1,4 @@ +import { refreshAccessTokenOrRelogin } from './refresh-token' import { API_PREFIX, IS_CE_EDITION, PUBLIC_API_PREFIX } from '@/config' import Toast from '@/app/components/base/toast' import type { AnnotationReply, MessageEnd, MessageReplace, ThoughtItem } from '@/app/components/base/chat/chat/type' @@ -356,39 +357,8 @@ const baseFetch = ( if (!/^(2|3)\d{2}$/.test(String(res.status))) { const bodyJson = res.json() switch (res.status) { - case 401: { - if (isPublicAPI) { - return bodyJson.then((data: ResponseError) => { - if (data.code === 'web_sso_auth_required') - requiredWebSSOLogin() - - if (data.code === 'unauthorized') { - removeAccessToken() - globalThis.location.reload() - } - - return Promise.reject(data) - }) - } - const loginUrl = `${globalThis.location.origin}/signin` - bodyJson.then((data: ResponseError) => { - if (data.code === 'init_validate_failed' && IS_CE_EDITION && !silent) - Toast.notify({ type: 'error', message: data.message, duration: 4000 }) - else if (data.code === 'not_init_validated' && IS_CE_EDITION) - globalThis.location.href = `${globalThis.location.origin}/init` - else if (data.code === 'not_setup' && IS_CE_EDITION) - globalThis.location.href = `${globalThis.location.origin}/install` - else if (location.pathname !== '/signin' || !IS_CE_EDITION) - globalThis.location.href = loginUrl - else if (!silent) - Toast.notify({ type: 'error', message: data.message }) - }).catch(() => { - // Handle any other errors - globalThis.location.href = loginUrl - }) - - break - } + case 401: + return Promise.reject(resClone) case 403: bodyJson.then((data: ResponseError) => { if (!silent) @@ -484,7 +454,9 @@ export const upload = (options: any, isPublicAPI?: boolean, url?: string, search export const ssePost = ( url: string, fetchOptions: FetchOptionType, - { + otherOptions: IOtherOptions, +) => { + const { isPublicAPI = false, onData, onCompleted, @@ -507,8 +479,7 @@ export const ssePost = ( onTextReplace, onError, getAbortController, - }: IOtherOptions, -) => { + } = otherOptions const abortController = new AbortController() const options = Object.assign({}, baseOptions, { @@ -532,21 +503,29 @@ export const ssePost = ( globalThis.fetch(urlWithPrefix, options as RequestInit) .then((res) => { if (!/^(2|3)\d{2}$/.test(String(res.status))) { - res.json().then((data: any) => { - if (isPublicAPI) { - if (data.code === 'web_sso_auth_required') - requiredWebSSOLogin() + if (res.status === 401) { + refreshAccessTokenOrRelogin(TIME_OUT).then(() => { + ssePost(url, fetchOptions, otherOptions) + }).catch(() => { + res.json().then((data: any) => { + if (isPublicAPI) { + if (data.code === 'web_sso_auth_required') + requiredWebSSOLogin() - if (data.code === 'unauthorized') { - removeAccessToken() - globalThis.location.reload() - } - if (res.status === 401) - return - } - Toast.notify({ type: 'error', message: data.message || 'Server Error' }) - }) - onError?.('Server Error') + if (data.code === 'unauthorized') { + removeAccessToken() + globalThis.location.reload() + } + } + }) + }) + } + else { + res.json().then((data) => { + Toast.notify({ type: 'error', message: data.message || 'Server Error' }) + }) + onError?.('Server Error') + } return } return handleStream(res, (str: string, isFirstMessage: boolean, moreInfo: IOnDataMoreInfo) => { @@ -568,7 +547,54 @@ export const ssePost = ( // base request export const request = (url: string, options = {}, otherOptions?: IOtherOptions) => { - return baseFetch(url, options, otherOptions || {}) + return new Promise((resolve, reject) => { + const otherOptionsForBaseFetch = otherOptions || {} + baseFetch(url, options, otherOptionsForBaseFetch).then(resolve).catch((errResp) => { + if (errResp?.status === 401) { + return refreshAccessTokenOrRelogin(TIME_OUT).then(() => { + baseFetch(url, options, otherOptionsForBaseFetch).then(resolve).catch(reject) + }).catch(() => { + const { + isPublicAPI = false, + silent, + } = otherOptionsForBaseFetch + const bodyJson = errResp.json() + if (isPublicAPI) { + return bodyJson.then((data: ResponseError) => { + if (data.code === 'web_sso_auth_required') + requiredWebSSOLogin() + + if (data.code === 'unauthorized') { + removeAccessToken() + globalThis.location.reload() + } + + return Promise.reject(data) + }) + } + const loginUrl = `${globalThis.location.origin}/signin` + bodyJson.then((data: ResponseError) => { + if (data.code === 'init_validate_failed' && IS_CE_EDITION && !silent) + Toast.notify({ type: 'error', message: data.message, duration: 4000 }) + else if (data.code === 'not_init_validated' && IS_CE_EDITION) + globalThis.location.href = `${globalThis.location.origin}/init` + else if (data.code === 'not_setup' && IS_CE_EDITION) + globalThis.location.href = `${globalThis.location.origin}/install` + else if (location.pathname !== '/signin' || !IS_CE_EDITION) + globalThis.location.href = loginUrl + else if (!silent) + Toast.notify({ type: 'error', message: data.message }) + }).catch(() => { + // Handle any other errors + globalThis.location.href = loginUrl + }) + }) + } + else { + reject(errResp) + } + }) + }) } // request methods diff --git a/web/service/refresh-token.ts b/web/service/refresh-token.ts new file mode 100644 index 0000000000..8bd2215041 --- /dev/null +++ b/web/service/refresh-token.ts @@ -0,0 +1,75 @@ +import { apiPrefix } from '@/config' +import { fetchWithRetry } from '@/utils' + +let isRefreshing = false +function waitUntilTokenRefreshed() { + return new Promise((resolve, reject) => { + function _check() { + const isRefreshingSign = localStorage.getItem('is_refreshing') + if ((isRefreshingSign && isRefreshingSign === '1') || isRefreshing) { + setTimeout(() => { + _check() + }, 1000) + } + else { + resolve() + } + } + _check() + }) +} + +// only one request can send +async function getNewAccessToken(): Promise { + try { + const isRefreshingSign = localStorage.getItem('is_refreshing') + if ((isRefreshingSign && isRefreshingSign === '1') || isRefreshing) { + await waitUntilTokenRefreshed() + } + else { + globalThis.localStorage.setItem('is_refreshing', '1') + isRefreshing = true + const refresh_token = globalThis.localStorage.getItem('refresh_token') + + // Do not use baseFetch to refresh tokens. + // If a 401 response occurs and baseFetch itself attempts to refresh the token, + // it can lead to an infinite loop if the refresh attempt also returns 401. + // To avoid this, handle token refresh separately in a dedicated function + // that does not call baseFetch and uses a single retry mechanism. + const [error, ret] = await fetchWithRetry(globalThis.fetch(`${apiPrefix}/refresh-token`, { + method: 'POST', + headers: { + 'Content-Type': 'application/json;utf-8', + }, + body: JSON.stringify({ refresh_token }), + })) + if (error) { + return Promise.reject(error) + } + else { + if (ret.status === 401) + return Promise.reject(ret) + + const { data } = await ret.json() + globalThis.localStorage.setItem('console_token', data.access_token) + globalThis.localStorage.setItem('refresh_token', data.refresh_token) + } + } + } + catch (error) { + console.error(error) + return Promise.reject(error) + } + finally { + isRefreshing = false + globalThis.localStorage.removeItem('is_refreshing') + } +} + +export async function refreshAccessTokenOrRelogin(timeout: number) { + return Promise.race([new Promise((resolve, reject) => setTimeout(() => { + isRefreshing = false + globalThis.localStorage.removeItem('is_refreshing') + reject(new Error('request timeout')) + }, timeout)), getNewAccessToken()]) +}