template.yaml

AWSTemplateFormatVersion: '2010-09-09'
Transform: AWS::Serverless-2016-10-31
Description: Serverless Chat App with Lambda Authorizer\



#This is the file that defines all the infrastructure behind this chat application
#Overview of Parts :
# 1. A DynamoDB with a Connections table (Logs if you are connected or not)
# 2. it's also got a Messages table with a searchable hash index (So we can search as well as secure the messages )
# 3. AWS Cognito Userpool : To store and authenticate user credentials it uses an obsured userpool with a UserPool CLient 
# 4. Handler files : The app runs via handlers to execute lambda functions (auth.py)

# 5.Lambda Functions , the heart of the execution : it has two funtions called AuthFunction and Auth Permission
# The AuthFunction Lambda invokes auth.py to generate and verify JWT Token.
# There's also functions to get and send messages and perform a search , connect and disconnect each with their own handler files 

# 6. Websocket and HTTP : These two handle diffeent things simultaneously, The HTTP Chat API takes care of the files like html , react , etc
# Meanwhile the Websocket works like a walkie talkie keeping an open connection between the chat people .

# 7. Routes and Integrations : These define the route API gateway uses 
# 8. Websocket URL Component 


#I chose defining my infra in a yaml file rather than terraform because SAM is purpose-built for serverless. It understands Lambda, API Gateway, and DynamoDB natively and generates the boilerplate CloudFormation for you. What would be 300 lines of Terraform is 80 lines of SAM YAML. so yeah it's easier to work with in this case 

Resources:

  # Maintains a connection ID and a hash key value , good for searches and finding people 
  ConnectionsTable:
    Type: AWS::DynamoDB::Table
    Properties:
      AttributeDefinitions:
        - {AttributeName: connectionId, AttributeType: S}
      KeySchema:
        - {AttributeName: connectionId, KeyType: HASH}
      BillingMode: PAY_PER_REQUEST
      TimeToLiveSpecification:
        AttributeName: ttl
        Enabled: true
#Holds messages with timestamps 
  MessagesTable:
    Type: AWS::DynamoDB::Table
    Properties:
      AttributeDefinitions:
        - {AttributeName: id, AttributeType: S}
        - {AttributeName: userId, AttributeType: S}
        - {AttributeName: timestamp, AttributeType: N}
      KeySchema:
        - {AttributeName: id, KeyType: HASH}
      BillingMode: PAY_PER_REQUEST
      GlobalSecondaryIndexes:
        - IndexName: userId-timestamp-index
          KeySchema:
            - {AttributeName: userId, KeyType: HASH}
            - {AttributeName: timestamp, KeyType: RANGE}
          Projection: {ProjectionType: ALL}
# User details live here 
  ChatUserPool:
    Type: AWS::Cognito::UserPool
    Properties:
      UserPoolName: ChatAppUsers
      UsernameAttributes: [email]
      AutoVerifiedAttributes: [email]

  ChatUserPoolClient:
    Type: AWS::Cognito::UserPoolClient
    Properties:
      ClientName: ChatAppClient
      UserPoolId: !Ref ChatUserPool
      ExplicitAuthFlows:
        - ALLOW_USER_PASSWORD_AUTH
        - ALLOW_USER_SRP_AUTH
        - ALLOW_REFRESH_TOKEN_AUTH
#This function calls the auth.py handler function to generate as well as authenticate or rather validate JWT
  AuthFunction:
    Type: AWS::Serverless::Function
    Properties:
      CodeUri: src/
      Handler: auth.lambda_handler
      Runtime: python3.12
      Environment:
        Variables:
          USER_POOL_ID: !Ref ChatUserPool
          APP_CLIENT_ID: !Ref ChatUserPoolClient

  AuthPermission:
    Type: AWS::Lambda::Permission
    Properties:
      Action: lambda:InvokeFunction
      FunctionName: !Ref AuthFunction
      Principal: apigateway.amazonaws.com
# The Websocket that acts as a walkie talkie allowing users to open chat 
  ChatWebSocketApi:
    Type: AWS::ApiGatewayV2::Api
    Properties:
      Name: ChatWebSocket
      ProtocolType: WEBSOCKET
      RouteSelectionExpression: "$request.body.action"
