Claude CodeでCognitoのカスタムクレームを書く — pre token generationを単体テスト

Amazon Cognitoのユーザープールが発行するトークンに、テナントIDや権限スコープを足したい。そんな場面で使うのが、pre token generation Lambdaトリガーです。トークンを発行する直前に関数が呼ばれ、クレームの追加・上書き・抑止ができます。

この記事では、そのLambdaをClaude Codeに書かせ、モックイベントで単体テストするまでを進めます。イベントの入出力形式はバージョンごとに違い、間違えると応答が反映されません。形式を先に渡し、テストで挙動を固定するのが要点です。

先に決めるのはイベントのバージョン

pre token generationトリガーのイベントには3つのバージョンがあります。どれを選ぶかで、応答の形も、触れるトークンも変わります。

ユーザープールは既定でV1_0を送ります。V2_0以降を受けるには、コンソールのトリガー設定で「Basic features + access token customization」を選ぶか、LambdaVersionに値を入れます。

aws cognito-idp update-user-pool --user-pool-id us-east-1_EXAMPLE \
  --lambda-config \
  'PreTokenGenerationConfig={LambdaArn=<関数のARN>,LambdaVersion=V2_0}'

注意点があります。update-user-poolは、指定しなかった設定を既定値に戻す更新です。既存のユーザープールに実行するときは、現在の設定を取得してから全項目を渡してください。Claude Codeに任せる場合も、この実行は人が確認してから流すのが安全です。

特定の関数バージョンを呼びたいときは、LambdaArnにバージョン付きの関数ARNを指定します。エイリアスや$LATESTを使い分ける運用をしているなら、テスト用と本番用のユーザープールで指すバージョンを分けておくと、検証中の関数が本番のトークンに混ざりません。

V2_0とV3_0の違いは、M2Mのアクセストークンに効くかどうかだけです。V3_0のTokenGeneration_ClientCredentialsイベントは、V3_0を選んだユーザープールにしか送られません。

利用できるプランにも条件があります。V1_0からV3_0のイベントは、EssentialsとPlusの機能プランで使えます。2024年11月22日18時(GMT)までに高度なセキュリティ機能を有効にしていたLiteプランのユーザープールはV1_0とV2_0、高度なセキュリティ機能がないLiteプランはV1_0のみです。V3_0はEssentialsとPlusだけです。

Claude Codeに渡す材料と指示の置き方

Claude Codeは、公式ドキュメントの形式を知っていても、バージョンの取り違えを起こしえます。V1_0の応答形を書いたままLambdaVersionをV2_0にする、といったずれです。防ぐには、動かしたいバージョンと形式をリポジトリに置いておきます。

CLAUDE.mdには、次のような断片を足します。

## Cognito pre token generation
- トリガーのイベントバージョンは V2_0 を使う
- 応答は event.response.claimsAndScopeOverrideDetails に入れる
  (idTokenGeneration / accessTokenGeneration / groupOverrideDetails)
- ハンドラーは event を書き換えて return する(新しいオブジェクトを返さない)
- 変更したら events/ のモックイベントで node --test を通す

そのうえで、公式ページの「request」と「response」のJSONをdocs/cognito-event-v2.mdなどに貼り、プロンプトから参照させます。

docs/cognito-event-v2.md の形式で pre token generation の Lambda を書いて。
custom:tenant_id を ID トークンとアクセストークンに tenant_id として足し、
admins グループの人にだけ MyAPI.admin スコープを足す。
先にテスト(node:test)を書いて、失敗を確認してから実装して。

Claude Codeの公式ガイドは、テストの足し方を「テストされていない関数の特定、足場の生成、境界条件の追加、実行して失敗を直す」の流れで示しています。テストを先に書かせると、イベントの形を取り違えたときに、実装の前に気づけます。

変更が大きいと感じるときは、プランモードで始める手もあります。claude --permission-mode planで起動すると、Claude Codeは読み取りと計画だけを行い、承認するまでファイルを書き換えません。

テストを先に置くと何が固まるか

Lambdaに載せる前に、ローカルでnode --testを回せるようにします。モックイベントは、公式ページのV2_0向けテストイベントを土台にします。triggerSource、callerContext.clientId、request.userAttributes、groupConfiguration、scopesを揃えた形です。

