AWSTemplateFormatVersion: "2010-09-09"
Description: Lambda@Edge CloudFront default document handler
Parameters:
  LambdaRuntime:
    Description: Runtime version for NodeJS
    Type: String
    Default: nodejs12.x
Resources:
  LambdaLogGroup:
    Type: AWS::Logs::LogGroup
    Properties: 
      LogGroupName: CloudFrontDefaultDocHandler
      RetentionInDays: 30
  LambdaExecutePolicy:
    Type: AWS::IAM::ManagedPolicy
    Properties:
      Description: "Lambda@Edge ExecutePolicy for prefix/folder redirection"
      Path: "/"
      PolicyDocument:
        Version: "2012-10-17"
        Statement:
          -
            Effect: "Allow"         
            Action: "logs:CreateLogGroup"
            Resource: !Join
              - ""
              - - "arn:aws:logs:*:"
                - !Ref "AWS::AccountId"
                - ":*"
          -
            Effect: "Allow"
            Action: 
              - "logs:CreateLogStream"
              - "logs:PutLogEvents"
            Resource: !Join
              - ""
              - - !GetAtt LambdaLogGroup.Arn
                - ":*"
  LambdaRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
        - Effect: "Allow"
          Principal: 
            Service:
            - edgelambda.amazonaws.com
            - lambda.amazonaws.com
          Action: "sts:AssumeRole"
      Path: "/"
      ManagedPolicyArns:
          - !Ref LambdaExecutePolicy
  LambdaCode:
    Type: AWS::Lambda::Function
    Properties: 
      Code:
        ZipFile: |
            const path = require('path')
            exports.handler = (event, context, callback) => {
                const { request } = event.Records[0].cf
                const olduri = request.uri;
                const extension = path.extname(olduri);
                if(extension && extension.length > 0){
                    return callback(null, request);
                }
                var newuri = ""
                if(olduri.slice(-1) === "/"){
                    newuri = olduri.replace(/\/$/, '\/index.html');
                } else {
                    newuri = olduri.replace(/$/, '\/index.html');
                }
                const redirect = {
                    status: '301',
                    statusDescription: 'Moved Permanently',
                    headers: {
                        location: [{
                            key: 'Location',
                            value: newuri,
                        }],
                    },
                };
                return callback(null, redirect);
            };
      Description: Handle default document for index.html
      FunctionName: !Join
        - '-'
        - - !Ref 'AWS::StackName'
          - "CloudFrontDefaultDocument"
      Handler: index.handler
      MemorySize: 128
      Role: !GetAtt LambdaRole.Arn
      Runtime: !Ref LambdaRuntime
      Timeout: 1
  LambdaCodeVersionPublish: 
    Type: AWS::Lambda::Version
    Properties: 
      FunctionName: !Ref LambdaCode
      Description: "A version of Default document handler"     
Outputs:
  FunctionARN:
    Description: ARN of the Function for CloudFront Distriutions to handle default documents (eg in prefixes/folders)
    Value: !Ref LambdaCodeVersionPublish