Implementing advanced authorization with AWS Lambda for endpoint-specific access

When handling more complex authorization patterns, we can implement the necessary logic using Cognito Lambda triggers and authorizer functions.

1. The scenario

I’m working on a maths contest application where participants move from one checkpoint to another, solving engaging maths puzzles. Each checkpoint has a task assigned to teams based on their year group.

Checkpoints are staffed by invigilators or supervisors who give the tasks to participants. Students step aside to solve the task, then share their - hopefully - correct answer with the supervisor, who logs it into the application.

Naturally, the application runs on AWS serverless infrastructure. 😄

One key requirement is that supervisors at any checkpoint can only access tasks assigned to their specific checkpoint. They should not see puzzles from other checkpoints.

The API Gateway includes an endpoint structured like this:

API_URL/checkpoints/:id

This endpoint returns data for the application to display about the corresponding checkpoint. The application has many endpoints, but, in this article, I will focus on the /checkpoints/:id endpoint to explain the solution.

So, supervisors at Checkpoint 1 can only call /checkpoints/1, those at Checkpoint 2 can call /checkpoints/2, and so forth.

The application uses a Cognito user pool to manage users, including the supervisors.

2. Challenges

As I planned the solution for this scenario, I encountered several challenges. I will address the following ones below.

2.1. Endpoint protection

Supervisors, regardless of the checkpoint, call the same endpoint, with the :id path parameter being the only difference. How can I implement endpoint authorization to reject calls when the path parameter does not match the supervisor’s assigned checkpoint?

2.2. Client-side

The client dynamically adds the path parameter to the URL based on the supervisor’s checkpoint. How does the client determine which path parameter to use when making the call?

2.3. Consistency

The client needs to call all endpoints, including /checkpoints/:id, using an access token as outlined by OAuth 2.0.

3. Pre-requisites

This post won't go into detail about how to create and configure the following:

• A Cognito user pool, groups and a pre-token generation trigger.
• An app client with resource servers and scopes.
• A REST-type API Gateway and its integrations.
• Lambda functions.

Links to the relevant documentation pages will be provided at the end of the post for those who need them.

4. Solution

Here’s one solution that works for my scenario.

4.1. Overview

The solution relies on two Lambda functions.

The first function triggers after user authentication, but before the user pool issues tokens. The second serves as the endpoint authorizer.

Before implementing these, though, we need to assign supervisors to the appropriate Cognito groups.

4.2. Groups

I opted for straightforward group names for supervisor assignments: checkpoint-1, checkpoint-2, and so on. Cognito allows up to 10,000 groups per user pool, which should be plenty unless our maths contest grows to an enormous scale.

Supervisors at Checkpoint 1 are added to the checkpoint-1 group, and the same pattern applies until all supervisors are assigned to their respective groups.

4.3. Resource servers

I configured a dedicated resource server and custom scopes in the user pool for supervisor authorization management.

I won’t explain the setup process here - I covered resource servers and custom scopes in a separate article.

The custom scopes follow this format: checkpoint/1, checkpoint/2, and so on. Supervisors at Checkpoint 1 receive the checkpoint/1 scope, those at Checkpoint 2 get checkpoint/2, and the pattern continues for all checkpoints.

4.4. Pre-token generation trigger

We should tackle two tasks: let the client know the correct path parameter and attach the matching scope to the access token.

A pre-token generation Lambda trigger can handle both. Once connected to the user pool, Cognito calls the Lambda function with a payload like this:

{
  // Some properties are omitted for brevity
  "region": "eu-central-1",
  "userPoolId": "USER_POOL_ID",
  "userName": "USER_NAME",
  "request": {
    "userAttributes": {
      "sub": "USER_ID_IN_COGNITO",
    },
    "groupConfiguration": {
      "groupsToOverride": [
        // WE NEED THIS:
        "checkpoint-1",
      ],
    },
    "scopes": [
      "openid"
    ]
  },
  "response": {
    "claimsAndScopeOverrideDetails": null
  }
}

As you can see, the pre-token generation Lambda trigger receives the supervisor’s assigned checkpoint! The function can then add this information to the ID and access tokens.

ID tokens carry details about the authenticated user. We can include a custom checkpoint_id variable to store the checkpoint ID. The front-end client can pull this ID from the token and dynamically insert it as a path parameter in the URL.

Access tokens contain authorization details in the scope property. The pre-token generation trigger can add the relevant scope to the access token.

Here’s what the function’s (pseudo) code might look like:

export async function handler(event) {
  // 1. Generate scopes
  const groups = event.request.groupConfiguration.groupsToOverride;
  const userScopes = generateUserScopesFromCognitoGroups(groups);

  // 2. Get the checkpoint ID for supervisors
  const checkpointId = extractCheckpointIdFromScope(userScopes);

  // 3. Add the info to the tokens
  return {
    ...event,
    response: {
      claimsAndScopeOverrideDetails: {
        // 3.a. Checkpoint ID to the ID token
        idTokenGeneration: {
          claimsToAddOrOverride: {
            checkpoint_id: checkpointId.toString(),
          },
        },
        // 3.b. Scopes to the access token
        accessTokenGeneration: {
          scopesToAdd: ['openid', ...userScopes],
        },
      },
    },
  };
}

