Deploy a Cognito Secured WebSocket API with AWS CDK
Tags: aws, cdk, apigateway, websocket
Published: 2024-03-08 07:12:09

[API Gateway][1] is great for knitting together serverless application, but have you ever wanted a little more? WebSocket APIs offer many benefits by allowing active two-way communication between your clients and backend. Consider a WebSocket API when building real-time apps.

This post will show you how to deploy a WebSocket API that is secured using [AWS Cognito][2]. Clients can connect to the WebSocket by using the `idToken` retrieved from `Cognito` after a successful login. This post is heavily inspired by [this AWS sample][3] and [this dev.to article][5]. I have simplified the code and refactored it a bit for my own purposes and understanding.

Let's start by creating our Auth backend. We create a `Cognito UserPool` and a `Cognito UserPool Client`:

```typescript
import {
  aws_cognito as cognito,
  RemovalPolicy,
} from "aws-cdk-lib";
import { Construct } from "constructs";

export class Auth extends Construct {
  userPool: cognito.UserPool;
  userPoolClient: cognito.UserPoolClient;

  constructor(scope: Construct, id: string) {
    super(scope, id);
    this.userPool = this._userPool();
    this.userPoolClient = this._userPoolClient(this.userPool);

  }
  _userPool(): cognito.UserPool {
    let userPool = new cognito.UserPool(this, "userPool", {
      removalPolicy: RemovalPolicy.DESTROY,
    });
    return userPool;
  }

  _userPoolClient(userPool: cognito.UserPool): cognito.UserPoolClient {
    let userPoolClient = new cognito.UserPoolClient(this, "userPoolClient", {
      userPool: userPool,
      authFlows: { userSrp: true, userPassword: true },

    });
    return userPoolClient;
  }
}
```

Now let's use these to associate the WebSocket `$connect` route with a [`Custom Authorizer Lambda`][4]. Follow along with the comments

```typescript
import {
  aws_cognito as cognito,
  aws_apigatewayv2 as apigatewayv2,
  aws_apigatewayv2_integrations as apigatewayv2_integrations,
  aws_apigatewayv2_authorizers as apigatewayv2_auth,
  aws_lambda_nodejs as lambdaNode,
  aws_lambda as lambda,
  aws_dynamodb as dynamo,
} from "aws-cdk-lib";
import { RemovalPolicy } from "aws-cdk-lib";
import { Construct } from "constructs";
import * as path from "path";

export class API extends Construct {
  apiv2: apigatewayv2.HttpApi;
  websocket: apigatewayv2.WebSocketApi;
  constructor(
    scope: Construct,
    id: string,
  // We get these from our `Auth` construct above
    userPool: cognito.UserPool,
    userPoolClient: cognito.UserPoolClient,
  ) {
    super(scope, id);
    this.websocket = this._webSocket(userPool, userPoolClient);

  _webSocket(
    userPool: cognito.UserPool,
    userPoolClient: cognito.UserPoolClient,
  ): apigatewayv2.WebSocketApi {
// This DynamoDB table will be used to save the `connectionId` when 
// clients connect.
    const connectionIdTable = new dynamo.Table(this, "ConnectionIdTable", {
      partitionKey: { name: "connectionId", type: dynamo.AttributeType.STRING },
      timeToLiveAttribute: "removedAt",
      billingMode: dynamo.BillingMode.PAY_PER_REQUEST,
      removalPolicy: RemovalPolicy.DESTROY,
    });

// The `userId` is used to uniquely ID the connected user associated 
// with the `Cognito UserPool`. We'll use this to clean-up the 
// database upon client `disconnect`
    connectionIdTable.addGlobalSecondaryIndex({
      partitionKey: { name: "userId", type: dynamo.AttributeType.STRING },
      indexName: "userIdIndex",
    });

// This lambda function will be used to authorize the connection. 
// Upon successful connection, it will save user information into 
// the DynamoDB table
    const authHandler = new lambdaNode.NodejsFunction(this, "AuthHandler", {
      runtime: lambda.Runtime.NODEJS_18_X,
      entry: path.join(__dirname, "../lambdas/websocket/authorizer/index.ts"),
      environment: {
        USER_POOL_ID: userPool.userPoolId,
        APP_CLIENT_ID: userPoolClient.userPoolClientId,
        CONNECTION_TABLE_NAME: connectionIdTable.tableName,
      },
    });

// This lambda function will handle all routes (in our simple example)
    const websocketHandler = new lambdaNode.NodejsFunction(
      this,
      "WebSocketHandler",
      {
        runtime: lambda.Runtime.NODEJS_18_X,
        entry: path.join(__dirname, "../lambdas/websocket/handler/index.ts"),
        environment: {
          CONNECTION_TABLE_NAME: connectionIdTable.tableName,
        },
      },
    );

    connectionIdTable.grantReadWriteData(websocketHandler);
    connectionIdTable.grantReadWriteData(authHandler);

// Create the authorizer. The `identitySource` is the information we 
// need to supply when connecting to the API. `authHandler` will 
// use the `idToken` to verify the user.
    const authorizer = new apigatewayv2_auth.WebSocketLambdaAuthorizer(
      "Authorizer",
      authHandler,
      {
        identitySource: ["route.request.querystring.idToken"],
      },
    );

// Now create our WebSocket.
    let apiv2Websocket = new apigatewayv2.WebSocketApi(
      this,
      "knowledgeBaseWebSocket",
      {
        description: "Used by your serverless app to integrate services",
// On the `$connect` route, specify our authorizer. When connecting, 
// traffic will go through the authorizer before being routed 
// to the `websocketHandler` integration.
        connectRouteOptions: {
          authorizer,
          integration: new apigatewayv2_integrations.WebSocketLambdaIntegration(
            "ConnectIntegration",
            websocketHandler,
          ),
        },
        disconnectRouteOptions: {
          integration: new apigatewayv2_integrations.WebSocketLambdaIntegration(
            "DisconnectIntegration",
            websocketHandler,
          ),
        },
// For this sample example, route all traffic through the `websocketHandler`
        defaultRouteOptions: {
          integration: new apigatewayv2_integrations.WebSocketLambdaIntegration(
            "DefaultIntegration",
            websocketHandler,
          ),
        },
      },
    );

// You need to manually add the stage.
    new apigatewayv2.WebSocketStage(this, "dev", {
      webSocketApi: apiv2Websocket,
      stageName: "dev",
      autoDeploy: true,
    });

// Be sure to grant the `websocketHandler` permission to manage 
// the WebSocket connection.
    apiv2Websocket.grantManageConnections(websocketHandler);

    return apiv2Websocket;
  }
}
```

