AWS Amplify Auth v6へのマイグレーションで対応した破壊的変更

ROUTE06 でソフトウェアエンジニアをしている @MH4GF です。私が関わるプロダクトでは認証や認可に Amazon Cognito を使っており、 React で実装したフロントエンドアプリケーションから Amazon Cognito との接続には AWS Amplify が提供している Auth ライブラリを利用しています。
AWS Amplify は 2023 年 11 月に v6 のメジャーバージョンがリリースされたのですが、Auth ライブラリにはいくつかの破壊的変更が含まれていたので、このブログで移行方法の一例を紹介します。

AWS Amplify JavaScript v6 での変更点

https://aws.amazon.com/jp/blogs/mobile/amplify-javascript-v6/

v6 の変更の目玉として以下の内容が取り上げられています。

  • ツリーシェイキングの改善とバンドルサイズの削減(55kb → 32kb (42%削減))
  • TypeScript 型定義の改善
  • Next.js App Router のサポート

開発者に影響の大きな変更としては、APIのデザインがメソッドチェーンベースから関数ベースに変更されたことです。 マイグレーションガイドはこちらで、マイグレーションを行う際にはまずこちらをご覧ください。

https://docs.amplify.aws/javascript/build-a-backend/auth/auth-migration-guide/

私たちの利用状況

v6 へのマイグレーションを説明する前に、私たちの利用状況を説明します。

  • Amplify のライブラリのうち Auth だけを利用(AppSync や Storage は使っていない)
  • JavaScript パッケージとして aws-amplify, @aws-amplify/ui-react を利用

そのため AWS Amplify の中でも Auth のみの紹介となることをご了承ください。

Amazon Cognito を React アプリケーションから利用する場合は aws-amplify だけ利用する方法もありますが、 @aws-amplify/ui-react は認証フローや認証状態を管理するためのカスタムフックや、サインイン・サインアップのためのフォームコンポーネントなどが提供されています。
これらの処理をライブラリ側に任せ開発効率を向上させたかったため @aws-amplify/ui-react も一緒に利用しています。

マイグレーションで対応した点

それでは v6 へのマイグレーションで対応した点を紹介していきます。

  • useAuthenticator で手に入る CognitoUser からトークンが取り出せなくなり、fetchAuthSession への移行が必要
    • React で非同期関数を扱うためのカスタムフックを作成
    • ID トークンの期限切れ判定を実装

useAuthenticator で手に入る CognitoUser からトークンが取り出せなくなり、 fetchAuthSession への移行が必要

Amplify Auth v5 までは CognitoUser インスタンスからトークンを取り出すことができましたが、v6 では関数をモジュールとして import する方式にリデザインされました。その中で fetchAuthSession がトークンを取り出すための関数として提供されています。

v5

import { Auth } from "aws-amplify";

const getAuthenticatedUser = async () => {
  const {
    username,
    signInUserSession: session,
    authenticationFlowType,
  } = await Auth.currentUserPoolUser();
};

v6

import { fetchAuthSession, getCurrentUser } from "aws-amplify/auth";

const getAuthenticatedUser = async () => {
  const { username, signInDetails } = await getCurrentUser();

  const { tokens: session } = await fetchAuthSession();

  return {
    username,
    session,
    authenticationFlowType: signInDetails.authFlowType,
  };
};

こちらから引用: https://docs.amplify.aws/javascript/build-a-backend/auth/auth-migration-guide/#authcurrentauthenticateduser-deprecated

fetchAuthSession は内部で以下の処理を行なっています。

  • ストレージ(デフォルトは localStorage)に保存されたトークンを取り出す
    • トークンが有効であれば返却して終了
  • 期限切れの場合やトークンが存在しない場合、また引数の forceRefresh が true の場合はリフレッシュを実行
  • リフレッシュが成功した場合はトークンを返却して終了、失敗した場合はサインアウト処理を実行

