[アップデート] Amazon Cognito のトークン生成前 Lambda トリガーのトリガーイベントバージョン V2 が使えるようになっていたので試してみた

いわさです。

Amazon Cognito では認証実行後に ID トークンとアクセストークンを取得することが出来ます。
このトークンクレームをカスタマイズしたいことがあるのですが、ID トークンについては以前からクレームの上書き・追加・抑制を行うことが出来ていました。

実は最近（2023 年 12 月 と 2024 年 5 月）になってこれらのトークンカスタマイズ機能でより高度なカスタマイズが行えるようになっています。

具体的には 2023 年 12 月には、従来の ID トークンのみだけでなくアクセストークンもカスタマイズ出来るようになっています。

そして 2024 年 5 月にはより複雑なデータ形式を設定することも出来るようになりました。

公式ドキュメントやマネジメントコンソール上では従来のカスタマイズを V1 あるいは基本、高度なカスタマイズを V2 なんて呼ばれていたりします。

高度なカスタマイズいいじゃないかとなりそうですが、一点前提条件があってこれらのトークンカスタマイズ機能 V2 を使うためには対象のユーザープールでアドバンスドセキュリティ機能を有効化する必要があります。
アドバンスドセキュリティを有効化すると Cognito ユーザー単価が変わるのでそこだけ注意が必要になっています。

今回私のほうでこれらのアップデートを試す機会がありまして、従来の V1 と V2 の違いを比較しながら簡単に検証してみましたので紹介したいと思います。

まずは普通に認証してトークンを取得してみる

まずは、特に Lambda トリガーを設定せずにトークンを取得してみます。
前提としてユーザプールにユーザーを作成し、初期パスワードから更新済みの状態でサインインに成功する状態です。

% aws cognito-idp initiate-auth --cli-input-json file://hoge.json
{
    "ChallengeParameters": {},
    "AuthenticationResult": {
        "AccessToken": "eyJraWQiOiJFa...v3DHBS-YsQ",
        "ExpiresIn": 3600,
        "TokenType": "Bearer",
        "RefreshToken": "eyJjdHki...DeX0c4Q",
        "IdToken": "eyJraWQiO...StAaXQ"
    }
}

ID トークンが取得出来ていますね。
jwt.io などでトークン検証を行いつつペイロード部分だけ抜きだしてみます。

{
  "sub": "47c42a58-1031-704d-38ff-c03ea785ca6a",
  "email_verified": true,
  "iss": "https://cognito-idp.ap-northeast-1.amazonaws.com/ap-northeast-1_KpV8GsQHn",
  "cognito:username": "47c42a58-1031-704d-38ff-c03ea785ca6a",
  "origin_jti": "7f0ab427-ddc0-497f-bcee-415b1defe5a9",
  "aud": "1gdidp3hgjlntu1dv6k5hn2uvb",
  "event_id": "52033edc-f439-4edc-a327-d8b9cc2855c8",
  "token_use": "id",
  "auth_time": 1718659347,
  "exp": 1718662947,
  "iat": 1718659347,
  "jti": "b5c1e424-c1e9-4b6e-8e3f-780695a8fd69",
  "email": "iwasa.takahito+hoge0618@example.com"
}

なるほど。これがデフォルトの ID トークンですね。

基本機能を使って ID トークンをカスタマイズしてみる

冒頭の北野さんの記事、あるいは次の公式ドキュメントを参考に、Lambda トリガーで使う関数を用意して、ユーザープールの Lambda トリガーを有効化してみましょう。

関数は公式ドキュメントと同様に次のような感じで実装しました。
claimOverrideDetailsで上書きするクレームと抑制するクレームを指定する感じですね。

で、これをユーザープールのトークン生成前 Lambda トリガーとして設定します。
非活性になってますがトリガーイベントバージョンに「基本機能 + アクセストークンのカスタマイズ」という選択肢があることがわかりますね。これは後で試すやつです。

上記 Lambda トリガーを設定した状態で再びinitiate-authコマンドでトークンを取得してみます。

{
  "sub": "47c42a58-1031-704d-38ff-c03ea785ca6a",
  "email_verified": true,
  "my_first_attribute": "first_value",
  "iss": "https://cognito-idp.ap-northeast-1.amazonaws.com/ap-northeast-1_KpV8GsQHn",
  "cognito:username": "47c42a58-1031-704d-38ff-c03ea785ca6a",
  "my_second_attribute": "second_value",
  "origin_jti": "d7351048-7b71-4bcd-972b-bd157cedb10a",
  "aud": "1gdidp3hgjlntu1dv6k5hn2uvb",
  "event_id": "cfbc8e5e-31b8-4b12-8e3d-a18fa65c836b",
  "token_use": "id",
  "auth_time": 1718659880,
  "exp": 1718663480,
  "iat": 1718659880,
  "jti": "4c7e0e65-adc1-4cd2-900a-fc8363420131"
}

