AuthProvider.tsx

import {
  useContext,
  useState,
  createContext,
  useRef,
  PropsWithChildren,
} from "react";
import { useInterval } from "../hooks/useInterval";
import { useApi } from "./ApiProvider";
import { AuthTokens } from "../api/types/AuthTokens";
import { SignInApi } from "../api/types/SignInApi";

/**
 * auth.ts provides all mechanisms requires to handle authentication in the app:
 * - Manual sign in
 * - Manual sign out
 * - Token refresh
 * - Client side token persistence.
 *
 * Use AuthProvider to set up auth, then call useAuth wherever you need access to AuthHelpers.
 */

/**
 * All data and functions required to manage auth
 */
export type AuthHelpers = {
  /**
   * Set to true when the user is signed in
   */
  isSignedIn: boolean;
  /**
   * User tokens (see api/createTokens.js)
   * Set to undefined when signed out.
   */
  tokens?: AuthTokens;
  /**
   * Signs you in and sets the user object (if successful).
   * Returns the success/error object received from the api.
   */
  signIn: SignInApi;
  /**
   * Signs you out
   */
  signOut: () => void;
};

/**
 * Context used to pass AuthHelpers around.
 * (Should not be used directly. Use AuthProvider instead)
 */
const AuthContext = createContext<AuthHelpers>({
  isSignedIn: false,
  signOut: () => {},
  signIn: async () => ({
    type: "error",
    error:
      "AuthContext not populated. Don't use AuthContext directly, use AuthProvider.",
  }),
});

export type AuthProviderProps = PropsWithChildren<{
  /**
   * How often to check if the current token has expired and refresh (in ms)
   */
  tryRefreshTokensInterval?: number;
}>;

/**
 * Sets up state and callbacks for the AuthContext object and renders an AuthContext.Provider.
 */
export const AuthProvider = ({
  children,
  tryRefreshTokensInterval = 1000,
}: AuthProviderProps) => {
  const [tokens, setTokens] = useState(getPersistedTokens());
  const api = useApi();
  useTokenRefresh(tokens, updateTokens, tryRefreshTokensInterval);

  function updateTokens(tokens?: AuthTokens) {
    setTokens(tokens);
    setPersistedTokens(tokens);
  }

  const signIn: SignInApi = async (...args) => {
    const response = await api.signIn(...args);
    if (response.type === "success") {
      updateTokens(response.tokens);
    }
    return response;
  };

  function signOut() {
    updateTokens(undefined);
  }

  const auth = { isSignedIn: !!tokens, tokens, signIn, signOut };

  return <AuthContext.Provider value={auth}>{children}</AuthContext.Provider>;
};

/**
 * Convenience hook to access auth object with less code.
 */
export const useAuth = () => useContext(AuthContext);

/**
 * Sets up an interval that calls the refreshApi function as soon as the current tokens expire
 * and passes the returned tokens to the updateTokens function. If token refresh fails,
 * undefined is passed to the updateTokens function, effectively signing the user out.
 */
const useTokenRefresh = (
  tokens: AuthTokens | undefined,
  updateTokens: (tokens?: AuthTokens) => void,
  refreshInterval: number
) => {
  const api = useApi();
  const isRefreshingTokensRef = useRef(false);

  const tryRefreshTokens = async () => {
    if (
      // Avoid concurrent refreshes
      !isRefreshingTokensRef.current &&
      tokens &&
      tokens.expiresAt <= new Date()
    ) {
      // Current tokens has expired, try to refresh
      isRefreshingTokensRef.current = true;
      const response = await api.refresh(tokens.refresh);
      isRefreshingTokensRef.current = false;

      if (response.type === "success") {
        // New tokens received
        updateTokens(response.tokens);
      } else {
        // Remove tokens (signs user out)
        updateTokens(undefined);
      }
    }
  };

  useInterval(
    tryRefreshTokens,
    tokens
      ? refreshInterval // Only try to refresh when we have tokens
      : undefined // Pause the interval whenever we're signed out
  );
};

const localStorageId = "tokens";

const setPersistedTokens = (tokens?: AuthTokens) => {
  if (tokens) {
    localStorage.setItem(localStorageId, JSON.stringify(tokens));
  } else {
    localStorage.removeItem(localStorageId);
  }
};

const getPersistedTokens = (): AuthTokens | undefined => {
  const jsonString = localStorage.getItem(localStorageId);
  const json: Record<keyof AuthTokens, any> = jsonString
    ? JSON.parse(jsonString)
    : undefined;
  if (!json) {
    return;
  }
  return {
    ...json,
    expiresAt: new Date(json.expiresAt), // Necessary since JSON.parse doesn't reinitialize Date objects
  };
};