v5 での CognitoUser.getSession() と似た処理を行なっています。amplify-js の利用者がリフレッシュを意識せずとも済むのは利点ですが、認証の仕様や挙動を理解するにはやはりソースコードを読んでおくことをおすすめします。

https://github.com/aws-amplify/amplify-js/blob/7e4eff3be3ec25520332330461ca2cefbf7bdbb7/packages/core/src/singleton/apis/fetchAuthSession.ts#L9-L11

React で非同期関数を扱うためのカスタムフックを作成

React での統合として私たちは @aws-amplify/ui-react を利用しており、そこで提供されている useAuthenticator フックから CognitoUser インスタンスを取り出すことができていました。しかし v6 からは CognitoUser インスタンスからトークンを取り出すことはできなくなり、fetchAuthSession への移行が必要となったため、関数をラップするカスタムフックを作成することにしました。

v5

import { useAuthenticator } from "@aws-amplify/ui-react";

export const Component = () => {
  const { user } = useAuthenticator();
  const token = user.getSignInUserSession()?.getIdToken().getJwtToken();
  // トークンを利用した処理
};

v6

import { useAuthenticator } from "@aws-amplify/ui-react";
import type { JWT, AuthSession } from "aws-amplify/auth";
import { fetchAuthSession } from "aws-amplify/auth";
import type { PropsWithChildren } from "react";
import { createContext, useCallback, useEffect, useMemo, useState, useContext } from "react";

type LoadSessionOptions = {
  forceRefresh?: boolean;
};

type ContextType = {
  /**
   * Cognito の ID トークン
   */
  idToken: JWT | undefined;
  /**
   * ID トークンの有効期限
   */
  expiration: Date | undefined;
  /**
   * サインイン処理の完了を示すフラグ
   */
  isConfigured: boolean;
  /**
   * サインイン判定を示すフラグ
   */
  isSignedIn: boolean;
  /**
   * 明示的にセッションをロードする関数
   */
  loadSession: (
    options?: LoadSessionOptions | undefined
  ) => Promise<AuthSession | undefined>;
};

export const SessionContext = createContext<ContextType>({
  idToken: undefined,
  expiration: undefined,
  isConfigured: false,
  isSignedIn: false,
  loadSession: async () => undefined,
});

/**
 * セッション情報を提供するProviderです。
 * サインイン状態に応じてセッション情報(idToken)を取得します。
 * Note:
 *  - サインイン済みかどうかは idToken が存在するかどうかで判断してください。
 *    - idToken が存在する => ブラウザにidTokenが存在するが、idTokenが有効期限内であることは保証されない
 *  - idTokenの有効期限が切れた状態で loadSession() を呼び出すと最新のidTokenに更新されます。
 */
export const SessionProvider = ({
  children,
}: PropsWithChildren): JSX.Element => {
  const { authStatus } = useAuthenticator((context) => [context.authStatus]);
  const [session, setSession] = useState<JWT>();
  const [isConfigured, setIsConfigured] = useState(false);
  const idToken = useMemo(() => session?.tokens?.idToken, [session]);
  const expiration = useMemo(() => session?.credentials?.expiration, [session]);
  const isSignedIn = useMemo(() => idToken != null, [idToken]);
  const loadSession = useCallback(
    async (options?: LoadSessionOptions) => {
      const session = await fetchAuthSession(options);
      setSession(session);
      return session;
    },
    [setSession]
  );

  useEffect(() => {
    switch (authStatus) {
      case "configuring":
        if (isConfigured) setIsConfigured(false);
        break;
      case "authenticated":
        loadSession();
        setIsConfigured(true);
        break;
      case "unauthenticated":
        setIsConfigured(true);
        break;
      default
        break;
    }
  }, [authStatus]);

  return (
    <SessionContext.Provider
      value={{ idToken, expiration, isConfigured, isSignedIn, loadSession }}
    >
      {children}
    </SessionContext.Provider>
  );
};

export const useSession = () => useContext(SessionContext);

