feat: implement tabby/auth (#969)

* feat: implement tabby/auth

* update

* [autofix.ci] apply automated fixes

---------

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>

ee/tabby-ui/components/providers.tsx

@@ -5,11 +5,14 @@ import { ThemeProvider as NextThemesProvider } from 'next-themes'
 import { ThemeProviderProps } from 'next-themes/dist/types'
 
 import { TooltipProvider } from '@/components/ui/tooltip'
+import { AuthProvider } from '@/lib/tabby/auth'
 
 export function Providers({ children, ...props }: ThemeProviderProps) {
   return (
     <NextThemesProvider {...props}>
-      <TooltipProvider>{children}</TooltipProvider>
+      <TooltipProvider>
+        <AuthProvider>{children}</AuthProvider>
+      </TooltipProvider>
     </NextThemesProvider>
   )
 }

ee/tabby-ui/lib/gql/generates/gql.ts

@@ -16,7 +16,9 @@ const documents = {
   '\n  query GetRegistrationToken {\n    registrationToken\n  }\n':
     types.GetRegistrationTokenDocument,
   '\n  query GetWorkers {\n    workers {\n      kind\n      name\n      addr\n      device\n      arch\n      cpuInfo\n      cpuCount\n      cudaDevices\n    }\n  }\n':
-    types.GetWorkersDocument
+    types.GetWorkersDocument,
+  '\n  mutation refreshToken($refreshToken: String!) {\n    refreshToken(refreshToken: $refreshToken) {\n      accessToken\n      refreshToken\n    }\n  }\n':
+    types.RefreshTokenDocument
 }
 
 /**
@@ -45,6 +47,12 @@ export function graphql(
 export function graphql(
   source: '\n  query GetWorkers {\n    workers {\n      kind\n      name\n      addr\n      device\n      arch\n      cpuInfo\n      cpuCount\n      cudaDevices\n    }\n  }\n'
 ): (typeof documents)['\n  query GetWorkers {\n    workers {\n      kind\n      name\n      addr\n      device\n      arch\n      cpuInfo\n      cpuCount\n      cudaDevices\n    }\n  }\n']
+/**
+ * The graphql function is used to parse GraphQL queries into a document that can be used by GraphQL clients.
+ */
+export function graphql(
+  source: '\n  mutation refreshToken($refreshToken: String!) {\n    refreshToken(refreshToken: $refreshToken) {\n      accessToken\n      refreshToken\n    }\n  }\n'
+): (typeof documents)['\n  mutation refreshToken($refreshToken: String!) {\n    refreshToken(refreshToken: $refreshToken) {\n      accessToken\n      refreshToken\n    }\n  }\n']
 
 export function graphql(source: string) {
   return (documents as any)[source] ?? {}

ee/tabby-ui/lib/gql/generates/graphql.ts

@@ -163,6 +163,19 @@ export type GetWorkersQuery = {
   }>
 }
 
+export type RefreshTokenMutationVariables = Exact<{
+  refreshToken: Scalars['String']['input']
+}>
+
+export type RefreshTokenMutation = {
+  __typename?: 'Mutation'
+  refreshToken: {
+    __typename?: 'RefreshTokenResponse'
+    accessToken: string
+    refreshToken: string
+  }
+}
+
 export const GetRegistrationTokenDocument = {
   kind: 'Document',
   definitions: [
@@ -214,3 +227,55 @@ export const GetWorkersDocument = {
     }
   ]
 } as unknown as DocumentNode<GetWorkersQuery, GetWorkersQueryVariables>
+export const RefreshTokenDocument = {
+  kind: 'Document',
+  definitions: [
+    {
+      kind: 'OperationDefinition',
+      operation: 'mutation',
+      name: { kind: 'Name', value: 'refreshToken' },
+      variableDefinitions: [
+        {
+          kind: 'VariableDefinition',
+          variable: {
+            kind: 'Variable',
+            name: { kind: 'Name', value: 'refreshToken' }
+          },
+          type: {
+            kind: 'NonNullType',
+            type: { kind: 'NamedType', name: { kind: 'Name', value: 'String' } }
+          }
+        }
+      ],
+      selectionSet: {
+        kind: 'SelectionSet',
+        selections: [
+          {
+            kind: 'Field',
+            name: { kind: 'Name', value: 'refreshToken' },
+            arguments: [
+              {
+                kind: 'Argument',
+                name: { kind: 'Name', value: 'refreshToken' },
+                value: {
+                  kind: 'Variable',
+                  name: { kind: 'Name', value: 'refreshToken' }
+                }
+              }
+            ],
+            selectionSet: {
+              kind: 'SelectionSet',
+              selections: [
+                { kind: 'Field', name: { kind: 'Name', value: 'accessToken' } },
+                { kind: 'Field', name: { kind: 'Name', value: 'refreshToken' } }
+              ]
+            }
+          }
+        ]
+      }
+    }
+  ]
+} as unknown as DocumentNode<
+  RefreshTokenMutation,
+  RefreshTokenMutationVariables
+>

ee/tabby-ui/lib/hooks/use-interval.ts

@@ -0,0 +1,31 @@
+import * as React from 'react'
+
+export default function useInterval(
+  callback: () => void,
+  delay: number | null
+): React.MutableRefObject<number | null> {
+  const savedCallback = React.useRef(callback)
+  const intervalRef = React.useRef<number | null>(null)
+
+  // Remember the latest callback if it changes.
+  React.useEffect(() => {
+    savedCallback.current = callback
+  }, [callback])
+
+  // Set up the interval.
+  React.useEffect(() => {
+    const tick = () => savedCallback.current()
+
+    if (typeof delay === 'number') {
+      intervalRef.current = window.setInterval(tick, delay * 60 * 1000)
+    }
+
+    return () => {
+      if (intervalRef.current) {
+        window.clearTimeout(intervalRef.current)
+      }
+    }
+  }, [delay])
+
+  return intervalRef
+}

ee/tabby-ui/lib/tabby/auth.tsx

@@ -0,0 +1,237 @@
+import * as React from 'react'
+import { graphql } from '@/lib/gql/generates'
+import useInterval from '@/lib/hooks/use-interval'
+import { gqlClient } from '@/lib/tabby-gql-client'
+import { jwtDecode } from 'jwt-decode'
+
+interface AuthData {
+  accessToken: string
+  refreshToken: string
+}
+
+type AuthState =
+  | {
+      status: 'authenticated'
+      data: AuthData
+    }
+  | {
+      status: 'loading' | 'unauthenticated'
+      data: null
+    }
+
+enum AuthActionType {
+  Init,
+  SignIn,
+  SignOut,
+  Refresh
+}
+
+interface InitAction {
+  type: AuthActionType.Init
+  data: AuthData | null
+}
+
+interface SignInAction {
+  type: AuthActionType.SignIn
+  data: AuthData
+}
+
+interface SignOutAction {
+  type: AuthActionType.SignOut
+}
+
+interface RefreshAction {
+  type: AuthActionType.Refresh
+  data: AuthData
+}
+
+type AuthActions = InitAction | SignInAction | SignOutAction | RefreshAction
+
+function authReducer(state: AuthState, action: AuthActions): AuthState {
+  switch (action.type) {
+    case AuthActionType.Init:
+    case AuthActionType.SignIn:
+    case AuthActionType.Refresh:
+      if (action.data) {
+        return {
+          status: 'authenticated',
+          data: action.data
+        }
+      } else {
+        return {
+          status: 'unauthenticated',
+          data: null
+        }
+      }
+    case AuthActionType.SignOut:
+      TokenStorage.reset()
+      return {
+        status: 'unauthenticated',
+        data: null
+      }
+  }
+}
+
+class TokenStorage {
+  static authName = '_tabby_auth'
+
+  initialState(): AuthData | null {
+    const authData = localStorage.getItem(TokenStorage.authName)
+    if (authData) {
+      return JSON.parse(authData)
+    } else {
+      return null
+    }
+  }
+
+  persist(state: AuthData) {
+    localStorage.setItem(TokenStorage.authName, JSON.stringify(state))
+  }
+
+  static reset() {
+    localStorage.removeItem(TokenStorage.authName)
+  }
+}
+
+interface AuthProviderProps {
+  children: React.ReactNode
+}
+
+interface AuthStore {
+  authState: AuthState | null
+  dispatch: React.Dispatch<AuthActions>
+}
+
+const AuthContext = React.createContext<AuthStore | null>(null)
+
+const refreshTokenMutation = graphql(/* GraphQL */ `
+  mutation refreshToken($refreshToken: String!) {
+    refreshToken(refreshToken: $refreshToken) {
+      accessToken
+      refreshToken
+    }
+  }
+`)
+
+async function doRefresh(token: string, dispatch: React.Dispatch<AuthActions>) {
+  let action: AuthActions
+  try {
+    action = {
+      type: AuthActionType.Refresh,
+      data: (
+        await gqlClient.request(refreshTokenMutation, { refreshToken: token })
+      ).refreshToken
+    }
+  } catch (err) {
+    console.error('Failed to refresh token', err)
+    action = {
+      type: AuthActionType.SignOut
+    }
+  }
+
+  dispatch(action)
+}
+
+const AuthProvider: React.FunctionComponent<AuthProviderProps> = ({
+  children
+}) => {
+  const storage = new TokenStorage()
+
+  const [authState, dispatch] = React.useReducer(authReducer, {
+    status: 'loading',
+    data: null
+  })
+
+  React.useEffect(() => {
+    const data = storage.initialState()
+    if (data?.refreshToken) {
+      doRefresh(data.refreshToken, dispatch)
+    } else {
+      dispatch({ type: AuthActionType.Init, data: null })
+    }
+  }, [])
+
+  React.useEffect(() => {
+    authState.data && storage.persist(authState.data)
+  }, [authState])
+
+  useInterval(async () => {
+    if (authState.status !== 'authenticated') {
+      return
+    }
+
+    await doRefresh(authState.data.refreshToken, dispatch)
+  }, 5)
+
+  return (
+    <AuthContext.Provider value={{ authState, dispatch }}>
+      {children}
+    </AuthContext.Provider>
+  )
+}
+
+function useAuthStore(): AuthStore {
+  const context = React.useContext(AuthContext)
+
+  if (!context) {
+    throw new Error(
+      'AuthProvider is missing. Please add the AuthProvider at root level'
+    )
+  }
+
+  return context
+}
+
+function useSignIn(): (params: AuthData) => Promise<boolean> {
+  const { dispatch } = useAuthStore()
+  return async data => {
+    dispatch({
+      type: AuthActionType.SignIn,
+      data
+    })
+
+    return true
+  }
+}
+
+function useSignOut(): () => Promise<void> {
+  const { dispatch } = useAuthStore()
+  return async () => {
+    dispatch({ type: AuthActionType.SignOut })
+  }
+}
+
+interface User {
+  email: string
+  isAdmin: boolean
+}
+
+type Session =
+  | {
+      data: null
+      status: 'loading' | 'unauthenticated'
+    }
+  | {
+      data: User
+      status: 'authenticated'
+    }
+
+function useSession(): Session {
+  const { authState } = useAuthStore()
+  if (authState?.status == 'authenticated') {
+    const { user } = jwtDecode<{ user: User }>(authState.data.accessToken)
+    return {
+      data: user,
+      status: authState.status
+    }
+  } else {
+    return {
+      status: authState?.status ?? 'loading',
+      data: null
+    }
+  }
+}
+
+export type { AuthStore, User, Session }
+
+export { AuthProvider, useSignIn, useSignOut, useSession }

ee/tabby-ui/package.json

@@ -40,6 +40,7 @@
     "focus-trap-react": "^10.1.1",
     "graphql": "^16.8.1",
     "graphql-request": "^6.1.0",
+    "jwt-decode": "^4.0.0",
     "lodash-es": "^4.17.21",
     "nanoid": "^4.0.2",
     "next": "^13.4.7",

ee/tabby-ui/yarn.lock

@@ -4422,6 +4422,11 @@ jsonify@^0.0.1:
     object.assign "^4.1.4"
     object.values "^1.1.6"
 
+jwt-decode@^4.0.0:
+  version "4.0.0"
+  resolved "https://registry.yarnpkg.com/jwt-decode/-/jwt-decode-4.0.0.tgz#2270352425fd413785b2faf11f6e755c5151bd4b"
+  integrity sha512-+KJGIyHgkGuIq3IEBNftfhW/LfWhXUIY6OmyVWjliu5KH1y0fw7VQ8YndE2O4qZdMSd9SqbnC8GOcZEy0Om7sA==
+
 katex@^0.16.0:
   version "0.16.8"
   resolved "https://registry.yarnpkg.com/katex/-/katex-0.16.8.tgz#89b453f40e8557f423f31a1009e9298dd99d5ceb"