How do I use InitiateAuth command (@aws-sdk/client-cognito-identity-provider) to trigger SRP Authentication followed by CUSTOM_CHALLANGE?

Question

Desired Flow:

1. Begin with Cognito SRP flow to verify user's username and password combination
2. If Username and password are correct, then move to CUSTOM_CHALLENGE
3. Upon completion of CUSTOM_CHALLENGE (OTP verification), then issue token

When logging in using Cognito with MFA, I want the OTP generated by Cognito to be of the format ABC-123456. However after investigation I believe that this configuration is not available out of the box. i.e. Cognito seems to only generate OTPs of format 123456 (6-digits)

Therefore I utilised the createAuth, defineAuth and verifyAuth lambda triggers to handle the OTP generation and verification

However, I still need Cognito to handle username and password verification before handing things off to my lambda triggers.

I read that Cognito allows SRP Authentication (not plaintext username and password) followed by CUSTOM_CHALLENGE

I'm using @aws-sdk/client-cognito-identity-provider library, but cannot seem to get the initiateAuth method to behave correctly. It skips the SRP Authentication and moves straight to my custom challanges.

Looking at the documentation here

For CUSTOM_AUTH: USERNAME (required), SECRET_HASH (if app client is configured with client secret), DEVICE_KEY. To start the authentication flow with password verification, include ChallengeName: SRP_A and SRP_A: (The SRP_A Value).

Note: I'm not adding code snippets here as I will be answering below

Answer 1

There are two parts that need to be tackled

1. Handling SRP authentication
2. Creating lambda triggers

Handling SRP Authentication

SRP authentication flow goes as such (NOTE this is to begin with SRP and then move to CUSTOM_CHALLENGE)

1. Generate SRP_A

  //npm install amazon-user-pool-srp-client
  import { SRPClient, calculateSignature, getNowString }from 'amazon-user-pool-srp-client';

  const userPoolId = this.userPoolId.split('_')[1]; //User pool id in env is in format of ap-southeast-1_9xxxxxx
  const srp = new SRPClient(userPoolId);
  const srpA = srp.calculateA();

2. InitateAuth Command with relevant parameters

   const initiateAuthParams = {
      AuthFlow: AuthFlowType.CUSTOM_AUTH,
      ClientId: this.clientId,
      AuthParameters: {
        USERNAME: phoneNumber, // My cognito is configured to allow phone number as an alias for username
        CHALLENGE_NAME: 'SRP_A',
        SRP_A: srpA,
      },
    };
    const command = new InitiateAuthCommand(initiateAuthParams);
    const initiateAuthResponse = await this.provider.send(command);

3. Response contains: SALT, SECRET_BLOCK, SRP_B, USER_ID_FOR_SRP, and Session (and more)

const userIdForSrp = initiateAuthResponse.ChallengeParameters.USER_ID_FOR_SRP;
const srpB = initiateAuthResponse.ChallengeParameters.SRP_B;
const salt = initiateAuthResponse.ChallengeParameters.SALT;
const secretBlock = initiateAuthResponse.ChallengeParameters.SECRET_BLOCK;
const session = initiateAuthResponse.Session;

4. Generate the PASSWORD_CLAIM_SIGNATURE and TIMESTAMP from user's password, userId, srpB and salt

   const hkdf = srp.getPasswordAuthenticationKey(
      userIdForSrp,
      pin, // This is the user's password
      srpB,
      salt,
    );
    const dateNow = getNowString();
    const signatureString = calculateSignature(
      hkdf,
      userPoolId,
      userIdForSrp,
      secretBlock,
      dateNow,
    );

5. RespondToAuthChallenge Command with relevant parameters

const respondToAuthParams: RespondToAuthChallengeCommandInput = {
      ClientId: this.clientId,
      ChallengeName: ChallengeNameType.PASSWORD_VERIFIER,
      ChallengeResponses: {
        PASSWORD_CLAIM_SIGNATURE: signatureString,
        PASSWORD_CLAIM_SECRET_BLOCK: secretBlock,
        TIMESTAMP: dateNow,
        USERNAME: userIdForSrp,
      },
      Session: session,
    };
    const respondToAuthCommand = new RespondToAuthChallengeCommand(
      respondToAuthParams,
    );
    const respondToAuthResponse = await this.provider.send(
      respondToAuthCommand,
    );

Lambda Triggers

I was referencing this blog post

Create Auth Trigger

import { SNSClient, PublishCommand } from '@aws-sdk/client-sns';
import { SESClient, SendEmailCommand } from '@aws-sdk/client-ses';

const region = 'ap-southeast-1';
const snsClient = new SNSClient({ region });
const sesClient = new SESClient({ region });

