CloudFormationTemplateForServerlessChatApp.yaml

Resources:
  ConnectionsTable:
    Type: AWS::DynamoDB::Table
    Properties:
      KeySchema:
        - AttributeName: connectionId
          KeyType: HASH
      AttributeDefinitions:
        - AttributeName: connectionId
          AttributeType: S
      ProvisionedThroughput: # Switching to provisioned throughput
        ReadCapacityUnits: 5
        WriteCapacityUnits: 5
      PointInTimeRecoverySpecification:
        PointInTimeRecoveryEnabled: true
      SSESpecification:
        SSEEnabled: true
        SSEType: KMS
        KMSMasterKeyId: "alias/<placeholder>" # Will insert the key as necessary
    UpdateReplacePolicy: Delete
    DeletionPolicy: Delete
  AutoScalingRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Action: sts:AssumeRole
            Effect: Allow
            Principal:
              Service: application-autoscaling.amazonaws.com
      Policies:
        - PolicyName: DynamoDBAutoScalingPolicy
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Action:
                  - dynamodb:DescribeTable
                  - dynamodb:UpdateTable
                Effect: Allow
                Resource: !GetAtt ConnectionsTable.Arn
  WriteAutoScalingTarget:
    Type: AWS::ApplicationAutoScaling::ScalableTarget
    Properties:
      MaxCapacity: 10
      MinCapacity: 5
      ResourceId: !Sub table/${ConnectionsTable}
      RoleARN: !GetAtt AutoScalingRole.Arn
      ScalableDimension: dynamodb:table:WriteCapacityUnits
      ServiceNamespace: dynamodb
  WriteAutoScalingPolicy:
    Type: AWS::ApplicationAutoScaling::ScalingPolicy
    Properties:
      PolicyName: DynamoDBWriteAutoScalingPolicy
      PolicyType: TargetTrackingScaling
      ScalingTargetId: !Ref WriteAutoScalingTarget
      TargetTrackingScalingPolicyConfiguration:
        PredefinedMetricSpecification:
          PredefinedMetricType: DynamoDBWriteCapacityUtilization
        TargetValue: 70
        ScaleInCooldown: 60
        ScaleOutCooldown: 60
  ReadAutoScalingPolicy:
    Type: AWS::ApplicationAutoScaling::ScalingPolicy
    Properties:
      PolicyName: ReadAutoScalingPolicy
      PolicyType: TargetTrackingScaling
      ScalingTargetId: !Ref ReadScalingTarget
      TargetTrackingScalingPolicyConfiguration:
        PredefinedMetricSpecification:
          PredefinedMetricType: DynamoDBReadCapacityUtilization
        ScaleInCooldown: 60
        ScaleOutCooldown: 60
        TargetValue: 70
  ReadScalingTarget:
    Type: AWS::ApplicationAutoScaling::ScalableTarget
    Properties:
      MaxCapacity: 10
      MinCapacity: 5
      ResourceId: !Sub table/${ConnectionsTable}
      RoleARN: !GetAtt AutoScalingRole.Arn
      ScalableDimension: dynamodb:table:ReadCapacityUnits
      ServiceNamespace: dynamodb
  ConnectHandlerServiceRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Action: sts:AssumeRole
            Effect: Allow
            Principal:
              Service: lambda.amazonaws.com
      ManagedPolicyArns:
        - !Join
          - ''
          - - 'arn:'
            - !Ref AWS::Partition
            - ':iam::aws:policy/service-role/AWSLambdaBasicExecutionRole'
  ConnectHandlerServiceRoleDefaultPolicy:
    Type: AWS::IAM::Policy
    Properties:
      PolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Action:
              - dynamodb:BatchWriteItem
              - dynamodb:PutItem
              - dynamodb:UpdateItem
              - dynamodb:DeleteItem
              - dynamodb:DescribeTable
            Effect: Allow
            Resource:
              - !GetAtt ConnectionsTable.Arn
              - !Ref AWS::NoValue
      PolicyName: ConnectHandlerServiceRoleDefaultPolicy
      Roles:
        - !Ref ConnectHandlerServiceRole
  ConnectHandler:
    Type: AWS::Lambda::Function
    Properties:
      Code:
        ZipFile: |-
          const AWS = require('aws-sdk');
          const ddb = new AWS.DynamoDB.DocumentClient({apiVersion: '2012-08-10'});

          const maxRetries = 3; // Maximum number of retries
          let retryCount = 0; // Current retry count

          exports.handler = async function (event, context) {
            try {
              const result = await attemptPutItem();
              return {
                statusCode: 200,
                body: JSON.stringify(result),
              };
            } catch (err) {
              console.error('Error in processing DynamoDB request:', err);
              return {
                statusCode: 500,
                body: 'Internal server error',
              };
            }
          };

          async function attemptPutItem() {
            try {
              return await ddb.put({
                TableName: process.env.table,
                Item: {
                  connectionId: event.requestContext.connectionId,
                },
              }).promise();
            } catch (error) {
              if (retryCount < maxRetries) {
                retryCount++;
                console.log(`Retry ${retryCount} for connection ID:`, event.requestContext.connectionId);
                await new Promise(resolve => setTimeout(resolve, 1000 * retryCount)); // Exponential back-off
                return attemptPutItem();
              } else {
                console.error('Max retries reached');
                throw new Error('Max retry attempts failed');
              }
            }
          }
      Role: !GetAtt ConnectHandlerServiceRole.Arn
      Environment:
        Variables:
          table: !Ref ConnectionsTable
      Handler: index.handler
      Runtime: nodejs16.x
    DependsOn:
      - ConnectHandlerServiceRoleDefaultPolicy
      - ConnectHandlerServiceRole
  ConnectHandlerVersion:
    Type: AWS::Lambda::Version
    Properties:
      FunctionName: !Ref ConnectHandler
  ConnectHandlerAlias:
    Type: AWS::Lambda::Alias
    Properties:
      FunctionName: !Ref ConnectHandler
      FunctionVersion: !GetAtt ConnectHandlerVersion.Version
      Name: Prod
      ProvisionedConcurrencyConfig:
        ProvisionedConcurrentExecutions: 5
  DisconnectHandlerServiceRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Action: sts:AssumeRole
            Effect: Allow
            Principal:
              Service: lambda.amazonaws.com
      ManagedPolicyArns:
        - !Join
          - ''
          - - 'arn:'
            - !Ref AWS::Partition
            - ':iam::aws:policy/service-role/AWSLambdaBasicExecutionRole'
  DisconnectHandlerServiceRoleDefaultPolicy:
    Type: AWS::IAM::Policy
    Properties:
      PolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Action:
              - dynamodb:BatchWriteItem
              - dynamodb:PutItem
              - dynamodb:UpdateItem
              - dynamodb:DeleteItem
              - dynamodb:DescribeTable
            Effect: Allow
            Resource:
              - !GetAtt ConnectionsTable.Arn  # Restrict access to specific table only
              - !Ref AWS::NoValue
      PolicyName: DisconnectHandlerServiceRoleDefaultPolicy
      Roles:
        - !Ref DisconnectHandlerServiceRole
  DisconnectHandler:
    Type: AWS::Lambda::Function
    Properties:
      Code:
        ZipFile: |-
          const AWS = require('aws-sdk');
                const ddb = new AWS.DynamoDB.DocumentClient();
                exports.handler = async function (event, context) {
                  await ddb
                    .delete({
                      TableName: process.env.table,
                      Key: {
                        connectionId: event.requestContext.connectionId,
                      },
                    })
                    .promise();
                  return {
                    statusCode: 200,
                  };
                };
      Role: !GetAtt DisconnectHandlerServiceRole.Arn
      Environment:
        Variables:
          table: !Ref ConnectionsTable
      Handler: index.handler
      Runtime: nodejs16.x
    DependsOn:
      - DisconnectHandlerServiceRoleDefaultPolicy
      - DisconnectHandlerServiceRole
  DisconnectHandlerVersion:
    Type: AWS::Lambda::Version
    Properties:
      FunctionName: !Ref DisconnectHandler
  DisconnectHandlerAlias:
    Type: AWS::Lambda::Alias
    Properties:
      FunctionName: !Ref DisconnectHandler
      FunctionVersion: !GetAtt DisconnectHandlerVersion.Version
      Name: Prod
      ProvisionedConcurrencyConfig:
        ProvisionedConcurrentExecutions: 5
  SendMessageHandlerServiceRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Statement:
          - Action: sts:AssumeRole
            Effect: Allow
            Principal:
              Service: lambda.amazonaws.com
        Version: '2012-10-17'
      ManagedPolicyArns:
        - !Join
          - ''
          - - 'arn:'
            - !Ref AWS::Partition
            - ':iam::aws:policy/service-role/AWSLambdaBasicExecutionRole'
  SendMessageHandlerServiceRoleDefaultPolicy:
    Type: AWS::IAM::Policy
    Properties:
      PolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Action:
              - dynamodb:BatchGetItem
              - dynamodb:GetRecords
              - dynamodb:GetShardIterator
              - dynamodb:Query
              - dynamodb:GetItem
              - dynamodb:Scan
              - dynamodb:ConditionCheckItem
              - dynamodb:DescribeTable
            Effect: Allow
            Resource:
              - !GetAtt ConnectionsTable.Arn  # Specify exact resource for better security
              - !Ref AWS::NoValue
      PolicyName: SendMessageHandlerServiceRoleDefaultPolicy
      Roles:
        - !Ref SendMessageHandlerServiceRole
  SendMessageHandler:
    Type: AWS::Lambda::Function
    Properties:
      Code:
        ZipFile: |-
          const AWS = require('aws-sdk');
                const ddb = new AWS.DynamoDB.DocumentClient();
                exports.handler = async function (event, context) {
                  let connections;
                  try {
                    connections = await ddb.scan({ TableName: process.env.table }).promise();
                  } catch (err) {
                    return {
                      statusCode: 500,
                    };
                  }
                  const callbackAPI = new AWS.ApiGatewayManagementApi({
                    apiVersion: '2018-11-29',
                    endpoint:
                      event.requestContext.domainName + '/' + event.requestContext.stage,
                  });
                  const message = JSON.parse(event.body).message;
                  const sendMessages = connections.Items.map(async ({ connectionId }) => {
                    if (connectionId !== event.requestContext.connectionId) {
                      try {
                        await callbackAPI
                          .postToConnection({ ConnectionId: connectionId, Data: message })
                          .promise();
                      } catch (e) {
                        console.log(e);
                      }
                    }
                  });
                  try {
                    await Promise.all(sendMessages);
                  } catch (e) {
                    console.log(e);
                    return {
                      statusCode: 500,
                    };
                  }
                  return { statusCode: 200 };
                };
      Role: !GetAtt SendMessageHandlerServiceRole.Arn
      Environment:
        Variables:
          table: !Ref ConnectionsTable
      Handler: index.handler
      Runtime: nodejs16.x
    DependsOn:
      - SendMessageHandlerServiceRoleDefaultPolicy
      - SendMessageHandlerServiceRole
  SendMessageHandlerVersion:
    Type: AWS::Lambda::Version
    Properties:
      FunctionName: !Ref SendMessageHandler
  SendMessageHandlerAlias:
    Type: AWS::Lambda::Alias
    Properties:
      FunctionName: !Ref SendMessageHandler
      FunctionVersion: !GetAtt SendMessageHandlerVersion.Version
      Name: Prod
      ProvisionedConcurrencyConfig:
        ProvisionedConcurrentExecutions: 5
  DefaultHandlerServiceRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Statement:
          - Action: sts:AssumeRole
            Effect: Allow
            Principal:
              Service: lambda.amazonaws.com
        Version: '2012-10-17'
      ManagedPolicyArns:
        - !Join
          - ''
          - - 'arn:'
            - !Ref AWS::Partition
            - ':iam::aws:policy/service-role/AWSLambdaBasicExecutionRole'
  DefaultHandlerServiceRoleDefaultPolicy:
    Type: AWS::IAM::Policy
    Properties:
      PolicyDocument:
        Statement:
          - Action: execute-api:ManageConnections
            Effect: Allow
            Resource:
              - !Join
                - ''
                - - 'arn:aws:execute-api:'
                  - !Ref AWS::Region
                  - ':'
                  - !Ref AWS::AccountId
                  - ':'
                  - '*/*/POST/@connections/*'
              - !Join
                - ''
                - - 'arn:aws:execute-api:'
                  - !Ref AWS::Region
                  - ':'
                  - !Ref AWS::AccountId
                  - ':'
                  - '*/*/GET/@connections/*'
        Version: '2012-10-17'
      PolicyName: DefaultHandlerServiceRoleDefaultPolicy
      Roles:
        - !Ref DefaultHandlerServiceRole
  DefaultHandler:
    Type: AWS::Lambda::Function
    Properties:
      Code:
        ZipFile: |-
          const AWS = require('aws-sdk');
          exports.handler = async function (event, context) {
            let connectionInfo;
            let connectionId = event.requestContext.connectionId;
            const callbackAPI = new AWS.ApiGatewayManagementApi({
              apiVersion: '2018-11-29',
              endpoint:
                event.requestContext.domainName + '/' + event.requestContext.stage,
            });
            try {
              connectionInfo = await callbackAPI
                .getConnection({ ConnectionId: event.requestContext.connectionId })
                .promise();
            } catch (e) {
              console.log(e);
            }
            connectionInfo.connectionID = connectionId;
            await callbackAPI
              .postToConnection({
                ConnectionId: event.requestContext.connectionId,
                Data:
                  'Use the sendmessage route to send a message. Your info:' +
                  JSON.stringify(connectionInfo),
              })
              .promise();
            return {
              statusCode: 200,
            };
          };
      Role: !GetAtt DefaultHandlerServiceRole.Arn
      Handler: index.handler
      Runtime: nodejs16.x
    DependsOn:
      - DefaultHandlerServiceRoleDefaultPolicy
      - DefaultHandlerServiceRole
  manageConnections:
    Type: AWS::IAM::Policy
    Properties:
      PolicyDocument:
        Statement:
          - Action: execute-api:ManageConnections
            Effect: Allow
            Resource: !Join
              - ''
              - - 'arn:aws:execute-api:'
                - !Ref AWS::Region
                - ':'
                - !Ref AWS::AccountId
                - ':'
                - '*/*/POST/@connections/*'
        Version: '2012-10-17'
      PolicyName: manageConnections
      Roles:
        - !Ref SendMessageHandlerServiceRole

  # This is Optional       
  LambdaWarmerFunction:
    Type: AWS::Lambda::Function
    Properties:
      Handler: lambdaWarmer.handler
      Role: !GetAtt LambdaWarmerRole.Arn
      Code:
        ZipFile: |
          import boto3
          client = boto3.client('lambda')
          
          def handler(event, context):
              functions_to_warm = [
                "arn:aws:lambda:region:account-id:function:ConnectHandler",
                "arn:aws:lambda:region:account-id:function:DisconnectHandler",
                "arn:aws:lambda:region:account-id:function:SendMessageHandler"
              ]
              for function_name in functions_to_warm:
                  client.invoke(FunctionName=function_name, InvocationType='RequestResponse')
                  
              return {
                  'statusCode': 200,
                  'body': 'Warmed up Lambda functions successfully.'
              }
      Runtime: python3.8
      Timeout: 120
      MemorySize: 128

  LambdaWarmerRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: [lambda.amazonaws.com]
            Action: ['sts:AssumeRole']
      Policies:
        - PolicyName: LambdaInvocationPolicy
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action: ['lambda:InvokeFunction']
                Resource: [
                  "arn:aws:lambda:region:account-id:function:ConnectHandler",
                  "arn:aws:lambda:region:account-id:function:DisconnectHandler",
                  "arn:aws:lambda:region:account-id:function:SendMessageHandler"
                ]

  LambdaWarmerTrigger:
    Type: AWS::Events::Rule
    Properties:
      ScheduleExpression: 'rate(5 minutes)'
      Targets:
        - Arn: !GetAtt LambdaWarmerFunction.Arn
          Id: "LambdaWarmerTarget"
          RoleArn: !GetAtt LambdaWarmerRole.Arn