feat: refresh-token (#9286)

Co-authored-by: NFish <douxc512@gmail.com>

web/app/components/swr-initor.tsx

@@ -4,6 +4,7 @@ import { SWRConfig } from 'swr'
 import { useEffect, useState } from 'react'
 import type { ReactNode } from 'react'
 import { useRouter, useSearchParams } from 'next/navigation'
+import useRefreshToken from '@/hooks/use-refresh-token'
 
 type SwrInitorProps = {
   children: ReactNode
@@ -13,18 +14,31 @@ const SwrInitor = ({
 }: SwrInitorProps) => {
   const router = useRouter()
   const searchParams = useSearchParams()
-  const consoleToken = searchParams.get('console_token')
+  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 || consoleTokenFromLocalStorage))
+    if (!(consoleToken || refreshToken || consoleTokenFromLocalStorage || refreshTokenFromLocalStorage)) {
       router.replace('/signin')
-
-    if (consoleToken) {
-      localStorage?.setItem('console_token', consoleToken!)
-      router.replace('/apps', { forceOptimisticNavigation: false } as any)
+      return
     }
+    if (consoleTokenFromLocalStorage && refreshTokenFromLocalStorage)
+      getNewAccessToken(consoleTokenFromLocalStorage, refreshTokenFromLocalStorage)
+
+    if (consoleToken && refreshToken) {
+      localStorage.setItem('console_token', consoleToken)
+      localStorage.setItem('refresh_token', refreshToken)
+      getNewAccessToken(consoleToken, refreshToken).then(() => {
+        router.replace('/apps', { forceOptimisticNavigation: false } as any)
+      }).catch(() => {
+        router.replace('/signin')
+      })
+    }
+
     setInit(true)
   }, [])
 

web/app/signin/normalForm.tsx

@@ -11,6 +11,7 @@ import { IS_CE_EDITION, SUPPORT_MAIL_LOGIN, apiPrefix, emailRegex } from '@/conf
 import Button from '@/app/components/base/button'
 import { login, oauth } from '@/service/common'
 import { getPurifyHref } from '@/utils'
+import useRefreshToken from '@/hooks/use-refresh-token'
 
 type IState = {
   formValid: boolean
@@ -61,6 +62,7 @@ function reducer(state: IState, action: IAction) {
 
 const NormalForm = () => {
   const { t } = useTranslation()
+  const { getNewAccessToken } = useRefreshToken()
   const useEmailLogin = IS_CE_EDITION || SUPPORT_MAIL_LOGIN
 
   const router = useRouter()
@@ -95,7 +97,9 @@ const NormalForm = () => {
         },
       })
       if (res.result === 'success') {
-        localStorage.setItem('console_token', res.data)
+        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)
         router.replace('/apps')
       }
       else {

web/app/signin/userSSOForm.tsx

@@ -7,6 +7,7 @@ import cn from '@/utils/classnames'
 import Toast from '@/app/components/base/toast'
 import { getUserOAuth2SSOUrl, getUserOIDCSSOUrl, getUserSAMLSSOUrl } from '@/service/sso'
 import Button from '@/app/components/base/button'
+import useRefreshToken from '@/hooks/use-refresh-token'
 
 type UserSSOFormProps = {
   protocol: string
@@ -15,8 +16,10 @@ type UserSSOFormProps = {
 const UserSSOForm: FC<UserSSOFormProps> = ({
   protocol,
 }) => {
+  const { getNewAccessToken } = useRefreshToken()
   const searchParams = useSearchParams()
-  const consoleToken = searchParams.get('console_token')
+  const consoleToken = searchParams.get('access_token')
+  const refreshToken = searchParams.get('refresh_token')
   const message = searchParams.get('message')
 
   const router = useRouter()
@@ -25,8 +28,10 @@ const UserSSOForm: FC<UserSSOFormProps> = ({
   const [isLoading, setIsLoading] = useState(false)
 
   useEffect(() => {
-    if (consoleToken) {
+    if (refreshToken && consoleToken) {
       localStorage.setItem('console_token', consoleToken)
+      localStorage.setItem('refresh_token', refreshToken)
+      getNewAccessToken(consoleToken, refreshToken)
       router.replace('/apps')
     }
 
@@ -36,7 +41,7 @@ const UserSSOForm: FC<UserSSOFormProps> = ({
         message,
       })
     }
-  }, [])
+  }, [consoleToken, refreshToken, message, router])
 
   const handleSSOLogin = () => {
     setIsLoading(true)

web/hooks/use-refresh-token.ts

@@ -0,0 +1,92 @@
+'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<NodeJS.Timeout>()
+  const advanceTime = useRef<number>(5 * 60 * 1000)
+  const interval = useRef<number>(55 * 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')
+    localStorage?.removeItem('last_refresh_time')
+    router.replace('/signin')
+  }, [])
+
+  const getNewAccessToken = useCallback(async (currentAccessToken: string, currentRefreshToken: string) => {
+    if (localStorage?.getItem('is_refreshing') === '1')
+      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()) {
+      localStorage?.setItem('is_refreshing', '1')
+      const [e, res] = await fetchWithRetry(fetchNewToken({
+        body: { refresh_token: currentRefreshToken },
+      }) as Promise<CommonResponse & { data: { access_token: string; refresh_token: string } }>)
+      if (e) {
+        handleError()
+        return e
+      }
+      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)
+      }, 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)
+      }, newTokenExpireTime - advanceTime.current - getCurrentTimeStamp())
+    }
+    return null
+  }, [getExpireTime, getCurrentTimeStamp, handleError])
+
+  useEffect(() => {
+    return () => {
+      clearTimeout(timer.current)
+      localStorage?.removeItem('is_refreshing')
+      localStorage?.removeItem('last_refresh_time')
+    }
+  }, [])
+
+  return {
+    getNewAccessToken,
+  }
+}
+
+export default useRefreshToken

