exanubes
Q&A

Passwordless authentication flow in Cognito User Pool

Following up on setting up a custom mailer in cognito we are going to configure and implement custom authentication flow for AWS Cognito User Pool. To do that we will use the cognito stack created in the previous article, change the authentication configuration and implement custom lambdas to handle the new authentication flow. Then we will have to go into the backend application and update the api to facilitate the new flow as well.

The authentication flow we’ll be implementing is a passwordless approach where users provide their username and in turn they will receive a login link on their email address.

You can find the finished code on GitHub or run the following command if you’d like to follow along

git clone -b start git@github.com:exanubes/cognito-passwordless-authentication-flow.git

Custom authentication flow

There are several steps Cognito has to take when using a custom authentication flow, many of which happen behind the scenes. However, it is useful to know the entire flow.

Naturally, it all starts with a user request to begin the authentication process. That leads Cognito to create a new session token for us – which expires after 3 minutes – and send it over to defineChallenge lambda. This is where we come in, as we’re going to have to implement the lambda handler to begin the authentication process. This lambda acts as a state machine for the entire authentication flow. As this is the first time the lambda is invoked in this session, we’re going to return one of the challenge types that have been predefined by AWS. You can find ChallengeNameType enum in the AWS SDK repository to see all possible challenge types.

Subsequently, Cognito takes that challenge type and triggers createChallenge lambda. Again, it will be up to us to implement handler logic and decide how we want the custom authentication to work. Once all the challenge parameters have been set we return public params to the client along with the session token.

We’re back to the user now. Once he enters the challenge response, we’re gonna take the username, the response and the session token and send the respondToAuthChallenge command. Cognito will then take our response and trigger the verifyChallenge lambda. Here, we’re gonna check if the provided challenge response matches the correct answer set by us in createChallenge lambda.

If successful, Cognito will invoke defineAuth lambda again, which will respond with authentication tokens this time. After all that Cognito sends tokens to the client and the user is authenticated.

Lambda resources

// stacks/cognito.stack.ts
  private createCustomChallengeLambdas(): Function[] {
    const defineChallengeLambda = new Function(
      this,
      "define-challenge-lambda",
      {
        runtime: Runtime.NODEJS_14_X,
        handler: "index.handler",
        code: Code.fromAsset(join(__dirname, "..", "lambdas/define-challenge")),
        environment: {
          CHALLENGE_NAME: ChallengeNameType.CUSTOM_CHALLENGE,
        },
      }
    );
    const createChallengeLambda = new Function(
      this,
      "create-challenge-lambda",
      {
        runtime: Runtime.NODEJS_14_X,
        handler: "index.handler",
        code: Code.fromAsset(join(__dirname, "..", "lambdas/create-challenge")),
        layers: [this.layer],
        environment: {
          SENDGRID_API_KEY: String(process.env.SENDGRID_API_KEY),
        },
      }
    );
    const verifyChallengeLambda = new Function(
      this,
      "verify-challenge-lambda",
      {
        runtime: Runtime.NODEJS_14_X,
        handler: "index.handler",
        code: Code.fromAsset(join(__dirname, "..", "lambdas/verify-challenge")),
      }
    );
    return [defineChallengeLambda, createChallengeLambda, verifyChallengeLambda]
  }

Picking up where we left off in Cognito with custom mailer repository . As described in Custom Authentication Flow , we need to create three lambdas for defining, creating and verifying the authentication challenge. We’re passing CHALLENGE_NAME as environment variable to defineChallenge lambda to determine what kind of challenge we want to use. Then we’re giving create lambda a variable with SENDGRID_API_KEY so that we can send an email to the user with the correct answer to the challenge.

Permissions

// stacks/cognito.stack.ts
  private static grantLambdaInvokePermission(lambda: Function, alias: string) {
    lambda.addPermission(`exanubes-cognito-invoke-permission-${alias}`, {
      principal: new ServicePrincipal("cognito-idp.amazonaws.com"),
      action: "lambda:InvokeFunction",
    });
  }

Cognito needs to have permission to invoke the lambdas so using this small helper method we can give it permission for all three authentication lambdas.

