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 and cannot be use as authorization.

Overview diagram:

AWS WSS API Architecture

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 is authorize yayy!!
        add_connection(connection_id, event)
        return {"statusCode": 200, "body": "connected"}
    else:
        # user is not authorize snap!
        return {"statusCode": 401, "body": "disconnected"}

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.

AWS WSS API Creation

Then create send route into your api.

AWS WSS API Creation

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.

$connect -> sample-echo-app-connect
$disconnect -> sample-echo-app-disconnect
$send -> sample-echo-app-send

That is all you need!

Let’s 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!"

Ritul Patel

How to create aws websocket api gateway?

Share