Setting Up Real-Time Communication with WebSockets Using AWS API Gateway v2 and Amazon CloudFront

Published in

Skillwell

12 min readMar 25, 2024

In association with our collaboration with JAMK University of Applied Sciences, our team offers mentorship to students undertaking thesis projects. Among the diverse topics this spring, one thesis is dedicated to exploring the integration of WebSocket communication with Flux architecture within web applications. To support this particular study, it became necessary to implement AWS API Gateway v2 for WebSocket communication as the backend infrastructure. To go beyond basic requirements and create a more practical learning environment, we implemented a setup with also CloudFront and Cognito. This blog post describes our approach to developing an effective WebSocket communication system.

Let’s start by highlighting the AWS services that collectively serve as the foundation of our solution. We chose a suite of AWS services beyond just AWS API Gateway v2 to craft a secure, scalable, and production-ready architecture.

AWS Services in The Solution

Our WebSocket communication system leverages a suite of AWS services, each contributing to a robust and scalable architecture:

AWS API Gateway v2 is a fully managed service that simplifies the process of creating, publishing, maintaining, monitoring, and securing both HTTP and WebSocket APIs at any scale.
Amazon Cognito User Pool reinforces our solution security by managing user authentication and access control through Cognito JWT authorizer in API Gateway for HTTP connections. For WebSocket connections, where direct integration with Cognito is not available, we architected a custom authentication mechanism within a Lambda function.
AWS Lambda functions serve as a powerful companion to AWS API Gateway v2, executing code in response to events like WebSocket messages and HTTP requests. This integration enables developers to efficiently handle real-time and standard web interactions without server management, facilitating the creation of responsive and scalable applications.
Amazon DynamoDB enhances our solution by offering a fully managed NoSQL database service, ideal for storing WebSocket connection information and application data.
Amazon CloudFront enhances our architecture by allowing a single custom domain name for both HTTP and WebSocket APIs via AWS API Gateway v2. Acting as a front-facing layer, CloudFront routes traffic to the correct API type, streamlining domain management while accelerating content delivery globally. It also simplifies SSL/TLS certificate management and bolsters security, enabling efficient and secure handling of diverse web traffic under one unified domain.
Amazon Certificate Manager secures our application by managing SSL/TLS certificates for our CloudFront distribution, automating the encryption process for our custom domain. Its seamless integration ensures continuous protection and hassle-free renewals.
Amazon Route 53 effectively directs traffic to CloudFront, managing our API’s domain.

Setting Up Base AWS Resources with CloudFormation

To simplify setting up AWS architectures, we mainly use CloudFormation. This tool helps us automatically set up and manage our AWS services by using code to define the infrastructure. With CloudFormation, we can easily create, duplicate, and manage our system’s setup, making it perfect for quickly deploying the AWS resources.

Preliminary Tasks: Setting Up a Hosted Zone and SSL/TLS Certificate

Before setting up our AWS resources with CloudFormation, we needed to take two important steps in the AWS Console.

First, we set up a hosted zone in Amazon Route 53 for our domain. This lets us control our domain’s DNS records.
Next, we needed to request an SSL/TLS certificate through AWS Certificate Manager, crucial for data encryption. Because we’re using Amazon CloudFront, we had to get this certificate in the US East (N. Virginia) region, which is a requirement for CloudFront custom domains.

These steps are standard procedures and part of common AWS knowledge, so I won’t go into further detail about them in this blog post.

Setting Up API Gateways with CloudFormation

Our API Gateways, acting as origin points for the CloudFront distribution, needed to be established first. We chose to initially define the core structure of the API Gateways using CloudFormation to set up the fundamental components. More detailed configurations were defined later using Serverless Application Model (SAM) templates. This is our normal approach and it allows us to methodically expand and refine the API Gateways.

Here are the CloudFormation resource definitions for the HTTP and WebSocket API Gateways:

