React jotai で localStorage を使って state を永続化するコンポーネント作ってみたのメモ

React で jotai を使って local storage に JTW を保存して永続化・リロード時に状態を復帰するのをやってみためも。
(※ サンプルなので local storage に JWT を保存するのはセキュリティ的によろしくないってのは今回考慮していません )

シナリオ

ログインの際に API から JWT token とユーザー情報が返される
JWT token は local storage に保存する
ユーザー情報は React の state にする

初回アクセス・リロード時にユーザー情報を永続化された JWT から復元するシナリオ

local storage に JWT token がある場合
API に JWT token を付けてアクセスする
JWT token が expire している場合はエラーが返される -> ログアウト状態にする
API から返されたユーザー情報を React の state にするセットする
ユーザー情報を復帰したら初期化完了とする

環境

react 18.2.0
typescript 5.0.4
jotai 2.0.4

1. JWT token とユーザー情報の state を jotai で定義する

JWT token を保持する atom `tokenAtom`

JWT token は localStorage に保存するので jotai の atomWithStorage を使う

// /atoms/token.atom.ts
import { atomWithStorage } from 'jotai/utils';

type HasTokenType = string;
type NoTokenType = null;
export type TokenType = HasTokenType | NoTokenType;

export const TOKEN_KEY = 'myToken' as const;
export const tokenAtom = atomWithStorage<TokenType>(TOKEN_KEY, undefined);

cf. React 👻 jotai を使うと localStorage を使った永続化が簡単だった件について - かもメモ

ユーザー情報を保持する atom `authAtom`

ユーザー情報の設定が完了したら初期化完了とするので合わせて isReady というフラグとなる state を作成する

// /atoms/auth.atom.ts
import { atom } from 'jotai';

export interface IAuthUser {
  id: string;
  email: string;
  displayName: string;
  role: number;
}

type noUserState = null;
const initialState = initialState;
export type AuthType = IAuthUser | noUserState | typeof initialState;

export const authAtom = atom<AuthType>(initialState);
// readonly atom
export const isReadyAtom = atom<boolean>((get) => {
  const auth = get(authAtom);
  return auth !== initialState;
});

User が存在する -> IAuthUser
User が存在しない -> null
User の存在確認前 -> undefined

User の存在確認が完了したら authAtom は undefined でなくなるので isReadyAtom が true になる

2. `tokenAtom` を扱う Hooks

// /hooks/useToken.ts
import { useCallback } from 'react';
import { useAtomValue, useSetAtom } from 'jotai';
import { TOKEN_KEY, TokenType, tokenAtom } from '@/atoms/token.atom';

// localStorage の値と比較する関数
const isAccessLocalStorage = (token: any) => {
  const storageValue = localStorage.getItem(TOKEN_KEY);

  if (storageValue === null) {
    if (storageValue === token) {
      return true;
    }
    throw new Error('Invalid token');
  } else {
    if (token === null) {
      return false;
    }
    if (JSON.parse(storageValue) === token) {
      return true;
    }
    throw new Error('Invalid token');
  }
};

export const useToken = () => {
  const token = useAtomValue(tokenAtom);
  const isReady = isAccessLocalStorage(token);

  return {
    token,
    isReady,
  };
};

export const useTokenMutators = () => {
  const setTokenAtom = useSetAtom(tokenAtom);

  const resetToken = useCallback(() => {
    setTokenAtom(null);
  }, [setTokenAtom]);

  const setToken = useCallback(
    (token: NoTokenType) => {
      setTokenAtom(token);
    },
    [setTokenAtom],
  );

  return {
    resetToken,
    setToken,
  };
};

token を取り扱うコンポーネントでは useToken() が返す isReady が true になるのを待って token を扱えば良い

💡 jotai の atom が localStorage にアクセス可能になっているかの判定

const isAccessLocalStorage = (token: any) => {
  const storageValue = localStorage.getItem(TOKEN_KEY);

  if (storageValue === null) {
    if (storageValue === token) {
      return true;
    }
    throw new Error('Invalid token');
  } else {
    if (token === null) {
      return false;
    }
    if (JSON.parse(storageValue) === token) {
      return true;
    }
    throw new Error('Invalid token');
  }
};

前提 - JavaScript で localStorage にアクセスしたときデータが存在しない場合は null が返される - jotai の atom は localStorage にアクセス可能になるまでタイムラグがあり、その間は初期値が返される (今回の場合は null)

