React & JWT Authentication - the right way!

React & JWT Authentication - the right way!

ยท

18 min read

In this series of posts, we create a secured end-to-end JWT-based authentication mechanism using NodeJS, Express, PassportJS, and React.

In this series I cover:

In this part, we focus on the client-side. If you haven't read Part 1, I strongly encourage you to do so before proceeding with this part, as it contains an important background that will help you grasp the concepts we use here.

What we want to achieve

  • A secured mechanism - we follow the rules described in the first part: access token is not stored in the local storage; utilize refresh tokens instead
  • User (and developer) friendly - automatic login & logout, multi-tabs support, automatic token refresh
  • State management - our app should know whether a user is authenticated

Let's start!

State Management: unstated-next

I really like Context API. It's a simple, integral part of React, and it works great for my use-cases. unstated-next is a tiny wrapper on top, that makes the use of Context API and hooks even better. You can, of course, use whatever state management solution you like (Redux ๐Ÿ™„, Mobx, etc). The concept remains the same.

HTTP Client: axios, axios-hooks

Axios is a really popular, open-source HTTP client for node and the browser. Axios has built-in support for request interceptors, which come handy when passing authorization headers. Since we use React hooks, we will add hooks support by integrating axios-hooks. If you don't like Axios, there are other great options out there (e.g SWR).

The useTokenExpiration hook

Let's start at the end. Since we use refresh tokens, we benefit from an increase in the level of the app's security by generating short-living access tokens. If access tokens expire often, we should refresh them to keep the users logged in. To achieve that, we set a timeout for each new access token we receive, and call the /refresh-token endpoint to receive a new one when it expires.

The useTokenExpiration gets the current token's expiration time, counts down until it expires, and lets the app know that a refresh is required, by calling the onTokenRefreshRequired callback.

// useTokenExpiration.ts

import { useEffect, useRef, useState } from 'react';

export function useTokenExpiration(onTokenRefreshRequired: Function) {
  const clearAutomaticRefresh = useRef<number>();
  const [tokenExpiration, setTokenExpiration] = useState<Date>();

  useEffect(() => {
    // get a new access token with the refresh token when it expires
    if (tokenExpiration instanceof Date && !isNaN(tokenExpiration.valueOf())) {
      const now = new Date();
      const triggerAfterMs = tokenExpiration.getTime() - now.getTime();

      clearAutomaticRefresh.current = window.setTimeout(async () => {
        onTokenRefreshRequired();
      }, triggerAfterMs);
    }

    return () => {
      window.clearTimeout(clearAutomaticRefresh.current);
    };
  }, [onTokenRefreshRequired, tokenExpiration]);

  const clearAutomaticTokenRefresh = () => {
    window.clearTimeout(clearAutomaticRefresh.current);
    setTokenExpiration(undefined);
  };

  return {
    clearAutomaticTokenRefresh,
    setTokenExpiration,
  };
}

The useToken hook

This is the heart of the authentication mechanism. It is responsible for the entire token management and lifecycle.

The token lifecycle

