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:

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.