How to create aws websocket api gateway?
In this post we will learn how to create AWS Websocket API Gateway, and how to secure and optimize it.
Things to note:
- AWS Websocket Gateway cannot be deployed as
private
gateway like REST API. $connect
and$disconnect
are pre-define routes.- All authorization need to happen on
$connect
event. - We will set
API Key
to manage the access from client, but this key will be visible andcannot be use as authorization
.
Overview diagram:

We’ll create a sample echo app to see how websocket messaging works. We’ll start by creating three python
lambdas.
sample-echo-app-connect:
def authorized(auth_header: str) -> bool:
"""
Function to check if user is authorized to connect.
"""
# add your authorization code
return True
def add_connection(connection_id: str, event: dict):
"""
Function to add connection info into backend for later use
"""
pass
def handler(event, context):
"""
This lambda will get invoked when user first tries to connect
"""
headers = event.get("headers", {})
request_context = event.get("requestContext", {})
connection_id = request_context.get("connectionId") # important info
if authorized(headers.get("Authorization")):
# user if authorize yayy!!
add_connection(connection_id, event)
return {"statusCode": 200, "body": "connected"}
else:
# user if not authorize snap!
return {"statusCode": 401, "body": "connected"}
sample-echo-app-disconnect:
def remove_connection(connection_id: str):
"""
Function to remove connection from backend
"""
pass
def handler(event, context):
"""
This lambda will get invoked when user disconnects
"""
request_context = event.get("requestContext", {})
connection_id = request_context.get("connectionId") # important info
remove_connection(connection_id)
return {"statusCode": 200, "body": "disconnected"}
sample-echo-app-send:
import boto3
import json
def echo_message(self, message: object, endpoint: str, connection_id: str):
"""
Echo message to client
"""
client = boto3.client("apigatewaymanagementapi", endpoint_url=endpoint)
client.post_to_connection(Data=json.dumps(message).encode("utf-8"), ConnectionId=connection_id)
def handler(event, context):
"""
This lambda will get invoked when user send message
"""
request_context = event.get("requestContext", {})
connection_id = request_context.get("connectionId") # important info
host_name = request_context.get("domainName")
message = json.loads(event.get("body"))
echo_message(message, f"https://{host_name}", connection_id)
return {"statusCode": 200, "body": "connected"}
Now, lets create Websocket API via AWS console.

Then create send
route into your api.

Time to map following lambdas to integration request
for each route. To do that, click on route -> Integration Request -> Lambda Function and select your lambda function.
sample-echo-app-connect -> $connect
sample-echo-app-disconnect -> $disconnect
sample-echo-app-send -> $send
That is all you need!
Let check that in action (I truncated some info):
wscat -c wss://blah.execute-api.amazonaws.com -H "Authorizatio: Bearer token" -H "X-API-Key: your-api-key"
> Connected
< {"message": "send", "body": "Hello World!"}
> "Hello World!"
Originally published at https://pritul95.github.io.