OKAY! Now what about the `authHandler` and `websocketHandler` functions?

```typescript
// `authHandler`
import { APIGatewayRequestAuthorizerHandler } from "aws-lambda";
import { CognitoJwtVerifier } from "aws-jwt-verify";
import { DynamoDBClient } from "@aws-sdk/client-dynamodb";
import { DynamoDBDocumentClient, PutCommand } from "@aws-sdk/lib-dynamodb";
import { CognitoIdTokenPayload } from "aws-jwt-verify/jwt-model";

const UserPoolId = process.env.USER_POOL_ID!;
const AppClientId = process.env.APP_CLIENT_ID!;

const client = DynamoDBDocumentClient.from(new DynamoDBClient({}));
const ConnectionTableName = process.env.CONNECTION_TABLE_NAME!;

export const handler: APIGatewayRequestAuthorizerHandler = async (
  event,
  _context,
) => {
  try {
// We will use the CognitoJwtVerifier to verify that the `idToken` 
// supplied by the client is indeed associated with our `UserPool`
    const verifier = CognitoJwtVerifier.create({
      userPoolId: UserPoolId,
      tokenUse: "id",
      clientId: AppClientId,
    });

    const connectionId = event.requestContext.connectionId!;
    const encodedToken = event.queryStringParameters!.idToken!;
    const payload = await verifier.verify(encodedToken);

// Successfully Authenticated!
    console.log("Token is valid. Payload:", payload);
// Save the session information to the DynamoDB table
    await saveSession(payload, connectionId);
    console.log("User information saved to database!");
// Now return a policy which allows the client to connect!
    return allowPolicy(event.methodArn, payload);
  } catch (error: any) {
    console.log(error.message);
    return denyAllPolicy();
  }
};

async function saveSession(
  payload: CognitoIdTokenPayload,
  connectionId: string,
) {
  try {
    await client.send(
      new PutCommand({
        TableName: ConnectionTableName,
        Item: {
          userName: payload["cognito:username"] ?? "",
          userGroups: payload["cognito:groups"] ?? "",
          userEmail: payload["email"] ?? "",
          userId: payload.sub,
          connectionId: connectionId,
          removedAt: Math.ceil(Date.now() / 1000) + 3600 * 3,
        },
      }),
    );
  } catch (err) {
    console.error(err);
  }
}

function allowPolicy(methodArn: string, idToken: any) {
  return {
    principalId: idToken.sub,
    policyDocument: {
      Version: "2012-10-17",
      Statement: [
        {
          Action: "execute-api:Invoke",
          Effect: "Allow",
          Resource: methodArn,
        },
      ],
    },
    context: {
      // set userId in the context
      userId: idToken.sub,
    },
  };
}

function denyAllPolicy() {
  return {
    principalId: "*",
    policyDocument: {
      Version: "2012-10-17",
      Statement: [
        {
          Action: "*",
          Effect: "Deny",
          Resource: "*",
        },
      ],
    },
  };
}
```