おぉなるほど。
my_first_attributeとmy_second_attributeが追加され、emailが出力されなくなりましたね。
Lambda トリガーによって期待どおりにカスタマイズされています。

文字列のみ設定できる

アップデートアナウンスを見てみると、今回の V2 では文字列だけではなくブール値や数値も設定出来るようになっているようです。
ということは V1 ではブールや文字列は設定出来ないということか。試してみましょう。

const handler = async (event) => {
  event.response = {
    claimsOverrideDetails: {
      claimsToAddOrOverride: {
        my_first_attribute: true,
        my_second_attribute: 111,
      },
      claimsToSuppress: ["email"],
    }
  };

  return event;
};

export { handler };

実行してみます。

{
  "sub": "47c42a58-1031-704d-38ff-c03ea785ca6a",
  "email_verified": true,
  "my_first_attribute": "true",
  "iss": "https://cognito-idp.ap-northeast-1.amazonaws.com/ap-northeast-1_KpV8GsQHn",
  "cognito:username": "47c42a58-1031-704d-38ff-c03ea785ca6a",
  "my_second_attribute": "111",
  "origin_jti": "00a342f1-acdf-4378-85bc-6877f5070134",
  "aud": "1gdidp3hgjlntu1dv6k5hn2uvb",
  "event_id": "0e9325a0-2995-4f51-848f-460f5b286cf5",
  "token_use": "id",
  "auth_time": 1718661092,
  "exp": 1718664692,
  "iat": 1718661092,
  "jti": "5c1ab8f9-de53-4603-9f8c-84dba1afa442"
}

なるほど。カスタマイズは出来ていますが、ブール値や数値を指定した部分が文字列として設定されていますね。

V2 (基本機能 + アクセストークンのカスタマイズ) を使ってみる

では続いて V2 を使ってみましょう。
まずは Lambda トリガーの設定をバージョンアップしたいのですが、前提として Cognito アドバンスドセキュリティ機能を有効化する必要があるのでそちらを先に有効化しておきましょう。

そうすると Lambda トリガーで「基本機能 + アクセストークンのカスタマイズ」が選択出来るようになりました。
Lambda トリガーの編集からトリガーイベントバージョンを更新します。

現在使われているトリガーイベントバージョンは以下からも確認が可能ですね。

さて、Lambda 関数は同じものを使うとして、コードの中身だけ修正してみましょう。

const handler = async (event) => {
  event.response = {
    // claimsOverrideDetails: {
    //   claimsToAddOrOverride: {
    //     my_first_attribute: "first_value",
    //     my_second_attribute: "second_value",
    //   },
    //   claimsToSuppress: ["email"],
    // },
    claimsAndScopeOverrideDetails: {
      idTokenGeneration: {
        claimsToAddOrOverride: {
          hoge1: "hoge1-id",
          hoge2: "hoge2-id",
        },
        claimsToSuppress: ["email"]
      },
      accessTokenGeneration: {
        claimsToAddOrOverride: {
          hoge1: "hoge1-access",
          hoge2: "hoge2-access",
        }
      }
    }
  };

  return event;
};

export { handler };

上記のコメント箇所は V1 で使えるコードで、V2 トリガーの場合は上記部分は動作しません。無視されます。
で、V2 の場合はまずclaimsAndScopeOverrideDetailsがあり、その中に ID トークンカスタマイズ設定内容であるidTokenGenerationと、アクセストークンカスタマイズ設定内容であるaccessTokenGenerationが存在しています。
V1 ではアクセストークンには触れなかったので、必要なデータ構造も変わっているわけですね。

上記コードでは ID トークンに 2 つのクレームを追加しつつ email をまた抑制しています。
そしてアクセストークンについても 2 つのカスタムクレームを追加しようとしています。

認証してみましょう。生成された ID トークンがこちらになります。

{
  "hoge2": "hoge2-id",
  "hoge1": "hoge1-id",
  "sub": "47c42a58-1031-704d-38ff-c03ea785ca6a",
  "email_verified": true,
  "iss": "https://cognito-idp.ap-northeast-1.amazonaws.com/ap-northeast-1_KpV8GsQHn",
  "cognito:username": "47c42a58-1031-704d-38ff-c03ea785ca6a",
  "origin_jti": "f074dbb8-4c66-40d9-b15a-48f656059f79",
  "aud": "1gdidp3hgjlntu1dv6k5hn2uvb",
  "event_id": "52fb404b-2d7b-4767-a5c7-04bd3369a2dd",
  "token_use": "id",
  "auth_time": 1718660612,
  "exp": 1718664212,
  "iat": 1718660612,
  "jti": "d148bf91-e9b7-43a0-aad4-ae6f93e14809"
}

