Amazon Cognito User Pools には AWS Lambda 関数を使って認証フローをカスタマイズできる「Lambda トリガー」という機能がある💡トークン生成前トリガー (Pre token generation Lambda trigger) を使うと，認証時に発行されるトークン (ID Token / Access Token) をカスタマイズできる👌

ちなみにトークン (Access Token) は2023年12月のリリースでカスタマイズできるようになっていたりする💡

Powertools for AWS Lambda (Python)

Powertools for AWS Lambda (Python) の Event Source Data Classes を使うと，シンプルに Lambda トリガーを実装できるようになる．今回は PreTokenGenerationTriggerEvent を使ってトークン (ID Token) のペイロードにクレームを追加／削除するトークン生成前トリガーを実装する❗️

docs.powertools.aws.dev

準備

まずは Amazon Cognito User Pools を作っておく❗️今回はサクッと認証を試すためにアプリケーションクライアントの認証フローとして ALLOW_USER_PASSWORD_AUTH も有効化する．あとカスタム属性として X アカウント名を表す custom:x を追加しておく．その他はデフォルト設定で良いかなと思う👌

トークンを取得する（トークン生成前トリガーなし）

まずは「トークン生成前トリガーなし」でトークン (ID Token) を取得する．AWS CLI で aws cognito-idp initiate-auth コマンドを実行する💡

awscli.amazonaws.com

$ CLIENT_ID=xxx
$ EMAIL=xxx
$ PASSWORD=xxx

$ aws cognito-idp initiate-auth \
  --client-id ${CLIENT_ID} \
  --auth-flow USER_PASSWORD_AUTH \
  --auth-parameters USERNAME=${EMAIL},PASSWORD=${PASSWORD} | jq .
{
  "ChallengeParameters": {},
  "AuthenticationResult": {
    "AccessToken": "xxx",
    "ExpiresIn": 3600,
    "TokenType": "Bearer",
    "RefreshToken": "xxx",
    "IdToken": "xxx"
  }
}

取得した IdToken の値から JWT Payload を抜き出すと以下のようになっていた❗️追加したカスタム属性 custom:x にも @kakakakakku という値が入っている👌

{
  "sub": "f7a4da98-a041-70e4-7954-8bbc5d3ca764",
  "email_verified": true,
  "iss": "https://cognito-idp.ap-northeast-1.amazonaws.com/ap-northeast-1_xxxxxxxxx",
  "cognito:username": "f7a4da98-a041-70e4-7954-8bbc5d3ca764",
  "custom:x": "@kakakakakku",
  "origin_jti": "85ed5e84-1346-41fa-8269-1e5aa0a35bab",
  "aud": "xxx",
  "event_id": "525e5170-54f5-450c-a053-258fde4ea50f",
  "token_use": "id",
  "auth_time": 1726876453,
  "exp": 1726880053,
  "iat": 1726876453,
  "jti": "db5441ca-8710-4fe7-b0cc-4b31f6fc77e0",
  "email": "y.yoshida22@gmail.com"
}

トークン (ID Token) の詳細は以下のドキュメントに載っている📝

docs.aws.amazon.com

トークン生成前トリガー

次に AWS Lambda 関数を実装してトークン生成前トリガーを設定する．

AWS Lambda 関数の実装は Powertools for AWS Lambda (Python) の Event Source Data Classes を使って，my_key クレームの追加と custom:x クレームの削除をする．実装自体は簡単にできる👌

👾 app.py（サンプル1）

シンプルに実装する場合

from aws_lambda_powertools.utilities.data_classes.cognito_user_pool_event import PreTokenGenerationTriggerEvent


def lambda_handler(event, context):
    event: PreTokenGenerationTriggerEvent = PreTokenGenerationTriggerEvent(event)
    event.response.claims_override_details.claims_to_add_or_override = {'my_key': 'my_value'}
    event.response.claims_override_details.claims_to_suppress = ['custom:x']

    return event.raw_event

👾 app.py（サンプル2）

ハンドラとロジックを分離してテストをしやすく実装する場合

import json

from aws_lambda_powertools.utilities.data_classes.cognito_user_pool_event import PreTokenGenerationTriggerEvent


def main(event):
    event: PreTokenGenerationTriggerEvent = PreTokenGenerationTriggerEvent(event)
    event.response.claims_override_details.claims_to_add_or_override = {'my_key': 'my_value'}
    event.response.claims_override_details.claims_to_suppress = ['custom:x']

    return event.raw_event


def lambda_handler(event, context):
    return main(event)


if __name__ == '__main__':
    with open('./events/event.json', 'r') as f:
        event = json.load(f)
    print(main(event))

トークンを取得する（トークン生成前トリガーあり）

もう一度 AWS CLI で aws cognito-idp initiate-auth コマンドを実行して，IdToken の値から JWT Payload を抜き出すと以下のようになっていた❗️期待通りに my_key が追加されていて，カスタム属性 custom:x は削除されていた👌

{
  "sub": "f7a4da98-a041-70e4-7954-8bbc5d3ca764",
  "email_verified": true,
  "iss": "https://cognito-idp.ap-northeast-1.amazonaws.com/ap-northeast-1_xxxxxxxxx",
  "cognito:username": "f7a4da98-a041-70e4-7954-8bbc5d3ca764",
  "origin_jti": "888ee4b8-4c3b-4a78-a34c-7e5f04c3788c",
  "aud": "xxx",
  "event_id": "bf25e9fe-beb5-4ad2-b1c7-188eda16ea75",
  "token_use": "id",
  "auth_time": 1726876789,
  "my_key": "my_value",
  "exp": 1726880389,
  "iat": 1726876789,
  "jti": "f75d7eaa-c826-4c40-9951-6dbe8bcc100a",
  "email": "y.yoshida22@gmail.com"
}