// stacks/cognito.stack.ts
CognitoStack.grantLambdaInvokePermission(
  defineChallengeLambda,
  'define-challenge-lambda'
);
CognitoStack.grantLambdaInvokePermission(
  createChallengeLambda,
  'create-challenge-lambda'
);
CognitoStack.grantLambdaInvokePermission(
  verifyChallengeLambda,
  'verify-challenge-lambda'
);

Custom authentication triggers

// stacks/cognito.stack.ts
const [defineChallenge, createChallenge, verifyChallenge] =
  this.createCustomChallengeLambdas();

userPool.addTrigger(UserPoolOperation.DEFINE_AUTH_CHALLENGE, defineChallenge);
userPool.addTrigger(UserPoolOperation.CREATE_AUTH_CHALLENGE, createChallenge);
userPool.addTrigger(
  UserPoolOperation.VERIFY_AUTH_CHALLENGE_RESPONSE,
  verifyChallenge
);

Here, we’re actually telling our User Pool which lambda should be invoked depending on which User Pool Operation is triggered.

Custom authentication user pool agent

// stacks/cognito.stack.ts
const client = userPool.addClient('exanubes-user-pool-client', {
  userPoolClientName: 'exanubes-cognito-app',
  authFlows: {
    custom: true,
  },
  accessTokenValidity: Duration.days(1),
  idTokenValidity: Duration.days(1),
  refreshTokenValidity: Duration.days(30),
  preventUserExistenceErrors: true,
});

One more thing we need to change in the previous infrastructure is the selected authentication flow in the user pool client. Previously it was userPassword now we want it to be custom.

Implementing Custom Lambdas

Define Challenge lambda

// lambdas/define-challenge/logger.ts
import { DefineAuthChallengeTriggerEvent } from 'aws-lambda';

const ALLOWED_ATTEMPTS = Infinity;
const challengeName = process.env.CHALLENGE_NAME || '';

exports.handler = async (event: DefineAuthChallengeTriggerEvent) => {
  const [challenge] = event.request.session.reverse();
  const challengeAttempts = event.request.session.length;
  if (challenge) {
    if (challengeAttempts >= ALLOWED_ATTEMPTS) {
      event.response.issueTokens = false;
      event.response.failAuthentication = true;
      return event;
    }
    if (challenge.challengeName === challengeName) {
      event.response.issueTokens = challenge.challengeResult;
      event.response.failAuthentication = !challenge.challengeResult;
      return event;
    }
  }
  event.response.issueTokens = false;
  event.response.failAuthentication = false;
  event.response.challengeName = challengeName;
  return event;
};

This is probably the most complicated one of the lambdas as it is supposed to work as a state machine for the entire authentication flow. First of all, Cognito gives us access to the current session in the request, which is an array of all challenge answers. This is why we reverse the array and destructure first element from it, as that will always be the latest attempt.

In order to create an appropriate response for user pool to handle, we have three properties on the response object we can use. issueTokens, failAuthentication and challengeName.

In the event where session is an empty array, we’re going to issue the challenge as that’s the very first time the lambda has been triggered. Otherwise, we’re checking the validity of the challenge. In this case we’re checking if the attempt threshold has been exceeded in which case we’re going to fail authentication and not issue tokens.

If, however, the attempts have not been exceeded and the challenge name matches the one we’ve configured the lambda for, we’re going to issue tokens if challengeResult is truthy and fail authentication if it’s falsy.

Create Challenge Lambda

// lambdas/create-challenge/logger.ts
import { CreateAuthChallengeTriggerEvent } from 'aws-lambda';

const { generateRandomString } = require('@exanubes/cdk-utils');
const emailClient = require('@sendgrid/mail');
const sendgridApiKey = String(process.env.SENDGRID_API_KEY);
emailClient.setApiKey(sendgridApiKey);

exports.handler = async (event: CreateAuthChallengeTriggerEvent) => {
  const code = generateRandomString();
  const user = event.request.userAttributes;
  const email = {
    to: `${event.userName} <${user.email}>`,
    from: 'Exanubes.com <example@email.com>',
    subject: `[${event.triggerSource}] Your login token`,
    text: `Use the link below to log in to exanubes.com\n http://localhost:4000/verify-login?code=${code}&username=${event.userName} \n this link will expire in three minutes`,
  };
  await emailClient.send(email);
  event.response.publicChallengeParameters = {
    email: user.email,
  };
  event.response.privateChallengeParameters = {
    code,
  };
  event.response.challengeMetadata = `EXANUBES_CHALLENGE`;
  return event;
};