おー、良いですね。動作しています。
アクセストークンも見てみましょう。

{
  "hoge2": "hoge2-access",
  "hoge1": "hoge1-access",
  "sub": "47c42a58-1031-704d-38ff-c03ea785ca6a",
  "iss": "https://cognito-idp.ap-northeast-1.amazonaws.com/ap-northeast-1_KpV8GsQHn",
  "client_id": "1gdidp3hgjlntu1dv6k5hn2uvb",
  "origin_jti": "f074dbb8-4c66-40d9-b15a-48f656059f79",
  "event_id": "52fb404b-2d7b-4767-a5c7-04bd3369a2dd",
  "token_use": "access",
  "scope": "aws.cognito.signin.user.admin",
  "auth_time": 1718660612,
  "exp": 1718664212,
  "iat": 1718660612,
  "jti": "8a188f8d-e3d2-4c0d-a024-009b0b638267",
  "username": "47c42a58-1031-704d-38ff-c03ea785ca6a"
}

こちらも良いですね！従来は出来なかったアクセストークンカスタマイズが出来るようになっています。
今回はクレームの追加と抑制だけ行っていますが、公式ドキュメントを見てみると、スコープの追加や抑制も新たに出来るようになっています。
各属性値の詳細な設定方法は公式ドキュメントのコード例も参考にしてください。

複雑なオブジェクトや文字列以外を設定してみる

V2 ではアクセストークンのカスタマイズに加えて文字列以外のデータ型も設定出来るようになっています。
今回は ID トークンに JSON 型、ブール型、数値型を指定してみましょう。コードは次のような感じになりました。

const handler = async (event) => {
  event.response = {
    claimsAndScopeOverrideDetails: {
      idTokenGeneration: {
        claimsToAddOrOverride: {
          hoge1: "hoge1-id",
          hoge2: "hoge2-id",
          hoge3: {
            hoge3a: {
              hoge3b1: "hogehoge",
              hoge3b2: "fugafuga"
            }
          },
          hoge4: true,
          hoge5: 111111
        },
        claimsToSuppress: ["email"]
      }
    }
  };

  return event;
};

export { handler };

トークンペイロードを確認してみましょう。

{
  "hoge2": "hoge2-id",
  "hoge1": "hoge1-id",
  "sub": "47c42a58-1031-704d-38ff-c03ea785ca6a",
  "hoge4": true,
  "hoge3": {
    "hoge3a": {
      "hoge3b1": "hogehoge",
      "hoge3b2": "fugafuga"
    }
  },
  "email_verified": true,
  "iss": "https://cognito-idp.ap-northeast-1.amazonaws.com/ap-northeast-1_KpV8GsQHn",
  "cognito:username": "47c42a58-1031-704d-38ff-c03ea785ca6a",
  "origin_jti": "d4daed5e-b680-4010-b6fd-e6890e1d779f",
  "aud": "1gdidp3hgjlntu1dv6k5hn2uvb",
  "event_id": "4f07fe96-f21b-4d9e-bca6-4e9735848831",
  "token_use": "id",
  "auth_time": 1718660922,
  "exp": 1718664521,
  "hoge5": 111111,
  "iat": 1718660922,
  "jti": "30f0a123-f0b3-4908-a6bd-817d987ba411"
}

おー、良いですね。
V1 では文字列として設定されてしまっていましたが複雑なデータ型や文字列以外の値も表現出来ていますね。

アドバンスドセキュリティを無効化する前にトリガーイベントバージョンの更新が必要

今回のトリガーイベント V2 はアドバンスドセキュリティが有効な場合のみ使うことが出来ます。
したがって、V2 の Lambda トリガーを設定している状態でアドバンスドセキュリティを無効化しようとすると次のようなダイアログが表示され、先にトリガーバージョンの変更が必要になりますのでご注意ください。

さいごに

本日は Amazon Cognito のトークン生成前 Lambda トリガーのトリガーイベントバージョン V2 が使えるようになっていたので試してみました。

アドバンスドセキュリティが必要ということで、導入前の場合は追加料金が発生しますが、アクセストークンのカスタマイズや基本機能よりもより柔軟なデータ型を設定した場合には選択肢になってきます。
まずはカスタマイズ出来るようになるということは知っておくと良いですね。