Cognito Starter Kit with Rust and Lambda

Welcome to the Cognito Starter Kit with a large helping of Rust seasoned with some CDK. I'm a big believer in Cognito and the power it gives builders to customize the various signup and authentication workflows. With Cognito, you get a managed service that has flexible usage-based pricing, numerous hooks and configurations and the ability to use OAuth and OIDC in your workflows. Let's dig in on the Cognito starter kit.

Components

I've written a good bit about Cognito in addition to customizing tokens and building authorizers. You can find those articles below:

With this Cognito Start Kit, I'm going to walk through building the below components.

Cognito User Pool with Advanced Security Features
A pre-authentication token customization Lambda written in Rust
An API Gateway Lambda authorizer written in Rust that verifies the JWT supplied by Cognito

Cognito User Pool

Defining a Cognito User Pool with AWS CDK is a straightforward effort. However, with the new Access Token customization features that were released in December 2023, the CDK L2 construct hasn't caught up yet. I'll walk through how to use the L1 to accomplish what is needed.

So which parts of the User Pool do we need to build? To customize the access token in addition to the ID token, the advanced security features need to be turned on. The Lambda pre-authentication hook needs to be enabled. This has a slight caveat that it's not yet supported by the normal L2 construct so I'll walk through that. Then, an often missed piece is to add a resource policy to the Lambda itself so that Cognito can invoke the Lambda as needed.

Congito Construct

For a quick point of clarity, when I say L1 construct, I mean the lowest-level construct that is generated from the CloudFormation resource specification. When I say L2, that's a higher-level construct with abstractions added for easier consumption. L1 constructs begin with Cfn. So let's look at the L1 User Pool construct.

const cfnUserPool = new CfnUserPool(this, "CfnUserPool", {
    userPoolName: `ExampleUserPool`,
    userPoolAddOns: {
        advancedSecurityMode: AdvancedSecurityMode.AUDIT,
    },
    lambdaConfig: {
        // @ts-ignore
        preTokenGenerationConfig: {
            lambdaArn: props.function.functionArn,
            lambdaVersion: "V2_0",
        },
    },
    policies: {},
});

The main 2 things to point out are that I'm enabling the advanced security mode to be able to customize the access token. As a fair word of warning, doing this will incur costs even below the 50,000-user free tier for non-advanced setups.

The second thing is to assign the token customizing Lambda to the configuration. And then to set the lambdaVerision to "V2_0" which enables that advanced payload.

Here's a small tip that I picked up from a GitHub discussion. There's nothing that says you can't convert that L1 INTO an L2 and then work with the higher-level API if you want to. Here's how I'm doing that for futher property settings. Pretty neat trick.

this._pool = UserPool.fromUserPoolId(scope, "RefdUserPool", cfnUserPool.ref);

Token Customization

The Cognito start kit includes a fully working Rust implementation of an access and ID token customizer all in one Lambda. One of the nice things about this workflow is that it frees me up from having to use Cognito attributes which can be super limiting. This freedom allows me to express my user profiles with whatever level of data that I desire. For this example, I'm going to use a DynamoDB table with a partition key.

DynamoDB Table

Nothing fancy with this setup. A pay-per-request table with just a partition key that is encrypted by AWS.

// dynamodb table
this._table = new dynamodb.Table(this, id, {
    billingMode: dynamodb.BillingMode.PAY_PER_REQUEST,
    removalPolicy: cdk.RemovalPolicy.DESTROY,
    partitionKey: {
        name: "id",
        type: dynamodb.AttributeType.STRING,
    },
    tableName: `SampleUserCustomized`,
    encryption: dynamodb.TableEncryption.AWS_MANAGED,
});

Token Customizer Function

I've been writing more and more about the benefits of Rust and Lambda quite a bit lately and I plan to take advantage of them in this customizer. Additionally, I've leveraged the Lambda Runtime project that includes data structures for the Lambda Events that I'll encounter while working with these payloads.

I'm going to leave the rest of the code in the repository for you to explore but I do want to walk through the logic and uniqueness of this handler.

As with any Lambda handler, I want to use long-lived APIs that were constructed outside of the function execution. I'm setting up an aws_sdk_dynamodb::Client to interact with my table to fetch the user profile when the request comes in.

Here's a quick peek at the record.

{
    "id": "USER#ben",
    "user_id": "ben",
    "first_name": "Ben",
    "last_name": "Pyle",
    "interesting_value": "Golf",
    "entity_type": "User"
}

This is the data that'll be used to add claims to both of the tokens.

Token Handler Code

The handler code itself will take the response from the DynamoDB query and use the User to add claims to the tokens.

