Using customized access tokens to set up authorization in API Gateway
Tags: aws, security, serverless
Published: 2024-01-11 08:12:14

We can easily add custom scopes to access tokens after the user has authenticated with a new Cognito user pools feature. Therefore, we can achieve fine-grained access control to various API endpoints with minimal work while delegating the authorization task to the API Gateway.

## 1. The scenario

Alice recently had to set up an **authorization** mechanism for the company's internal application, which has hundreds of users. The application consists of microservices and an <a href="https://docs.aws.amazon.com/apigateway/latest/developerguide/apigateway-rest-api.html" class="link" target="_blank" rel="noopener noreferrer">Amazon API Gateway</a> protects the backend resources running the business logic.

Alice's main problem is that the API Gateway contains dozens of endpoints, and almost every endpoint should have its permission!

The following list contains part of the requirements:

| Endpoint   | Description   | Permission      |
|------------|---------------|-----------------|
| GET /pets  | List all pets | ReadPetsAccess  |
| POST /pets | Add a pet     | WritePetsAccess |

The list contains a few dozen more endpoints and permissions.

She must ensure that only users with the relevant permission can access each endpoint.

Alice <a href="https://dev.to/aws-builders/implementing-access-control-on-api-gateway-endpoints-with-id-tokens-k9p" class="link" target="_blank" rel="noopener noreferrer">is aware</a> that API Gateway <a href="https://docs.aws.amazon.com/apigateway/latest/developerguide/apigateway-integrate-with-cognito.html" class="link" target="_blank" rel="noopener noreferrer">supports authorization</a> with Cognito <a href="https://docs.aws.amazon.com/cognito/latest/developerguide/cognito-user-identity-pools.html" class="link" target="_blank" rel="noopener noreferrer">user pool</a> **tokens**. She also knows that an ideal solution would be to define a **required scope** which maps to a specific permission in the given endpoint's configuration and have the **access token** contain that scope. In this case, API Gateway would handle authorization, and the backend code would only take care of what it's supposed to, which is running the business logic.

## 2. Current solution

Until recently, this solution was only partially possible to do. Sure, we could define the required scope in the **Method request setting** in the API Gateway configuration. However, the `scope` property in the access token was not customizable based on the specific authenticated user. The authentication request must have claimed all defined scopes so that they are present in the access token. It made the user- or permission-based access control hard to implement because the token contained all scopes for everyone.

So instead, Alice decided to <a href="https://dev.to/aws-builders/controlling-access-to-the-api-with-lambda-authorizers-5bbj" class="link" target="_blank" rel="noopener noreferrer">create</a> a <a href="https://docs.aws.amazon.com/cognito/latest/developerguide/cognito-user-identity-pools.html" class="link" target="_blank" rel="noopener noreferrer">Lambda authorizer</a> that validates the authorization token and allows or denies the request accordingly based on the assigned Cognito groups in the token.

## 3. A new Cognito feature to the rescue

But it's a lot of work with some extra dependencies in the code package. <a href="https://aws.amazon.com/lambda" class="link" target="_blank" rel="noopener noreferrer">Lambda</a> authorizer functions have their well-deserved place in our AWS developer toolbox, but we can now simplify the authorization process. The <a href="https://aws.amazon.com/about-aws/whats-new/2023/12/amazon-cognito-user-pools-customize-access-tokens/" class="link" target="_blank" rel="noopener noreferrer">new Cognito user pools feature</a> makes it possible to customize the access token's `scope` property. With this, we can finally have API Gateway's built-in feature to authorize requests in a secure way with access tokens.

It means that we can now **remove** the Lambda authorizer from the architecture, add the required scope to the access token and configure API Gateway to look for that scope in the token for each endpoint!

## 4. Main steps

This post won't go over the entire process. Instead, I will highlight some key points.

### 4.1. Create a custom scope for each permission

We can create a custom OAuth2.0 scope for each permission in the **Resource server** section in Cognito user pools. So we can modify our list accordingly:

| Endpoint   | Permission      | Scope         |
|------------|-----------------|---------------|
| GET /pets  | ReadPetsAccess  | petsapi/read  |
| POST /pets | WritePetsAccess | petsapi/write |

### 4.2. Configure the API Gateway

Next, we configure API Gateway to use the access token for authorization.

![Configure the required scope in API Gateway](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/6zpuzv45qdphyu0ecjd1.png)

We go to the **Method request** of the `GET /pets` endpoint and define the `petsapi/read` authorization scope. API Gateway will validate the token then and look for the `petsapi/read` value in the `scope` property. If the access token contains this scope, API Gateway will let the request proceed, otherwise it will deny it.