localStorage の値が null の時 -> token がそもそも存在していない状態
1. jotai の atom は localStorage へのアクセスに関わらず null になるの筈である
localStorage の値が strings の時 -> token が存在する状態
1. jotai の atom が null -> まだ localStorage へアクセスできてない状態
2. jotai の atom が string -> localStorage の値と atom の値が等しければ正しい
  ※ jotai の atomWithStorage は保存時に JSON.sanitaize() するので文字列でも JSON.parse() して比較する必要がある (localStorage から取得される値には " コーテーションが含まれるため)

3. `authAtom` を扱う Hooks

// /hooks/useAuth.ts
import { useCallback } from 'react';
import { useAtomValue, useSetAtom } from 'jotai';
import { authAtom } from '@/atoms/auth.atom';
import { TokenType } from '@/atoms/token.atom';
import { useTokenMutators } from '@/hooks/useToken';

export const useAuth = () => {
  const user = useAtomValue(authAtom, { store });
  const isReady = useAtomValue(isAuthReadyAtom);

  return {
    user,
    isReady,
  };
};

const fetchUser = async <T>(url: string, payload: any): Promise<T> => {
  const method = 'POST';
  const headers = {
    Accept: 'application/json',
    'Content-Type': 'application/json',
  };

  return await fetch(url, { method, headers, body: payload }).then(
    async (res) => {
      // status code 200 番台以外はエラーにする
      if ( !res.ok ) {
        throw new Error(res.statusText);
      }
      return await res.json<T>();
    },
  );
};

export const useAuthMutators = () => {
  const setAuth = useSetAtom(authAtom, { store });
  const { resetToken,  setToken } = useTokenMutators();

  const signOut = useCallback(() => {
    setAuth(null);
    resetToken();
  }, [setAuth, resetToken]);

  const signIn = useCallback(async (payload: any) => {
    try {
      const { user, token } = await fetchUser<{user: IAuthUser, token: string }>(API_SIGN_IN, payload);
      setAuth(user);
      setToken(token);
    } catch (error) {
      // signin 失敗
      setAuth(null);
      resetToken();
    }
  }, [setAuth, setToken]);

  // signup 割愛

  const signInByToken = useCallback(async (token: AuthType) => {
    if (!token) {
      signOut();
      return;
    }

    try {
      const { user, token: newToken } = await fetchUser<{user: IAuthUser, token: string}>(API_GET_USER, { authorization: `Bearer ${token}` });
      setAuth(user);
      setToken(newToken);
    } catch (error) {
      // token が expire などで user データが取得できなかった場合
      signOut();
    }
  }, [setAuth, signOut]);

  return {
    signOut,
    signIn,
    signInByToken,
  };
};

ユーザー情報を扱うコンポーネントでは useAuth() から取得される isReady が true になるまでまって user オブジェクトを扱えば良い
localStorage に永続化された JWT token から user state を再構築するには useAuthMutators() の signInByToken 関数に JWT token を渡せば良い

4. 初回ロード時に永続化された token から user state を設定するコンポーネント

AuthObserver.tsx

import { useEffect } from 'react';
import { useToken } from '@/hooke/useToken';
import { useAuthMutators } from '@/hooks/useAuth';

function AuthObserver(): null {
  const { token, isReady } = useToken();
  const { signInByToken } = useAuthMutators();

  useEffect(() => {
    if (!isReady) { return; }
    signInByToken(token);
  }, [isReady]);
  return null;
}
export { AuthObserver };

localStorage にアクセス可能になったら token を signInByToken に渡す
signInByToken 内で token が null なら sign out され、 token がある場合は API からユーザー情報の取得を試みる
ユーザー情報が取得できたら state に設定、ユーザー情報の取得に失敗したら sign out 扱いになる

App.tsx

import { AuthObserver } from './AuthObserver';
function App(): JSX.Element {
  return (
    <>
      <MyComponent />
      <AuthObserver />
    </>
  )
}

5. ログインユーザー情報を扱うコンポーネント

import { useAuth } from '@/hooks/useAuth';

function MyComponent(): JSX.Element {
  const { user, isReady } = useAuth();
  if (!isReady) { return <div>Loading...</div> }
  if (!user) { return <Signin /> }
  return {
    <div>Hello, {user.displaName} !</div>
  };
}

所感

一度 Zenn のストックにメモ書きしておいた React, Next.js でアプリを作っているとよく作る初回ロード時に永続化された情報からログイン情報を復元する処理を改めてまとめてみました。
toknAtom と authAtom の 2つを作成したのですが依存関係があるので、read-write atom などを使えば jotai の設定内で依存関係を作ってしまうこともできそうですが、今回はそこまでやっていません。
このままでも初期化の処理は <AuthObserver /> に任せてしまって、ユーザー情報を扱うコンポーネントでは useAuth さえ使えば、初期化中・非ログイン / ログインが判別できるので割といい感じなったのではないかと思います。(コードについては作成しておいたものを調整して blog にしたのでコピペで動かなかったらスミマセン…)
フロー図的なのを先に考えるといい感じに実装できる事に今さら気づきました

おわり₍ ᐢ. ̫ .ᐢ ₎

[参考]