refactor(ui): use web locks api to manage refresh requests (#1905)

* refactor(ui): use web locks api to manage refresh requests refreshPromise update * update
TabbyML · May 6, 2024 · 6eadc5c · 6eadc5c
1 parent d0ddd33
commit 6eadc5c
Show file tree

Hide file tree

Showing 3 changed files with 72 additions and 107 deletions.
diff --git a/ee/tabby-ui/lib/tabby/fetcher.ts b/ee/tabby-ui/lib/tabby/fetcher.ts
@@ -1,13 +1,8 @@
-import { createRequest } from '@urql/core'
+import { Client, createRequest, fetchExchange } from '@urql/core'
 import { jwtDecode } from 'jwt-decode'
 
 import { refreshTokenMutation } from './auth'
-import { client } from './gql'
-import {
-  getAuthToken,
-  isTokenExpired,
-  tokenManagerInstance
-} from './token-management'
+import { getAuthToken, isTokenExpired, tokenManager } from './token-management'
 
 interface FetcherOptions extends RequestInit {
   responseFormat?: 'json' | 'blob'
@@ -26,7 +21,7 @@ export default async function authEnhancedFetch(
   const currentFetcher = options?.customFetch ?? window.fetch
 
   if (willAuthError(url)) {
-    return tokenManagerInstance.refreshToken(doRefreshToken).then(res => {
+    return tokenManager.refreshToken(doRefreshToken).then(res => {
       return requestWithAuth(url, options)
     })
   }
@@ -37,7 +32,7 @@ export default async function authEnhancedFetch(
   )
 
   if (response.status === 401) {
-    return tokenManagerInstance.refreshToken(doRefreshToken).then(res => {
+    return tokenManager.refreshToken(doRefreshToken).then(res => {
       return requestWithAuth(url, options)
     })
   } else {
@@ -54,7 +49,7 @@ function willAuthError(url: string) {
     // Check whether `token` JWT is expired
     try {
       const { exp } = jwtDecode(accessToken)
-      return exp ? isTokenExpired(exp) : true
+      return isTokenExpired(exp)
     } catch (e) {
       return true
     }
@@ -85,6 +80,11 @@ function addAuthToRequest(options?: FetcherOptions): FetcherOptions {
 }
 
 async function refreshAuth(refreshToken: string) {
+  const client = new Client({
+    url: `/graphql`,
+    requestPolicy: 'network-only',
+    exchanges: [fetchExchange]
+  })
   const refreshAuth = client.createRequestOperation(
     'mutation',
     createRequest(refreshTokenMutation, { refreshToken })

diff --git a/ee/tabby-ui/lib/tabby/gql.ts b/ee/tabby-ui/lib/tabby/gql.ts
@@ -28,12 +28,7 @@ import {
   listInvitations,
   listRepositories
 } from './query'
-import {
-  clearAuthToken,
-  getAuthToken,
-  isTokenExpired,
-  tokenManagerInstance
-} from './token-management'
+import { getAuthToken, isTokenExpired, tokenManager } from './token-management'
 
 interface ValidationError {
   path: string
@@ -265,13 +260,26 @@ const client = new Client({
           accessToken = authData?.accessToken
           refreshToken = authData?.refreshToken
 
+          if (
+            operation.kind === 'query' &&
+            operation.query.definitions.some(definition => {
+              return (
+                definition.kind === 'OperationDefinition' &&
+                definition.name?.value &&
+                ['GetServerInfo'].includes(definition.name.value)
+              )
+            })
+          ) {
+            return false
+          }
+
           if (
             operation.kind === 'mutation' &&
             operation.query.definitions.some(definition => {
               return (
                 definition.kind === 'OperationDefinition' &&
                 definition.name?.value &&
-                ['tokenAuth', 'registerUser'].includes(definition.name.value)
+                ['tokenAuth', 'register'].includes(definition.name.value)
               )
             })
           ) {
@@ -292,10 +300,10 @@ const client = new Client({
           }
 
           if (accessToken) {
-            // Check whether `token` JWT is expired
             try {
               const { exp } = jwtDecode(accessToken)
-              return exp ? isTokenExpired(exp) : true
+              // Check whether `token` JWT is expired
+              return isTokenExpired(exp)
             } catch (e) {
               return true
             }
@@ -304,18 +312,16 @@ const client = new Client({
           }
         },
         async refreshAuth() {
-          if (refreshToken) {
-            return tokenManagerInstance.refreshToken(() =>
-              utils
-                .mutate(refreshTokenMutation, {
-                  refreshToken: refreshToken as string
-                })
-                .then(res => res?.data?.refreshToken)
-            )
-          } else {
-            // This is where auth has gone wrong and we need to clean up and redirect to a login page
-            clearAuthToken()
-          }
+          return tokenManager.refreshToken(async () => {
+            const refreshToken = getAuthToken()?.refreshToken
+            if (!refreshToken) return undefined
+
+            return utils
+              .mutate(refreshTokenMutation, {
+                refreshToken
+              })
+              .then(res => res?.data?.refreshToken)
+          })
         }
       }
     }),

diff --git a/ee/tabby-ui/lib/tabby/token-management.ts b/ee/tabby-ui/lib/tabby/token-management.ts
@@ -1,9 +1,11 @@
+import { jwtDecode } from 'jwt-decode'
+import { isNil } from 'lodash-es'
+
 import { isClientSide } from '../utils'
 import { AuthData } from './auth'
 
 export const AUTH_TOKEN_KEY = '_tabby_auth'
 export const AUTH_LOCK_KEY = '_tabby_auth_lock'
-export const AUTH_LOCK_EXP = 1000 * 10
 
 const getAuthToken = (): AuthData | undefined => {
   if (isClientSide()) {
@@ -36,99 +38,56 @@ const clearAuthToken = () => {
   )
 }
 
-const isTokenExpired = (exp: number) => {
-  return Date.now() > exp * 1000
+const isTokenExpired = (exp: number | undefined): boolean => {
+  return isNil(exp) ? true : Date.now() > exp * 1000
 }
 
-class TokenManager {
-  private retryQueue: Array<(success: boolean, error?: Error) => void>
-
-  constructor() {
-    this.retryQueue = []
-
-    if (typeof window !== 'undefined') {
-      window.addEventListener('storage', this.handleStorageChange)
-    }
-  }
+// Checks if the JWT token's issued-at time (iat) is within the last minute.
+const isTokenRecentlyIssued = (iat: number | undefined): boolean => {
+  return isNil(iat) ? false : Date.now() - iat * 1000 < 60 * 1000
+}
 
-  private handleStorageChange = (event: StorageEvent) => {
+class TokenManager {
+  async refreshToken(doRefreshToken: () => Promise<AuthData | undefined>) {
     try {
-      if (
-        event.key === AUTH_LOCK_KEY &&
-        event.newValue === null &&
-        this.retryQueue?.length
-      ) {
-        this.processQueue()
+      if (typeof navigator?.locks === 'undefined') {
+        console.error(
+          'The Web Locks API is not supported in your browser. Please upgrade to a newer browser version.'
+        )
+        throw new Error()
       }
-    } catch (e) {}
-  }
-
-  tryGetRefreshLock() {
-    const currentLock = localStorage.getItem(AUTH_LOCK_KEY)
-    const lockTimestamp = currentLock ? parseInt(currentLock, 10) : null
-    const now = Date.now()
-    if (
-      !currentLock ||
-      (lockTimestamp && now - lockTimestamp > AUTH_LOCK_EXP)
-    ) {
-      localStorage.setItem(AUTH_LOCK_KEY, now.toString())
-      return true
-    }
-    return false
-  }
-
-  releaseRefreshLock() {
-    localStorage.removeItem(AUTH_LOCK_KEY)
-  }
 
-  enqueueRetryRequest(
-    retryCallback: (success: boolean, error?: Error) => void
-  ) {
-    this.retryQueue.push(retryCallback)
-  }
-
-  processQueue() {
-    this.retryQueue.forEach(retryCallback => retryCallback(true))
-    this.retryQueue = []
-    this.releaseRefreshLock()
-  }
+      await navigator.locks.request(AUTH_LOCK_KEY, async () => {
+        const authToken = getAuthToken()
+        const accessToken = getAuthToken()?.accessToken
 
-  rejectQueue(error?: Error) {
-    this.retryQueue.forEach(retryCallback => retryCallback(false, error))
-    this.retryQueue = []
-    this.releaseRefreshLock()
-  }
+        let newAuthToken: AuthData | undefined
 
-  async refreshToken(doRefreshToken: () => Promise<AuthData | undefined>) {
-    if (!this.tryGetRefreshLock()) {
-      // refreshing
-      return new Promise<void>((resolve, reject) => {
-        this.enqueueRetryRequest((success: boolean, error?: Error) => {
-          if (!success || error) {
-            reject(error ?? 'Failed to refresh token')
+        if (accessToken) {
+          const { iat } = jwtDecode(accessToken)
+          if (isTokenRecentlyIssued(iat)) {
+            newAuthToken = authToken
           } else {
-            resolve()
+            newAuthToken = await doRefreshToken()
           }
-        })
-      })
-    }
+        }
 
-    const newToken = await doRefreshToken()
-    if (newToken) {
-      await saveAuthToken(newToken)
-      this.processQueue()
-    } else {
-      this.rejectQueue()
+        if (newAuthToken) {
+          saveAuthToken(newAuthToken)
+        } else {
+          clearAuthToken()
+        }
+      })
+    } catch (e) {
       clearAuthToken()
-      throw new Error('Failed to refresh token')
     }
   }
 }
 
-const tokenManagerInstance = new TokenManager()
+const tokenManager = new TokenManager()
 
 export {
-  tokenManagerInstance,
+  tokenManager,
   getAuthToken,
   saveAuthToken,
   clearAuthToken,