Cognito UserPools - Hosted UI Customizations (including logo upload)

As per the https://github.com/aws/aws-cdk/issues/6765 tracking issue, CDK doesn’t yet have construct support for all Cognito things. We can use escape hatches, but it would be good to have native support to be able to apply Cognito UserPools Hosted UI Customisations (CSS, logo upload, etc)

Use Case

To apply Cognito UserPool Hosted UI customizations as part of my CDK stack, without having to resort to escape hatches/workarounds.

Proposed Solution

Implement some new constructs that support the UI customisations. As Cloudformation doesn’t currently support uploading the logo, this would probably be achieved by a custom resource that calls the AWS SDK or similar.

The following code snippets are my initial attempts to workaround this with the escape hatch, but while they seemed to deploy, I don’t know that they actually worked in the end. In particular, I’m not sure the latter AWS SDK call was working particularly well for the image upload. Originally I intended to use the CloudFormation part for the CSS, and just do the logo via the SDK, but I chopped and changed the code a little in trying to get things working.

You require a domain set on the UserPool before you are able to apply customisations.

The logo can apparently only be JPG/PNG.

Refs:

• CDK
  • https://docs.aws.amazon.com/cdk/api/latest/docs/@aws-cdk_aws-cognito.CfnUserPoolUICustomizationAttachment.html
  • https://docs.aws.amazon.com/cdk/api/latest/docs/@aws-cdk_custom-resources.AwsCustomResource.html
• https://docs.aws.amazon.com/AWSCloudFormation/latest/UserGuide/aws-resource-cognito-userpooluicustomizationattachment.html
• https://docs.aws.amazon.com/cognito/latest/developerguide/cognito-user-pools-app-ui-customization.html
• https://docs.aws.amazon.com/cognito-user-identity-pools/latest/APIReference/API_SetUICustomization.html
• https://docs.aws.amazon.com/AWSJavaScriptSDK/latest/AWS/CognitoIdentityServiceProvider.html#setUICustomization-property

Warning, this code may not actually work as it is currently:

const fs = require('fs')

/**
 * Cognito Hosted UI customisations.
 *
 * @type {CfnUserPoolUICustomizationAttachment}
 *
 * @see https://docs.aws.amazon.com/cdk/api/latest/docs/@aws-cdk_aws-cognito.CfnUserPoolUICustomizationAttachment.html
 * @see https://docs.aws.amazon.com/AWSCloudFormation/latest/UserGuide/aws-resource-cognito-userpooluicustomizationattachment.html
 * @see https://docs.aws.amazon.com/cognito/latest/developerguide/cognito-user-pools-app-ui-customization.html
 */
const userPoolHostedUICustomisation = new cognito.CfnUserPoolUICustomizationAttachment(
  this,
  'UserPoolHostedUICustomisation',
  {
    userPoolId: userPool.userPoolId,
    clientId: 'ALL',
    css: fs.readFileSync('./assets/cognito-hosted-ui.css').toString('utf-8'),
  }
)
userPoolHostedUICustomisation.node.addDependency(userPool)
userPoolHostedUICustomisation.node.addDependency(userPoolDomain)

/**
 * Use the AWS SDK to upload a custom logo + CSS for the Cognito Hosted UI customisations.
 *
 * @type {AwsCustomResource}
 *
 * @see https://docs.aws.amazon.com/cdk/api/latest/docs/@aws-cdk_custom-resources.AwsCustomResource.html
 * @see https://docs.aws.amazon.com/cognito-user-identity-pools/latest/APIReference/API_SetUICustomization.html
 * @see https://docs.aws.amazon.com/AWSJavaScriptSDK/latest/AWS/CognitoIdentityServiceProvider.html#setUICustomization-property
 */