export const handler = async (event) => {
  console.log('event: ');
  console.log(event);

  console.log('session');
  console.log(event.request.session);
  let otp;
  if (event.request.session.length === 2) {
    // Username password auth complete, generate OTP
    otp = generateOTP();
    await sendSMS(event.request.userAttributes.phone_number, otp);
    if (event.request.userAttributes.email) {
      await sendEmail(event.request.userAttributes.email, otp);
    }
  } else {
    // There's an existing session. Don't generate new digits but
    // re-use the code from the current session. This allows the user to
    // make a mistake when keying in the code and to then retry, rather
    // the needing to e-mail the user an all new code again.
    const previousChallenge = event.request.session.slice(-1)[0];
    otp = previousChallenge.challengeMetadata.match(/CODE-([A-Z]*-\d*)/)[1];
  }

  // This is sent back to the client app
  const otpPrefix = otp.split('-')[0];
  event.response.publicChallengeParameters = {
    otpPrefix,
  };

  // Add the secret login code to the private challenge parameters
  // so it can be verified by the "Verify Auth Challenge Response" trigger
  event.response.privateChallengeParameters = { secretLoginCode: otp };

  // Add the secret login code to the session so it is available
  // in a next invocation of the "Create Auth Challenge" trigger
  event.response.challengeMetadata = `CODE-${otp}`;

  return event;
};

function generateOTP() {
  const alphabet = 'abcdefghijklmnopqrstuvwxyz'.toUpperCase();
  let letters = '';
  for (let i = 0; i < 3; i++) {
    letters += alphabet[Math.floor(Math.random() * alphabet.length)];
  }
  const numbers = Math.floor(Math.random() * 1000000)
    .toString()
    .padStart(6, '0');

  const output = `${letters}-${numbers}`;
  return output;
}

async function sendEmail(emailAddress, otp) {
  const params = {
    Destination: { ToAddresses: [emailAddress] },
    Message: {
      Body: {
        Html: {
          Charset: 'UTF-8',
          Data: `<html><body><p>Your OTP is:</p>
                           <h3>${otp}</h3></body></html>`,
        },
        Text: {
          Charset: 'UTF-8',
          Data: `Your secret login code: ${otp}`,
        },
      },
      Subject: {
        Charset: 'UTF-8',
        Data: 'Your One Time Password',
      },
    },
    Source: 'abc@abc.com',
  };
  const command = new SendEmailCommand(params);
  const response = await sesClient.send(command);
  return response;
}

async function sendSMS(phoneNumber, otp) {
  const otpMessage = 'Your OTP is: ' + otp;
  const params = {
    PhoneNumber: phoneNumber,
    Message: otpMessage,
  };

  try {
    const command = new PublishCommand(params);
    const response = await snsClient.send(command);
    console.log('Success. SMS Send Response: ', response);
    return response; // For unit tests.
  } catch (err) {
    console.log(err, err.stack);
  }
}

Define Auth Trigger

export const handler = async (event) => {
  console.log('event: ');
  console.log(event);
  if (
    event.request.session &&
    event.request.session.length === 1 &&
    event.request.session[0].challengeName === 'SRP_A' &&
    event.request.session[0].challengeResult === true
  ) {
    //SRP_A is the first challenge, this will be implemented by cognito. Set next challenge as PASSWORD_VERIFIER.
    event.response.issueTokens = false;
    event.response.failAuthentication = false;
    event.response.challengeName = 'PASSWORD_VERIFIER';
  } else if (
    event.request.session &&
    event.request.session.length === 2 &&
    event.request.session[1].challengeName === 'PASSWORD_VERIFIER' &&
    event.request.session[1].challengeResult === true
  ) {
    //If password verification is successful then set next challenge as CUSTOM_CHALLENGE.
    event.response.issueTokens = false;
    event.response.failAuthentication = false;
    event.response.challengeName = 'CUSTOM_CHALLENGE';
  } else if (
    event.request.session &&
    //first session is password verification, after that 3 tries for OTP
    event.request.session.length >= 5 &&
    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.slice(-1)[0].challengeName === 'CUSTOM_CHALLENGE' &&
    event.request.session.slice(-1)[0].challengeResult === true
  ) {
    // The user provided the right answer; succeed auth
    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;
};

And lastly Verify Auth Trigger

export const handler = async (event) => {
  const expectedAnswer =
    event.request.privateChallengeParameters.secretLoginCode;
  if (event.request.challengeAnswer === expectedAnswer) {
    event.response.answerCorrect = true;
  } else {
    event.response.answerCorrect = false;
  }
  return event;
};