Resources:
  HTTPApi:
    Type: AWS::ApiGatewayV2::Api
    Properties:
      Description: HTTP API
      Name: !Sub ${ApplicationName}-${EnvironmentName}-api-http
      ProtocolType: HTTP

  HttpApiDefaultStage:
    Type: AWS::ApiGatewayV2::Stage
    Properties:
      ApiId: !Ref HTTPApi
      AutoDeploy: true
      StageName: $default

  WSApi:
    Type: AWS::ApiGatewayV2::Api
    Properties:
      Description: WEBSOCKET API
      Name: !Sub ${ApplicationName}-${EnvironmentName}-api-ws
      ProtocolType: WEBSOCKET
      RouteSelectionExpression: "$request.body.action"

Please note how similar the definitions of HTTP and WebSocket API Gateways are. And despite WebSocket communication starting with an HTTP request that upgrades to a WebSocket connection, a single AWS API Gateway instance cannot function as both an HTTP and WebSocket API.

Setting Up CloudFront Distribution and Route 53 RecordSet with CloudFormation

Following the setup of the API Gateways, we crafted a CloudFormation template for configuring the CloudFront distribution and associated DNS record.

Here’s the CloudFormation template that we used:

Resources:
  APIOriginRequestPolicy:
    Type: AWS::CloudFront::OriginRequestPolicy
    Properties:
      OriginRequestPolicyConfig:
        Name: APIOriginRequestPolicy
        Comment: Policy to forward all headers to API
        HeadersConfig:
          HeaderBehavior: allExcept
          Headers:
            - Host
        CookiesConfig:
          CookieBehavior: none
        QueryStringsConfig:
          QueryStringBehavior: all

  Distribution:
    Type: AWS::CloudFront::Distribution
    Properties:
      DistributionConfig:
        Origins:
          - Id: !Join
              - ""
              - - "S3-Website-"
                - !FindInMap [
                    Environment,
                    !Ref "EnvironmentName",
                    "OriginBucket",
                  ]
                - !Sub ".s3-${AWS::Region}.amazonaws.com"
            DomainName: !Join
              - ""
              - - !FindInMap [
                    Environment,
                    !Ref "EnvironmentName",
                    "OriginBucket",
                  ]
                - !Sub ".s3-${AWS::Region}.amazonaws.com"
            S3OriginConfig:
              OriginAccessIdentity: !Join
                - ""
                - - "origin-access-identity/cloudfront/"
                  - Fn::ImportValue: !Sub "${ApplicationName}-${EnvironmentName}-cloudfront-oai"
          - Id: !Join
              - ""
              - - "HTTP-API-"
                - !FindInMap [Environment, !Ref "EnvironmentName", "HTTPApiId"]
                - !Sub ".execute-api.${AWS::Region}.amazonaws.com"
            DomainName: !Join
              - ""
              - - !FindInMap [Environment, !Ref "EnvironmentName", "HTTPApiId"]
                - !Sub ".execute-api.${AWS::Region}.amazonaws.com"
            CustomOriginConfig:
              HTTPSPort: 443
              OriginKeepaliveTimeout: 60
              OriginProtocolPolicy: https-only
              OriginReadTimeout: 60
              OriginSSLProtocols:
                - TLSv1.2
          - Id: !Join
              - ""
              - - "WS-API-"
                - !FindInMap [Environment, !Ref "EnvironmentName", "WSApiId"]
                - !Sub ".execute-api.${AWS::Region}.amazonaws.com"
            DomainName: !Join
              - ""
              - - !FindInMap [Environment, !Ref "EnvironmentName", "WSApiId"]
                - !Sub ".execute-api.${AWS::Region}.amazonaws.com"
            CustomOriginConfig:
              HTTPSPort: 443
              OriginKeepaliveTimeout: 60
              OriginProtocolPolicy: https-only
              OriginReadTimeout: 60
              OriginSSLProtocols:
                - TLSv1.2
        Enabled: "true"
        Aliases:
          - !FindInMap [Environment, !Ref "EnvironmentName", "Alias"]
        CacheBehaviors:
          - PathPattern: "/http/*"
            TargetOriginId: !Join
              - ""
              - - "HTTP-API-"
                - !FindInMap [Environment, !Ref "EnvironmentName", "HTTPApiId"]
                - !Sub ".execute-api.${AWS::Region}.amazonaws.com"
            ViewerProtocolPolicy: redirect-to-https
            CachePolicyId: "4135ea2d-6df8-44a3-9df3-4b5a84be39ad"
            OriginRequestPolicyId: !Ref APIOriginRequestPolicy
            AllowedMethods:
              - GET
              - HEAD
              - OPTIONS
              - PUT
              - POST
              - PATCH
              - DELETE
            CachedMethods:
              - GET
              - HEAD
              - OPTIONS
          - PathPattern: "/ws/*"
            TargetOriginId: !Join
              - ""
              - - "WS-API-"
                - !FindInMap [Environment, !Ref "EnvironmentName", "WSApiId"]
                - !Sub ".execute-api.${AWS::Region}.amazonaws.com"
            ViewerProtocolPolicy: redirect-to-https
            CachePolicyId: "4135ea2d-6df8-44a3-9df3-4b5a84be39ad"
            OriginRequestPolicyId: !Ref APIOriginRequestPolicy
            AllowedMethods:
              - GET
              - HEAD
              - OPTIONS
            CachedMethods:
              - GET
              - HEAD
              - OPTIONS
        DefaultCacheBehavior:
          CachedMethods:
            - GET
            - HEAD
            - OPTIONS
          AllowedMethods:
            - GET
            - HEAD
            - OPTIONS
          Compress: true
          CachePolicyId: "4135ea2d-6df8-44a3-9df3-4b5a84be39ad"
          TargetOriginId: !Join
            - ""
            - - "S3-Website-"
              - !FindInMap [Environment, !Ref "EnvironmentName", "OriginBucket"]
              - !Sub ".s3-${AWS::Region}.amazonaws.com"
          ForwardedValues:
            QueryString: "false"
            Cookies:
              Forward: none
          ViewerProtocolPolicy: redirect-to-https
          FunctionAssociations:
            - EventType: viewer-request
              FunctionARN: !GetAtt ViewerRequestCloudFrontFunction.FunctionMetadata.FunctionARN
        PriceClass: PriceClass_100
        ViewerCertificate:
          AcmCertificateArn:
            !FindInMap [Environment, !Ref "EnvironmentName", "CertificateArn"]
          MinimumProtocolVersion: TLSv1.2_2018
          SslSupportMethod: sni-only

  DnsRecord:
    Type: AWS::Route53::RecordSet
    DependsOn: Distribution
    Properties:
      AliasTarget:
        DNSName: !GetAtt Distribution.DomainName
        # Constant value for CloudFront: Z2FDTNDATAQYW2
        HostedZoneId: "Z2FDTNDATAQYW2"
      HostedZoneId:
        !FindInMap [Environment, !Ref "EnvironmentName", "HostedZoneId"]
      Name: !FindInMap [Environment, !Ref "EnvironmentName", "Alias"]
      Type: A

