Use Amazon Cognito to add claims to an identity token for fine-grained authorization

August 2, 2023: Amazon Verified Permissions now offers a direct integration with Amazon Cognito to add fine-grained authorization within your applications. Learn more.

With Amazon Cognito, you can quickly add user sign-up, sign-in, and access control to your web and mobile applications. After a user signs in successfully, Cognito generates an identity token for user authorization. The service provides a pre token generation trigger, which you can use to customize identity token claims before token generation. In this blog post, we’ll demonstrate how to perform fine-grained authorization, which provides additional details about an authenticated user by using claims that are added to the identity token. The solution uses a pre token generation trigger to add these claims to the identity token.

Scenario

Imagine a web application that is used by a construction company, where engineers log in to review information related to multiple projects. We’ll look at two different ways of designing the architecture for this scenario: a standard design and a more optimized design.

Standard architecture

A sample standard architecture for such an application is shown in Figure 1, with labels for the various workflow steps:

1. The user interface is implemented by using ReactJS (a JavaScript library for building user interfaces).
2. The user pool is configured in Amazon Cognito.
3. The back end is implemented by using Amazon API Gateway.
4. AWS Lambda functions exist to implement business logic.
5. The AWS Lambda CheckUserAccess function (5) checks whether the user has authorization to call the AWS Lambda functions (4).
6. The project information is stored in an Amazon DynamoDB database.

Figure 1: Lambda functions that need the user’s projectID call the GetProjectID Lambda function

In this scenario, because the user has access to information from several projects, several backend functions use calls to the CheckUserAccess Lambda function (step 5 in Figure 1) in order to serve the information that was requested. This will result in multiple calls to the function for the same user, which introduces latency into the system.

Optimized architecture

This blog post introduces a new optimized design, shown in Figure 2, which substantially reduces calls to the CheckUserAccess API endpoint:

1. The user logs in.
2. Amazon Cognito makes a single call to the PretokenGenerationLambdaFunction-pretokenCognito function.
3. The PretokenGenerationLambdaFunction-pretokenCognito function queries the Project ID from the DynamoDB table and adds that information to the Identity token.
4. DynamoDB delivers the query result to the PretokenGenerationLambdaFunction-pretokenCognito function.
5. This Identity token is passed in the authorization header for making calls to the Amazon API Gateway endpoint.
6. Information in the identity token claims is used by the Lambda functions that contain business logic, for additional fine-grained authorization. Therefore, the CheckUserAccess function (7) need not be called.

The improved architecture is shown in Figure 2.

Figure 2. Get the projectID and inset it in a custom claim in the Identity token

The benefits of this approach are:

1. The number of calls to get the Project ID from the DynamoDB table are reduced, which in turn reduces overall latency.
2. The dependency on the CheckUserAccess Lambda function is removed from the business logic. This reduces coupling in the architecture, as depicted in the diagram.

In the code sample provided in this post, the user interface is run locally from the user’s computer, for simplicity.

Code sample

You can download a zip file that contains the code and the AWS CloudFormation template to implement this solution. The code that we provide to illustrate this solution is described in the following sections.

Prerequisites

Before you deploy this solution, you must first do the following:

1. Download and install Python 3.7 or later.
2. Download the AWS SDK for Python (Boto3) library by using the following pip command.
   pip install boto3
3. Install the argparse package by using the following pip command.
   pip install argparse
4. Install the AWS Command Line Interface (AWS CLI).
5. Configure the AWS CLI.
6. Download a code editor for Python. We used Visual Studio Code for this post.
7. Install Node.js.

Description of infrastructure

The code provided with this post installs the following infrastructure in your AWS account.

Deploy the solution

In this section, we describe how to deploy the infrastructure, save the trigger configuration, add users to the Cognito user pool, and run the web application.

To deploy the solution infrastructure

1. Download the zip file to your machine. The readme.md file in the addclaimstoidtoken folder includes a table that describes the key files in the code.
2. Change the directory to addclaimstoidtoken.
   cd addclaimstoidtoken
3. Review stackInputs.json. Change the value of the userPoolDomainName parameter to a random unique value of your choice. This example uses pretokendomainname as the Amazon Cognito domain name; you should change it to a unique domain name of your choice.
4. Deploy the infrastructure by running the following Python script.
   python3 setup_pretoken.py

   After the CloudFormation stack creation is complete, you should see the details of the infrastructure created as depicted in Figure 3.

   

   Figure 3: Details of infrastructure

Now you’re ready to add users to your Amazon Cognito user pool.

To add users to your Cognito user pool

1. To add users to the Cognito user pool and configure the DynamoDB store, run the Python script from the addclaimstoidtoken directory.
   python3 add_user_info.py
2. This script adds one user. It will prompt you to provide a username, email, and password for the user.
   

   Note: Because this is sample code, advanced features of Cognito, like multi-factor authentication, are not enabled. We recommend enabling these features for a production application.

   The addUserInfo.py script performs two actions:

   • Adds the user to the Cognito user pool.
     

     Figure 4: User added to the Cognito user pool

   • Adds sample data to the DynamoDB table.
     

     Figure 5: Sample data added to the DynamoDB table named UserInfoTable

Now you’re ready to run the application to verify the custom claim addition.

To run the web application

1. Change the directory to the pre-token-web-app directory and run the following command.
   cd pre-token-web-app
2. This directory contains a ReactJS web application that displays details of the identity token. On the terminal, run the following commands to run the ReactJS application.
   npm install
   npm start

   This should open http://localhost:8081 in your default browser window that shows the Login button.

   

   Figure 6: Browser opens to URL http://localhost:8081

3. Choose the Login button. After you do so, the Cognito-hosted login screen is displayed. Log in to the website with the user identity you created by using the addUserInfo.py script in step 1 of the To add users to your Cognito user pool procedure.
   

   Figure 7: Input credentials in the Cognito-hosted login screen

4. When the login is successful, the next screen displays the identity and access tokens in the URL. You can reveal the token details to verify that the custom claim has been added to the token by choosing the Show Token Detail button.
   

   Figure 8: Token details displayed in the browser

What happened behind the scenes?

In this web application, the following steps happened behind the scenes:

1. When you ran the npm start command on the terminal command line, that ran the react-scripts start command from package.json. The port number (8081) was configured in the pre-token-web-app/.env file. This opened the web application that was defined in app.js in the default browser at the URL http://localhost:8081.
2. The Login button is configured to navigate to the URL that was defined in the constants.js file. The constants.js file was generated during the running of the setup_pretoken.py script. This URL points to the Cognito-hosted default login user interface.
3. When you provided the login information (username and password), Amazon Cognito authenticated the user. Before generating the set of tokens (identity token and access token), Cognito first called the pre-token-generation Lambda trigger. This Lambda function has the code to connect to the DynamoDB database. The Lambda function can then access the project information for the user that is stored in the userInfo table. The Lambda function read this project information and added it to the identity token that was delivered to the web application.

   Lambda function code

   const AWS = require("aws-sdk");
   
   // Create the DynamoDB service object
   var ddb = new AWS.DynamoDB({ apiVersion: "2012-08-10" });
   
   // PretokenGeneration Lambda
   exports.handler = async function (event, context) {
       var eventUserName = "";
       var projects = "";
   
       if (!event.userName) {
           return event;
       }
   
       var params = {
           ExpressionAttributeValues: {
               ":v1": {
                   S: event.userName
               }
           },
           KeyConditionExpression: "userName = :v1",
           ProjectionExpression: "projects",
           TableName: "UserInfoTable"
       };
   
       event.response = {
           "claimsOverrideDetails": {
               "claimsToAddOrOverride": {
                   "userName": event.userName,
                   "projects": null
               },
           }
       };
   
       try {
           let result = await ddb.query(params).promise();
           if (result.Items.length > 0) {
               const projects = result.Items[0]["projects"]["S"];
               console.log("projects = " + projects);
               event.response.claimsOverrideDetails.claimsToAddOrOverride.projects = projects;
           }
       }
       catch (error) {
           console.log(error);
       }
   
       return event;
   };

   The code for the Lambda function is as follows.

4. After a successful login, Amazon Cognito redirected to the URL that was specified in the App Client Settings section, and added the token to the URL.
5. The webpage detected the token in the URL and displayed the Show Token Detail button. When you selected the button, the webpage read the token in the URL, decoded the token, and displayed the information in the relevant text boxes.
6. Notice that the Decoded ID Token box shows the custom claim named projects that displays the projectID that was added by the PretokenGenerationLambdaFunction-pretokenCognito trigger.

How to use the sample code in your application

We recommend that you use this sample code with the following modifications:

1. The code provided does not implement the API Gateway and Lambda functions that consume the custom claim information. You should implement the necessary Lambda functions and read the custom claim for the event object. This event object is a JSON-formatted object that contains authorization data.
2. The ReactJS-based user interface should be hosted on an Amazon Simple Storage Service (Amazon S3) bucket.
3. The projectId of the user is available in the token. Therefore, when the token is passed by the Authorization trigger to the back end, this custom claim information can be used to perform actions specific to the project for that user. For example, getting all of that user’s work items that are related to the project.
4. Because the token is valid for one hour, the information in the custom claim information is available to the user interface during that time.
5. You can use the AWS Amplify library to simplify the communication between your web application and Amazon Cognito. AWS Amplify can handle the token retention and refresh token mechanism for the web application. This also removes the need for the token to be displayed in the URL.
6. If you’re using Amazon Cognito to manage your users and authenticate them, using the Amazon Cognito user pool to control access to your API is easier, because you don’t have to write the authentication code in your authorizer.
7. If you decide to use Lambda authorizers, note the following important information from the topic Steps to create an API Gateway Lambda authorizer: “In production code, you may need to authenticate the user before granting authorization. If so, you can add authentication logic in the Lambda function as well by calling an authentication provider as directed in the documentation for that provider.”
8. Lambda authorizer is recommended if the final authorization (not just token validity) decision is made based on custom claims.

Conclusion

In this blog post, we demonstrated how to implement fine-grained authorization based on data stored in the back end, by using claims stored in an identity token that is generated by the Amazon Cognito pre token generation trigger. This solution can help you achieve a reduction in latency and improvement in performance.

For more information on the pre token generation Lambda trigger, refer to the Amazon Cognito Developer Guide.

 
If you have feedback about this post, submit comments in the Comments section below. If you have questions about this post, contact AWS Support.

Want more AWS Security news? Follow us on Twitter.