Multi-Tenant AWS Amplify: Method 1: Cognito Custom Attributes

7 min readJan 5, 2020

Return to Multi-Tenant AWS Amplify: Overview

This is method 1 of 3 for creating multi-tenant AWS Amplify mobile apps in React Native. In this method, the tenant information is stored in an AWS Cognito custom attribute.

The example code for this post uses React Native 61.5 and AWS Amplify 2.2.1, and is at https://github.com/dantasfiles/AmplifyMultiTenant1

We use an AWS Cognito Post Confirmation Lambda Trigger to store the tenant information in an AWS Cognito custom attribute.
Recall that a Post Confirmation Lambda Trigger is invoked “after a new user is confirmed, allowing you to… add custom logic.”

Custom attributes, including our tenant custom attribute, only appear in the ID token, not the access token. The AWS Amplify API passes the access token by default for use in access control, as seen in aws-amplify/amplify-js/packages/api/src/API.ts

// aws-amplify/amplify-js/packages/api/src/API.ts
case 'AMAZON_COGNITO_USER_POOLS':
  const session = await Auth.currentSession();    
  headers = {     
    Authorization: session.getAccessToken().getJwtToken(),    
};`

However, the AWS Amplify API can be overridden to pass the ID token, which does contains the tenant custom user attribute, instead of the access token.

Note: as specified in Why You Should Always Use Access Tokens to Secure APIs, this has downsides.

This tenant custom attribute can then be used by Custom Claims in the API for access control checks.

Upsides of this method include that it is less expensive: the lambda function is only invoked once upon confirmation of a user, limiting AWS Lambda costs (as opposed using the Pre Token Generation Lambda Trigger in Method 3: Virtual Cognito Groups).

Downsides of this method include that it’s more complicated to setup than the other methods, and that it’s not preferable to use ID tokens to secure APIs as described above.

Setup Authentication

Perform the initial steps as described in Multi-Tenant AWS Amplify: Overview

Add AWS Amplify authentication. Make sure to enable Add User to Group which creates a Post Confirmation Lambda Trigger, and enter any value for the name of the group. We don’t actually want to add a user to a group (unlike in Method 2: Cognito Groups), but we want a Post Confirmation Lambda Trigger and all of its associated resources and roles to be created so we can edit them later.

> amplify add auth  
... 
 Do you want to enable any of the following capabilities?
 ( ) Add Google reCaptcha Challenge
 ( ) Email Verification Link with Redirect
>(*) Add User to Group
 ( ) Email Domain Filtering (blacklist)
 ( ) Email Domain Filtering (whitelist)
 ( ) Custom Auth Challenge Flow (basic scaffolding - not for production)
 ( ) Override ID Token Claims
 Do you want to enable any of the following capabilities? Add User to Group
? Enter the name of the group to which users will be added. ENTER ANYTHING HERE
Succesfully added the Lambda function locally

Push the authentication setup into the cloud, which will create a user pool to which you can add the tenant custom attribute in the next section.
> amplify push

Add the `tenant` custom attribute in AWS Cognito

In Manage User Pools in the AWS Cognito console, choose amplifymultitenantXXXX_userpool_4bc30edc-env, then click Choose custom attributes…

Click Add Custom attribute

Add the tenant custom attribute, and save changes

Setup the Post Confirmation Lambda Trigger

Modify amplify\backend\function\amplifymultitenantXXXXPostConfirmation\src\add-to-group.js

// amplify\backend\function\amplifymultitenantXXXXPostConfirmation\src\add-to-group.js
/* eslint-disable-line */ const aws = require('aws-sdk');
exports.handler = async (event, context, callback) => {
  const cisp = new aws.CognitoIdentityServiceProvider({
    apiVersion: '2016-04-18',
  });
  const updateParams = {
    UserAttributes: [
      {
        Name: 'custom:tenant',
        Value: '', // ADD YOUR TENANT LOGIC HERE
       },
    ],
    UserPoolId: event.userPoolId,
    Username: event.userName,
  };
  try {
    await cisp.adminUpdateUserAttributes(updateParams).promise();
    callback(null, event);
  } catch (e) {
    callback(e);
  }
};

This function is run after a user is confirmed, and it uses the CognitoIdentityServiceProvider.adminUpdateUserAttributes function to add the tenant information to the tenant custom attribute.
Your tenant selection logic can read and use any Post Confirmation Lambda Trigger parameters, including AWS Cognito user attributes, which are passed to the add-to-group.js Lambda function in the event.request.userAttributes parameter.

The Lambda function needs permission to modify user attributes, but the function was originally set up with an Add User to Group policy.
Overwrite the previous Add User to Group policy with cognito-idp:AdminUpdateUserAttributes permissions by modifying amplify\backend\auth\amplifymultitenantXXXX\amplifymultitenantXXXX-cloudformation-template.yml

# amplify\backend\auth\amplifymultitenantXXXX\amplifymultitenantXXXX-cloudformation-template.yml
Resources:
  ...
  amplifymultitenantXXXXPostConfirmationAddToGroupCognito:
    Type: AWS::IAM::Policy
    Properties:
      PolicyName: 
        amplifymultitenantXXXXPostConfirmationAddToGroupCognito
      PolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Action:
              - cognito-idp:AdminUpdateUserAttributes
            Resource: !GetAtt
              - UserPool
              - Arn
      Roles:
        - !Join [
            '',
         ['amplifymultitenantXXXXPostConfirmation', '-', !Ref env],
          ]

Note: Instead of directly modifying the cloud formation template, it would be preferable to modify the permissions field in amplify\backend\auth\amplifymultitenantXXXX\parameters.json, but this does not appear to work.

Push the function setup into the cloud
> amplify push

The ID token of a user will now contain the required tenant information in its tenant custom attribute.

Viewing the tenant

We can extract the tenant information from the tenant custom attribute in the ID token with the following example code in App.js

// App.js
async function fetchUserInfo(setTenant) {
  // get the id token of the signed in user
  const {idToken} = await Auth.currentSession();
  // get the tenant custom attribute from the id token
  const tenant = idToken.payload['custom:tenant'];
  setTenant(tenant);
}
const App = withAuthenticator(() => {
  const [tenant, setTenant] = useState('');
  useEffect(() => {
    fetchTenant(setTenant);
  }, []);
  return (
    ...
    <Text style={styles.sectionDescription}>
      Your tenant is {tenant}
    </Text>
    ...
  );});

Note: This is a good opportunity to test if your tenant custom attribute is working before moving on to setting up the API.

Setup AWS Amplify API to use the ID token instead of the access token

Custom attributes only appear in the ID token, not the access token. We must modify the API so that it uses the ID instead of API token.
Add to index.js as documented in Set Custom Request Headers for GraphQL

// index.js
import Amplify, {Auth} from 'aws-amplify';
import config from './aws-exports';
Amplify.configure(config);
// Use ID token instead of access token in API calls
Amplify.configure({
  API: {
    graphql_headers: async () => {
      const session = await Auth.currentSession();
      return {
        Authorization: session.getIdToken().getJwtToken(),
      };
    },
  },
});

Note: as specified in Why You Should Always Use Access Tokens to Secure APIs, this has downsides.

Securing an API

We can now use the tenant custom attribute in the ID token to secure our API.

Add an AWS Amplify API
> amplify add api

Edit amplify\backend\api\amplifymultitenant\schema.graphql

type Todo
  @model(subscriptions: null)
  @auth(rules: [{
    allow: owner, 
    ownerField: "tenant", 
    identityClaim: "custom:tenant"}]) {
  id: ID!
  tenant: ID!
  name: String!
  description: String
}

First,tenant: ID! stores the tenant associated with the Todo item.

Second, the Custom Claims in @auth(rules: [{allow: owner, ownerField: "tenant", identityClaim: "custom:tenant"}]) looks in the tenant field (specified by the ownerField) of the Todo and matches it against the tenant custom attribute (specified by the identityClaim) in the ID token. Recall that in a previous step, we added the tenant custom attribute to the ID token.

For queries, the @auth rule will only allow a user to view Todos that are associated with their tenant. For example, here is a simplified version of the access control code that AWS Amplify generates for the getTodo GraphQL operation in the Query.getTodo.res.vtl resolver.
It makes sure that a Todo associated with a tenant (stored in $ctx.result.tenant) can only be read by a user associated with that tenant (stored in $ctx.identity.claims.get("custom:tenant")).

#if($ctx.result.tenant == $ctx.identity.claims.get("custom:tenant"))
  #set( $isOwnerAuthorized = true )
#end

For mutations, the @auth rule will only allow a user to add Todos that have the tenant field set to the user’s tenant.
In addition, if no tenant is specified as an input to a createTodo mutation, it’ll automatically insert the tenant custom attribute of the user into the tenant field of the Todo. Here is a simplified version of the access control code in the generated Mutation.createTodo.req.vtl resolver:

#set( $allowedOwners0 = 
        $util.defaultIfNull($ctx.args.input.tenant, null) )
#if( $allowedOwners0 == $ctx.identity.claims.get("custom:tenant") )
  #set( $isOwnerAuthorized = true )
#end
#if( $util.isNull($allowedOwners0) && 
     (! $ctx.args.input.containsKey("tenant")) )
  $util.qr($ctx.args.input.put("tenant", 
                        $ctx.identity.claims.get("custom:tenant")))
  #set( $isOwnerAuthorized = true )
#end

Dynamic Group Authorization

Instead of owner authorization, you can use Dynamic Group Authorization with the rule @auth(rules: [{allow: groups, groupsField: "tenant", groupClaim: "custom:tenant"}])

However, AWS Amplify expects the groupClaim to contain a list, not a single string. So to use Dynamic Group Authorization, you need to change your Post Confirmation Lambda Trigger code slightly. If your tenant’s name is ClientInc then you need to place ["ClientInc"] in the tenant custom attribute instead of just ClientInc, so the Dynamic Group Authorization interprets the tenant custom attribute as a list with one string element.

For more information on the difference between owner and group authorization, see Owner vs. Group Access Control in AWS Amplify API