How to allow or restrict users of specific domains in AWS Cognito

In this article, we're going to discuss how to allow or restrict users of specific domains in AWS Cognito

Why you may want to restrict based on domains?

You may want to build a web application for your company event where you want only the employees of your company to register. One way to identify the employees of your company, he or she will have an email id with your company domain name.

For example, if I'm creating such an application, I would allow anyone whose email address ends with  @cloudtechsimplified.com  

Infrastructure

We're going to use AWS CDK for creating the necessary infrastructure. It's an open-source software development framework that lets you define cloud infrastructure. AWS CDK supports many languages including TypeScript, Python, C#, Java, and others. You can learn more about AWS CDK from a beginner's guide here. We're going to use TypeScript in this article.

Cognito Pool

Let's create a Cognito Pool as shown below

const cognitoDomainPool = new cognito.UserPool(this, "cognitoDomainPool", {
      userPoolName: "cognito-domain-pool",
      signInCaseSensitive: false, // case insensitive is preferred in most situations
      selfSignUpEnabled: true,
      accountRecovery: cognito.AccountRecovery.EMAIL_ONLY,
      mfa: cognito.Mfa.OFF,
      userVerification: {
        emailSubject: "Verify your email for Cogntio Domains App",
        emailBody:
          "Hey, Thanks for signing up to Cognito Domains App. Your verification code is {####}",
        emailStyle: cognito.VerificationEmailStyle.CODE,
      },
      signInAliases: {
        email: true,
        username: false,
      },
      standardAttributes: {
        email: {
          required: true,
        },
      },
    });

Let's add an app client for the pool

const appClient = cognitoDomainPool.addClient("cognito-domain-app-client", {
      userPoolClientName: "cognito-domain-app-client",
      generateSecret: true,
      authFlows: {
        userPassword: true,
      },
      oAuth: {
        flows: {
          authorizationCodeGrant: true,
        },
        callbackUrls: ["http://localhost:3000/api/auth/callback/cognito"],
        logoutUrls: ["http://localhost:3000"],
        scopes: [
          cognito.OAuthScope.EMAIL,
          cognito.OAuthScope.OPENID,
          cognito.OAuthScope.PROFILE,
        ],
      },
    });

Then, we can create a custom domain so that users of our application may recognize the app. Otherwise, it may create subdomain with random characters.

  cognitoDomainPool.addDomain("CognitoDomain", {
      cognitoDomain: {
        domainPrefix: "cts-domain-app-2706",
      },
    });

Lambda function

Let's create a lambda function with pre signup authentication trigger of Cognito. This lambda function will be called when the user tries to sign-up.

 const nodeJsFunctionProps: NodejsFunctionProps = {
      bundling: {
        externalModules: [
          "aws-sdk", // Use the 'aws-sdk' available in the Lambda runtime
        ],
      },
      runtime: Runtime.NODEJS_18_X,
      timeout: Duration.seconds(30), // Api Gateway timeout
    };

    const preSignupAuthLambda = new NodejsFunction(
      this,
      "preSignupAuthLambda",
      {
        entry: "lambdas/pre-signup-auth-lambda.ts",
        ...nodeJsFunctionProps,
        functionName: "cognito-domain-app-pre-signup-auth-lambda",
      }
    );

Configure the pre-sign-up trigger of AWS Cognito to the Lambda function

In the below code snippet, we're going to configure AWS Cognito to add a trigger for the pre-sign-up to the lambda function.  This configuration will make cognito to call our lambda function when a user tries to signup.

  cognitoDomainPool.addTrigger(
      cognito.UserPoolOperation.PRE_SIGN_UP,
      preSignupAuthLambda
    );

Lambda function code

In the actual lambda function code, we're getting the email information from user attributes. And, we're doing a simple check to check whether our email address contains the allowed domain.

export async function handler(event: any, context: any) {
  const userEmailDomain = event.request.userAttributes.email.split("@")[1];
  const allowedDomains = ["cloudtechsimplified.com"];
  if (!allowedDomains.includes(userEmailDomain)) {
    throw new Error("Invalid email domain");
  }
  return event;
}

You can implement any additional logic here.

Deployment & Testing

You can deploy the infrastructure (Cognito & Lambda) by executing cdk deploy command. All the necessary resources would be created.

For testing, I'm creating a simple NextJS application with next-auth library to test our signup process.

I've tried to sign-up using the email someuser@gmail.com - as it does not end with @cloudtechsimplified.com , the error is thrown "Invalid email domain" as shown in the below screenshot.

When you valid email domain - the signup process will complete.

Please let me know your thoughts in the comments.