A CloudFormation template for deploying a serverless websocket API on AWS with a CloudFront function to convert path parameters to query string parameters

cloudfront-template.yaml

AWSTemplateFormatVersion: "2010-09-09"
Transform: "AWS::Serverless-2016-10-31"

Resources:
  ApiGateway:
    Type: "AWS::ApiGatewayV2::Api"
    Properties:
      Name: !Sub "${AWS::StackName}-wss-api"
      ProtocolType: "WEBSOCKET"
      RouteSelectionExpression: "\\$default"

  Stage:
    Type: "AWS::ApiGatewayV2::Stage"
    Properties:
      StageName: "Prod"
      AutoDeploy: true
      ApiId: !Ref "ApiGateway"

  LambdaRole:
    Type: "AWS::IAM::Role"
    Properties:
      RoleName: !Sub "${AWS::StackName}-lambda-role"
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: "Allow"
            Principal:
              Service:
                - "lambda.amazonaws.com"
            Action:
              - "sts:AssumeRole"

  LambdaPolicy:
    Type: "AWS::IAM::Policy"
    Properties:
      PolicyName: !Sub "${AWS::StackName}-lambda-policy"
      PolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: "Allow"
            Action:
              - "logs:CreateLogGroup"
              - "logs:CreateLogStream"
              - "logs:PutLogEvents"
            Resource: "*"
      Roles:
        - !Ref "LambdaRole"

  ConnectFunction:
    Type: "AWS::Serverless::Function"
    Properties:
      Runtime: "python3.10"
      Timeout: 30
      Architectures:
        - "arm64"
      MemorySize: 256
      Role: !GetAtt "LambdaRole.Arn"
      Handler: "index.handler"
      InlineCode: |
        def handler(event, context):
          print(event)
          return {'statusCode': 200}

  ConnectFunctionLogGroup:
    Type: "AWS::Logs::LogGroup"
    Properties:
      LogGroupName: !Sub "/aws/lambda/${ConnectFunction}"
      RetentionInDays: 30

  ConnectInvokePermission:
    Type: "AWS::Lambda::Permission"
    DependsOn:
      - "ApiGateway"
    Properties:
      Action: "lambda:InvokeFunction"
      FunctionName: !Ref "ConnectFunction"
      Principal: "apigateway.amazonaws.com"

  ConnectRoute:
    Type: "AWS::ApiGatewayV2::Route"
    Properties:
      ApiId: !Ref "ApiGateway"
      RouteKey: "$connect"
      OperationName: "ConnectRoute"
      Target: !Sub "integrations/${ConnectIntegration}"

  ConnectIntegration:
    Type: "AWS::ApiGatewayV2::Integration"
    Properties:
      ApiId: !Ref "ApiGateway"
      IntegrationType: "AWS_PROXY"
      IntegrationUri: !Sub "arn:aws:apigateway:${AWS::Region}:lambda:path/2015-03-31/functions/${ConnectFunction.Arn}/invocations"

  CloudFrontDist:
    Type: "AWS::CloudFront::Distribution"
    Properties:
      DistributionConfig:
        Origins:
          - Id: !Sub "${AWS::StackName}-cloudfront-origin"
            DomainName: !Sub "${ApiGateway}.execute-api.${AWS::Region}.amazonaws.com"
            OriginPath: !Sub "/${Stage}"
            CustomOriginConfig:
              HTTPSPort: 443
              OriginProtocolPolicy: "https-only"
        DefaultCacheBehavior:
          ViewerProtocolPolicy: "https-only"
          TargetOriginId: !Sub "${AWS::StackName}-cloudfront-origin" # must be the same as the origin defined above
          CachePolicyId: "4135ea2d-6df8-44a3-9df3-4b5a84be39ad" # Managed-CachingDisabled
          OriginRequestPolicyId: !Ref "CloudFrontOriginRequestPolicy"
          FunctionAssociations:
            - EventType: "viewer-request"
              FunctionARN: !GetAtt "CloudFrontFunction.FunctionMetadata.FunctionARN"
        Enabled: true
        IPV6Enabled: false

  CloudFrontOriginRequestPolicy:
    Type: "AWS::CloudFront::OriginRequestPolicy"
    Properties:
      OriginRequestPolicyConfig:
        Name: !Sub "${AWS::StackName}-cloudfront-orp"
        HeadersConfig:
          HeaderBehavior: "whitelist"
          Headers:
            - "Sec-WebSocket-Key"
            - "Sec-WebSocket-Version"
            - "Sec-WebSocket-Protocol"
            - "Sec-WebSocket-Accept"
        QueryStringsConfig:
          QueryStringBehavior: "all"
        CookiesConfig:
          CookieBehavior: "none"

  CloudFrontFunction:
    Type: "AWS::CloudFront::Function"
    Properties:
      Name: !Sub "${AWS::StackName}-cloudfront-function"
      AutoPublish: true
      FunctionCode: |
        function handler(event) {
          var request = event.request;

          var re = /^(.*?\/)([^.]+)$/;
          var match = re.exec(request.uri);

          if (match) {
            request.uri = match[1];

            request.querystring.path = {
                'multiValue': match[2].split('/').map(p => { return {'value': p} })
            };
          }

          return request;
        }
      FunctionConfig:
        Comment: "Change path parameters to query string"
        Runtime: "cloudfront-js-1.0"

Outputs:
  ServerApi:
    Description: "Api Gateway endpoint URL"
    Value: !Sub "${ApiGateway.ApiEndpoint}/${Stage}"
  CloudFrontUrl:
    Description: "Cloudfront URL"
    Value: !Sub "wss://${CloudFrontDist.DomainName}"