How to forward Cognito custom claims during token refresh?

Yes, you can use the Cognito "pre token generation lambda trigger" in combination with client metadata to add custom claims to identity or access tokens. However, you're correct that when the tokens are refreshed, the client metadata is no longer available, which makes it challenging to forward the custom claims to the new tokens. One potential solution is to store the custom claims as user attributes in the Cognito User Pool when you first issue the tokens. Then, during the token refresh process, the pre-token generation Lambda trigger can retrieve these custom claims from the user attributes and add them to the new tokens.

Here's a high-level overview of the steps:

1. When a user first authenticates and you issue the initial tokens, use the pre-token generation Lambda trigger to add the custom claims based on the client metadata.
2. In the same pre-token generation Lambda trigger, store these custom claims as user attributes in the Cognito User Pool using the adminUpdateUserAttributes API.
3. During the token refresh process, the pre-token generation Lambda trigger is invoked again. In this trigger, you can retrieve the custom claims from the user attributes using the adminGetUser API.
4. Add the retrieved custom claims to the new tokens being issued during the refresh process.

Here's some sample code in Node.js to illustrate this approach:

// Initial token issuance
exports.handler = async (event) => {
  if (event.request.clientMetadata) {
    const customClaims = getCustomClaimsFromClientMetadata(event.request.clientMetadata);

    // Add custom claims to tokens
    event.response = {
      ...event.response,
      claimsOverrideDetails: {
        claimsToAddOrOverride: customClaims
      }
    };

    // Store custom claims as user attributes
    const userAttributes = Object.entries(customClaims).map(([key, value]) => ({
      Name: key,
      Value: value
    }));
    await adminUpdateUserAttributes(event.userName, userAttributes);
  }

  return event;
};

// Token refresh
exports.handler = async (event) => {
  const user = await adminGetUser(event.userName);
  const customClaims = user.UserAttributes.reduce((claims, attr) => {
    claims[attr.Name] = attr.Value;
    return claims;
  }, {});

  event.response = {
    ...event.response,
    claimsOverrideDetails: {
      claimsToAddOrOverride: customClaims
    }
  };

  return event;
};

This approach has a few caveats:

• It assumes that the custom claims can be stored as user attributes (within the attribute name and value limits).
• It requires additional logic to handle user attribute updates if the custom claims change over time.
• It may not work if the user is deleted and recreated, as the user attributes would be lost.

Another potential solution is to store the client metadata or a reference to it in a separate data store (like DynamoDB) during the initial token issuance. Then, during the refresh process, you can retrieve this metadata from the data store and use it to add the custom claims to the new tokens. Ultimately, the best approach depends on your specific requirements and the nature of the custom claims you're adding to the tokens.

Useful resources:

1. https://docs.aws.amazon.com/cognito/latest/developerguide/user-pool-lambda-pre-token-generation.html
2. https://aws.amazon.com/blogs/security/how-to-customize-access-tokens-in-amazon-cognito-user-pools/