Once the client is authenticated, the connection request is routed to the `websocketHandler` function:

```typescript
// `websocketHandler`
import { APIGatewayProxyHandler } from "aws-lambda";
import { DynamoDBClient } from "@aws-sdk/client-dynamodb";
import { DeleteCommand, DynamoDBDocumentClient } from "@aws-sdk/lib-dynamodb";
import {
  ApiGatewayManagementApiClient,
  PostToConnectionCommand,
} from "@aws-sdk/client-apigatewaymanagementapi";

const client = DynamoDBDocumentClient.from(new DynamoDBClient({}));
const ConnectionTableName = process.env.CONNECTION_TABLE_NAME!;

export const handler: APIGatewayProxyHandler = async (event, _context) => {
  console.log(event);
  const routeKey = event.requestContext.routeKey!;
  const connectionId = event.requestContext.connectionId!;

// If the route is `$connect` go ahead and connect
  if (routeKey == "$connect") {
    return { statusCode: 200, body: "Connected." };
  }
// If `$disconnect`
  if (routeKey == "$disconnect") {
    try {
// This function will delete the connection information from the 
// DynamoDB table when you disconnect.
      await removeConnectionId(connectionId);
      return { statusCode: 200, body: "Disconnected." };
    } catch (err) {
      console.error(err);
      return { statusCode: 500, body: "Disconnection failed." };
    }
// Default route for any other requests
  } else {
// Just echo back messages in other route than connect, disconnect 
// (for testing purpose)
    const domainName = event.requestContext.domainName!;
// When we use a custom domain, we don't need to append a stage name
    const endpoint = domainName.endsWith("amazonaws.com")
      ? `https://${event.requestContext.domainName}/${event.requestContext.stage}`
      : `https://${event.requestContext.domainName}`;

// Now let's send a message from the WebSocket back to the client!
    const managementApi = new ApiGatewayManagementApiClient({
      endpoint,
    });

    try {
// This is where the message gets sent back to client. This example 
// just echos the event body. So if we connect and just send RAW text, 
// the WebSocket will send back the same RAW text.
      await managementApi.send(
        new PostToConnectionCommand({
          ConnectionId: connectionId,
          Data: Buffer.from(
            JSON.stringify({
              message: event.body,
            }),
            "utf-8",
          ),
        }),
      );
    } catch (e: any) {
      if (e.statusCode == 410) {
        await removeConnectionId(connectionId);
      } else {
        console.log(e);
        throw e;
      }
    }

    return { statusCode: 200, body: "Received." };
  }
};

const removeConnectionId = async (connectionId: string) => {
  return await client.send(
    new DeleteCommand({
      TableName: ConnectionTableName,
      Key: {
        connectionId,
      },
    }),
  );
};
```

Hope somebody found this useful! In the end probably using the already established IAM authorizer is more simple, but if you are looking to use your `Cognito` login credentials, then this example should help!


[1]: https://aws.amazon.com/api-gateway/
[2]: https://aws.amazon.com/cognito/
[3]: https://github.com/aws-samples/websocket-api-cognito-auth-sample/blob/main/README_en.md#implementation
[4]: https://docs.aws.amazon.com/apigateway/latest/developerguide/apigateway-websocket-api-lambda-auth.html
[5]: https://dev.to/aws-builders/building-websocket-apis-with-aws-api-gateway-and-cdk-1i27