Implementing both Password-less Email OTP based Sign In and Authenticator App/SOFTWARE_TOKEN_MFA based MFA in AWS Cognito and Amplify

Question

We are trying to implement MultiFactor Authentication using Authenticator App for a custom auth flow using passwordless email authentication. We are currently using the password-less email authentication flow from aws blog "Implementing passwordless email authentication with Amazon Cognito".

We are required to provide MFA to our users using an Authenticator App. but are unable to implement SOFTWARE_TOKEN_MFA on top of this CUSTOM_CHALLENGE.

We have successfully implemented the enabling MFA steps and now MFA is enabled for the user and the preferredMFA is set to "SOFTWARE_TOKEN_MFA". We made few modifications to the code in "Define Auth Challenge" function so that if a user has enabled MFA then upon successfully completing the "CUSTOM_CHALLENGE" the users challengeName is set to "SOFTWARE_TOKEN_MFA".

We were expecting that after the CUSTOM_CHALLENGE is successfully completed and the challengeName is changed to "SOFTWARE_TOKEN_MFA" we can call Auth.confirmSignIn function to validate the code from the authenticator app, but this throws "CodeMismatchException: Invalid code or auth state for the user.". We have throughly checked that code entered is correct so the issue is in the auth state.

Are we doing this right ?

Modified code of DefineAuthChallenge

/**
 * @type {import('@types/aws-lambda').DefineAuthChallengeTriggerHandler}
 */

const getUserPreferredMFA = require('./functions/getUserPreferredMFA');

//define-challenge
exports.handler = async event => {

  try {
    console.log("event: ", JSON.stringify(event, null, 2));
  }
  catch (err) {
    console.log("error: ", err);
  }

  if (event.request.session &&
    event.request.session.find(attempt => attempt.challengeName !== 'CUSTOM_CHALLENGE')) {
    // We only accept custom challenges; fail auth
    event.response.issueTokens = false;
    event.response.failAuthentication = true;
  }
  else if (event.request.session &&
    event.request.session.length >= 3 &&
    event.request.session.slice(-1)[0].challengeResult === false) {
    // The user provided a wrong answer 3 times; fail auth
    event.response.issueTokens = false;
    event.response.failAuthentication = true;
  }
  else if (event.request.session &&
    event.request.session.length &&
    event.request.session.slice(-1)[0].challengeName === 'CUSTOM_CHALLENGE' && // Doubly stitched, holds better
    event.request.session.slice(-1)[0].challengeResult === true) {

    //check if user has enabled MFA by checking preferredMFA is "SOFTWARE_TOKEN_MFA"
    if (await getUserPreferredMFA(event.userPoolId, event.userName) === "SOFTWARE_TOKEN_MFA") {
      event.response.challengeName = 'SOFTWARE_TOKEN_MFA';
      event.response.issueTokens = false;
      event.response.failAuthentication = false;
    }
    else {
      event.response.issueTokens = true;
      event.response.failAuthentication = false;
    }
  }
  else {
    // The user did not provide a correct answer yet; present challenge
    event.response.issueTokens = false;
    event.response.failAuthentication = false;
    event.response.challengeName = 'CUSTOM_CHALLENGE';
  }
  return event;
};

The code for CreateAuthChallenge and Verify AuthChallenge is same as in the AWS Blog.

To Sum up what we are trying to do

1. User enters Email Address.
2. Auth.signIn function is called.
3. An email with login code is sent to users email address and cognito/amplify returns the following:

{
    "challengeName": "CUSTOM_CHALLENGE",
    "challengeParam": {
        "USERNAME": "-------------------",
        "email": "user@example.com"
    },
    "signInUserSession": null,
    "authenticationFlowType": "CUSTOM_AUTH" 
    ....
}

4. User enter the code in FE and calls Auth.sendCustomChallengeAnswer.
5. The code is validated by verifyAuthChallenge and passes the result to DefineAuthChallenge which will check if user has setup MFA and set preferredMFA as "SOFTWARE_TOKEN_MFA". if yes sets the challengeName to "SOFTWARE_TOKEN_MFA". This returns the following

{
    "challengeName": "SOFTWARE_TOKEN_MFA",
    "challengeParam": {
        "USERNAME": "-------------------",
        "email": "user@example.com"
    },
    "signInUserSession": null,
    "authenticationFlowType": "CUSTOM_AUTH" .....
} 

6. User enters the code from the authenticator app and calls Auth.confirmSignIn(user, mfaCode, "SOFTWARE_TOKEN_MFA").
7. This call fails and returns "CodeMismatchException: Invalid code or auth state for the user."