const userPoolHostedUICustomisations = new cr.AwsCustomResource(
  this,
  'UserPoolHostedUILogo',
  {
    resourceType: 'Custom::SetCognitoUserPoolHostedUILogo',
    onCreate: {
      service: 'CognitoIdentityServiceProvider',
      action: 'setUICustomization',
      parameters: {
        UserPoolId: userPool.userPoolId,
        ClientId: 'ALL',
        // Note: Wrap with Buffer because the SDK automatically Base64 encodes string, which double encodes the image
        ImageFile: Buffer.from(fs.readFileSync('./assets/logo.png').toString('base64')),
        // ImageFile: fs.readFileSync('./assets/logo.png'),
        CSS: fs
          .readFileSync('./assets/cognito-hosted-ui.css')
          .toString('utf-8'),
      },
      physicalResourceId: cr.PhysicalResourceId.of(
        `cognito-ui-logo-all-clients-${userPoolDomain.domain}`
      ),
    },
    // TODO: can we restrict this policy more? Get the ARN for the user pool domain? Or the user pool maybe?
    policy: cr.AwsCustomResourcePolicy.fromSdkCalls({
      resources: cr.AwsCustomResourcePolicy.ANY_RESOURCE,
      // TODO: ?? resources: [userPool.userPoolArn],
    }),
  }
)

Other

• 👋 I may be able to implement this feature request
• ⚠️ This feature might incur a breaking change

---

This is a 🚀 Feature Request

[skrud-dt]

Hi all, this is my hacky solution, I hope you all find it helpful.

1. Create a CDK construct that uploads the Logo and CSS file as s3_asset.Assets.
2. Create a Custom Resource using the Provider framework, whose inputs are the userPoolId, as well as the bucketName and objectKey for each of the logo/css file.
3. And here is the “weird” part – the onEvent handler is actually a no-op (see explanation below). Instead the logic lives in the isComplete handler, which downloads the files from S3 and calls SetUICustomization.

OK, so why is the onEvent lambda a no-op? Well,

• Cognito requires that you have a domain in order to call SetUiCustomization.
• The domain must be ACTIVE when call SetUiCustomization, otherwise the call will fail with this error: The domain associated with this user pool is currently not Active.
• The domain is initially in the CREATED state. So when you’re creating the userPoolDomain via CDK, _you’ll need to wait until the domain is “active” until you can do the customization.

IMO, the construct that creates the UserPoolDomain should do the awaiting – since you can’t really login using that domain until it’s Active anyway. But since it doesn’t, this is the approach I took.

Here are some files that may help:

This is the construct, which assumes that the assets (logo and css) are in ../static/. Note that the permissions for DescribeUserPoolDomain are for ['*']…

import {
  Construct,
  CustomResource,
  Duration,
  RemovalPolicy,
} from '@aws-cdk/core';
import * as s3_asset from '@aws-cdk/aws-s3-assets';
import * as cr from '@aws-cdk/custom-resources';
import * as lambda_node from '@aws-cdk/aws-lambda-nodejs';
import * as iam from '@aws-cdk/aws-iam';
import * as cognito from '@aws-cdk/aws-cognito';
import * as logs from '@aws-cdk/aws-logs';
import path from 'path';
import {UpdateCognitoUiProperties} from './cognitoUiStyle.setUiStyle';

export interface CognitoUiStyleProps {
  userPool: cognito.IUserPool;
  userPoolClient: cognito.IUserPoolClient;
}