The key highlights of this template include:

We configured it to forward all headers except Host to both HTTP and WebSocket APIs, utilizing a custom origin request policy for this purpose.
We applied the cache policy ID 4135ea2d-6df8–44a3–9df3–4b5a84be39ad, a managed policy designed to disable caching, ensuring that requests to both API types are not stored in cache.
A default cache behavior is required for CloudFront distributions. Our default cache behavior returns web application that is served from S3 bucket.
The CloudFront distribution routes traffic based on request paths: /http/ prefixes direct to the HTTP API, while /ws/ prefixes route to the WebSocket API.
Lastly, we configured our chosen domain name to serve as an alias for the CloudFront distribution, facilitating domain-based access to the distributed content.

Creating Cognito User Pool and User Pool Client

To secure our HTTP API and manage user authentication efficiently, we integrated Amazon Cognito into our architecture. Specifically, we defined very basic versions of Cognito User Pool and User Pool Client with the following CloudFormation template:

Resources:
  LambdaExecutionRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          Effect: Allow
          Principal:
            Service: lambda.amazonaws.com
          Action: sts:AssumeRole
      Path: "/"
      Policies:
        - PolicyName: root
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: Allow
                Action:
                  - logs:CreateLogGroup
                  - logs:CreateLogStream
                  - logs:PutLogEvents
                Resource: arn:aws:logs:*:*:*

  PreSignUpLambda:
    Type: AWS::Serverless::Function
    Properties:
      CodeUri: ../build
      Handler: pre-signup-lambda.handler
      MemorySize: 128
      Role: !GetAtt LambdaExecutionRole.Arn
      Runtime: nodejs20.x
      Timeout: 5

  CognitoUserPool:
    Type: AWS::Cognito::UserPool
    Properties:
      UserPoolName:
        !FindInMap [Environment, !Ref "EnvironmentName", "AppUserPoolName"]
      Policies:
        PasswordPolicy:
          MinimumLength: 6
          RequireLowercase: false
          RequireNumbers: false
          RequireSymbols: false
          RequireUppercase: false
          TemporaryPasswordValidityDays: 7
      LambdaConfig:
        PreSignUp: !GetAtt PreSignUpLambda.Arn

  CognitoUserPoolClient:
    Type: "AWS::Cognito::UserPoolClient"
    Properties:
      ClientName:
        !FindInMap [Environment, !Ref "EnvironmentName", "AppUserPoolClient"]
      GenerateSecret: "false"
      RefreshTokenValidity: 30
      UserPoolId: !Ref CognitoUserPool

  PreSignUpLambdaInvocationPermission:
    Type: AWS::Lambda::Permission
    Properties:
      Action: lambda:InvokeFunction
      FunctionName: !GetAtt PreSignUpLambda.Arn
      Principal: cognito-idp.amazonaws.com
      SourceArn: !GetAtt CognitoUserPool.Arn

