How to create aws websocket api gateway?

Ritul Patel
3 min readSep 3, 2020

--

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.

Sign up to discover human stories that deepen your understanding of the world.

Free

Distraction-free reading. No ads.

Organize your knowledge with lists and highlights.

Tell your story. Find your audience.

Membership

Read member-only stories

Support writers you read most

Earn money for your writing

Listen to audio narrations

Read offline with the Medium app

--

--

Ritul Patel
Ritul Patel

Written by Ritul Patel

Software Engineer with passion to learn new tech!

No responses yet

Write a response