Creates a really basic websocket service and testing tools

websocket-cloudformation.yaml

# 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>