Implementing access control on API Gateway endpoints with ID tokens

1. The scenario

Say that we have an internal application used by most employees at the company. Some users can add data to the Tasks applications, while others can only read them. We should find a way to validate the user’s permission before the request reaches the backend.

For the sake of simplicity, the application will have two endpoints, GET /tasks and POST /tasks. Users wanting to add tasks to the application must have the write.task permission. Everyone can read the data, so we should assign read.tasks to all users. That is, elevated users can access both the read and write endpoints.

2. Architecture overview

In this solution, we will store user identities (username, password, email, etc.) in a Cognito user pool.

We will also have a DynamoDB table dedicated to each user’s permissions. An item’s partition key is the user’s Cognito User pools ID called sub. Among others, each item has a permissions attribute where we - surprise - store the permissions assigned to the given user. For example, if Alice is an elevated user, her item’s permissions attribute in the table will look something like this:

[{ "S": "read.tasks" }, { "S": "write.tasks" }]

Bob is a regular user, so his permissions attribute will only have [{ "S": "read.tasks" }].

As mentioned above, the application will have two REST APIs set up in API Gateway. The endpoint handlers are Lambda functions.

3. Pre-token generation

Cognito User pools will return some tokens after the user has signed in to the application. One of them is the ID token that contains information about the user, like username or email address.

If we could add the user’s permission to the token, we could validate it when the requests hit the API. This way, we could decide if the user has the authorization to access the endpoint.

Cognito has a feature called Lambda triggers, which is available under User pool properties:

We can generate Lambda triggers for various workflows here (more on these in future posts). One is the Pre token generation trigger under the Authentication block. The target function will run between the user’s login activity and Cognito’s token generation, so it seems the perfect time and place to add some custom claims to the ID token.

3.1. Lambda code

We need to create a regular Lambda function and assign it as the target to the trigger.

The function’s code can look like this:

import { PreTokenGenerationTriggerEvent } from 'aws-lambda';
import { DynamoDBClient } from '@aws-sdk/client-dynamodb';
import { DynamoDBDocumentClient, GetCommand, GetCommandInput } from '@aws-sdk/lib-dynamodb';

const { TABLE_NAME } = process.env;

const ddbClient = new DynamoDBClient();
const docClient = DynamoDBDocumentClient.from(ddbClient);

export const handler = async (event: PreTokenGenerationTriggerEvent): Promise<PreTokenGenerationTriggerEvent> => {
  const sub = event.request.userAttributes.sub;
  const input: GetCommandInput = {
    TableName: TABLE_NAME,
    Key: {
      user_id: sub,
    },
  };
  const command = new GetCommand(input);
  let permissionAttribute: string[];

  try {
    const response = await docClient.send(command);
    permissionAttribute = response.Item?.permissions ?? [];
  } catch (error) {
    throw error;
  }

  // The claim must be a string
  const permissions = permissionAttribute.join(' ');

  return {
    ...event,
    response: {
      claimsOverrideDetails: {
        claimsToAddOrOverride: {
          permissions,
        },
      },
    },
  };
};

The code uses the AWS SDK for JavaScript v3 with DynamoDBDocumentClient, where we don’t have to specify DynamoDB data types. As a result, we will have simpler, easier-to-read code.

3.2. Key points

There are two main points to highlight in the code.

First, we get the sub value from the input event. It’s inside the request.userAttributes object, accompanied by some other properties:

{
  // other properties
  "request": {
    "userAttributes": {
      "sub": "49615ae7-bd00-4206-93df-f7a4140ecc9e",
      "cognito:user_status": "CONFIRMED",
      "email_verified": "true",
      "email": "USER_EMAIL_ADDRESS"
    },
  }
}

The second point is that the function must return the same event type with the permissions claim we want to add:

{
  // request properties here
  response: {
    claimsOverrideDetails: {
      claimsToAddOrOverride: {
        permissions,
      },
    },
  }
},

We should add the extra token field to the response object. The custom claim should have a string value, so we join the array items into a string.

3.3. Final ID token

The ID token returned by Cognito will now contain the permissions claim:

{
  "sub": "49615ae7-bd00-4206-93df-f7a4140ecc9e",
  "email_verified": true,
  "iss": "https://cognito-idp.eu-central-1.amazonaws.com/USER_POOL_ID",
  "cognito:username": "alice",
  "origin_jti": "d3bce7b0-2b4b-4eeb-b8b5-f68831cd38df",
  "aud": "APP_CLIENT_ID",
  "event_id": "9968f619-caa4-4abc-9b2d-cc0f83618fed",
  "token_use": "id",
  "permissions": "read.tasks write.task", // <-- here it is
  "auth_time": 1694591667,
  "exp": 1694595267,
  "iat": 1694591668,
  "jti": "21d0ee60-15c9-44c6-9cff-8366909c3eb7",
  "email": "USER_EMAIL_ADDRESS"
}

3.4. IAM permission

For this solution to work, the Lambda function’s execution role must have dynamodb:GetItem permission.

4. Validation

We can validate the ID token in a couple of ways. This example will present two of them.

4.1. Cognito authorizer

We can create a Cognito authorizer and let User pools validate the token. The token source will be the Authorization header:

Because we use the ID token, we should leave the Authorization Scopes field blank. If we add anything here, API Gateway will think the token is an access token and will look for its scope claim. ID tokens don’t have this claim. (We can’t deceive API Gateway by adding a scope claim to the ID token with the pre-token generator function.) After we have assigned the authorizer to the endpoints, API Gateway will check if the token originates from the user pool. This way, we restricted access to our endpoint to the legitimate users of our company application.

But this is just the first part of the story because non-elevated users can still access the write endpoint. The authorizer won’t look at our custom permissions claim, so the backend code should do some checks. Assuming that we have a proxy integration to the backend, we can get the permissions claim from the event object like this:

// other properties
{
  "requestContext": {
    "authorizer": {
      "claims": {
        "permissions": "read.tasks write.task"
      }
    }
  }
}

Alternatively, if we don’t want a proxy integration, we can create a mapping template in the Integration request. We can extract the claim from the request and map it to a friendly property name in the event object the function will receive:

#set($inputRoot = $input.path('$'))
{
  "permissions": "$context.authorizer.claims.permissions"
}

The function can access the permissions directly from the input event object as event.permissions and perform the necessary check.

4.2. Lambda authorizer

Alternatively, we can create a Lambda authorizer instead of using Cognito. In this case, we won’t let the request reach the backend, but we are responsible for the token validation logic. I wrote a post about token-based Lambda authorizers a while ago, which describes how to do it in detail.

5. Considerations

The solution described in this post is not the only answer that solves the problem described in the scenario. I already wrote about some, for example, Cognito groups can address this architectural issue too. As for other possible options, I’m planning to discuss them in future posts.

6. Summary

If we need more granular access control on our API endpoints, we can add custom claims that refer to the permissions to the ID token with a pre-token generation trigger. Its target is a regular Lambda function that returns the request event object extended by the new claim we want to add to the token.

We can then use a Cognito authorizer in API Gateway to validate the token. In this case, we should check the custom permission on the backend. Alternatively, we can create a Lambda authorizer, which performs the validation and the permission check at the API Gateway level.

7. References, further reading

Tutorial: Creating a user pool - How to create a user pool

Creating a REST API in Amazon API Gateway - How to create a REST API in API Gateway

Create a DynamoDB table - How to create a DynamoDB table

Pre token generation Lambda trigger - With more examples

Control access to a REST API using Amazon Cognito user pools as authorizer - How to create a Cognito authorizer

Controlling access to HTTP APIs with JWT authorizers - And the little brother