Using PostgreSQL Row Level Security to isolate Tenant data in generated GraphQL CRUD APIs

This article is based on the recipe to Generate GraphQL CRUD APIs for AppSync from existing Aurora PostgreSQL DB.

Related Demo Video

Goal

Add conditional checks into all CRUD operations to isolate Tenant data.

Requirements

• PostGraphile - builds a GraphQL API from a PostgreSQL schema in seconds.
• AWS Aurora PostgreSQL - instance with existing DB schema.

Applicable use-cases

• GraphQL APIs for СRUD operations on existing PostgreSQL DB schema.
• Database tables shared between multiple Tenants.

Implementation

As an example, I will use the following simple table with Tenant ID stored in customer_id column:

create table team
(
    id serial not null constraint team_pk primary key,
    customer_id INTEGER not null,
    name varchar(256) not null,
    description varchar(256)
);

create index idx_team_customer on team (customer_id);

Enable Cognito Authentication

const pool = new UserPool(this, 'PortalUserPool', {
    ...
        customAttributes: {
            customer_id: new NumberAttribute({
                max: Number.MAX_SAFE_INTEGER,
                mutable: false,
            })
        }
})

new CfnGraphQLApi(this, 'PorltalApi', {
    name: `${Stack.of(this).stackName}-api`,
    authenticationType: 'AMAZON_COGNITO_USER_POOLS',
    userPoolConfig: { 
                awsRegion: pool.stack.region,
        userPoolId: pool.userPoolId,
    },
});

• the custom attribute customer_id will be used to store the Tenant ID user belongs to. It will be automatically passed into Direct Resolver Lambda as a value of the payload key identity.claims['custom:customer_id']

Create the RLS (row-level security) check

The following stored function is the actual implementation of a Tenant isolation check:

CREATE OR REPLACE FUNCTION tenant_check(input_tenant INTEGER) RETURNS BOOLEAN AS
$$
DECLARE
    current_tenant INTEGER;
BEGIN
    current_tenant := current_setting('jwt.claims.customer_id', TRUE)::INTEGER;
    IF current_tenant IS NOT NULL AND current_tenant != 0 THEN
        RETURN (input_tenant = current_tenant);
    END IF;
    RETURN TRUE;
END
$$ LANGUAGE plpgsql;

Few notes:

• input_tenant INTEGER - is an input variable of the same type (INTEGER) as the tenant key column (our case it is customer_id).
• current_setting('jwt.claims.customer_id', TRUE) - reads Tenant ID of the currently authenticated user (see more details see below).
• 'jwt.claims.customer_id' - custom connection config attribute used to pass the user's Tenant ID from identity.claims['custom:customer_id'] received by Lambda Resolver.
• IF current_tenant IS NOT NULL AND current_tenant != 0 - used to bypass check for the "root" (or "super-user") access where a logged-in user that has tenant ID equal to 0 permitted to access all tenants.

Enable RLS on tables

DROP POLICY IF EXISTS tenant_isolation ON public.team;
CREATE POLICY tenant_isolation ON public.team TO PUBLIC
    USING (tenant_check(customer_id)) 
    WITH CHECK (tenant_check(customer_id));
ALTER TABLE public.team ENABLE ROW LEVEL SECURITY;
ALTER TABLE public.broadcast_group FORCE ROW LEVEL SECURITY;

• USING (tenant_check(customer_id)) - will check SELECT queries (GraphQL Query fields).
• WITH CHECK (tenant_check(customer_id)) - will check INSERT/UPDATE/DELETE operations (GraphQL Mutation fields).

Propagate Authenticated Identity Claims

Postgraphile connection configuration used in Direct Lambda Resolver should be updated in the following way:

class UnauthorizedException extends Error {
    constructor() {
        super('You are not authorized to make this call.');
        this.name = 'UnauthorizedException';
    }
}

export const handler = async (event: AppSyncPayload, context: Context, callback: Callback): Promise<void> => {

    // Ensure call is authenticated
    const claims = event.identity?.claims;
    if (!claims) {
        throw new UnauthorizedException();
    }
    
    // Postgraphile connection attributes. 
    const pgSettings = {
        'jwt.claims.customer_id': claims['custom:customer_id'];
    };

    ...
}

Testing

Add a few users to Cognito User Pool:

export POOL_ID=<USER_POLL_ID> 

// Admin with super-user access (all tenants)
aws cognito-idp admin-create-user \
--user-pool-id $POOL_ID \
--username admin \
--temporary-password 'S3cure=S3cret' \
--message-action SUPPRESS \
--user-attributes \
Name=custom:customer_id,Value=0 --

// Tenant user with customer_id = 1
aws cognito-idp admin-create-user \
--user-pool-id $POOL_ID \
--username customer-user1 \
--temporary-password 'S3cure=S3cret' \
--message-action SUPPRESS \
--user-attributes \
Name=custom:customer_id,Value=1 --

Add a few records to DB table:

INSERT INTO public.team (customer_id, name, description) VALUES (1, 'Private Team', 'Customer 1 private team')
(2, 'Internal Team', 'Customer 2 internal team');

And Open Queries tab in AWS AppSync API console

Scenario 1: admin query for all records in the table

Login with UserPool as admin / S3cure=S3cret then execute the GraphQL query:

teamsConnection {
  nodes {
    customerId
    name
  }
  totalCount
}

The expected response will contain all records because tenant check bypassed for admin:

{
  "data": {
    "teamsConnection": {
      "nodes": [
        {
          "customerId": 1,
          "name": "Private Team"
        },
        {
          "customerId": 2,
          "name": "Internal Team"
        }
      ],
      "totalCount": 2
    }
  }
}

Scenario 2: user query for all records

Login with UserPool as customer-user1 / S3cure=S3cret then execute the GraphQL query:

teamsConnection {
  nodes {
    customerId
    name
  }
  totalCount
}

The expected response will contain only records having customerId = 1:

{
  "data": {
    "teamsConnection": {
      "nodes": [
        {
          "customerId": 1,
          "name": "Private Team"
        },
      ],
      "totalCount": 1
    }
  }
}

Scenario 3: User of Customer 1 tries to create a record with Tenant ID of another Customer 2 (negative test)

mutation {
  createTeam(input: { 
    team: { 
        customerId: 2, 
        name: "hacked-by-kid"
      }
  }) {
    team {
      id
    }
  }
}

The expected response will contain an error:

{
  "data": {
    "createTeam": null
  },
  "errors": [
    {
      "path": [
        "createTeam"
      ],
      "data": null,
      "errorType": "GraphQLError",
      "errorInfo": null,
      "locations": [
        {
          "line": 2,
          "column": 3,
          "sourceName": null
        }
      ],
      "message": "new row violates row-level security policy for table \"team\""
    }
  ]
}