// index.test.mjs(抜粋)
const baseEvent = (overrides = {}) => ({
  version: "2",
  triggerSource: "TokenGeneration_Authentication",
  region: "us-east-1",
  userPoolId: "us-east-1_EXAMPLE",
  userName: "JaneDoe",
  callerContext: { clientId: "1example23456789" },
  request: {
    userAttributes: {
      sub: "a1b2c3d4-5678-90ab-cdef-EXAMPLE11111",
      email: "jane@example.com",
      phone_number: "+12065551212",
      "custom:tenant_id": "t-001",
    },
    groupConfiguration: {
      groupsToOverride: ["admins"],
      iamRolesToOverride: [],
      preferredRole: [],
    },
    scopes: ["openid"],
    ...overrides,
  },
  response: { claimsAndScopeOverrideDetails: [] },
});

テストでは、次の4点を固めました。テナントIDが両方のトークンに入る、adminsにだけスコープが付く、custom:tenant_idが欠けても落ちない、グループ構成が壊れない、です。

test("admins グループにだけ MyAPI.admin スコープを足す", async () => {
  const admin = await handler(baseEvent());
  assert.deepEqual(
    admin.response.claimsAndScopeOverrideDetails
      .accessTokenGeneration.scopesToAdd,
    ["MyAPI.admin"],
  );
  const user = await handler(
    baseEvent({ groupConfiguration: { groupsToOverride: [] } }),
  );
  assert.deepEqual(
    user.response.claimsAndScopeOverrideDetails
      .accessTokenGeneration.scopesToAdd,
    [],
  );
});

ハンドラーの形と、グループを消さない一行

テストを通すハンドラーは次のとおりです。公式のV1_0のサンプルと同じく、event.responseを組み立ててeventを返す形にしています。

export const handler = async (event) => {
  const attrs = event.request.userAttributes ?? {};
  const groups = event.request.groupConfiguration?.groupsToOverride ?? [];
 
  const claims = {
    tenant_id: attrs["custom:tenant_id"] ?? "unknown",
    is_admin: groups.includes("admins"),
  };
 
  event.response = {
    claimsAndScopeOverrideDetails: {
      idTokenGeneration: {
        claimsToAddOrOverride: claims,
        claimsToSuppress: ["phone_number"],
      },
      accessTokenGeneration: {
        claimsToAddOrOverride: claims,
        scopesToAdd: groups.includes("admins") ? ["MyAPI.admin"] : [],
      },
      groupOverrideDetails: event.request.groupConfiguration,
    },
  };
  return event;
};

groupOverrideDetailsの一行が、見落とされやすい箇所です。この項目に空のオブジェクトやnullを返すと、Cognitoはグループを抑止します。既存の構成を保つには、リクエストのgroupConfigurationをそのままgroupOverrideDetailsに入れて返します。クレームを足すだけのつもりで応答を作り直すと、cognito:groupsが消える可能性があります。

手元のNode.js v20.18.0で、この4つのテストを実行した結果は次のとおりです。

$ npm test
...
1..4
# tests 4
# pass 4
# fail 0

テストは、ハンドラーが返すオブジェクトの形を確かめる単体テストです。Cognito本体がその応答をどうトークンに反映するかまでは検証できません。最終確認は、開発用のユーザープールに関数を関連付け、実際にサインインして、IDトークンとアクセストークンをデコードして行います。

公式の制限表を、テストに写しておく

クレームには、足せるもの、上書きできるもの、抑止できるものの制限があります。Claude Codeに「好きなクレームを足して」と頼むと、効かないクレームを書いてしまうことがあるため、制限をテストに写すと安全です。

アクセストークンのaudは、足せる場合でも値がリクエストのクライアントIDと一致しなければなりません。値はイベントのcallerContext.clientIdから取れます。scopesToAddには、空白を含むスコープを入れられません。

cognito:preferred_roleとcognito:rolesを抑止したいときは、claimsToSuppressにcognito:groupsを足す必要があります。同じクレームを追加と抑止の両方に指定した場合は、抑止が勝ちます。

スコープの抑止にも条件があります。scopesToSuppressにaws.cognito.signin.user.adminを入れると、ユーザーがCognitoのユーザープールAPIで自分のプロフィールを読み書きする操作を止められます。phone_numberスコープの削除で電話番号の取得を防げるのは、残りのスコープにopenidと標準スコープが1つ以上あるときだけです。テストでは、スコープ配列の全体を比べておくと、この条件の崩れに気づけます。

見落としやすい条件

