Skip to content

Instantly share code, notes, and snippets.

@ca0abinary
Last active July 31, 2021 21:23
Show Gist options
  • Save ca0abinary/c3c5ec94186db8df1f7b02b324a37eb0 to your computer and use it in GitHub Desktop.
Save ca0abinary/c3c5ec94186db8df1f7b02b324a37eb0 to your computer and use it in GitHub Desktop.
Creates a really basic websocket service and testing tools
# sample deployment command
# aws cloudformation deploy --template-file template.yaml --capabilities CAPABILITY_NAMED_IAM --stack-name ws-test --profile dev
# sample testing uri
# aws cloudformation describe-stacks --query "Stacks[?contains(@.StackName,'ws-test')].Outputs[0][?OutputKey=='WebSocketURI'].OutputValue" --profile dev | ConvertFrom-Json
# You can test with the test.html produced by the TestHtmlPage output
# or
# npm install -g wscat
# wscat -c $(aws cloudformation describe-stacks --query "Stacks[?contains(@.StackName,'ws-test')].Outputs[0][?OutputKey=='WebSocketURI'].OutputValue" --profile dev | ConvertFrom-Json)
# and paste this test message:
# {"action":"message","message":"Hello, world!"}
AWSTemplateFormatVersion: 2010-09-09
Description: Simple WebSockets example
Mappings:
static:
webSocket:
deployStageName: deploy
Resources:
WebSocket:
Type: AWS::ApiGatewayV2::Api
Properties:
Name: !Sub ${AWS::StackName}
ProtocolType: WEBSOCKET
RouteSelectionExpression: $request.body.action
WebSocketDeployment:
Type: AWS::ApiGatewayV2::Deployment
DependsOn: [ConnectRoute, DisconnectRoute, MessageRoute]
Properties:
ApiId: !Ref WebSocket
WebSocketStage:
Type: AWS::ApiGatewayV2::Stage
Properties:
ApiId: !Ref WebSocket
DeploymentId: !Ref WebSocketDeployment
StageName: deploy
ConnectRoute:
Type: AWS::ApiGatewayV2::Route
Properties:
ApiId: !Ref WebSocket
AuthorizationType: NONE
RouteKey: $connect
Target: !Sub integrations/${ConnectIntegration}
ConnectIntegration:
Type: AWS::ApiGatewayV2::Integration
Properties:
ApiId: !Ref WebSocket
IntegrationType: AWS_PROXY
IntegrationUri: !Sub arn:aws:apigateway:${AWS::Region}:lambda:path/2015-03-31/functions/${Function.Arn}/invocations
DisconnectRoute:
Type: AWS::ApiGatewayV2::Route
Properties:
ApiId: !Ref WebSocket
AuthorizationType: NONE
RouteKey: $disconnect
Target: !Sub integrations/${ConnectIntegration}
DisconnectIntegration:
Type: AWS::ApiGatewayV2::Integration
Properties:
ApiId: !Ref WebSocket
IntegrationType: AWS_PROXY
IntegrationUri: !Sub arn:aws:apigateway:${AWS::Region}:lambda:path/2015-03-31/functions/${Function.Arn}/invocations
MessageRoute:
Type: AWS::ApiGatewayV2::Route
Properties:
ApiId: !Ref WebSocket
AuthorizationType: NONE
RouteKey: message
Target: !Sub integrations/${ConnectIntegration}
MessageIntegration:
Type: AWS::ApiGatewayV2::Integration
Properties:
ApiId: !Ref WebSocket
IntegrationType: AWS_PROXY
IntegrationUri: !Sub arn:aws:apigateway:${AWS::Region}:lambda:path/2015-03-31/functions/${Function.Arn}/invocations
Function:
Type: AWS::Lambda::Function
Properties:
Code:
ZipFile: !Sub
- |
# -*- coding: utf-8 -*-
import boto3, json, logging, traceback
logger = logging.getLogger()
logger.setLevel(logging.DEBUG)
def handler(event, context):
logger.debug(json.dumps(event))
ip_address = event['requestContext']['identity']['sourceIp']
connection_id = event['requestContext']['connectionId']
route_key = event['requestContext']['routeKey']
body = json.loads(event['body']) if 'body' in event else ''
ws = boto3.client('apigatewaymanagementapi',
endpoint_url = 'https://${WebSocket}.execute-api.${AWS::Region}.amazonaws.com/${stageName}')
try:
if route_key == '$connect':
return {'statusCode': 200, 'body': 'Connected.'}
elif route_key == 'message':
reply_to_message(ws, json.dumps({
'webservice': '🦜',
'message': f'**squawk** [{ip_address}] said {body["message"]}'
}), connection_id)
return {'statusCode': 200, 'body': 'Got the message'}
elif route_key == '$disconnect':
return {'statusCode': 200, 'body': 'Disconnected.'}
else:
return {'statusCode': 500, 'body': 'Failed to process message.'}
except Exception:
traceback.print_exc()
return {'statusCode': 500, 'body': 'Failed to process message.'}
def reply_to_message(ws, response, connection_id):
try:
ws.post_to_connection(
Data=response,
ConnectionId=connection_id
)
except Exception as e:
logger.error(e)
- { stageName: !FindInMap [static, webSocket, deployStageName] }
FunctionName: !Sub ${AWS::StackName}-Function
Handler: index.handler
MemorySize: 128
Role: !GetAtt FunctionRole.Arn
Runtime: python3.7
FunctionRole:
Type: AWS::IAM::Role
Properties:
AssumeRolePolicyDocument:
Version: 2012-10-17
Statement:
- Effect: Allow
Principal: { Service: [lambda.amazonaws.com] }
Action: ['sts:AssumeRole']
ManagedPolicyArns:
- arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole
Policies:
- PolicyName: !Sub ${AWS::StackName}-FunctionRolePolicy
PolicyDocument:
Statement:
- Effect: Allow
Action: execute-api:*
Resource: !Sub
- arn:aws:execute-api:${AWS::Region}:${AWS::AccountId}:${WebSocket}/${stageName}/POST/@connections/*
- { stageName: !FindInMap [static, webSocket, deployStageName] }
ConnectionFunctionPermission:
Type: AWS::Lambda::Permission
DependsOn: [WebSocket]
Properties:
Action: lambda:InvokeFunction
FunctionName: !Ref Function
Principal: apigateway.amazonaws.com
ConnectionFunctionLogGroup:
Type: AWS::Logs::LogGroup
Properties:
RetentionInDays: 30
LogGroupName: !Sub /aws/lambda/${AWS::StackName}-Function
Outputs:
WebSocketURI:
Description: "The endpoint used to connect to the websocket."
Value: !Sub wss://${WebSocket}.execute-api.${AWS::Region}.amazonaws.com/${WebSocketStage}
TestHtmlPage:
Description: "Paste this into test.html and run it to test"
Value: !Sub |
<!DOCTYPE html>
<meta charset="utf-8" />
<title>WebSocket Test</title>
<script language="javascript" type="text/javascript">
var wsUri = "wss://${WebSocket}.execute-api.${AWS::Region}.amazonaws.com/${WebSocketStage}";
var output;
function init()
{
output = document.getElementById("output");
testWebSocket();
}
function testWebSocket()
{
websocket = new WebSocket(wsUri);
websocket.onopen = function(evt) { onOpen(evt) };
websocket.onclose = function(evt) { onClose(evt) };
websocket.onmessage = function(evt) { onMessage(evt) };
websocket.onerror = function(evt) { onError(evt) };
}
function onOpen(evt)
{
writeToScreen("CONNECTED");
var message = "Hello, world!";
doSend(JSON.stringify({ "action": "message", message }));
}
function onClose(evt)
{
writeToScreen("DISCONNECTED");
console.info(evt);
}
function onMessage(evt)
{
const response = JSON.parse(evt.data);
writeToScreen('<span style="color: blue;">RESPONSE: ' + response.message + '</span>');
websocket.close();
}
function onError(evt)
{
writeToScreen('<span style="color: red;">ERROR:</span> ' + evt.data);
console.error(evt);
}
function doSend(message)
{
writeToScreen("SENT: " + message);
websocket.send(message);
}
function writeToScreen(message)
{
var pre = document.createElement("p");
pre.style.wordWrap = "break-word";
pre.innerHTML = message;
output.appendChild(pre);
}
window.addEventListener("load", init, false);
</script>
<h2>WebSocket Test</h2>
<div id="output"></div>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment