diff --git a/ee/tabby-ui/components/providers.tsx b/ee/tabby-ui/components/providers.tsx index 8922264..f26ea99 100644 --- a/ee/tabby-ui/components/providers.tsx +++ b/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 ( - {children} + + {children} + ) } diff --git a/ee/tabby-ui/lib/gql/generates/gql.ts b/ee/tabby-ui/lib/gql/generates/gql.ts index 4c6b1d5..2267b05 100644 --- a/ee/tabby-ui/lib/gql/generates/gql.ts +++ b/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] ?? {} diff --git a/ee/tabby-ui/lib/gql/generates/graphql.ts b/ee/tabby-ui/lib/gql/generates/graphql.ts index c1233b4..6293c3b 100644 --- a/ee/tabby-ui/lib/gql/generates/graphql.ts +++ b/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 +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 +> diff --git a/ee/tabby-ui/lib/hooks/use-interval.ts b/ee/tabby-ui/lib/hooks/use-interval.ts new file mode 100644 index 0000000..08f284b --- /dev/null +++ b/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 { + const savedCallback = React.useRef(callback) + const intervalRef = React.useRef(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 +} diff --git a/ee/tabby-ui/lib/tabby/auth.tsx b/ee/tabby-ui/lib/tabby/auth.tsx new file mode 100644 index 0000000..16d4321 --- /dev/null +++ b/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 +} + +const AuthContext = React.createContext(null) + +const refreshTokenMutation = graphql(/* GraphQL */ ` + mutation refreshToken($refreshToken: String!) { + refreshToken(refreshToken: $refreshToken) { + accessToken + refreshToken + } + } +`) + +async function doRefresh(token: string, dispatch: React.Dispatch) { + 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 = ({ + 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 ( + + {children} + + ) +} + +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 { + const { dispatch } = useAuthStore() + return async data => { + dispatch({ + type: AuthActionType.SignIn, + data + }) + + return true + } +} + +function useSignOut(): () => Promise { + 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 } diff --git a/ee/tabby-ui/package.json b/ee/tabby-ui/package.json index eacbe8a..50f702f 100644 --- a/ee/tabby-ui/package.json +++ b/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", diff --git a/ee/tabby-ui/yarn.lock b/ee/tabby-ui/yarn.lock index 7f4ca9b..7445eb0 100644 --- a/ee/tabby-ui/yarn.lock +++ b/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"