Cognito Workshop > 8. Setup Cognito Extension > 8.1. Build Lambda pre-generation tokens trigger

8.1. Build Lambda pre-generation tokens trigger

Build Lambda Trigger

We will build a Lambda trigger for both TypeScript and Python. This trigger will update claims using the claimsAndScopeOverrideDetails property in event.response. Specifically:

idTokenGeneration.claimsToAddOrOverride: update ID Token claims
accessTokenGeneration.claimsToAddOrOverride: update Access Token claims

TypeScript

In your local project, go to src/runtimes/lambda-functions.
Create a folder authorizer and inside it, create pre-token-generation-trigger.ts.
Add your code in this file.

/**
 * Hàm này sẽ dùng để tổng hợp thông tin và tạo claims của user
 * và bỏ vào trong token trước khi tạo token. Hoạt động với cả
 * khi sign in và refresh tokens.
 *
 * @param event - lambda event
 * @param context - lambda context
 *
 * @returns
 */
export async function handler(event: any, context: any) {
  // Lấy thông tin từ userAttributes
  const userAttrs = event.request.userAttributes;

  const role = userAttrs["custom:role"] || "";
  const fullName = userAttrs["name"];

  // Gán lại response với claims và scopes
  event.response = {
    claimsAndScopeOverrideDetails: {
      idTokenGeneration: {
        claimsToAddOrOverride: {
          "custom:role": role,
          full_name: fullName,
        },
      },
      accessTokenGeneration: {
        claimsToAddOrOverride: {
          "custom:role": role,
        },
        scopesToAdd: role ? [`role:${role}`] : [],
      },
    },
  };

  return event;
}

8.1.1

Python

In your local project, go to src/runtimes/lambda_functions.
Create a folder authorizer and inside it, create pre_token_generation_trigger.py.
Add your code in this file.

def handler(event, context):
    """
    Hàm này sẽ dùng để tổng hợp thông tin và tạo claims của user
    và bỏ vào trong token trước khi tạo token. Hoạt động với cả
    khi sign in và refresh tokens.

    :param event: lambda event
    :param context: lambda context
    :return: event với response được gán thêm claims
    """

    # Lấy thông tin từ userAttributes
    user_attrs = event["request"]["userAttributes"]

    role = user_attrs.get("custom:role", "")
    full_name = user_attrs.get("name", "")

    # Gán lại response với claims và scopes
    event["response"] = {
        "claimsAndScopeOverrideDetails": {
            "idTokenGeneration": {
                "claimsToAddOrOverride": {
                    "custom:role": role,
                    "full_name": full_name,
                },
            },
            "accessTokenGeneration": {
                "claimsToAddOrOverride": {
                    "custom:role": role,
                },
                "scopesToAdd": [f"role:{role}"] if role else [],
            },
        },
    }

    return event