Federated User With authorization code fail first time retrieving the token

Hello,

From the case description, I understand that you are connecting Azure B2C as an IDP to your Cognito user pool. You have a Post Confirmation Lambda Trigger[1] in the user pool which links IDP users to your Cognito user pool users and deletes the IDP user. When you are hitting the token endpoint using the code obtained during the first login of IdP user, you are getting 400 error. But in the second attempt (with a new code) you are able to get the tokens.

I replicated this scenario at my end and was able to reproduce the issue. Please allow me to explain what is causing the issue:

During the first login of the IdP user, the Post Confirmation lambda would get triggered which would link the users and would delete the IdP user. As the IdP user gets deleted by the Lambda function during the first login so, the generated code is not associated with any existing user. Hence, if you try to hit the token endpoint to get the tokens using that code then you would get an error as the user is already deleted.

But in the second attempt, as the users (IDP user and native user) are already linked so, the code you receive now is for an existing user. Hence, when you hit the token endpoint you get the required tokens.

In order to avoid this scenario, you can use the PreSignup Lambda Trigger[2] instead of Post Confirmation Lambda Trigger. Shortly before Amazon Cognito signs up a new user, it activates the PreSignup AWS Lambda function. That is, the Pre Signup Lambda gets triggered even before the user is created in the pool so, there is no need to delete the IdP user using the Lambda function. You can just link IdP user to Cognito native user in the Pre Signup Lambda function. In this way, the code generated during the first login would be for an existing linked user. Hence, you would be able to get the token using the code obtained during first sign-in as well. I tried implementing this workaround at my end and was able to get the tokens for the IdP user using the code for each login attempt.

Below is a sample code that I used in the Pre Signup Lambda function at my end:

client = boto3.client('cognito-idp')

response = client.admin_link_provider_for_user(
UserPoolId='<your user pool id>',
DestinationUser={
    'ProviderName': 'Cognito',
    'ProviderAttributeValue': '<username>'
},
SourceUser={
    'ProviderName': '<IDP provider name>',
    'ProviderAttributeName': 'Cognito_Subject',
    'ProviderAttributeValue': '<user-id of idp>'
}
)

return response

[1]. https://docs.aws.amazon.com/cognito/latest/developerguide/user-pool-lambda-post-confirmation.html

[2]. https://docs.aws.amazon.com/cognito/latest/developerguide/user-pool-lambda-pre-sign-up.html