The generateUserScopesFromCognitoGroups utility converts a Cognito group name like checkpoint-1 into a scope format, checkpoint/1.

The extractCheckpointIdFromScope utility pulls the checkpoint ID 1 from checkpoint-1.

Once the function runs, the ID token includes this property:

"checkpoint_id": "1"
// other ID token properties

The access token reflects the corresponding checkpoint scope:

"scope": "openid checkpoint/1"
// other access token properties

With these enriched tokens, the client can extract the checkpoint_id from the ID token, append it to the URL path, and call the endpoint with the access token.

4.5. Lambda authorizer

Now, the client calls the correct URL with the checkpoint ID in the path and the access token in the Authorization header. The final step is to confirm that the checkpoint ID in the endpoint matches the scope in the token.

Cognito authorizers are too rigid for this case. Using one would require attaching all checkpoint scopes to the endpoint, allowing any supervisor to call it with any checkpoint ID.

Lambda authorizers work well for custom authorization logic. They are standard Lambda functions that receive specific authorizer payloads from API Gateway.

Of the two payload types, TOKEN and REQUEST, I chose REQUEST for this scenario. I have covered TOKEN-type payloads elsewhere, so I will focus on REQUEST payloads here.

A Lambda authorizer with a REQUEST payload gets an event object like this:

{
  // properties are omitted for brevity
  "type": "REQUEST",
  "methodArn": "arn:aws:execute-api:eu-central-1:ACCOUNT_ID:API_ID/STAGE_NAME/GET/checkpoints/1",
  "resource": "/checkpoints/{id}",
  "path": "/checkpoints/1",
  "httpMethod": "GET",
  "headers": {
    // WE NEED THIS
    "Authorization": "Bearer ACCESS_TOKEN",
  },
  "pathParameters": {
    // AND THIS
    "id": "1"
  },
  "stageVariables": {},
  "requestContext": {
    // lots of properties
  }
}

The event includes everything the authorizer needs to compare the :id path parameter with the access token’s scope.

Here’s a snippet of the authorizer’s (pseudo) code:

export async function handler(event) {
  const { headers, pathParameters, methodArn: resource } = event;

  // 1. Get the token from the Authorization header
  const token = extractAuthorizationToken(headers);
  // 2. Verify and decode the token
  const tokenPayload = verifyAccessToken(token);

  // 3. Get the "sub" and "scope" properties from the token
  const { sub: userId, scope } = payload;
  // 4. Get the checkpoint ID from the invoked URL path
  const checkPointIdInPath =
    extractCheckpointIdFromPathParameters(pathParameters);

  // 5. Compare the checkpoint IDs from the path and the token
  const matchingCheckpointIds =
    matchPathParameterIdToScope(checkPointIdInPath, scope);
  // 6. Return IAM policies as per the result
  return matchingCheckpointIds
    ? generatePolicy({ resource, effect: Effect.ALLOW, principalId: userId })
    : generatePolicy({ resource, effect: Effect.DENY, principalId: userId });
}

// Util function to generate the policy
function generatePolicy({ resource, effect, principalId }) {
  return {
    policyDocument: {
      Version: '2012-10-17',
      Statement: [
        {
          Action: 'execute-api:Invoke',
          Resource: resource,
          Effect: effect,
        },
      ],
    },
    principalId,
  };
}

These steps outline the logic described earlier.

The authorizer function must return an IAM policy. If the checkpoint IDs match, the policy will Allow the request to proceed. If not, it will Deny it.

The principalId that identifies the user making the call is a mandatory attribute. The sub property from the token, the user’s unique ID in the user pool, serves this purpose and identifies the supervisor.

5. Considerations

As noted, Cognito groups are a practical way to organize supervisors by their checkpoint assignments.

If we wanted to adapt this approach for an application with a larger user base, creating and managing groups might become impractical.

In that case, we could store user details and their properties in a separate database like DynamoDB. The pre-token trigger function could then fetch the user’s profile from the database.

This setup might be easier to manage in a rapidly changing environment, though it adds an extra API call to the database, which could slow down the response time.

6. Summary

Combining a pre-token generation trigger with a custom API Gateway authorizer function allows for more sophisticated authorization flows.

The pre-token generation function adds custom details to the ID and access tokens. Then, a Lambda authorizer can enforce the custom authorization logic.

7. Further reading

Getting started with user pools - How to create a user pool with an app client

Adding groups to a user pool - The title says it all

Pre token generation Lambda trigger - Detailed reference for adding a trigger

Creating an app client - How to create an app client for a user pool

Get started with API Gateway - How to create an API Gateway resource

Create your first function - AWS Lambda "Getting started"