For full-scale production, we usually add more specific user details, set strong password rules, and use custom Lambda triggers for extra features. But for this case we wanted to have the simplest possible setup. We also used a pre-signup lambda function to automatically confirm all users, avoiding the need of email or phone number verification.

Setting Up a DynamoDB Table for WebSocket Connections

To manage WebSocket connections efficiently, we set up a DynamoDB table tailored for tracking active connections and supporting real-time communication. DynamoDB’s capability to scale and its high-performance nature suit the dynamic requirements of WebSocket connections perfectly, ensuring rapid reads and writes.

This is the CloudFormation template that we used for the connections table:

Resources:
  WebSocketConnections:
    Type: AWS::DynamoDB::Table
    Properties:
      AttributeDefinitions:
        - AttributeName: "username"
          AttributeType: "S"
        - AttributeName: "sessionKey"
          AttributeType: "S"
        - AttributeName: "connectionId"
          AttributeType: "S"
      KeySchema:
        - AttributeName: "username"
          KeyType: "HASH"
        - AttributeName: "sessionKey"
          KeyType: "RANGE"
      BillingMode: PAY_PER_REQUEST
      GlobalSecondaryIndexes:
        - IndexName: !Sub "${ApplicationName}-${EnvironmentName}-websocket-connections-index-connection"
          KeySchema:
            - AttributeName: "connectionId"
              KeyType: "HASH"
          Projection:
            ProjectionType: "ALL"
      SSESpecification:
        SSEEnabled: True
      TableName: !Sub "${ApplicationName}-${EnvironmentName}-websocket-connections"
      TimeToLiveSpecification:
        AttributeName: "expires"
        Enabled: True

The main features of this template are:

A single user can maintain multiple concurrent connections, one for each active session.
Through a global secondary index, the table allows for querying data by connection ID.
A Time to Live (TTL) setting is configured to automatically purge data, maintaining table efficiency and relevance.

Integrating AWS Lambda function with the HTTP API Gateway

