Refreshing Apollo subscription websocket connections with updated auth tokens

.useApolloClient.md

useApolloClient

The primary idea of this hook is to re-create the underlying ApolloClient being used by all of the Subscription and Query components. There is lots of talk throughout the Apollo community about how to best refresh a WebSocket connection when a new auth token is available, and this is my attempt at providing something useful, apart from middleware that seems to have issues.

Related Links

• apollographql/subscriptions-transport-ws#171
• apollographql/apollo-link#197
• apollographql/subscriptions-transport-ws#102
• https://www.prisma.io/forum/t/authentication-issues-with-subscriptions/45?u=nilan

Philosophy

Architecturally, this solution keeps data flowing in the same direction, by passing a new apolloClient down to the components that consume it whenever the auth token changes (after calling .stop() on the current client). When downstream components receive this new prop, they reinstate their subscription connections as a part of their re-render, this time with the new token.

What's Happening in useApolloClient

useApolloClient is essentially a wrapper around useState, which is holding a client instance. When the dispatch function is called with the new token (setApolloClient(token)), a new client is created and the old client is replaced in State. The data then flows down through the Provider's components, triggering re-renders of all consumers.

In Auth0Provider, the idToken (use any token you like in your own implementation) is also tracked in State, and on each re-render of the Auth0Provider, the current idToken is compared with the token currently in use in the current apolloClient, by reading from the cache (see Apollo Docs on managing local state):

  const { idToken: cachedIdToken } = apolloClient.readQuery({ query: TOKEN_QUERY })

  if (idToken !== cachedIdToken) {
    apolloClient.stop()
    setApolloClient(idToken)
  }

In the event they do not match (because a new token is available, or the token has been cleared during a logout operation), the old client is stopped, cache is cleared (if that is a concern for security or whatever), and a new client is created.

Performance Note

I haven't yet profiled the memory or network overhead of this approach, but because this should be happening very rarely (once every couple hours at most, in the case of a refresh; and ideally when someone logs out/logs in, all of this should be rebuilt anyway, right?), I'm not overly concerned about performance in this regard. Unless the token changes, none of this re-rendering takes place.

App.js

import React from 'react'
import { BrowserRouter as Router, Switch, Route, NavLink } from 'react-router-dom'
import { Home, Profile } from './screens'
import { ApolloProvider } from 'react-apollo'
import { useAuth0 } from './Auth0Provider'
import PrivateRoute from './PrivateRoute'

export default props => {
  const { isAuthenticated, loginWithPopup, logout, apolloClient } = useAuth0()

  return <ApolloProvider client={apolloClient}>
    <Router>
      <nav className='flex bg-gray-800 text-gray-500 items-center'> {/* tailwind utility classes... 😎 */}
        <NavLink activeClassName='text-white' className='p-4 hover:bg-gray-700' exact to='/'>Home</NavLink>
        {isAuthenticated && <>
          <NavLink activeClassName='text-white' className='p-4 hover:bg-gray-700' to='/profile'>Profile</NavLink>
          <button className='p-4 hover:bg-gray-700' onClick={() => logout()}>Logout</button>
        </>}
        {!isAuthenticated && <button className='p-4 hover:bg-gray-700' onClick={() => loginWithPopup({})}>Login</button>}
      </nav>
      <Switch>
        <PrivateRoute path='/profile' component={Profile} />
        <Route exact path='/' component={Home} />
      </Switch>
    </Router>
  </ApolloProvider>
}

Auth0Provider.js

import React, { useState, useEffect, useContext } from 'react'
import createAuth0Client from '@auth0/auth0-spa-js'
import { useApolloClient } from '../hooks'
import gql from 'graphql-tag'

const TOKEN_QUERY = gql`
  query getIdToken {
    idToken @client
  }
`

export const Auth0Context = React.createContext()
export const useAuth0 = () => useContext(Auth0Context)