Create challenge lambda is responsible for actually creating the temporary password that has to be entered in order to log in. This is also the place where we’re going to notify the user of the correct password whether it’s email, a text message or a push notification. I’ve opted for just sending an email with a login link to my frontend application which will take the code and username values from the query parameter and call my backend API with those values to respond to the challenge. We can also set public and private challenge parameters. The public params will be sent back to the client so it probably shouldn’t include the answer to the challenge, but the private params will only be accessible withing the auth flow. The challengeMetada prop is used for assigning a custom name to the challenge if we want to do so.

Verify Challenge Lambda

// lambdas/verify-challenge/logger.ts
import { VerifyAuthChallengeResponseTriggerEvent } from 'aws-lambda';

exports.handler = async (event: VerifyAuthChallengeResponseTriggerEvent) => {
  const validCode = event.request.privateChallengeParameters.code;
  event.response.answerCorrect = validCode === event.request.challengeAnswer;
  return event;
};

The most straight-forward of the lambdas, here we just want to compare the user input with the correct answer saved in privateChallengeParameters in create challenge lambda and assign the value to the answerCorrect param on the response object.

Recap

So far we’ve gone over the overall custom authentication flow. We’ve created the lambdas, added user pool triggers and gave the necessary permissions to Cognito for invoking the lambda functions. Then we covered the implementation of each one of the functions and now with the AWS CDK part finished we have to move to the server code and make the necessary adjustments to migrate from userPassword to the custom authentication flow.

AWS SDK

Login

// auth/auth.service.ts
async login({ username }: LoginDto) {
  const input: InitiateAuthCommandInput = {
    ClientId: this.awsConfigService.userPoolClientId,
    AuthFlow: AuthFlowType.CUSTOM_AUTH,
    AuthParameters: {
      USERNAME: username,
    },
  };
  const command = new InitiateAuthCommand(input);
  return this.client.send(command);
}

This time the AuthFlow is set to CUSTOM_AUTH so that cognito knows to trigger our define challenge handler as well as not expect a PASSWORD field in the AuthParameters. Those are the only changes to the login method.

Verify login

// auth/auth.service.ts
async verifyLogin({ code, username, session }: VerifyLoginDto) {
  const input: RespondToAuthChallengeCommandInput = {
    ClientId: this.awsConfigService.userPoolClientId,
    ChallengeName: ChallengeNameType.CUSTOM_CHALLENGE,
    Session: session,
    ChallengeResponses: {
      ANSWER: code,
      USERNAME: username,
    },
  };
  const command = new RespondToAuthChallengeCommand(input);
  const response = await this.client.send(command);
  return response.AuthenticationResult;
}

Similar to verify signup, we now have to verify login as we’ve split the authentication process rather than providing username and password at once. With the response we need the session token which was returned from the initialAuth command in login method. The challenge name should be the same as in the defineChallenge lambda and lastly username and answer to the challenge. If successful, the command will respond with authentication tokens.

Signup

// auth/auth.service.ts
  async signup({ username, email }: SignUpDto) {
    const input: SignUpCommandInput = {
      Password: generateRandomString(),
      Username: username,
      ClientId: this.awsConfigService.userPoolClientId,
      UserAttributes: [{ Name: 'email', Value: email }],
    };
    const command = new SignUpCommand(input);
    return this.client.send(command);
  }

Last but not least, we have to slightly adjust the signup input for the sign up command as there’s no need for the user to provide us with a password, however, the command still expects to receive a password nonetheless. That’s why we’re gonna generate a random string to keep it satisfied.

Summary

In this article we’ve learned how to setup an infrastructure using the AWS CDK for custom authentication flow in AWS Cognito. We’ve gone over setting up lambda resources, adding user pool triggers and giving Cognito permission to invoke the lambdas. Then we’ve covered the implementation of each of the lambdas necessary for custom authentication flow. Last but not least we needed to adjust the backend API to facilitate the different auth flow. In the signup process user no longer provides a password so we generate one for him as that’s a requirement in the SDK API. We also needed to change the login setup as now we no longer receive tokens from that API call and we needed to change the authentication flow it’s supposed to initiate. Finally, we had to implement a new API endpoint for answering the Cognito challenge and receiving authentication tokens.