To define the business logic of the HTTP API, we deployed an AWS Lambda function, integrating it with API Gateway via AWS CloudFormation and the Serverless Application Model (SAM). In our approach, we specifically utilized SAM for the Lambda function definition due to its streamlined deployment capabilities.

SAM’s package and deploy commands significantly simplify the deployment process, automating the manual steps typically required to deploy a Lambda function. This focused use of SAM, combined with CloudFormation for other needed resources, offered a balanced and effective strategy for managing our serverless application deployment.

Below is the template we used for deploying the lambda function and integrating it with the HTTP API Gateway.

Resources:
  LambdaExecutionRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          Effect: Allow
          Principal:
            Service: lambda.amazonaws.com
          Action: sts:AssumeRole
      Path: "/"
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole

  APIFunction:
    Type: AWS::Serverless::Function
    Properties:
      CodeUri: ../build
      Handler: index.handler
      MemorySize: 512
      Role: !GetAtt LambdaExecutionRole.Arn
      Runtime: nodejs20.x
      Timeout: 30
      Environment:
        Variables:
          APPLICATION_NAME: !Ref ApplicationName
          ENVIRONMENT_NAME: !Ref EnvironmentName

  APIFunctionIntegration:
    Type: AWS::ApiGatewayV2::Integration
    Properties:
      ApiId: !FindInMap [Environment, !Ref "EnvironmentName", "HTTPApiId"]
      IntegrationType: AWS_PROXY
      IntegrationUri: !Join
        - ""
        - - "arn:"
          - !Ref "AWS::Partition"
          - ":apigateway:"
          - !Ref "AWS::Region"
          - ":lambda:path/2015-03-31/functions/"
          - !GetAtt APIFunction.Arn
          - /invocations
      PayloadFormatVersion: "2.0"

  JWTAuthorizer:
    Type: AWS::ApiGatewayV2::Authorizer
    Properties:
      Name: !Sub ${ApplicationName}-${EnvironmentName}-jwt-authorizer
      ApiId: !FindInMap [Environment, !Ref "EnvironmentName", "HTTPApiId"]
      AuthorizerType: JWT
      IdentitySource:
        - "$request.header.Authorization"
      JwtConfiguration:
        Audience:
          - !FindInMap [
              Environment,
              !Ref "EnvironmentName",
              "CognitoAppClientID",
            ]
        Issuer: !Sub
          - "https://cognito-idp.${AWS::Region}.amazonaws.com/${UserPool}"
          - UserPool:
              !FindInMap [Environment, !Ref "EnvironmentName", "UserPool"]

  APIDefaultRoute:
    Type: AWS::ApiGatewayV2::Route
    Properties:
      ApiId: !FindInMap [Environment, !Ref "EnvironmentName", "HTTPApiId"]
      AuthorizationType: JWT
      AuthorizerId: !Ref JWTAuthorizer
      RouteKey: $default
      Target: !Join
        - /
        - - integrations
          - !Ref APIFunctionIntegration

  APIFunctionPermission:
    Type: AWS::Lambda::Permission
    Properties:
      Action: lambda:InvokeFunction
      FunctionName: !Ref APIFunction
      Principal: apigateway.amazonaws.com
      SourceArn: !Join
        - ""
        - - !Sub "arn:${AWS::Partition}:execute-api:${AWS::Region}:${AWS::AccountId}:"
          - !FindInMap [Environment, !Ref "EnvironmentName", "HTTPApiId"]
          - "/*"

The main features of this template are:

IAM Role for Lambda: Essential for granting execution permissions to the Lambda function, including logging and accessing other AWS services like DynamoDB. Please note that the policies that grant access to DynamoDB and other AWS resources are omitted for simplicity.
Lambda Function Configuration: Describes the setup for the serverless function, including its code, runtime environment, and operational parameters like memory and timeout settings. Environment variables are used to tailor the function’s behavior for different deployment stages.
API Gateway Integration: This component connects the Lambda function to the API Gateway, ensuring smooth request forwarding and responses. It’s key for making the function accessible online as an API, leveraging proxy integration to enable straightforward communication.
JWT Authorization: Establishes security for the API by integrating JWT-based user authentication, leveraging a user pool for token validation. This layer of security ensures that only authenticated requests can access the API, protecting it from unauthorized use.
Route Configuration: Establishes how API requests are directed to the Lambda function, creating a structured flow for handling incoming traffic.
Invocation Permission: Grants the API Gateway authority to trigger the Lambda function.

Integrating AWS Lambda functions with the WebSocket API Gateway

While we typically employ a single Lambda function to handle all routes of an HTTP API through a proxy-type deployment, our approach with WebSocket APIs differs. For WebSocket functionality, we designate individual Lambda functions for various purposes within the WebSocket API.

In the structure of WebSocket APIs utilizing AWS API Gateway v2, three indispensable routes form the core framework for managing and facilitating connections effectively:

$connect route: Initiated with an HTTP request, this route is the first point of contact for clients attempting to establish a WebSocket connection.
$disconnect route: Triggered when a WebSocket connection comes to an end, this route is essential for performing clean-up operations and managing resources.
$default route: Acts as a catch-all for any messages sent over the WebSocket that do not correspond to other predefined routes.

Beyond the three mandatory routes, we define an extra additional route specifically designed for ping/pong communication, which is initiated by the clients. Since CloudFront has a maximum idle connection timeout as low as 60 seconds, this heartbeat mechanism is important for keeping the WebSocket connection alive. Please note that while WebSocket API Gateway has a 10 minutes idle connection timeout, CloudFront has only maximum of 1 minute.

The deployment and configuration of these routes are efficiently managed through the Serverless Application Model (SAM) and AWS CloudFormation, similarly to HTTP API.

Below is the template we used for deploying the lambda functions and integrating those with the WebSocket API Gateway.

