This past week, I was tackling a problem where I need to authenticate a React client with a GraphQL server using Amazon Cognito and had to handle auto refreshing the expired tokens.

Problem

What can we do when an API token has expired?

  1. Just show it as an error to the user, which might be a bad user experience (maybe ask them to re-log in)
  2. Check whether our token is valid before every request, which might not be a scalable option if we have many requests.
  3. ”Assume” that the token is valid, and if it turns out otherwise, we will handle refreshing the token and retrying the request.

In this article, I will go with the third option.

TL; DR

This what my final Apollo Link stack looks like

import { ApolloClient } from 'apollo-client';
import { InMemoryCache } from 'apollo-cache-inmemory';
import { HttpLink } from 'apollo-link-http';
import { onError } from 'apollo-link-error';
import { ApolloLink, Observable } from 'apollo-link';
import { setContext } from "apollo-link-context";
import { refreshSession, getAccessToken } from '../lib/auth';
 
// cached in memory access token
let inMemoryToken;
 
const unauthenticatedErrorLink = onError(
  ({ networkError, operation, forward }
  ) => {
  if ( networkError && 
    networkError.name ==='ServerError' && 
    networkError.statusCode === 401) {
    // unauthenticated. token expired. Refetch token using refresh token
    console.warn("Refreshing token and trying again");
    // get new token and set inMemoryToken
    return promiseToObservable(
      new Promise((resolve, reject) => refreshSession(
        (err, session) => {
          if (err) {
            console.error("Failed refreshing session");
            reject(err);
          }
          // set inMemoryToken to refreshed token
          inMemoryToken = session.accessToken.jwtToken;
          resolve();
        }
      ))
    ).flatMap(() => forward(operation));
  }
});
 
const setAuthorizationLink = setContext((request, previousContext) => {
  // if you have a cached value return immediately
  if (inMemoryToken) return _buildAuthHeader(inMemoryToken);
 
  return getAccessToken().then(userToken => {
    inMemoryToken = userToken;
    return _buildAuthHeader(inMemoryToken)
  });
});
 
const apolloClient = new ApolloClient({
  link: ApolloLink.from([
    // handle auth error + auth header middleware
    ApolloLink.from([unauthenticatedErrorLink, setAuthorizationLink ]),
    // http request middleware
    new HttpLink({ uri: process.env.REACT_APP_YOUR_GRAPHQL_API_ENDPOINT }),
  ]),
  cache: new InMemoryCache(),
});
 
export default apolloClient;
 
// private functions
 
const _buildAuthHeader = (token) => {
  return {
    headers: {
      Authorization: `Bearer ${token}`
    }
  }
}
 
const promiseToObservable = promise =>
  new Observable((subscriber) => {
    promise.then(
      (value) => {
        if (subscriber.closed) return;
        subscriber.next(value);
        subscriber.complete();
      },
      err => subscriber.error(err)
    );
  });

Solution

Amazon Cognito

Cognito is AWS’s user authentication solution, similar to other authentication provider such as Auth0, Okta, or Firebase.

When a user is logged in, it will return 3 JWT tokens: Access, Identity and Refresh tokens. Identity token is used to get user information such as username and email. Access token is used to get access to resources. Both tokens have short expiry, which is set to 1 hour by default. The other token, the refresh token, is different. It has a longer expiry date, and is used to “refresh” access and identity token. Learn more about JWT tokens here.

First things first, we need to set up an Apollo Client. Here we opt to use the individual packages provided by Apollo instead of apollo-boost package because we need flexibility to configure our Apollo Client later on.

import React from 'react';
import { ApolloProvider } from '@apollo/react-hooks';
import { ApolloClient } from 'apollo-client';
import { InMemoryCache } from 'apollo-cache-inmemory';
import { HttpLink } from 'apollo-link-http';
import { ApolloLink, Observable } from 'apollo-link';
import { setContext } from "apollo-link-context";
import { refreshSession, getAccessToken } from '../lib/auth';
 
export default function App() {
  const client = new ApolloClient({
    cache: new InMemoryCache(),
    link: new HttpLink({
      uri: process.env.REACT_APP_YOUR_GRAPHQL_API_ENDPOINT
    })
  })
  return (
    <ApolloProvider client={client}>
      ...
    </ApolloProvider>
  )
}

Before we go further, I want to explain the concepts of link in Apollo. While Apollo have built in defaults in its request, it also gives us the ability to modify how our request should look like by exposing its network layer through a family of apollo-link-* packages. We have seen one in our code above, which uses the HttpLink provided by apollo-link-http package.

a single Apollo Link

Observable

observable explanation

You can learn more about observables from the difference with promises or its documentation

We will now focus more on building the Apollo Link stack. As discussed before, each link worked as a middleware and is able to modify the request or handle errors. Links can also be chained together. Apollo provides the ApolloLink.from([ Link1, Link2, ... , LinkN ]) constructor to chain multiple links.

