Controlling user auth flow with Lambda & Cognito

Disclaimer: the hero image of this post was the result of the following prompt AWS lambda and AWS cognito logos into a Renaissance paint. Use full logos and a less known painting. I think I still have much to learn into AI image prompts 😅😅

---

Authentication is a common topic between many kinds of systems. There are different ways to handle it and my preferred ones make usage of managed services. I found AWS Cognito a really great solution to handle authentication speacially if you are later connecting the authenticated app with other hosted services. Cognito will provide you built in ways to manage and cross validate users against services and recently I've been using it's hooks to build even more complex auth features

Cognito triggers

Cognito user pools have a feature named 'Lambda triggers' which let's you use previously created Lambdas to perform custom actions during four types of flow:

1. Sign up
2. Authentication
3. Custom authentication (such as CAPTCHA or security questions)
4. Messaging

Each of these flows have different triggers that will execute lambda code in between specific steps of the flow. Sign up for instance has Pre sign-up trigger, Post confirmation trigger and Migrate user trigger that can be attached to a Lambda function.

To test the capacities of Lambda triggers, we will develop a system that prevents login after 5 consecutive failures.

Coding the lambdas

We're gonna need two lambdas to make the flow controll, one of them would take care of updating the user data so we may count how many times the user tryed login. This one will also block the user if the number of attempts exceeds the maximumm. The second one, will be used reset our counter, so in the future the user will still have the maximumm number of attempts left.

The first lambda trigger would be like this

module.exports.preAuthTrigger = async (event) => {
    if (!(await this.isUserEnabled(event))) throw new Error('Usuário Bloqueado')

    const attempts = await this.getLoginAttempts(event)
    if (attempts > 4) {
        await this.disableUser(event)
        throw new Error('Usuário Bloqueado')
    }

    await this.updateLoginAttempts(event, attempts)
    return event
}

Our first step is to check whether the user is already blocked by the amount of attempts. We can do it with a separate fn:

exports.isUserEnabled = async (event) => {
    const getParams = {
        UserPoolId: event.userPoolId,
        Username: event.userName,
    }
    const userData = await cognitoService.adminGetUser(getParams).promise()
    return userData.Enabled
}

With this we are accessing the properties of the user on the cognito user pool and checking out the Enable property that dictates if the user is able to user it's username and password to login. A disabled user can't login into a cognito pool and that's exactly we want here.

For the second step, we need to check if the number of attempts is greater than the max permitted.

exports.getLoginAttempts = async (event) => {
    const getParams = {
        UserPoolId: event.userPoolId,
        Username: event.userName,
    }
    const userData = await cognitoService.adminGetUser(getParams).promise()
    const attribute = userData.UserAttributes.find(
        (att) => att.Name === 'custom:attempts'
    )
    if (attribute !== undefined && attribute !== null)
        return parseInt(attribute.Value)
    else return 0
}

It is a very simmilar process to the previous fn, but now we're looking for a custom attribute named custom:attempts that we will create into our user pool in the next steps. If the user has more than 5 attempts (we start counting at 0), then we should block the user. Piece of cake:

exports.disableUser = async (event) => {
    await cognitoService
        .adminDisableUser({
            UserPoolId: event.userPoolId,
            Username: event.userName,
        })
        .promise()
}

We also have to throw an Error and stop executing the lambda since this will make the login process fail as we want. Now that we are able to block the user, we just need to update the number of attempts if it isn't blocked:

exports.updateLoginAttempts = async (event, attempts) => {
    const updateParams = {
        UserAttributes: [
            {
                Name: 'custom:login_attempts',
                Value: (attempts + 1).toString(),
            },
        ],
        UserPoolId: event.userPoolId,
        Username: event.userName,
    }

    await cognitoService.adminUpdateUserAttributes(updateParams).promise()
}

This last function sets everything for the first lambda trigger. Now we are able to perform all the actions from our main lambda function. The final code with all functions will be like this:

module.exports.preAuthTrigger = async (event) => {
    if (!(await this.isUserEnabled(event))) throw new Error('Usuário Bloqueado')

    const attempts = await this.getLoginAttempts(event)
    if (attempts > 4) {
        await this.disableUser(event)
        throw new Error('Usuário Bloqueado')
    }

    await this.updateLoginAttempts(event, attempts)
    return event
}

exports.isUserEnabled = async (event) => {
    const getParams = {
        UserPoolId: event.userPoolId,
        Username: event.userName,
    }
    const userData = await cognitoService.adminGetUser(getParams).promise()
    return userData.Enabled
}

exports.getLoginAttempts = async (event) => {
    const getParams = {
        UserPoolId: event.userPoolId,
        Username: event.userName,
    }
    const userData = await cognitoService.adminGetUser(getParams).promise()
    const attribute = userData.UserAttributes.find(
        (att) => att.Name === 'custom:login_attempts'
    )
    if (attribute !== undefined && attribute !== null)
        return parseInt(attribute.Value)
    else return 0
}


exports.disableUser = async (event) => {
    await cognitoService
        .adminDisableUser({
            UserPoolId: event.userPoolId,
            Username: event.userName,
        })
        .promise()
}


exports.updateLoginAttempts = async (event, attempts) => {
    const updateParams = {
        UserAttributes: [
            {
                Name: 'custom:login_attempts',
                Value: (attempts + 1).toString(),
            },
        ],
        UserPoolId: event.userPoolId,
        Username: event.userName,
    }

    await cognitoService.adminUpdateUserAttributes(updateParams).promise()
}

In my next post we will write the code for the PostAuth lambda trigger and see how can we setup cognito to use both lambdas!