web/package.json

@@ -55,6 +55,7 @@
     "immer": "^9.0.19",
     "js-audio-recorder": "^1.0.7",
     "js-cookie": "^3.0.1",
+    "jwt-decode": "^4.0.0",
     "katex": "^0.16.10",
     "lamejs": "^1.2.1",
     "lexical": "^0.16.0",

web/service/common.ts

@@ -38,8 +38,21 @@ import type {
 import type { RETRIEVE_METHOD } from '@/types/app'
 import type { SystemFeatures } from '@/types/feature'
 
-export const login: Fetcher<CommonResponse & { data: string }, { url: string; body: Record<string, any> }> = ({ url, body }) => {
-  return post(url, { body }) as Promise<CommonResponse & { data: string }>
+type LoginSuccess = {
+  result: 'success'
+  data: { access_token: string;refresh_token: string }
+}
+type LoginFail = {
+  result: 'fail'
+  data: string
+}
+type LoginResponse = LoginSuccess | LoginFail
+export const login: Fetcher<LoginResponse, { url: string; body: Record<string, any> }> = ({ url, body }) => {
+  return post(url, { body }) as Promise<LoginResponse>
+}
+
+export const fetchNewToken: Fetcher<CommonResponse & { data: { access_token: string; refresh_token: string } }, { body: Record<string, any> }> = ({ body }) => {
+  return post('/refresh-token', { body }) as Promise<CommonResponse & { data: { access_token: string; refresh_token: string } }>
 }
 
 export const setup: Fetcher<CommonResponse, { body: Record<string, any> }> = ({ body }) => {

web/utils/index.ts

@@ -39,3 +39,21 @@ export const getPurifyHref = (href: string) => {
 
   return escape(href)
 }
+
+export async function fetchWithRetry<T = any>(fn: Promise<T>, retries = 3): Promise<[Error] | [null, T]> {
+  const [error, res] = await asyncRunSafe(fn)
+  if (error) {
+    if (retries > 0) {
+      const res = await fetchWithRetry(fn, retries - 1)
+      return res
+    }
+    else {
+      if (error instanceof Error)
+        return [error]
+      return [new Error('unknown error')]
+    }
+  }
+  else {
+    return [null, res]
+  }
+}

web/yarn.lock

@@ -6205,6 +6205,11 @@ jsonc-eslint-parser@^2.0.4, jsonc-eslint-parser@^2.1.0:
     array-includes "^3.1.5"
     object.assign "^4.1.3"
 
+jwt-decode@^4.0.0:
+  version "4.0.0"
+  resolved "https://registry.npmmirror.com/jwt-decode/-/jwt-decode-4.0.0.tgz#2270352425fd413785b2faf11f6e755c5151bd4b"
+  integrity sha512-+KJGIyHgkGuIq3IEBNftfhW/LfWhXUIY6OmyVWjliu5KH1y0fw7VQ8YndE2O4qZdMSd9SqbnC8GOcZEy0Om7sA==
+
 katex@^0.16.0, katex@^0.16.10:
   version "0.16.10"
   resolved "https://registry.npmjs.org/katex/-/katex-0.16.10.tgz"