React Apollo + Firebase authorization

App.tsx

import Dashboard from 'pages/Dashboard'
import Layout from 'components/Layout'
import Login from 'pages/Login'
import React from 'react'
import { useRoutes, Navigate, useLocation } from 'react-router-dom'
import { useAuth, userContext } from 'hooks/useAuth'

function App() {
  const storedToken = localStorage.getItem('token')
  const location = useLocation()
  const { user, initializing } = useAuth()

  return useRoutes([
    { path: '/login', element: <Login /> },
    {
      path: '/',
      element: storedToken ? (
        <userContext.Provider value={{ user, initializing }}>
          <Layout />
        </userContext.Provider>
      ) : (
        <Navigate
          to="/login"
          replace={true}
          state={{ from: location.pathname }}
        />
      ),
      children: [
        { path: 'dashboard', element: <Dashboard /> }
        { path: '*', element: <Navigate to="/dashboard" /> },
      ],
    },
  ])
}

export default App

Layout.tsx

import { ApolloProvider } from '@apollo/client'
import { userContext } from 'hooks/useAuth'
import { useClient } from 'hooks/useClient'
import React, { useContext } from 'react'
import { Outlet } from 'react-router'

interface Props {}

const Layout: React.FC<Props> = () => {
  const client = useClient()
  const { user, initializing } = useContext(userContext)

  if (initializing) {
    return <span>Loading</span>
  }
  
  return (
    <ApolloProvider client={client}>
      <Outlet />
    </ApolloProvider>
  )
}

export default Layout

useAuth.ts

import * as firebase from 'firebase/app'
import { useEffect, useState, createContext } from 'react'

export type AuthState = {
  initializing: boolean
  user: firebase.User | null
}

export const userContext = createContext<AuthState>({
  initializing: true,
  user: null,
})

export const useAuth = () => {
  const [state, setState] = useState<AuthState>({
    initializing: true,
    user: null,
  })

  useEffect(() => {
    const unsubscribe = firebase.auth().onIdTokenChanged(async (user) => {
      if (user) {
        const token = await user.getIdToken()
        localStorage.setItem('token', token)
      }
    })
    return () => unsubscribe()
  }, [])

  useEffect(() => {
    const unsubscribe = firebase.auth().onAuthStateChanged(async (user) => {
      if (user) {
        const token = await user.getIdToken()
        localStorage.setItem('token', token)
      }

      setState({ initializing: false, user })
    })

    return () => unsubscribe()
  }, [])

  return state
}

useClient.ts

import {
  ApolloClient,
  ApolloLink,
  HttpLink,
  InMemoryCache,
} from '@apollo/client'
import { getMainDefinition } from '@apollo/client/utilities'
import { setContext } from '@apollo/link-context'
import { onError } from '@apollo/link-error'
import { RetryLink } from '@apollo/link-retry'
import { WebSocketLink } from '@apollo/link-ws'
import { useNavigate } from 'react-router'

export const useAuthLink = () => {
  return setContext((_, { headers }) => {
    const token = localStorage.getItem('token')
    return { headers: { Authorization: `Bearer ${token}` } }
  })
}

export const useErrorLink = () => {
  const navigate = useNavigate()

  return onError(({ graphQLErrors, networkError, operation, forward }) => {
    if (graphQLErrors) {
      for (let err of graphQLErrors) {
        switch (err.extensions?.code) {
          case 'access-denied':
            navigate('/login')
            break
          case 'invalid-jwt':
            localStorage.removeItem('token')
            navigate('/login')
            break
        }
      }
    }
  })
}

export const useRetryLink = () => {
  return new RetryLink({
    delay: {
      initial: 500,
      max: Infinity,
      jitter: true,
    },
    attempts: {
      max: 5,
      retryIf: error => !!error,
    },
  })
}

export const useLink = () => {
  const token = localStorage.getItem('token')

  const httpLink = new HttpLink({
    uri: process.env.REACT_APP_HASURA_API_URL || '/graphql',
    headers: { Authorization: `Bearer ${token}` },
  })

  const wsLink = new WebSocketLink({
    uri: process.env.REACT_APP_HASURA_WSS_URL || '/graphql',
    options: {
      reconnect: true,
      lazy: true,
      connectionParams: () => {
        const token = localStorage.getItem('token')
        return { headers: { Authorization: `Bearer ${token}` } }
      },
      connectionCallback: console.log,
    },
  })

  return ApolloLink.split(
    ({ query }) => {
      const definition = getMainDefinition(query)
      return (
        definition.kind === 'OperationDefinition' &&
        definition.operation === 'subscription'
      )
    },
    wsLink,
    httpLink
  )
}

export const useClient = () => {
  const authLink = useAuthLink()
  const errorLink = useErrorLink()
  const retryLink = useRetryLink()
  const link = useLink()
  return new ApolloClient({
    cache: new InMemoryCache(),
    link: ApolloLink.from([authLink, retryLink, errorLink, link]),
    connectToDevTools: true,
  })
}