import { useSession } from "./useSession";

export const Component = () => {
  const { idToken } = useSession();
  // トークンを利用した処理
};

fetchAuthSession は非同期関数なので、React18.2.0 では useEffect で実行することになりました。React Canary の use APIが安定版で使えるようになればもう少し楽になるかもしれません。

isSignedIn の判定方法は idToken が存在するかどうかで行なっています。ドキュメントにはサインイン可否を判断したいだけであれば useAuthenticator で取得できる authStatus だけ見れば良いと記載されていますが1、それでは「サインイン済みだがトークンの取得がまだ完了していない」という状況が起こってしまい、トークンを利用したい箇所で不都合があったためです。

ID トークンの期限切れ判定を実装

私たちは ID トークンを Ruby on Rails 製バックエンドの GraphQL API の認可にも使用しており、HTTP リクエストの Authorization ヘッダに Bearer トークンとして付与しています。バックエンドでの検証時にトークンが期限切れになってしまっていることを防ぐため、HTTP リクエストを送信する前にトークンの有効期限を確認、失効していればリフレッシュを行いたいです。2
v5 では CognitoUser インスタンスからトークンの検証をするための isValid() が実行できましたが、 v6 では代替するメソッドが提供されなくなりました。そのためトークンの検証処理を自作する必要がありました。

v5

import { useAuthenticator } from "@aws-amplify/ui-react";

export const useAuthExchange = () => {
  const { user } = useAuthenticator();
  const isValid = user.getSignInUserSession()?.isValid();
};

v6

import type { JWT } from "aws-amplify/auth";

/**
 *
 * AmplifyのJWTトークンが有効期限切れかどうかを判定します。
 * 参考:
 * - https://github.com/aws-amplify/amplify-js/blob/0aba4fbcf3f343d92700d4e62e8147ee9c29c880/packages/core/src/singleton/Auth/index.ts#L12-L22
 * - https://github.com/aws-amplify/amplify-js/blob/d0911f5f9108f224627491738c2eb7461af0639c/packages/auth/src/providers/cognito/tokenProvider/TokenOrchestrator.ts#L84-L89
 *
 * ※ clockDrift は取り出せないため考慮していません。
 */
export const isTokenExpired = (token: JWT) => {
  const exp = token.payload.exp;
  if (!exp) return true;

  const now = Date.now() / 1000;
  return exp < now;
};

import { useSession } from "./useSession";
import { isTokenExpired } from "./isTokenExpired";

export const useAuthExchange = () => {
  const { idToken } = useSession(); // 前述で作成したカスタムフック
  const isValid = idToken && !isTokenExpired(idToken);
};

fetchAuthSession の内部で利用されている isTokenExpired を参考に判定関数を自作しました。
コードコメントに記載の通り、v6 では clockDrift にアクセスできないため正確な判定ができていないことを許容しています。clockDrift はサーバーとクライアントの時刻のズレを表す値で、例えば利用者の端末の時間設定が変更されている時にズレを吸収するために利用されます。
clockDrift は amplify-js から取得できなくなったものの、一応 localStorage から取り出すことはできます。ただそれは内部実装の詳細に踏み込みすぎていると判断し今回は省略しています。サーバー側での検証もあるため、この程度の精度で問題ないと判断しました。

まとめ

本記事では AWS Amplify v6 へのマイグレーションに伴う私たちが行なった対応を紹介しました。v6 へのマイグレーションは fetchAuthSession への変更が中心であり、それに伴う変更点を説明しました。
これらの対応はできればライブラリ側で吸収して欲しかったというのが正直な感想ではありますが、マイグレーションをする方の参考になれば幸いです。

  1. https://ui.docs.amplify.aws/react/connected-components/authenticator/advanced#authentication-check↩
  2. 私たちは GraphQL クライアントとして urql を利用しており、これらの処理は urql のミドルウェアである @urql/exchange-auth を利用して実装しています。ここでは本筋から外れるため割愛します。↩