export const Auth0Provider = ({
  children,
  ...initOptions
}) => {
  const [isAuthenticated, setIsAuthenticated] = useState()
  const [user, setUser] = useState()
  const [idToken, setIdToken] = useState('') // either empty string or null
  const [authAppState, setAuthAppState] = useState() // in redirect mode, this is useful
  const [auth0Client, setAuth0] = useState() // "hydrated" auth0 class (aware of your auth0 client details)
  const [loading, setLoading] = useState(true) // if auth0 is still working, let the app know
  const [popupOpen, setPopupOpen] = useState(false) // useful for knowing when to display/hide stuff
  const [apolloClient, setApolloClient] = useApolloClient()

  const { idToken: cachedIdToken } = apolloClient.readQuery({ query: TOKEN_QUERY })

  if (idToken !== cachedIdToken) {
    apolloClient.stop()
    setApolloClient(idToken)
  }

  useEffect(() => {
    const initAuth0 = async () => {
      const auth0FromHook = await createAuth0Client(initOptions)
      setAuth0(auth0FromHook)

      if (window.location.search.includes('code=')) {
        const { appState } = await auth0FromHook.handleRedirectCallback()

        setAuthAppState(appState)

        await auth0FromHook.getTokenSilently()
      }

      const isAuthenticated = await auth0FromHook.isAuthenticated()

      setIsAuthenticated(isAuthenticated)

      if (isAuthenticated) {
        const [user, fetchedToken] = await Promise.all([auth0FromHook.getUser(), auth0FromHook.getIdToken()])
        // getIdToken method is not yet publicly available: https://github.com/auth0/auth0-spa-js/pull/54
        // use getTokenSilently() or another method to obtain access token, and configure an Auth0 server
        // so it returns in JWT format (if you need that format).

        setUser(user)
        setIdToken(fetchedToken)
        window.localStorage.setItem('idToken', fetchedToken)
      }

      setLoading(false)
    }

    initAuth0()
    // eslint-disable-next-line
  }, []);

  const loginWithPopup = async (params = {}) => {
    if (!auth0Client) return

    setPopupOpen(true)

    try {
      await auth0Client.loginWithPopup(params)
    } catch (error) {
      console.error('popup error', error)
    } finally {
      setPopupOpen(false)
    }

    const [user, token] = await Promise.all([auth0Client.getUser(), auth0Client.getIdToken()])

    setIdToken(token)
    window.localStorage.setItem('idToken', token)
    setUser(user)
    setIsAuthenticated(true)
  }

  const handleRedirectCallback = async () => {
    setLoading(true)

    const { appState } = await auth0Client.handleRedirectCallback()
    const [user, token] = await Promise.all([auth0Client.getUser(), auth0Client.getIdToken()])

    setLoading(false)
    setIdToken(token)
    window.localStorage.setItem('idToken', token)
    setAuthAppState(appState)
    setIsAuthenticated(true)
    setUser(user)
  }

  return (
    <Auth0Context.Provider
      value={{
        isAuthenticated,
        user,
        idToken,
        apolloClient,
        loading,
        popupOpen,
        loginWithPopup,
        handleRedirectCallback,
        authAppState,
        setAuthAppState,
        getIdTokenClaims: async (...p) => auth0Client && auth0Client.getIdTokenClaims(...p),
        loginWithRedirect: (...p) => auth0Client.loginWithRedirect(...p),
        getTokenSilently: async (...p) => auth0Client && auth0Client.getTokenSilently(...p),
        getTokenWithPopup: (...p) => auth0Client.getTokenWithPopup(...p),
        logout: (...p) => {
          window.localStorage.removeItem('idToken')
          auth0Client.logout(...p)
        }
      }}
    >
      {children}
    </Auth0Context.Provider>
  )
}

index.js

// bootstrap everything and wrap it all in Auth0Provider

import React from 'react'
import ReactDOM from 'react-dom'
import App from './src/App'
import { Auth0Provider } from './src/auth/Auth'

ReactDOM.render(
  <Auth0Provider
    domain='{your-domain}.auth0.com'
    client_id='1234'
    redirect_uri={window.location.origin}
  >
    <App />
  </Auth0Provider>,
  document.getElementById('app')
)

PrivateRoute.js

import React, { useEffect, useState } from 'react'
import { Route } from 'react-router-dom'
import { useAuth0 } from './Auth'
import { PleaseLogin } from './PleaseLogin'

const PrivateRoute = ({ component: Component, path, ...rest }) => {
  const [authenticating, setAuthenticating] = useState(false) // use this to prevent nasty loops
  const { isAuthenticated, loginWithPopup, loading } = useAuth0()

  useEffect(() => {
    const fn = async () => {
      if (!isAuthenticated && !loading && !authenticating) {
        setAuthenticating(true)
        await loginWithPopup({})

        // now that the process is completely done, we can allow
        // the function to run again, if for some reason it needs to
        setAuthenticating(false)
      }
    }
    fn()
  }, [isAuthenticated, loginWithPopup, path])

  // render function provides dummy component (or a redirect if you choose) when not authenticated
  const render = isAuthenticated ? props => <Component {...props} /> : props => <PleaseLogin {...props} />

  return <Route path={path} render={render} {...rest} />
}

export default PrivateRoute

useApolloClient.js

import { useState } from 'react'
import ApolloClient from 'apollo-client'
import { InMemoryCache } from 'apollo-cache-inmemory'
import { split } from 'apollo-link'
import { HttpLink } from 'apollo-link-http'
import { WebSocketLink } from 'apollo-link-ws'
import { getMainDefinition } from 'apollo-utilities'

// I opted to use global cache, which can be reset after a new token is
// received, but probably preferable to just grab a new cache for each client.
const cache = new InMemoryCache()

// get any token that might already exist when the app starts up
// returns null if nothing found, which is good, because we can't
// put `undefined` in the cache.
const initialToken = window.localStorage.getItem('idToken')

// function to dynamically generate a new client
// when a new token is available
const createClient = token => {
  // Create an http link:
  const httpLink = new HttpLink({
    uri: 'https://the-mind.herokuapp.com/v1/graphql'
  })

  // Create a WebSocket link:
  const wsLink = new WebSocketLink({
    uri: `wss://the-mind.herokuapp.com/v1/graphql`,
    options: {
      reconnect: true,
      connectionParams: () => ({
        headers: {
          Authorization: `Bearer ${token}`
        }
      }),
      lazy: true
    }
  })

  // split the links based on type of request
  const link = split(
    ({ query }) => {
      const definition = getMainDefinition(query)

      return (
        definition.kind === 'OperationDefinition' &&
        definition.operation === 'subscription'
      )
    },
    wsLink,
    httpLink
  )

  return new ApolloClient({
    link,
    cache, // pulling the same cache in each time a new client
           // is created. Or use a new one each time, if you choose
    resolvers: {}
  })
}

// create an initial client with any existing tokens (or null)
const firstClient = createClient(initialToken)

// we will be querying for the token later, so let's
// make sure there's something there up front
firstClient.writeData({
  data: {
    idToken: initialToken
  }
})

// custom hook
export const useApolloClient = () => {
  // store the initial client in State
  const [client, setClient] = useState(firstClient)

  // return the current client from State and a function
  // for dynamically updating State with a new client
  return [client, newToken => {
    const newClient = createClient(newToken)

    // having created the new client, update the cache
    // with the new idToken value.
    newClient.writeData({
      data: {
        idToken: newToken
      }
    })

    // lock the new client into State for use throughout the app!
    return setClient(newClient)
  }]
}

We've moved to using server sent events rather than a websocket. This means we can have a simple link which will restart the operation, rather than having to tear down the whole websocket.

This example uses Okta, but anything with an event emitter could follow a similar pattern.

class RetryableOperation {
  private subscription: ObservableSubscription | null = null;

  constructor(
    private observer: SubscriptionObserver<FetchResult>,
    private operation: Operation,
    private forward: NextLink,
  ) {
    this.try();

    oktaAuth.tokenManager.on("renewed", this.tokenHandler);
  }

  private tokenHandler: TokenManagerRenewEventHandler = (key) => {
    if (key === "accessToken") {
      this.cancel();
      this.try();
    }
  };

  private cancel = () => {
    if (this.subscription) {
      this.subscription.unsubscribe();
    }
    this.subscription = null;
  };

  public close = () => {
    oktaAuth.tokenManager.off("renewed", this.tokenHandler);
    this.cancel();
  };

  private try = () => {
    this.subscription = this.forward(this.operation).subscribe({
      next: this.observer.next.bind(this.observer),
      error: this.observer.error.bind(this.observer),
      complete: this.observer.complete.bind(this.operation),
    });
  };
}

export class SubscriptionRetryLink extends ApolloLink {
  public request(operation: Operation, nextLink: NextLink): Observable<FetchResult> {
    const definition = getMainDefinition(operation.query);

    // for non-subscription operations, move onto the next link
    const subscriptionOperation =
      definition.kind === Kind.OPERATION_DEFINITION && definition.operation === "subscription";

    if (!subscriptionOperation) return nextLink(operation);

    return new Observable((observer) => {
      const retryable = new RetryableOperation(observer, operation, nextLink);

      return () => {
        retryable.close();
      };
    });
  }
}

I wonder if you could achieve something similar with Websocket by closing the connection, then all open operations could retry.

Are you saying that you use SSE with Apollo Client's Subscription API?

Are you saying that you use SSE with Apollo Client's Subscription API?

Yeah, we use graphql-sse to achieve this.

You can also use multi-part HTTP responses which is what comes out of the box with the Apollo HTTP link, but I've never tried it, not sure what it would look like on the server.