The useToken hook holds the access token in-memory (we explained why it's important, remember?). Therefore we store the access token in a useRef hook - more on that soon. The hook is also responsible for managing the entire token lifecycle: initialization, removal, expiration, and refresh. The token expiration is handled by using the useTokenExpiration hook we created.

A word about clearing the token - since the refresh token is saved in an httpOnly cookie, we can't access it or modify it from the browser using Javascript. This raises a problem: we can't really complete the action on the client-side alone as the browser will keep sending the refresh token cookie and a new access token will be received on every refresh. To solve that, we call the /logout endpoint, which returns a new empty cookie, and in practice, deletes the refresh token from the browser.

Let's create the lifecycle methods:

// useToken.ts

// .... 

export function useToken(onTokenInvalid: Function, onRefreshRequired: Function) {
 const accessToken = useRef<string>();
  const { clearAutomaticTokenRefresh, setTokenExpiration } = useTokenExpiration(onRefreshRequired);

  const setToken = useCallback(
    ({ token_expiration, access_token }: TokenResponse) => {
      accessToken.current = access_token;
      const expirationDate = new Date(token_expiration);
      setTokenExpiration(expirationDate);
    },
    [setTokenExpiration],
  );

  const isAuthenticated = useCallback(() => {
    return !!accessToken.current;
  }, []);

  const clearToken = useCallback(
    (shouldClearCookie = true) => {
      // if we came from a different tab, we should not clear the cookie again
      const clearRefreshTokenCookie = shouldClearCookie ? axios.get('logout') : Promise.resolve();

      // clear refresh token
      return clearRefreshTokenCookie.finally(() => {
        // clear token
        accessToken.current = '';

        // clear auto refresh interval
        clearAutomaticTokenRefresh();
      });
    },
    [clearAutomaticTokenRefresh],
  );

  return {
    clearToken,
    setToken,
    isAuthenticated,
  };
}

Initiating Axios and setting interceptors

Once we have the token, we want to pass it to all of the requests automatically. For that we use interceptors. Interceptors are pieces of code that run before and after we fire a request and allow us to intervene in the process. Since we should do the initiating process just once on the application startup, we store the access token in a useRef hook. By doing that, the useEffect containing the definitions of our interceptors will not redefine them over and over again when the access token changes.

The first interceptor attaches the access token to the authorization header before we send it to the backend. The second interceptor handles the 401 response from the backend by clearing the token's data and notifying the app that the current user isn't logged in anymore.

// useToken.ts

export const axios = Axios.create({
  baseURL: BASE_URL,
});

export function useToken(onTokenInvalid: Function, onRefreshRequired: Function) {
// .....

    useEffect(() => {
    // add authorization token to each request
    axios.interceptors.request.use(
      (config: AxiosRequestConfig): AxiosRequestConfig => {
        config.headers.authorization = `Bearer ${accessToken.current}`;
        return config;
      },
    );

    // if the current token is expired or invalid, logout the user
    axios.interceptors.response.use(
      (response) => response,
      (error) => {
        if (error.response.status === 401 && accessToken.current) {
          clearToken();

          // let the app know that the current token was cleared
          onTokenInvalid();
        }
        return Promise.reject(error);
      },
    );

    // configure axios-hooks to use this instance of axios
    configure({ axios });
  }, [clearToken, onTokenInvalid]);

And finally, the entire hook:

// useToken.ts

import Axios, { AxiosRequestConfig } from 'axios';
import { useCallback, useEffect, useRef } from 'react';
import { configure } from 'axios-hooks';
import { BASE_URL } from '../config/config';
import { User } from '../Containers/AuthContainer';
import { useTokenExpiration } from './useTokenExpiration';

interface TokenResponse {
  access_token: string;
  token_expiration: string;
}

export interface UserAndTokenResponse extends TokenResponse {
  user: User;
}

export const axios = Axios.create({
  baseURL: BASE_URL,
});

export function useToken(onTokenInvalid: Function, onRefreshRequired: Function) {
  const accessToken = useRef<string>();
  const { clearAutomaticTokenRefresh, setTokenExpiration } = useTokenExpiration(onRefreshRequired);

  const setToken = useCallback(
    ({ token_expiration, access_token }: TokenResponse) => {
      accessToken.current = access_token;
      const expirationDate = new Date(token_expiration);
      setTokenExpiration(expirationDate);
    },
    [setTokenExpiration],
  );

  const isAuthenticated = useCallback(() => {
    return !!accessToken.current;
  }, []);

  const clearToken = useCallback(
    (shouldClearCookie = true) => {
      // if we came from a different tab, we should not clear the cookie again
      const clearRefreshTokenCookie = shouldClearCookie ? axios.get('logout') : Promise.resolve();

      // clear refresh token
      return clearRefreshTokenCookie.finally(() => {
        // clear token
        accessToken.current = '';

        // clear auto refresh interval
        clearAutomaticTokenRefresh();
      });
    },
    [clearAutomaticTokenRefresh],
  );

  useEffect(() => {
    // add authorization token to each request
    axios.interceptors.request.use(
      (config: AxiosRequestConfig): AxiosRequestConfig => {
        config.headers.authorization = `Bearer ${accessToken.current}`;
        return config;
      },
    );

    // if the current token is expired or invalid, logout the user
    axios.interceptors.response.use(
      (response) => response,
      (error) => {
        if (error.response.status === 401 && accessToken.current) {
          clearToken();

          // let the app know that the current token was cleared
          onTokenInvalid();
        }
        return Promise.reject(error);
      },
    );

    // configure axios-hooks to use this instance of axios
    configure({ axios });
  }, [clearToken, onTokenInvalid]);

  return {
    clearToken,
    setToken,
    isAuthenticated,
  };
}

The Authentication Container

The Authentication Container holds the logged-in user state, and exposes the user authentication lifecycle methods: register, login, logout, and refresh token. The container is also responsible for cross-tabs login and logout support.

Initialization

The access token is saved in memory and therefore isn't available in the initialization phase of the app, however, the refresh token that is saved in an httpOnly cookie is available. When our app starts up, we try to fetch a new access token using the /refresh-token endpoint and we attempt to initialize our app with it. If the request is successful, we update the state with the received access token and the user object. In case the refresh token is not available or expired, the user will have to login again. Let's create the container, and the useEffect hook that is responsible for the initialization.

// AuthContainer.ts

//...

function useAuth() {
  const history = useHistory();
  const [user, setUser] = useState<User | null>(null);
  const refreshToken = useCallback(refresh, []);

  const onTokenInvalid = useCallback(() => setUser(null), []);
  const { setToken, clearToken, isAuthenticated } = useToken(onTokenInvalid, refreshToken);

  useEffect(() => {
    // try to get new token on first render using refresh token
    refreshToken();
  }, [refreshToken]);

  async function refresh() {
    const {
      data: { user, ...rest },
    } = await axios.get<UserAndTokenResponse>('refresh-token');

    setUser(user);
    setToken(rest);
  }

//.....

The user lifecycle methods

The lifecycle methods have several roles:

  • Update the user state on different lifecycle events: login, logout, register, refresh token
  • Set a new token data when user registers or logs in
  • Clear token data when the user logs out
  • Notify the other open applications tabs on login and logout events
// AuthContainer.ts

//...


const logout = useCallback(() => {
    clearToken().finally(() => {
      setUser(null);
      history.push('/');

      // fire an event to logout from all tabs
      window.localStorage.setItem(AuthEvents.LOGOUT, new Date().toISOString());
    });
  }, [history, clearToken]);

  const register = useCallback(
    async (userToRegister: UserBase) => {
      const {
        data: { user, ...rest },
      } = await axios.post<UserAndTokenResponse>('register', userToRegister);
      setUser(user);
      setToken(rest);
    },
    [setToken],
  );

  const login = useCallback(
    async (email: string, password: string) => {
      const {
        data: { user, ...rest },
      } = await axios.post<UserAndTokenResponse>('login', {
        email,
        password,
      });
      setUser(user);
      setToken(rest);

      // fire an event to let all tabs know they should login
      window.localStorage.setItem(AuthEvents.LOGIN, new Date().toISOString());
    },
    [setToken],
  );

  async function refresh() {
    const {
      data: { user, ...rest },
    } = await axios.get<UserAndTokenResponse>('refresh-token');

    setUser(user);
    setToken(rest);
  }
//...

Multiple tabs support for login and logout

If the user has multiple tabs open, when they log in or out, we should notify the other tabs that the status has changed. For that, we use storage events. When we set a new variable on the local storage, an event is fired across all tabs, not including the tab that initiated the change. This is exactly what we are looking for!

Let's first add the listener that listens to storage events fired by other tabs:

  useEffect(() => {
    // add listener for login or logout from other tabs
    window.addEventListener('storage', async (event: WindowEventMap['storage']) => {
      if (event.key === AuthEvents.LOGOUT && isAuthenticated()) {
        await clearToken(false);
        setUser(null);
      } else if (event.key === AuthEvents.LOGIN) {
        refreshToken();
      }
    });
  }, [clearToken, isAuthenticated, refreshToken]);

Let's combine everything together:

// AuthContainer.ts

import { useCallback, useEffect, useState } from 'react';
import { createContainer } from 'unstated-next';
import { useHistory } from 'react-router';
import { axios, UserAndTokenResponse, useToken } from '../Hooks/useToken';
import { AuthEvents } from '../services/AuthEvents';

export interface UserBase {
  name: string;
  email: string;
  password: string;
}

export interface User extends UserBase {
  _id: string;
  role: 'user' | 'admin';
}

function useAuth() {
  const history = useHistory();
  const [user, setUser] = useState<User | null>(null);
  const refreshToken = useCallback(refresh, []);

  const onTokenInvalid = () => setUser(null);
  const { setToken, clearToken, isAuthenticated } = useToken(onTokenInvalid, refreshToken);

  useEffect(() => {
    // try to get new token on first render using refresh token
    refreshToken();
  }, [refreshToken]);

  useEffect(() => {
    // add listener for login or logout from other tabs
    window.addEventListener('storage', async (event: WindowEventMap['storage']) => {
      if (event.key === AuthEvents.LOGOUT && isAuthenticated()) {
        await clearToken(false);
        setUser(null);
      } else if (event.key === AuthEvents.LOGIN) {
        refreshToken();
      }
    });
  }, [clearToken, isAuthenticated, refreshToken]);

  const logout = useCallback(() => {
    clearToken().finally(() => {
      setUser(null);
      history.push('/');

      // fire an event to logout from all tabs
      window.localStorage.setItem(AuthEvents.LOGOUT, new Date().toISOString());
    });
  }, [history, clearToken]);

  const register = useCallback(
    async (userToRegister: UserBase) => {
      const {
        data: { user, ...rest },
      } = await axios.post<UserAndTokenResponse>('register', userToRegister);
      setUser(user);
      setToken(rest);
    },
    [setToken],
  );

  const login = useCallback(
    async (email: string, password: string) => {
      const {
        data: { user, ...rest },
      } = await axios.post<UserAndTokenResponse>('login', {
        email,
        password,
      });
      setUser(user);
      setToken(rest);

      // fire an event to let all tabs know they should login
      window.localStorage.setItem(AuthEvents.LOGIN, new Date().toISOString());
    },
    [setToken],
  );

  async function refresh() {
    const {
      data: { user, ...rest },
    } = await axios.get<UserAndTokenResponse>('refresh-token');

    setUser(user);
    setToken(rest);
  }

  return {
    user,
    setUser,
    register,
    login,
    logout,
    refreshToken,
  };
}

export const AuthContainer = createContainer(useAuth);

That's it! Now we just need to wrap everything with the AuthContainer to make the authentication state available in the entire app.

// App.tsx

export default () => {
  return (
    <ThemeProvider theme={theme}>
        <Router>
          <AuthContainer.Provider>
            <App />
          </AuthContainer.Provider>
        </Router>
    </ThemeProvider>
  );
};

And now we can use our new and shiny AuthContainer to protect the application routes

import React, { ComponentType } from 'react';
import { BrowserRouter as Router, Redirect, Route } from 'react-router-dom';

interface ProtectedRouteProps { 
  component: ComponentType<any>;
  path: string
}

const ProtectedRoute: FC<ProtectedRouteProps> = ({ component: Component, path = '' }) => {
  const { user } = AuthContainer.useContainer();

  return (
    <Route
      path={path}
      render={(props) =>
        user ? (
          <Component {...props} />
        ) : (
          <Redirect
            to={{
              pathname: '/login',
              state: { from: props.location },
            }}
          />
        )
      }
    />
  );
};

That's all for this part. In the next part, we'll integrate multiple social logins that work seamlessly using the infrastructure we created so far. I hope you find this post helpful, let me know what you think!