トリガーが呼ばれる契機はtriggerSourceで分かれます。マネージドログイン経由のTokenGeneration_HostedAuth、認証完了後のTokenGeneration_Authentication、管理者作成ユーザーの初回パスワード変更後のTokenGeneration_NewPasswordChallenge、デバイス認証のTokenGeneration_AuthenticateDevice、トークン更新のTokenGeneration_RefreshTokens、そしてV3_0だけのTokenGeneration_ClientCredentialsです。

リフレッシュトークンで更新するときも関数は呼ばれます。計算が重い処理やネットワーク呼び出しを置くと、更新のたびに走ります。外部のAPIから属性を引く関数を書かせるときは、タイムアウトと失敗時の挙動をテストに含めてください。

clientMetadataにも条件があります。ClientMetadataの値は、RespondToAuthChallengeとAdminRespondToAuthChallengeの呼び出しでは関数に届きますが、InitiateAuthとAdminInitiateAuthでは届きません。呼び出しの経路によってテスト用のイベントが変わるため、clientMetadataを使う分岐は経路別に分けてモックを用意すると迷いません。

IDトークンの複雑なオブジェクトは、phone_number_verified、email_verified、updated_at、addressのクレームには入れられません。

V1_0の既存関数をV2_0へ移すとき

すでにV1_0で動く関数があるなら、Claude Codeに「応答の形だけを移す」改修として頼みます。V1_0ではclaimsOverrideDetails直下にあったclaimsToAddOrOverrideとclaimsToSuppressが、V2_0ではIDトークン用のidTokenGenerationの下に移ります。アクセストークンに触れる項目はaccessTokenGenerationの下に新しく置きます。groupOverrideDetailsは、コンテナの直下に残ります。

// V1_0(IDトークンだけ)
event.response = {
  claimsOverrideDetails: {
    claimsToAddOrOverride: { tenant_id: "t-001" },
    claimsToSuppress: ["phone_number"],
  },
};
 
// V2_0(IDトークンとアクセストークン)
event.response = {
  claimsAndScopeOverrideDetails: {
    idTokenGeneration: {
      claimsToAddOrOverride: { tenant_id: "t-001" },
      claimsToSuppress: ["phone_number"],
    },
    accessTokenGeneration: {
      claimsToAddOrOverride: { tenant_id: "t-001" },
    },
  },
};

IDトークンとアクセストークンの両方を変えたいときは、関数をV2_0以降に更新し、同じ応答で両方のデータを返す必要があります。移行の差分は、テストの期待値を先に書き換えて失敗させ、そのあとで実装を直す順にすると、取りこぼしを減らせます。

V3_0でM2Mのトークンも変えるなら、クライアント認証情報を使う呼び出し側にも注意点があります。クライアントメタデータは、トークンエンドポイントへのPOSTボディにaws_client_metadataとしてURLエンコードしたJSON文字列を載せます。関数にはclientMetadataとして届くので、モックイベントでも同じキーで用意しておきます。

あわせて、デバッグ用に応答を丸ごとログへ出す場合は、応答のサイズにも気を配ってください。公式のサンプルは、応答を文字列にした長さまでログに出しています。Claude Codeに関数を書かせるときは、本番で個人情報を含むuserAttributesをそのままログに出さないよう、指示に一行足しておくと安全です。

CIでモックイベントを回す

モックイベントをファイルに置いておけば、claude -pでイベントを流し込む使い方もできます。公式ガイドは、標準入力と標準出力をUnixツールと同じに扱えると説明しています。

cat events/hosted-auth.json | claude -p \
  "このイベントを handler に通したときの応答を要約して"

ただし、これは挙動の要約を読むだけで、合否の判定ではありません。判定はnode --testに任せ、モックイベントの追加(新しいtriggerSourceへの対応など)だけをClaude Codeに頼む分担が扱いやすい形です。

Lambda側の運用には、別の記事が手がかりになります。メモリ設定はClaude CodeでLambdaのメモリを最適化、本番でのエラー急増の調べ方はClaudeにCloudWatch Logsのクエリを書かせてLambdaのエラー急増を調べるで扱っています。

まとめ

pre token generationのLambdaは、イベントのバージョンと応答の形を取り違えないことが全てです。CLAUDE.mdにバージョンと応答のコンテナ名を置き、公式のリクエスト形式をプロンプトから参照させ、モックイベントのテストを先に書かせます。そうすれば、groupOverrideDetailsの取りこぼしのような、トークンを発行して初めて気づく不具合を減らせます。

関数の関連付け(update-user-pool)と、実際のサインインでのトークン確認は、人が握る部分です。テストが通ったことと、本番のトークンが期待どおりであることは別の確認として扱ってください。