export class CognitoUiStyle extends Construct {
  constructor(scope: Construct, id: string, props: CognitoUiStyleProps) {
    super(scope, id);

    const eventFn = new lambda_node.NodejsFunction(this, 'noop', {
      description: 'No op waiting for domain to come up',
    });
    const completeFn = new lambda_node.NodejsFunction(this, 'setUiStyle', {
      description: 'Update Cognito Hosted UI Style',
    });

    const policy = new iam.PolicyStatement({
      actions: [
        'cognito-idp:SetUICustomization',
        'cognito-idp:DescribeUserPool',
      ],
      effect: iam.Effect.ALLOW,
      resources: [props.userPool.userPoolArn],
    });
    completeFn.addToRolePolicy(policy);

    // Note that the resource for DescribeUserPoolDomain needs to be "*" since we can't get an ARN for the cognitoDomain.
    completeFn.addToRolePolicy(
      new iam.PolicyStatement({
        actions: ['cognito-idp:DescribeUserPoolDomain'],
        effect: iam.Effect.ALLOW,
        resources: ['*'],
      })
    );

    const cssFileAsset = new s3_asset.Asset(this, 'Css', {
      path: path.resolve(__dirname, '..', 'static', 'hosted-ui.css'),
      readers: [completeFn],
    });
    const logoAsset = new s3_asset.Asset(this, 'Logo', {
      path: path.resolve(__dirname, '..', 'static', 'logo.png'),
      readers: [completeFn],
    });

    const provider = new cr.Provider(this, 'CrProvider', {
      onEventHandler: eventFn,
      isCompleteHandler: completeFn,
      queryInterval: Duration.seconds(30),
      totalTimeout: Duration.hours(1),
      logRetention: logs.RetentionDays.ONE_WEEK,
    });

    const crProperties: UpdateCognitoUiProperties = {
      cssLocator: {
        bucketName: cssFileAsset.s3BucketName,
        objectKey: cssFileAsset.s3ObjectKey,
      },
      logoLocator: {
        bucketName: logoAsset.s3BucketName,
        objectKey: logoAsset.s3ObjectKey,
      },
      userPoolClientId: props.userPoolClient.userPoolClientId,
      userPoolId: props.userPool.userPoolId,
    };
    new CustomResource(this, 'CustomResource', {
      serviceToken: provider.serviceToken,
      removalPolicy: RemovalPolicy.DESTROY,
      properties: crProperties,
      resourceType: 'Custom::CognitoUiCustomization',
    });
  }
}

This is the noop:

import {CdkCustomResourceHandler} from 'aws-lambda';

export const handler: CdkCustomResourceHandler = async event => {
    const {userPoolId} = event.ResourceProperties;
    switch (event.RequestType) {
      case 'Create':
      case 'Update':
        return {
          PhysicalResourceId: userPoolId + '-' + new Date().toISOString(),
        };

      case 'Delete':
        return {
          PhysicalResourceId: event.PhysicalResourceId,
        };
    }
  }

This is the lambda that does the customization (I’m using runtypes to type-check the ResourceProperties and stream-buffers to read the S3 Readable stream into a Buffer).

import {
  CognitoIdentityProviderClient,
  DescribeUserPoolCommand,
  DescribeUserPoolDomainCommand,
  DomainDescriptionType,
  SetUICustomizationCommand,
  UserPoolType,
} from '@aws-sdk/client-cognito-identity-provider';
import {GetObjectCommand, S3Client} from '@aws-sdk/client-s3';
import {
  CdkCustomResourceIsCompleteHandler,
  CdkCustomResourceIsCompleteResponse,
  CloudFormationCustomResourceCreateEvent,
  CloudFormationCustomResourceUpdateEvent,
} from 'aws-lambda';
import * as runtypes from 'runtypes';
import {Readable} from 'stream';
import streamBuffers from 'stream-buffers';

const s3ResourceLocatorRuntype = runtypes.Record({
  bucketName: runtypes.String,
  objectKey: runtypes.String,
});

export const updateCognitoUiPropertiesRuntype = runtypes.Record({
  userPoolId: runtypes.String,
  userPoolClientId: runtypes.String,
  cssLocator: s3ResourceLocatorRuntype,
  logoLocator: s3ResourceLocatorRuntype,
});

export type S3ResourceLocator = runtypes.Static<
  typeof s3ResourceLocatorRuntype
>;

export type UpdateCognitoUiProperties = runtypes.Static<
  typeof updateCognitoUiPropertiesRuntype
>;

const s3Client = new S3Client({region: process.env.AWS_REGION});
const cognitoClient = new CognitoIdentityProviderClient({
  region: process.env.AWS_REGION,
});

