Architecture Graph of APIGateway-Lambda-DDB

Create a Simple Chat WebSocket Lambda Service with CloudFormation

Recently, I have been working on setting up WebSocket connection service with AWS ApiGateway/Lambda/CloudFormation, and found it hard to accomplish only with http://docs.aws.amazon.com/.

Although AWS has given some nicely written blogs (links here, here and here) about how to set up a WebSocket API, but it is done with AWS Console UI and mouse clicking. Since it is not based on CloudFormation, we could not pipeline it or maintain it in a professional developer way.

This article will introduce how to create a WebSocket ApiGateway Chat service, and explain the CloudFormation template files so that you can create your own chat service with a full understanding what is going on on your CloudFormation stacks.

Github link: https://github.com/zijing07/simple-websockets-chat-app.

Understand the Project Objective

We will set up a chat service to allow any people to connect, and send messages to each other. For example in the below image, two sessions use wscat to connect to the WebSocket service, and send messages to each other:

Design the Components

1. Database to store all the connected sessions
   This article will utilize AWS DynamoDB to store the connection information.
2. OnConnect Lambda Function
   Required by WebSocket API Documentation
   When a new session is established, this function will be invoked. In this function, the connectionId should be persisted in DynamoDB to broadcast coming messages.
3. OnDisconnect Lambda Function
   Required by WebSocket API Documentation
   When a session disconnects, this function will be invoked. Once executed, the connectionId should be removed from DynamoDB to avoid unnecessary message delivery.
4. OnSendMessage Lambda Function
   We define this function. This function gathers all the connected sessions and broadcast messages to every one.

Directory Structure

I forked from this github repo by aws-lab, and changed the architecture from a single stack into master/nested stack pattern:

The template.yaml is the master stack. The AWS lambda resources are placed in the nested stack lambdas/lambdas.yaml. With the nested stack set up, the architecture becomes scalable when we need to add other stacks like Frontend Page, Message Encryption, Authorizer and so on.

Define the DynamoDB

ConnectionDDBTable:
  Type: AWS::DynamoDB::Table
  Properties:
    AttributeDefinitions:
      - AttributeName: "connectionId"
        AttributeType: "S"
    KeySchema:
      - AttributeName: "connectionId"
        KeyType: "HASH"
    BillingMode: PAY_PER_REQUEST

This is a simple table with one functionality: storing the active connectionIds.

Define the WebSocket

The CloudFormation definition of WebSocket API is like:

SimpleChatWebSocket:
  Type: AWS::ApiGatewayV2::Api
  Properties:
    Name: !Sub ${AWS::StackName}-SimpleChatWebSocket
    ProtocolType: WEBSOCKET
    RouteSelectionExpression: "$request.body.action"
    ApiKeySelectionExpression: $request.header.x-api-key

This is a standard template from AWS docs site, please remember the $request.body.action, we will revisit it later.

Lambdas

CloudFormation scripts for each lambda consists of four resources:

Route:
  Type: AWS::ApiGatewayV2::Route
Integration:
  Type: AWS::ApiGatewayV2::Integration
LambdaFunction:
  Type: AWS::Serverless::Function
LambdaPermission:
  Type: AWS::Lambda::Permission

• Route: Entry point in the ApiGateway
• Integration: Link the entry point to the lambda function
• LambdaFunction: Where the code lives
• LambdaPermission: Define what capability the lambda function has, such as DynamoDB access in our case

This article will not paste the full CloudFormation template of each function, templates can be found at: github:zijing07/../lambdas.yaml.

OnConnect Lambda

CloudFormation route configuration:

ConnectRoute:
  Type: AWS::ApiGatewayV2::Route
  Properties:
    ApiId: !Ref SimpleChatWebSocket
    RouteKey: $connect

The route key $connect is default supported by ApiGateway, every time a new connection comes, $connect will be the entry point, then OnConnect Lambda stores the connectionId to the DynamoDB:

const putParams = {
  TableName: process.env.TABLE_NAME,
  Item: { connectionId: event.requestContext.connectionId }
};
await ddb.put(putParams).promise();

Full code could be found at: github:zijing07/../onConnectHandler.

OnDisconnect Lambda

CloudFormation route configuration:

DisconnectRoute:
  Type: AWS::ApiGatewayV2::Route
  Properties:
    ApiId: !Ref SimpleChatWebSocket
    RouteKey: $disconnect

Like $connect, $disconnect is also a builtin route in ApiGateway. Once the connection is closed, $disconnect will be the entry point, then OnDisconnect Lambda is responsible to remove the connectionId from the DynamoDB:

const deleteParams = {
  TableName: process.env.TABLE_NAME,
  Key: { connectionId: event.requestContext.connectionId }
};
await ddb.delete(deleteParams).promise();

Full code could be found at: github:zijing07/../onDisconnectHandler.

If you find this article useful, please follow this account for future updates. Thanks for the support!OnSendMessage Lambda

CloudFormation route configuration:

SendRoute:
  Type: AWS::ApiGatewayV2::Route
  Properties:
    ApiId: !Ref SimpleChatWebSocket
    RouteKey: sendmessage

Please notice the route key used here. Remember the $request.body.action configured in SimpleChatWebSocket above? In the ApiGateway/WebSocket resource, other than $connect and $disconnect, we could define our own route keys. ApiGateway will extract the route key from the request body, with the RouteSelectionExpression as the key. In our case, if we send a message:

{ "action": "sendmessage", "data": "live long" }

ApiGateway will know we are targeting the entry point configured with sendmessage.

Then the above message will be passed to the OnSendMessage Lambda, the lambda will execute the following operations to broadcast the message:

// 1. Extract the message from the request body
const postData = JSON.parse(event.body).data;

// 2. Scan the DynamoDB for all connectionIds
const connectionData = await ddb.scan(...).promise();

// 3. Send the message to all connections
const postCalls = connectionData.Items
    .map(async ({ connectionId }) => 
        await apigwManagementApi.postToConnection(...).promise());
await Promise.all(postCalls);

Full code could found at: github:zijing07/…/onSendMessageHandler.

Package and Build

This project utilizes Makefile to manage the packing and building flow. Yeah, I hear you, make is too old, but I have to say that make is really really wonderful for managing simple build flows.

There are three steps to deploy the stack:

1. Create S3 bucket to store the template file
2. Package local template files and upload them to S3
3. Deploy the stack on CloudFormation

To deploy the stack, just type make in the command line:

$ make

The Makefile could be found at: github:zijing07/../Makefile.

Deployed stacks will look like:

Demo Time!

Next Steps

1. Create a frontend page for the chat service
   That article will introduce how to use CloudFront to host your website without a server!
2. Add more functionalities
   - Chat room support
   - End-to-End Encryption

Please reply if you have any question or suggestion.

Thanks for reading.

If you find this article useful, please follow this account for future updates. Thanks for the support!