AWS Cognito Pre-Token Generation not adding custom claims to ID Token (with ALB setup + Auth Code flow) - aws-lambda

Question

I'm adding custom claims to Cognito's ID token using the "Pre Token Generation" trigger.
Problem
The lambda is triggered, but the issued ID Token doesn't include the claims I added. Am I missing something?
My setup
Using OAuth 2.0 with authorization code flow
My client app sits behind a load balancer (alb). The alb interacts with Cognito to get the Access + ID Tokens in the form of a ALBSessionCookie. Very similar to [0]
To get the ID Token, the client calls a custom endpoint to my backend with the ALBSessionCookie. The backend uses that cookie to return a decoded ID Token to the user. This is the ID Token that I expect should have the custom claims included.
[0] https://www.exampleloadbalancer.com/auth_detail.html
Lambda function (pre-token generation trigger)
Format taken from https://docs.aws.amazon.com/cognito/latest/developerguide/user-pool-lambda-pre-token-generation.html#aws-lambda-triggers-pre-token-generation-example-1
exports.handler = (event, context, callback) => {
event.response = {
"claimsOverrideDetails": {
"claimsToAddOrOverride": {
"my-custom-claims-namespace": JSON.stringify({
"custom-claim-1": "hello",
"custom-claim-2": "hello",
})
}
}
}
callback(null, event)
}
If I can't make this work with ALB, what are some workarounds? Some ideas:
Call Cognito directly for an ID Token (somehow), hoping that will trigger the lambda to issue a JWT with the custom claims
Call Cognito via AmplifyJS

Answer 1

I have a feeling this is expected behavior, though seems like a limitation. Looking here:
https://www.exampleloadbalancer.com/auth_detail.html
We can see that the following steps occur:
ALB receives JWT (ID token, Access Token)
ALB to send access token
ALB receives user info(claims)
I believe the ALB is then not sending the contents of the Decoded ID token (That were manipulated by the Lambda trigger) back to the backend but instead sends the 'user info(claims)' (returned from the UserInfo endpoint) which are not effected by the Cognito trigger.

Answer 2

Yeah the ALB doesn't work that way, the ID Token that Lambda trigger customizes is the one you get when a user Authenticates. There are a couple of options.
Custom User Attributes
The least invasive IMO if instead of adding these attributes in the Lambda trigger, you could have them as custom attributes in Cognito, these I do believe will be in this token. You can sync these attributes at each successful Authorization. That may meet your requirements.
API GW
You could put an API GW either between your LB and your APP or infront of your LB. The API GW does give you a layer in which you can do all this stuff and more with customizing headers, tokens etc. For example you could have a Lambda Authorizer which reads this access token, and returns a context which you can reference in your integration requests back to your backend. It's a bit more involved and will add at least some latency to your app, although you can safely have a large TTL on your auth response because your LB is already doing Auth and you only want some extra attributes. You could also do a re-design and put this all in API GW and get all the bells and whistles it has but you might not need them.
But yeah probably easiest to use the first option if possible as that won't require you to do a redesign and you will just need to change your attribute names to custom:....