export const handler: CdkCustomResourceIsCompleteHandler = async event => {
    switch (event.RequestType) {
      case 'Create':
      case 'Update':
        return createCognitoUiSettings(event);

      case 'Delete':
        return {
          IsComplete: true,
        };
    }
  }
);

async function createCognitoUiSettings(
  event: CdkCustomResourceIsCompleteEvent
): Promise<CdkCustomResourceIsCompleteResponse> {
  const {cssLocator, logoLocator, userPoolClientId, userPoolId} =
    updateCognitoUiPropertiesRuntype.check(event.ResourceProperties);

  console.log('Checking userpool domain');
  const userPool = await getUserPool(userPoolId);
  if (!userPool.Domain && !userPool.CustomDomain) {
    return {
      IsComplete: false,
    };
  }

  const domain = await getUserPoolDomain(userPool);
  if (domain.Status?.toLowerCase() !== 'active') {
    console.log('Domain is not yet active');
    return {
      IsComplete: false,
    };
  }

  console.log('Updating cognito settings');
  // Load resources from S3
  const [cssResource, logoResource] = await Promise.all([
    s3Client
      .send(
        new GetObjectCommand({
          Bucket: cssLocator.bucketName,
          Key: cssLocator.objectKey,
        })
      )
      .then(async ({Body}) => getFileContents(Body as Readable)),
    s3Client
      .send(
        new GetObjectCommand({
          Bucket: logoLocator.bucketName,
          Key: logoLocator.objectKey,
        })
      )
      .then(async ({Body}) => getFileContents(Body as Readable)),
  ]);

  const cssContents = cssResource.toString('utf-8');

  const res = await cognitoClient.send(
    new SetUICustomizationCommand({
      UserPoolId: userPoolId,
      ClientId: userPoolClientId,
      CSS: cssContents,
      ImageFile: logoResource,
    })
  );

  return {
    IsComplete: true,
    Data: res.UICustomization,
  };
}

async function getFileContents(readable: Readable): Promise<Buffer> {
  return new Promise((resolve, reject) => {
    const writable = new streamBuffers.WritableStreamBuffer();
    writable.on('error', err => reject(err));
    writable.on('finish', () => {
      const buf = writable.getContents();
      if (!buf) {
        reject(new Error('Failed to load data'));
        return;
      }
      resolve(buf);
    });
    readable.pipe(writable);
  });
}

async function getUserPool(userPoolId: string): Promise<UserPoolType> {
  return cognitoClient
    .send(
      new DescribeUserPoolCommand({
        UserPoolId: userPoolId,
      })
    )
    .then(({UserPool}) => {
      if (!UserPool) {
        throw new Error('Failed to get userPool');
      }
      return UserPool;
    });
}

async function getUserPoolDomain(
  userPool: UserPoolType
): Promise<DomainDescriptionType> {
  // use the CustomDomain if there is one, otherwise the 'prefix'
  const domain = userPool.CustomDomain ?? userPool.Domain;
  if (!domain) {
    throw new Error('No domain!');
  }

  return cognitoClient
    .send(
      new DescribeUserPoolDomainCommand({
        Domain: domain,
      })
    )
    .then(({DomainDescription}) => {
      if (!DomainDescription) {
        throw new Error('No domain description');
      }
      console.log('Domain', {DomainDescription});

      return DomainDescription;
    });
}

[dcarrion87]

I’ve moved bits over to managing using SDK and JS scripts directly that I run after to get around this as I ran out of time. E.g.:

import * as AWS from 'aws-sdk';
const cognito = new AWS.CognitoIdentityServiceProvider();
await cognito.setUICustomization({
            ClientId: c.ClientId,
            CSS: Buffer.from(fs.readFileSync(cssPath)).toString('utf-8'),
            ImageFile: Buffer.from(fs.readFileSync(logoPath)),
            UserPoolId: c.UserPoolId
}).promise()