Once the token Hashmaps have been adjusted, I need to work with the two data structures that are used by the Cognito API to apply the claims to the tokens. CognitoAccessTokenGenerationV2 and CognitoIdTokenGenerationV2

The last piece of the handler code is to build a response and then return that to the invoker.

async fn function_handler(
    client: &Client,
    table_name: &String,
    mut event: LambdaEvent<CognitoEventUserPoolsPreTokenGenV2>,
) -> Result<CognitoEventUserPoolsPreTokenGenV2, Error> {
    let mut access = HashMap::new();
    let mut id = HashMap::new();
    match event.payload.cognito_event_user_pools_header.user_name {
        Some(ref user_name) => {
            // fetch the user from DynamoDB
            let user = data::fetch_item(client, &table_name, user_name).await?;
            // insert interesting_value into the access token
            access.insert("interesting_value".to_string(), user.interesting_value);
            // insert first_name and last_name into the ID token
            id.insert("first_name".to_string(), user.first_name);
            id.insert("last_name".to_string(), user.last_name);
        }
        None => {
            event
                .payload
                .response
                .claims_and_scope_override_details
                .as_mut()
                .unwrap()
                .group_override_details
                .groups_to_override = vec![];
        }
    }

    // access token customize struct
    let access_token = CognitoAccessTokenGenerationV2 {
        claims_to_add_or_override: access,
        claims_to_suppress: vec![],
        scopes_to_add: vec![],
        scopes_to_suppress: vec![],
    };

    // ID token customize struct
    let id_token = CognitoIdTokenGenerationV2 {
        claims_to_add_or_override: id,
        claims_to_suppress: vec![],
    };

    let ovr = ClaimsAndScopeOverrideDetailsV2 {
        access_token_generation: Some(access_token),
        group_override_details: GroupConfiguration {
            ..Default::default()
        },
        id_token_generation: Some(id_token),
    };

    event.payload.response = CognitoEventUserPoolsPreTokenGenResponseV2 {
        claims_and_scope_override_details: Some(ovr),
    };

    Ok(event.payload)
}

Rust Token Handler Thoughts

What I like most about these new payloads is that I can customize both the ID and access token at one time. For a quick refresher, adding some claims to the access token can speed up the authorization process by not having to fetch key pieces of data needed to perform those authorizations. By sprinkling in key claims to the ID token, the client or UI can alter the user's experience without needing to fetch those additional details.

Those are both upfront costs that I can save in user experience down the line later for 60 minutes or however long I choose to set the expiration on these two tokens.

The last thought I want to explore that you'll see in the repository is that my Rust skills are improving. I'm using the ? operator in the handler code quite a bit to reduce the need to have nested match blocks. I'm doing that by implementing the From trait and converting the errors into a custom Error that I defined. This code is in the models.rs file.

API Authorizer Function

API Gateway offers a few options for authorizers. There's IAM authorization, Cognito authorization and a custom Lambda authorization. Why would I be building a Lambda authorizer IF I can make use of native Cognito authorization? That's a really good question and here's the answer. With custom Lambda authorizers I get the ability to add context to my payloads that can be forwarded down to the lower API requests.

Take this example for instance. I'm building a multi-tenant application that uses the user's location or tenant in the queries to the database. I could fetch that data based on the token supplied OR I could parse that data out of the customized access token and then forward those elements to the API. Again, time savings. All of these hops and CPU cycles add up to user experience and if not for the users, what's the point!?

Authorizer Initialization Code

I tend to not show init code in articles, but this is worthwhile. I had someone ask me the other day how to handle fetching the keyset needed to verify the signature of the JWT. I do this type of logic in the main function because I can reuse the output over and over while not having to hit the jwks URL again.

That code looks like this:

let keyset = jsonwebtokens_cognito::KeySet::new(region_id, user_pool_id).unwrap();
let _ = keyset.prefetch_jwks().await;

Authorizer Handler Code

The Cognito starter kit comes together nicely in this authorizer code. By using Cognito, customizing the tokens, doing the authorization and then forwarding context to resulting API calls, I get a fully serverless workflow with the performance of the Rust programming language.

I found this crate while working on something else a few weeks back and it makes working with Cognito and JWT super simple.

The handler code is not complex. Let's walk through the steps real quick.

By default, I'm returning Allow to the authorizer. I intend to disprove that the request should be allowed. I start by building a verifier from the client_idsupplied to the handler. This is sourced from an environment variable in the main function.

I second pull the token from the payload. I'm OK with unwrap() here because if the unwrap fails, the function should fail. API Gateway will always give me the payload I expect.