With this step, we automatically set up API Gateway to expect an **access token** instead of an ID token (ID tokens don't have the `scope` property). After we have **deployed** the API, the endpoint will reject requests with no tokens, ID tokens or access tokens without `scope: petsapi/read`.

### 4.3. Enable Advanced security

The new access token customization will only work if we enable **Advanced security** in the Cognito user pool. It's a **paid feature** which currently costs $0.05 per Monthly Active User in the Frankfurt region.

### 4.4. Customize the access token with the pre token generator

We can now build a **pre token generation** Lambda function to modify the access token.

![Create pre token generation trigger](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/bqphmgo0s8bxdj9grtzn.png)

The function will run after the user has authenticated (so we know who it is) but before Cognito generates the tokens. We can use the function to add and remove scopes from the access token or modify the ID token. We should select the **Basic features + access token customization** option here.

![Configure access token customization](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/bv9jp06p0leuwm03xhoi.png)

The function's logic will add the `petsapi/read` and `petsapi/write` values to the access token's `scope` property if the user has `ReadPetsAccess` and `WritePetsAccess` permissions, respectively.

## 5. Code example

The Lambda function code can look like this:

```ts
import { PreTokenGenerationV2TriggerEvent } from 'aws-lambda';

// 1. Map the scopes to the permissions
const permissionToScopeMap: Record<string, string> = {
  ReadPetsAccess: 'petsapi/read',
  WritePetsAccess: 'petsapi/write',
  ...
};

export const handler = async (event: PreTokenGenerationV2TriggerEvent): Promise<PreTokenGenerationV2TriggerEvent> => {
  // 2. Get the permissions from the token
  const userGroups = event.request.groupConfiguration.groupsToOverride;
  if (!userGroups) {
    console.log('No group is configured to the user');
    return event;
  }

  // 3. Retrieve the scopes corresponding to the user's permissions
  const userScopes = userGroups.map((group) => {
    return permissionToScopeMap[group];
  });

  // 4. Add the customized scope to the token
  const response: PreTokenGenerationV2TriggerEvent = {
    ...event,
    response: {
      claimsAndScopeOverrideDetails: {
        accessTokenGeneration: {
          scopesToAdd: ['openid', 'profile', ...userScopes],
        },
      },
    },
  };

  return response;
};
```

First, we map the permissions to the custom scopes. In this case, I have an object for the sake of simplicity. But we can store the mappings in a separate config file or externally, for example, in <a href="https://docs.aws.amazon.com/systems-manager/latest/userguide/systems-manager-parameter-store.html" class="link" target="_blank" rel="noopener noreferrer">Parameter Store</a>.

Next, we retrieve the permission we have already configured in Cognito user pools as **groups**. For this example, it doesn't matter what IAM permissions we assign to the groups because we don't use the role. It can be a blank role with no permissions. We use the groups to know which user needs which scopes. A user can belong to multiple groups in Cognito, so the user's access token will contain all corresponding scopes.

The function's `return` value must be the original <a href="https://docs.aws.amazon.com/cognito/latest/developerguide/user-pool-lambda-pre-token-generation.html#cognito-user-pools-lambda-trigger-syntax-pre-token-generation" class="link" target="_blank" rel="noopener noreferrer">`event` object</a> with all changes we want to make on the token. In this case, we should write our code to expect the **Version 2** of the pre token generation trigger parameters as input.

## 6. Considerations

I believe the new access token customization feature will simplify how we control access to API endpoints. There's no need for workarounds as it's plain and simple. We don't have to set up Lambda authorizers since we can now use API Gateway's built-in token validation feature. The trade-offs are that it's a paid feature, and we will still have a Lambda function running as the pre token generator trigger.

## 7. Summary

Cognito access tokens are now customizable! It means that we can add or remove claims and scopes from them. The new feature simplifies fine-grained access control to our endpoints.

With the help of a pre token generator Lambda trigger function, we can add the scope the authenticated user needs to the access token. Then, the API Gateway will validate the token, look for the required scope, and allow or deny the request based on the `scope` property.

## 8. Further reading

<a href="https://docs.aws.amazon.com/cognito/latest/developerguide/cognito-user-pools-define-resource-servers.html" class="link" target="_blank" rel="noopener noreferrer">OAuth 2.0 scopes and API authorization with resource servers</a> - How to create custom scopes in Cognito user pools

<a href="https://docs.aws.amazon.com/cognito/latest/developerguide/cognito-user-pool-settings-advanced-security.html" class="link" target="_blank" rel="noopener noreferrer">Adding advanced security to a user pool</a> - The title says it all

<a href="https://aws.amazon.com/cognito/pricing/" class="link" target="_blank" rel="noopener noreferrer">Amazon Cognito Pricing</a> - Cognito pricing page

<a href="https://docs.aws.amazon.com/cognito/latest/developerguide/getting-started-with-cognito-user-pools.html" class="link" target="_blank" rel="noopener noreferrer">Getting started with user pools</a> - How to create a user pool with an app client

<a href="https://docs.aws.amazon.com/cognito/latest/developerguide/cognito-user-pools-user-groups.html" class="link" target="_blank" rel="noopener noreferrer">Adding groups to a user pool</a> - Straighforward title from AWS

<a href="https://docs.aws.amazon.com/apigateway/latest/developerguide/how-to-create-api.html" class="link" target="_blank" rel="noopener noreferrer">Creating a REST API in Amazon API Gateway</a> - How to create a REST API in API Gateway

<a href="https://docs.aws.amazon.com/cognito/latest/developerguide/user-pool-lambda-pre-token-generation.html" class="link" target="_blank" rel="noopener noreferrer">Pre token generation Lambda trigger</a> - With more examples