Fine-Grained Authorization with Amazon Cognito

As of a week ago, you can now modify the contents of the access token, making it possible to add fine-grained authorization decisions to your application that is leveraging Amazon Cognito for user authentication. Previously it has only been possible to modify the contents of the id token, which should solely be used to authenticate the user, not to authorize operations. This is what the access token is for. Modifying the properties of the access token requires that the User Pool is configured to enforce Advanced Security.

Amazon Cognito Lambda Triggers

Amazon Cognito has support for different Lambda triggers that can be used to customize the behavior of a user pool at various points in the authentication and authorization process. In my previous post, I explained how some of these triggers can be used to implement the delegation (or on-behalf-of) flow using Amazon Cognito, one of the triggers used in that solution is the "Pre Token Generation" trigger. This trigger can be used to add or remove properties of the id and access token, which is exactly what we need to be able to implement fine-grained authorization in our application.

Structuing the Permissions

To keep the id and access token relatively slim, we need to ensure that we only add the permissions valid for the current authorization context, based on the client that is requesting the token. In this example we will work with a booking application, the solution will contain two different domains, one that handles booking and one that handles reviews. When our front-end application requests an access token for the end user, we know that it will only communicate with the booking services, not with the review service, therefore we only need to add permissions relevant to that domain.

We also know that there will likely be a lot of users using our application over time, therefore we need to ensure that we have an easy way to manage the permissions for our user base, we can't add and remove permissions on every user, instead, we need a role-concept. We bind permissions to roles rather than to the users directly. Amazon Cognito has a user group concept that we can utilize for this use case, users of the group "User" get a certain list of permissions, and users of the group "Admin" get another.

Based on this, this is how we will structure our permissions for this solution:

Client	Role	Permissions
BookingClient	Admin	booking:read, booking:write, booking:delete
BookingClient	User	booking:read
ReviewClient	Admin	review:read, review:write, review:delete
ReviewClient	User	review:read

Implementation

We will use CDK to define our stack, I will focus on the parts essential for this solution, the complete example can be found here.

Let's start with creating a new CDK project by running cdk init app --language typescript in an empty folder called "permissions" and edit the permissions-stack.ts-file to add a UserPool and enable advanced security mode:

// Create user pool
const userPool = new UserPool(this, "UserPool", {
  userPoolName: "permissions",
  customAttributes: {
    permission: new StringAttribute({ mutable: true }),
  },
});

userPool.addDomain("UserPoolDomain", {
  cognitoDomain: {
    domainPrefix: "permissions",
  },
});

// Enable advanced security
const cfnUserPool = userPool.node.defaultChild as CfnUserPool;
cfnUserPool.userPoolAddOns = {
  advancedSecurityMode: "ENFORCED",
};

// Store user pool id in SSM, to be used by Lambda integrations
new StringParameter(this, "UserPoolIdParameter", {
  parameterName: "/permissions/userpool/id",
  stringValue: userPool.userPoolId,
});

Then we will add a resource server to manage the scopes required to access our different services, as mentioned, in this example, I will add a booking service and a review service.

// Add resource server
const resourceServer = userPool.addResourceServer("ResourceServer", {
  identifier: "resources",
  scopes: [
    {
      scopeName: "booking-service",
      scopeDescription: "Access booking service",
    },
    {
      scopeName: "review-service",
      scopeDescription: "Access review service",
    },
  ],
});

We will use DynamoDB to store the client permissions, let's define the table next:

// Add dynamodb table to store permissions in
const table = new Table(this, "Permissions", {
  partitionKey: { name: "pk", type: AttributeType.STRING },
  sortKey: { name: "sk", type: AttributeType.STRING },
  tableName: "Permissions",
  billingMode: BillingMode.PAY_PER_REQUEST,
});

We will also add the Admin and User groups that we will use to tie permissions to:

// Add Admin group to user pool
new CfnUserPoolGroup(this, "AdminGroup", {
  groupName: "Admin",
  userPoolId: userPool.userPoolId,
  description: "Admin group",
});

// Add User group to user pool
new CfnUserPoolGroup(this, "UserGroup", {
  groupName: "User",
  userPoolId: userPool.userPoolId,
  description: "User group",
});

Next, we will define the Pre Token Generation Lambda (pre-token-generation.ts), which we store in a new folder called handlers in the lib-folder. The type definitions for Pre Token Generation V2 have not yet been included in the aws-lambda-package, so it must be manually defined, they can be found here. I have excluded them in the below code:

import { BatchGetItemCommand, DynamoDBClient } from "@aws-sdk/client-dynamodb";

// Type definitions

const client = new DynamoDBClient({ region: "eu-north-1" });
const tableName = process.env.TABLE_NAME!;

export const handler = async (
  event: PreTokenGenerationTriggerHandler
) => {
  const claims = event.request.userAttributes;
  const groups = event.request.groupConfiguration.groupsToOverride;

  if ((groups?.length ?? 0) > 0) {
    const command = new BatchGetItemCommand({
      RequestItems: {
        [tableName]: {
          Keys: groups.map((group) => ({
            pk: { S: event.callerContext.clientId },
            sk: { S: group },
          })),
        },
      },
    });
    const data = await client.send(command);

    const permissions = [...new Set(data.Responses![tableName].map(
      (item) => item.permissions.SS ?? []
    ).flat())];

    claims["permissions"] = permissions.join(",");
  }

  event.response = {
    claimsAndScopeOverrideDetails: {
      idTokenGeneration: {
        claimsToAddOrOverride: claims,
      },
      accessTokenGeneration: {
        claimsToAddOrOverride: claims,
      },
    },
  };

  return event;
};

In the code above we look up permissions based on the calling client and the end-user groups, the permissions are added to the id and access tokens as a permissions claim, separated with a space. Now we will continue to update our stack and add the Pre-Token Generation trigger:

// Define pre token generation lambda
const preTokenGeneration = new NodejsFunction(this, "PreTokenGeneration", {
  runtime: Runtime.NODEJS_20_X,
  handler: "handler",
  entry: `${__dirname}/handlers/pre-token-generation.ts`,
  memorySize: 1769,
  environment: {
    TABLE_NAME: table.tableName,
  },
});

preTokenGeneration.addToRolePolicy(
  new PolicyStatement({
    actions: ["ssm:GetParameter"],
    resources: [
      Stack.of(this).formatArn({
        service: "ssm",
        resource: "parameter/permissions/userpool/id",
      }),
    ],
  }),
);

preTokenGeneration.addToRolePolicy(
  new PolicyStatement({
    actions: ["cognito-idp:DescribeUserPoolClient"],
    resources: ["*"],
  })
);

userPool.addTrigger(
  UserPoolOperation.PRE_TOKEN_GENERATION,
  preTokenGeneration
); // Needs to be changed to V2_0 in console until CDK supports it

// Grant read access to table
table.grantReadData(preTokenGeneration);

That's it, we have now added the resources required for our fine-grained authorization solution, the next steps are to add our clients to test the solution with. We will add two clients, "BookingClient" and "ReviewClient".

// Create booking client
const bookingClient = new UserPoolClient(this, "BookingClient", {
  userPool,
  generateSecret: false,
  userPoolClientName: "Booking",
  authFlows: {
    userPassword: false,
    userSrp: true,
    custom: false,
    adminUserPassword: false,
  },
  oAuth: {
    flows: {
      authorizationCodeGrant: true,
    },
    scopes: [
      OAuthScope.EMAIL,
      OAuthScope.OPENID,
      OAuthScope.PROFILE,
      OAuthScope.custom(
        `${resourceServer.userPoolResourceServerId}/booking-service`
      ),
    ],
    callbackUrls: ["http://localhost:3000"],
    logoutUrls: ["http://localhost:3000"],
  },
});

// Create review client
const reviewClient = new UserPoolClient(this, "ReviewClient", {
  userPool,
  generateSecret: false,
  userPoolClientName: "Review",
  authFlows: {
    userPassword: false,
    userSrp: true,
    custom: false,
    adminUserPassword: false,
  },
  oAuth: {
    flows: {
      authorizationCodeGrant: true,
    },
    scopes: [
      OAuthScope.EMAIL,
      OAuthScope.OPENID,
      OAuthScope.PROFILE,
      OAuthScope.custom(
        `${resourceServer.userPoolResourceServerId}/review-service`
      ),
    ],
    callbackUrls: ["http://localhost:3000"],
    logoutUrls: ["http://localhost:3000"],
  },
});

We will also add a get-booking Lambda handler, which we will use as a proof-of-concept to verify that the permissions is present and can be used to determine whether or not the user is authorized to read bookings. We'll add the handler, get-booking.ts, in the handlers-folder that we created before.

import { APIGatewayProxyEvent } from "aws-lambda";