However, when I try and verify the token, if the value is something garbage that could come from the client, then I'm going to set allowance to Deny.

If the token is valid, I'm then going to pull claims from the serde_json::Value and attach them as context into the response.

async fn function_handler(
    client_id: &str,
    keyset: &jsonwebtokens_cognito::KeySet,
    event: LambdaEvent<ApiGatewayCustomAuthorizerRequest>,
) -> Result<ApiGatewayCustomAuthorizerResponse, claims::AuthorizerError> {
    let mut allowance = "Allow";
    let mut ctx = serde_json::Value::default();

    let verifier = keyset.new_access_token_verifier(&[client_id]).build()?;
    let token = event.payload.authorization_token.unwrap();
    let claims: Result<serde_json::Value, jsonwebtokens_cognito::Error> =
        keyset.try_verify(token.as_str(), &verifier);

    match claims {
        Ok(c) => ctx = dump_claims(&c)?,
        Err(_) => {
            allowance = "Deny";
        }
    }

    let response = new_response(allowance, ctx);
    Ok(response)
}

Adding Claims to the Context

I want to show this for completeness. Above I mentioned a function called dump_claims. Its sole purpose is to take the claims from the token and then build a small subset that can be passed into the request context for downstream API requests.

pub fn dump_claims(value: &serde_json::Value) -> Result<serde_json::Value, serde_json::Error> {
    let claim: Result<Claim, serde_json::Error> = serde_json::from_value(value.clone());
    tracing::debug!("(Claim_JSON): {}", value);
    tracing::debug!("(Claim_Struct): {:?}", claim);

    match claim {
        Ok(c) => {
            let pc = PrivateClaim {
                user_name: c.username,
                location_id: c.interesting_value,
            };
            tracing::debug!("(PrivateClaim): {:?}", pc);
            let pc_v = serde_json::to_value(pc)?;
            Ok(pc_v)
        }
        Err(e) => {
            tracing::error!("(Claim_Struct): {:?}", e);
            Err(e)
        }
    }
}

Rust API Authorizer Thoughts

Just like with the customizer, I'm able to take advantage of the Rust language's performance and ergonomics to build simple functions that perform amazingly well. Here are a few sample runs when testing from the AWS Console. Those aren't fake and one of them is < 1ms. I'm not kidding about Rust and Lambda's performance.

Another thing that I often forget and I'll admit is a bias is to check the optimal memory, speed and cost tuning output. I naturally think, oh, this function needed 24MB of memory, let's give it 128MB. However, the optimal setting for Rust Lambda's is 256MB per my experience.

The AWS Lambda Power Tuning State Machine agrees.

Putting it All Together

What does this look like when all put together? Let's take a look.

Cognito

When the User Pool is configured with the pre-authentication Lambda with the V2 payload, it should look like this.

Customized Tokens

With the handler in place, when a user logs in and authenticates via this User Pool, I'm going to get the ID and access token with the claims added from the handler code reviewed above.

ID Token

Note the first_name and last_name claims that were added in the code.

And the same thing holds for the access token. Look in the image for the interesting_value claim that should say Golf.

Access Token

Testing the API Authorizer

I didn't include a working API Gateway example in this repository but it would be easy to extend this by using a CloudFormation output and then including the authorizer in another project. I did run a few samples with an API Gateway I already had and here are what the screenshots look like for that.

If I supply a valid and non-expired token, here's the output that I'm going to receive.

Note the "Allow" in the Effect property.

And if I supply anything invalid, it'll look the same minus the "Effect" value.

Note the "Deny" in the Effect property.

Wrapping Up

First off, here's a link to the GitHub repository that is referenced in this article.

Second, I'm starting to feel more productive with Rust. I'm reaching for the docs and Google less and less. The patterns for using enums, errors and traits are starting to be easier to implement. And the borrower checker is starting to be more friend than my annoying door bouncer. Those are all good things in my opinion.

I'm also close to declaring VSCode as my primary editor over Rust Rover. I'm not all the way there yet, but close. I might write something on my Rust development setup at some point.

As for Cognito, it is a super powerful, robust and highly scalable managed service from AWS that gives a builder so much power out of the box. However, I hope I've shown you a nice Cognito starter kit that brings together the power of Rust and Lambda to jump-start your use cases. I've used the product for years and reached for it first as it's AWS native and has all of the features I'm looking for at a price that makes sense. Just remember, customizing access tokens and enabling the audit features do come at a cost. So use it wisely.

I hope this has been helpful and I'm looking forward to more Rust and Serverless content this year! Stay tuned!

Thanks for reading and happy building!

Blog