# Allows HTTP Web Content to be served 
  ChatHttpApi:
    Type: AWS::ApiGatewayV2::Api
    Properties:
      Name: ChatHttpApi
      ProtocolType: HTTP


  ChatApiStage:
    Type: AWS::ApiGatewayV2::Stage
    Properties:
      ApiId: !Ref ChatWebSocketApi
      StageName: Prod
      AutoDeploy: true

  ChatApiAuthorizer:
    Type: AWS::ApiGatewayV2::Authorizer
    Properties:
      Name: LambdaAuthorizer
      ApiId: !Ref ChatWebSocketApi
      AuthorizerType: REQUEST
      IdentitySource:
        - "route.request.querystring.token"
      AuthorizerUri: !Sub "arn:aws:apigateway:${AWS::Region}:lambda:path/2015-03-31/functions/${AuthFunction.Arn}/invocations"

  OnConnectRoute:
    Type: AWS::ApiGatewayV2::Route
    Properties:
      ApiId: !Ref ChatWebSocketApi
      RouteKey: $connect
      Target: !Sub "integrations/${OnConnectIntegration}"
      AuthorizationType: CUSTOM
      AuthorizerId: !Ref ChatApiAuthorizer

  SendMessageRoute:
    Type: AWS::ApiGatewayV2::Route
    Properties:
      ApiId: !Ref ChatWebSocketApi
      RouteKey: sendMessage
      Target: !Sub "integrations/${SendMessageIntegration}"

  OnDisconnectRoute:
    Type: AWS::ApiGatewayV2::Route
    Properties:
      ApiId: !Ref ChatWebSocketApi
      RouteKey: $disconnect
      Target: !Sub "integrations/${OnDisconnectIntegration}"

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

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

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


  OnConnectFunction:
    Type: AWS::Serverless::Function
    Properties:
      CodeUri: src/
      Handler: app.lambda_handler
      Runtime: python3.12
      Environment:
        Variables:
          TABLE_NAME: !Ref ConnectionsTable
      Policies:
        - DynamoDBCrudPolicy: {TableName: !Ref ConnectionsTable}
# Fetches the messages ealier 
  GetMessagesFunction:
    Type: AWS::Serverless::Function
    Properties:
      CodeUri: src/
      Handler: get_messages.lambda_handler
      Runtime: python3.12
      Environment:
        Variables:
          MESSAGES_TABLE: !Ref MessagesTable
      Policies:
        - DynamoDBReadPolicy: {TableName: !Ref MessagesTable}  
      Events:
        GetMessages:
          Type: HttpApi
          Properties:
            Path: /messages
            Method: get
            ApiId: !Ref ChatHttpApi    

  OnConnectPermission:
    Type: AWS::Lambda::Permission
    Properties:
      Action: lambda:InvokeFunction
      FunctionName: !Ref OnConnectFunction
      Principal: apigateway.amazonaws.com
# This bad boy takes care of sending the actual message to the person on the other end + the DB for persistent storage 
  SendMessageFunction:
    Type: AWS::Serverless::Function
    Properties:
      CodeUri: src/
      Handler: send_message.lambda_handler
      Runtime: python3.12
      Environment:
        Variables:
          TABLE_NAME: !Ref ConnectionsTable
          MESSAGES_TABLE: !Ref MessagesTable
          USER_POOL_ID: !Ref ChatUserPool
      Policies:
        - DynamoDBCrudPolicy: {TableName: !Ref ConnectionsTable}
        - DynamoDBCrudPolicy: {TableName: !Ref MessagesTable}
        - Statement:
            - Effect: Allow
              Action: execute-api:ManageConnections
              Resource: !Sub "arn:aws:execute-api:${AWS::Region}:${AWS::AccountId}:${ChatWebSocketApi}/*"

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

  OnDisconnectFunction:
    Type: AWS::Serverless::Function
    Properties:
      CodeUri: src/
      Handler: disconnect.lambda_handler
      Runtime: python3.12
      Environment:
        Variables:
          TABLE_NAME: !Ref ConnectionsTable
      Policies:
        - DynamoDBCrudPolicy: {TableName: !Ref ConnectionsTable}
# Allows the disconnect function through the API gateway
  OnDisconnectPermission:
    Type: AWS::Lambda::Permission
    Properties:
      Action: lambda:InvokeFunction
      FunctionName: !Ref OnDisconnectFunction
      Principal: apigateway.amazonaws.com

Outputs:
  WebSocketURL:
    Value: !Sub "wss://${ChatWebSocketApi}.execute-api.${AWS::Region}.amazonaws.com/Prod"
  UserPoolId:
    Value: !Ref ChatUserPool
  ClientId:
    Value: !Ref ChatUserPoolClient