export const handler = async (event: APIGatewayProxyEvent) => {
  const claims = event.requestContext.authorizer?.claims;

  // Verify that the permissions is set and that the permissions contain booking:read
  if (!Boolean(claims?.permissions) || !claims.permissions.includes("booking:read")) {
    return {
      statusCode: 403,
      body: JSON.stringify({
        message: "Forbidden",
      }),
    };
  }

  return {
    statusCode: 200,
    body: JSON.stringify({
      id: "456",
      date: new Date().toISOString(),
      name: "Hotel California",
    }),
  };
};

And add the handler to the stack with an API Gateway and Cognito authorizer infront:

// Crate REST API
const api = new RestApi(this, "PermissionsApi", {
  restApiName: "permissions-api",
});

// Create Booking Lambda
const getBookingLambda = new NodejsFunction(this, "GetBooking", {
  runtime: Runtime.NODEJS_20_X,
  handler: "handler",
  entry: `${__dirname}/handlers/get-booking.ts`,
  memorySize: MEMORY_SIZE,
});
const getBookingIntegration = new LambdaIntegration(getBookingLambda);

// Create Cognito authorizer
const authorizer = new CfnAuthorizer(this, "CognitoAuthorizer", {
  restApiId: api.restApiId,
  type: "COGNITO_USER_POOLS",
  identitySource: "method.request.header.Authorization",
  providerArns: [userPool.userPoolArn],
  name: "CognitoAuthorizer",
});

// Add booking resource
const bookingResource = api.root.addResource("booking");
bookingResource.addMethod("GET", getBookingIntegration, {
  authorizationType: AuthorizationType.COGNITO,
  authorizer: {
    authorizerId: authorizer.ref,
  },
  authorizationScopes: [
    `${resourceServer.userPoolResourceServerId}/booking-service`,
  ],
});
bookingResource.addCorsPreflight({
  allowOrigins: ["http://localhost:3000"],
  allowMethods: ["GET", "OPTIONS", "POST", "PUT", "DELETE"],
  allowHeaders: ["*"],
});

That's it, now we have all the resources required to also test the concept, let's deploy them and test it out!

Testing the Solution

To test the fine-grained authorization we need first to create a user that can login to the user pool and add the user to one of the two groups, "User" or "Admin", this can be done through the AWS Console or the AWS CLI:

aws cognito-idp admin-create-user \
  --user-pool-id <USER_POOL_ID> \
  --username <USERNAME> \
  --user-attributes Name=email,Value=<EMAIL> Name=email_verified,Value=true \
  --temporary-password <TEMPORARY_PASSWORD> \
  --message-action SUPPRESS

We also need to populate the Permissions-table, this can be done with AWS CLI by running aws dynamodb batch-write-item --request-items file://sample-data.json where "sample-data.json" contains the following (make sure to replace the client id's):

{
  "Permissions": [
    {
      "PutRequest": {
        "Item": {
          "pk": { "S": "<BOOKING_CLIENT_ID>" },
          "sk": { "S": "Admin" },
          "permissions": { "SS": ["booking:write", "booking:read", "booking:delete"] }
        }
      }
    },
    {
      "PutRequest": {
        "Item": {
          "pk": { "S": "<BOOKING_CLIENT_ID>" },
          "sk": { "S": "User" },
          "permissions": { "SS": ["booking:read"] }
        }
      }
    },
    {
      "PutRequest": {
        "Item": {
          "pk": { "S": "<REVIEW_CLIENT_ID>" },
          "sk": { "S": "Admin" },
          "permissions": { "SS": ["review:write", "review:read", "review:delete"] }
        }
      }
    },
    {
      "PutRequest": {
        "Item": {
          "pk": { "S": "<REVIEW_CLIENT_ID>" },
          "sk": { "S": "User" },
          "permissions": { "SS": ["review:read"] }
        }
      }
    }
  ]
}

Now it is time to authenticate using the "BookingClient" and call the API using the end-user credentials. This can be done through any API client such as Insomnia or Postman, or you could use the simple client defined in the companion repository of this post. Then you can call the booking endpoint to verify that you have permission to call it, also try to remove yourself from the groups that give the permissions and call the API again to see that you no longer have access to call it.

I recommend that you combine this fine-grained authorization method with the delegation (or on-behalf-of) flow mentioned in my previous post when you need to perform machine-to-machine calls, this to ensure that the permissions are updated between the requests, based on the Cognito client.

<- Back