Resources:
  LambdaExecutionRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          Effect: Allow
          Principal:
            Service: lambda.amazonaws.com
          Action: sts:AssumeRole
      Path: "/"
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole
      Policies:
        - PolicyName: root
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: Allow
                Action:
                  - "execute-api:ManageConnections"
                Resource:
                  - !Sub "arn:aws:execute-api:${AWS::Region}:${AWS::AccountId}:${WebSocketApi}/*"

  CustomAuthorizerFunction:
    Type: AWS::Serverless::Function
    Properties:
      FunctionName: !Sub "${ApplicationName}-${EnvironmentName}-websocket-authorizer"
      CodeUri: ../build
      Handler: authorizer.handler
      Runtime: nodejs20.x
      MemorySize: 256
      Role: !GetAtt LambdaExecutionRole.Arn

  WebSocketApiAuthorizer:
    Type: AWS::ApiGatewayV2::Authorizer
    Properties:
      Name: !Sub "${ApplicationName}-${EnvironmentName}-websocket-authorizer"
      ApiId: !FindInMap [Environment, !Ref "EnvironmentName", "WSApiId"]
      AuthorizerType: REQUEST
      AuthorizerUri:
        Fn::Sub: arn:aws:apigateway:${AWS::Region}:lambda:path/2015-03-31/functions/${CustomAuthorizerFunction.Arn}/invocations

  OnConnectFunction:
    Type: AWS::Serverless::Function
    Properties:
      FunctionName: !Sub "${ApplicationName}-${EnvironmentName}-websockets-on-connect"
      CodeUri: ../build
      Handler: onconnect.handler
      Runtime: nodejs20.x
      MemorySize: 256
      Role: !GetAtt LambdaExecutionRole.Arn

  OnConnectInteg:
    Type: AWS::ApiGatewayV2::Integration
    Properties:
      ApiId: !FindInMap [Environment, !Ref "EnvironmentName", "WSApiId"]
      Description: Connect Integration
      IntegrationType: AWS_PROXY
      IntegrationUri:
        Fn::Sub: arn:aws:apigateway:${AWS::Region}:lambda:path/2015-03-31/functions/${OnConnectFunction.Arn}/invocations

  OnConnectRoute:
    Type: AWS::ApiGatewayV2::Route
    Properties:
      ApiId: !FindInMap [Environment, !Ref "EnvironmentName", "WSApiId"]
      RouteKey: $connect
      OperationName: ConnectRoute
      AuthorizationType: CUSTOM
      AuthorizerId: !Ref WebSocketApiAuthorizer
      Target: !Join
        - "/"
        - - "integrations"
          - !Ref OnConnectInteg

  OnConnectPermission:
    Type: AWS::Lambda::Permission
    Properties:
      Action: lambda:InvokeFunction
      FunctionName: !Ref OnConnectFunction
      Principal: apigateway.amazonaws.com

  CustomAuthorizerFunctionInvokePermission:
    Type: AWS::Lambda::Permission
    Properties:
      Action: lambda:InvokeFunction
      FunctionName: !Ref CustomAuthorizerFunction
      Principal: apigateway.amazonaws.com

  OnDisconnectFunction:
    Type: AWS::Serverless::Function
    Properties:
      FunctionName: !Sub "${ApplicationName}-${EnvironmentName}-websockets-on-disconnect"
      CodeUri: ../build/
      Handler: ondisconnect.handler
      MemorySize: 256
      Role: !GetAtt LambdaExecutionRole.Arn
      Runtime: nodejs20.x

  OnDisconnectInteg:
    Type: AWS::ApiGatewayV2::Integration
    Properties:
      ApiId: !FindInMap [Environment, !Ref "EnvironmentName", "WSApiId"]
      Description: Disconnect Integration
      IntegrationType: AWS_PROXY
      IntegrationUri:
        Fn::Sub: arn:aws:apigateway:${AWS::Region}:lambda:path/2015-03-31/functions/${OnDisconnectFunction.Arn}/invocations

  OnDisconnectRoute:
    Type: AWS::ApiGatewayV2::Route
    Properties:
      ApiId: !FindInMap [Environment, !Ref "EnvironmentName", "WSApiId"]
      RouteKey: $disconnect
      AuthorizationType: NONE
      OperationName: DisconnectRoute
      Target: !Join
        - "/"
        - - "integrations"
          - !Ref OnDisconnectInteg

  OnDisconnectPermission:
    Type: AWS::Lambda::Permission
    Properties:
      Action: lambda:InvokeFunction
      FunctionName: !Ref OnDisconnectFunction
      Principal: apigateway.amazonaws.com

  OnDefaultFunction:
    Type: AWS::Serverless::Function
    Properties:
      FunctionName: !Sub "${ApplicationName}-${EnvironmentName}-websockets-on-default"
      CodeUri: ../build/
      Handler: ondefault.handler
      MemorySize: 256
      Role: !GetAtt LambdaExecutionRole.Arn
      Runtime: nodejs20.x

  OnDefaultInteg:
    Type: AWS::ApiGatewayV2::Integration
    Properties:
      ApiId: !FindInMap [Environment, !Ref "EnvironmentName", "WSApiId"]
      Description: Default Integration
      IntegrationType: AWS_PROXY
      IntegrationUri:
        Fn::Sub: arn:aws:apigateway:${AWS::Region}:lambda:path/2015-03-31/functions/${OnDefaultFunction.Arn}/invocations

  OnDefaultRoute:
    Type: AWS::ApiGatewayV2::Route
    Properties:
      ApiId: !FindInMap [Environment, !Ref "EnvironmentName", "WSApiId"]
      RouteKey: $default
      AuthorizationType: NONE
      OperationName: DefaultRoute
      Target: !Join
        - "/"
        - - "integrations"
          - !Ref OnDefaultInteg

  OnDefaultPermission:
    Type: AWS::Lambda::Permission
    Properties:
      Action: lambda:InvokeFunction
      FunctionName: !Ref OnDefaultFunction
      Principal: apigateway.amazonaws.com

  OnPingFunction:
    Type: AWS::Serverless::Function
    Properties:
      FunctionName: !Sub "${ApplicationName}-${EnvironmentName}-websockets-on-ping"
      CodeUri: ../build/
      Handler: onping.handler
      MemorySize: 256
      Role: !GetAtt LambdaExecutionRole.Arn
      Runtime: nodejs20.x

  OnPingInteg:
    Type: AWS::ApiGatewayV2::Integration
    Properties:
      ApiId: !FindInMap [Environment, !Ref "EnvironmentName", "WSApiId"]
      Description: Ping Pong Integration
      IntegrationType: AWS_PROXY
      IntegrationUri:
        Fn::Sub: arn:aws:apigateway:${AWS::Region}:lambda:path/2015-03-31/functions/${OnPingFunction.Arn}/invocations

  OnPingRoute:
    Type: AWS::ApiGatewayV2::Route
    Properties:
      ApiId: !FindInMap [Environment, !Ref "EnvironmentName", "WSApiId"]
      RouteKey: ping
      AuthorizationType: NONE
      OperationName: PingRoute
      Target: !Join
        - "/"
        - - "integrations"
          - !Ref OnPingInteg

  OnPingPermission:
    Type: AWS::Lambda::Permission
    Properties:
      Action: lambda:InvokeFunction
      FunctionName: !Ref OnPingFunction
      Principal: apigateway.amazonaws.com

  Deployment:
    Type: AWS::ApiGatewayV2::Deployment
    DependsOn:
      - OnConnectPermission
      - OnDisconnectPermission
      - OnDefaultPermission
      - OnPingPermission
    Properties:
      ApiId: !FindInMap [Environment, !Ref "EnvironmentName", "WSApiId"]
      Description: !Sub "${EnvironmentName} Deployment ${DeploymentTimestamp}"

  Stage:
    Type: AWS::ApiGatewayV2::Stage
    Properties:
      StageName: !FindInMap [Environment, !Ref "EnvironmentName", "StageName"]
      Description: !Sub "${EnvironmentName} Stage"
      DeploymentId: !Ref Deployment
      ApiId: !FindInMap [Environment, !Ref "EnvironmentName", "WSApiId"]

