AWS API Gateway Websockets + AWS Lambda + CDK + IAM Security with Cognito

***Recomendation: Use aws-cdk v2***

Recently I had to create a websocket using API Gateway, but using the new tool for infrastructure as code called AWS CDK. Also it was important to secure the websocket. In our case we were using Role-Based Access Control (https://docs.aws.amazon.com/cognito/latest/developerguide/role-based-access-control.html) for authenticating our APIs.

The first thing I guess you already know, is that a websocket is a bi-directional communication protocol, so we will create a Lambda for receiving messages, and another one for sending to the websocket.

I will use Node + Typescript in my examples but you can translate those to any other language supported by CDK (TypeScript, JavaScript, Python, Java, and C#).

Just for doing an easy example, we will use DynamoDB to store the websocket connections (you can use thedatabase that fits best for your case)

Code

First, we create the lambda that will receive data, the one that’s going to send data and the API Gateway websocket. Also we will create a policy for the sender lambda (the one that sends messages back to the user to be able to send those messages):

Our second step is to allow the cognito IAM Role to execute $connect in the websocket and deny the access to ManageConnections and secrets:

In the third step we will create the DynamoDB Table to store the connections. You can skip this step if you are changing the database. What I’m doing here is to store the connectionId and the userId as keys because I want to be able to send a message back to an specific person:

Our step four is creating a policy for the API to be able to invoke the lambda (the receiver lambda):

Step five, create the routes for the api gateway. As you can see the only one that needs to be secured is the connect route, once you are connected you should be able to disconnect or send messages without validating IAM again.

Step six, deployment configuration:

Extra step, create and configure a custom domain:

And that’s all you need. I will share the code of the lambdas so you can see how to get and store the connections, and how to send messages.

Receiver

Sender

For sending I’m using an SNS topic, so I can, from any part of my application, send an SNS notification to an user and this lambda will notify the websocket and have real live interaction with the user.

The interface of the lambda is a destination that could be undefined (send a broadcast message), one specific user or an array of users, and a message.

If you enjoy reading, follow me for more articles.

Keep learning.