multiple Apollo Links

We currently have one link, which is the HttpLink. This will always be in the link stack since it is the terminating link (the last link) that will fetch GraphQL results from a GraphQL endpoint over http connection.

Authentication

Most API endpoints are secured with authentication, so we need to somehow pass in a token or an API key to the request headers. We can do this by adding the setContext function from apollo-link-context package. The getAccessToken is a asynchronous function that return a promise to return the token.

...
 
// cached storage for the user token
let inMemoryToken;
 
const setAuthorizationLink = setContext((request, previousContext) => {
  // if you have a cached value return immediately
  if (inMemoryToken) return _buildAuthHeader(inMemoryToken);
 
  return getAccessToken().then(userToken => {
    inMemoryToken = userToken;
    return _buildAuthHeader(inMemoryToken)
  });
});
 
const apolloClient = new ApolloClient({
  link: ApolloLink.from([
    // auth header middleware
    setAuthorizationLink,
    // http request middleware
    new HttpLink({ uri: process.env.REACT_APP_YOUR_GRAPHQL_API_ENDPOINT }),
  ]),
  cache: new InMemoryCache(),
});
 
export default apolloClient;
 
// private functions
 
const _buildAuthHeader = (token) => {
  return {
    headers: {
      Authorization: `Bearer ${token}`
    }
  }
}
 
...

Error handling

Okay, so our request now have the Authorization header. How do we handle Unauthenticated errors? We can use the onError function provided by the apollo-link-error package. The onError link will only run if there is an error.

So our request will look like this :

onSuccess: onSuccess

onError: onError

Modify your code to include this part of the code. It creates a new link that will log whenever we get a 401 error, which indicates that our token has expired. A thing to note is the order in which the links are placed.

...
 
const unauthenticatedErrorLink = onError(
  ({ networkError, operation, forward }) => 
  { if (networkError && 
    networkError.name ==='ServerError' && 
    networkError.statusCode === 401) {
    // unauthenticated. token expired. Refetch token using refresh token
    console.warn("Refreshing token and trying again");
  }
});
 
const apolloClient = new ApolloClient({
  link: ApolloLink.from([
    // handle auth error + auth header middleware
    unauthenticatedErrorLink, 
    setAuthorizationLink,
    // http request middleware
    new HttpLink({ uri: process.env.REACT_APP_YOUR_GRAPHQL_API_ENDPOINT }),
  ]),
  cache: new InMemoryCache(),
});
 
export default apolloClient;
 
...

Handling retries

Currently, we only logged the error whenever we get an unauthenticated error. Now, we want to handle the refresh token and retry the request. As shown in the docs, we can retry a request using forward(operation). So we can try something like

if (networkError && 
    networkError.name ==='ServerError' && 
    networkError.statusCode === 401) {
    // unauthenticated. token expired. Refetch token using refresh token
    console.warn("Refreshing token and trying again");
 
    // an asynchronous function that will run the callback function after 
    // refreshing the token. then run the 
    refreshSession(
      (err, session) => {
        if (err) {
          console.error("Failed refreshing session");
        }
        // set inMemoryToken to refreshed token
        inMemoryToken = session.accessToken.jwtToken;
      }
    )
    
    return forward(operation);
  }
 

Unfortunately,onError function provided to handle error does not support asynchronous request and only returns an Observable or a void. So we created a helper function to wrap the promise inside an Observable and changed the code to use our helper function. A thing to note here is that we need to wrap the refreshSession(callback) inside a promise which will only resolve after the callback has been executed. I need to do this because the API for AWS Amplify (which I am using as a client-side library to access AWS Cognito), is returning a function with a callback, and not a function that returns a promise. If your refreshSession returns a promise, then you should not need to wrap it again.

...
 
  if ( networkError && 
    networkError.name ==='ServerError' && 
    networkError.statusCode === 401) {
    // unauthenticated. token expired. Refetch token using refresh token
    console.warn("Refreshing token and trying again");
    // get new token and set inMemoryToken
    return promiseToObservable(
      new Promise((resolve, reject) => refreshSession(
        (err, session) => {
          if (err) {
            console.error("Failed refreshing session");
            reject(err);
          }
          // set inMemoryToken to refreshed token
          inMemoryToken = session.accessToken.jwtToken;
          resolve();
        }
      ))
    ).flatMap(() => forward(operation));
  }
 
...
 
// private functions
...
 
const promiseToObservable = promise =>
  new Observable((subscriber) => {
    promise.then(
      (value) => {
        if (subscriber.closed) return;
        subscriber.next(value);
        subscriber.complete();
      },
      err => subscriber.error(err)
    );
  });
 

The final code is here

Note: AWS Cognito stores the token inside localStorage, so you can actually hack the problem by accessing the token(s) directly and synchronously without the hassle of wrapping promises into observables. However, I believe that it is more maintainable to access the token via the API provided.