The main features of this template are:

Lambda functions are deployed for the needed WebSocket events ($connect, $disconnect, $default, ping), each with specific roles and integrated securely with API Gateway.
For WebSocket APIs, directly using JWT as the authorizer type is not supported. However, similar levels of security can be achieved by employing a custom Lambda authorizer to validate the JWT. I’ll write more of that in my follow-up blog post.

Conclusion

This blog post has focused primarily on setting up the foundational AWS resources necessary for enabling real-time communication through WebSockets using AWS API Gateway v2, alongside Amazon CloudFront.

The emphasis has been on the infrastructure setup — leveraging CloudFormation and the Serverless Application Model (SAM) for a streamlined deployment process. This groundwork lays the foundation for developing responsive and scalable applications, capable of handling real-time communication with ease.

Looking ahead, a subsequent blog post will delve into the Lambda functions’ logic as well as details and noticeable caveats to take into account. Thank you for reading, and stay tuned for more in-depth articles on leveraging AWS services for your SaaS solutions.

Setting Up Real-Time Communication with WebSockets Using AWS API Gateway v2 and Amazon CloudFront

AWS Services in The Solution

Setting Up Base AWS Resources with CloudFormation

Preliminary Tasks: Setting Up a Hosted Zone and SSL/TLS Certificate

Setting Up API Gateways with CloudFormation

Setting Up CloudFront Distribution and Route 53 RecordSet with CloudFormation

Creating Cognito User Pool and User Pool Client

Setting Up a DynamoDB Table for WebSocket Connections

Integrating AWS Lambda function with the HTTP API Gateway

Integrating AWS Lambda functions with the WebSocket API Gateway

Conclusion

Written by Jari Ikävalko