Shows how to use CloudFormation to attach a Lambda@Edge function to a CloudFront distribution to add HSTS and CSP custom headers. NOTE: The stack must be updated twice: once with the condition set to false, and once with it set to true.

lambda-at-edge.yaml

AWSTemplateFormatVersion: 2010-09-09
Parameters:
  RootDomainName:
    Type: String
  IncludeLambdaEdge:
    Type: String
    AllowedValues: ['true', 'false']
Conditions:
  IncludeLambdaEdge:
    !Equals ['true', !Ref IncludeLambdaEdge]
Mappings:
  RegionMap:
    us-east-1:
      S3HostedZoneID: Z3AQBSTGFYJSTF
      S3WebsiteEndpoint: s3-website-us-east-1.amazonaws.com
    us-west-1:
      S3HostedZoneID: Z2F56UZL2M1ACD
      S3WebsiteEndpoint: s3-website-us-west-1.amazonaws.com
    us-west-2:
      S3HostedZoneID: Z3BJ6K6RIION7M
      S3WebsiteEndpoint: s3-website-us-west-2.amazonaws.com
    eu-west-1:
      S3HostedZoneID: Z1BKCTXD74EZPE
      S3WebsiteEndpoint: s3-website-eu-west-1.amazonaws.com
    ap-southeast-1:
      S3HostedZoneID: Z3O0J2DXBE1FTB
      S3WebsiteEndpoint: s3-website-ap-southeast-1.amazonaws.com
    ap-southeast-2:
      S3HostedZoneID: Z1WCIGYICN2BYD
      S3WebsiteEndpoint: s3-website-ap-southeast-2.amazonaws.com
    ap-northeast-1:
      S3HostedZoneID: Z2M4EHUR26P7ZW
      S3WebsiteEndpoint: s3-website-ap-northeast-1.amazonaws.com
    sa-east-1:
      S3HostedZoneID: Z31GFT0UA1I2HV
      S3WebsiteEndpoint: s3-website-sa-east-1.amazonaws.com

Resources:

  RootCertificate:
    Type: 'AWS::CertificateManager::Certificate'
    Properties: 
      DomainName: !Ref RootDomainName

  SubdomainCertificate:
    Type: 'AWS::CertificateManager::Certificate'
    Properties: 
      DomainName: !Sub
        - '*.${Domain}'
        - Domain: !Ref RootDomainName

  PublicWebsiteRootBucket:
    Type: 'AWS::S3::Bucket'
    Properties:
      BucketName: !Ref RootDomainName
      AccessControl: PublicRead
      WebsiteConfiguration:
        RedirectAllRequestsTo:
          HostName: !Ref PublicWebsiteWwwBucket

  PublicWebsiteWwwBucket:
    Type: 'AWS::S3::Bucket'
    Properties:
      BucketName: !Sub
        - www.${Domain}
        - Domain: !Ref RootDomainName
      AccessControl: PublicRead
      WebsiteConfiguration:
        IndexDocument: index.html
        ErrorDocument: 404.html

  PublicWwwBucketPolicy:
      Type: 'AWS::S3::BucketPolicy'
      Properties:
        PolicyDocument:
          Id: PublicWebsitePolicy
          Version: 2012-10-17
          Statement:
            - Sid: PublicReadForGetBucketObjects
              Effect: Allow
              Principal: '*'
              Action: 's3:GetObject'
              Resource: !Join 
                - ''
                - - 'arn:aws:s3:::'
                  - !Ref PublicWebsiteWwwBucket
                  - /*
        Bucket: !Ref PublicWebsiteWwwBucket

  PublicRootBucketPolicy:
      Type: 'AWS::S3::BucketPolicy'
      Properties:
        PolicyDocument:
          Id: PublicWebsitePolicy
          Version: 2012-10-17
          Statement:
            - Sid: PublicReadForGetBucketObjects
              Effect: Allow
              Principal: '*'
              Action: 's3:GetObject'
              Resource: !Join 
                - ''
                - - 'arn:aws:s3:::'
                  - !Ref PublicWebsiteRootBucket
                  - /*
        Bucket: !Ref PublicWebsiteRootBucket

  InternalTracingBucket:
    Type: 'AWS::S3::Bucket'
    Properties:
      BucketName: !Join
        - '-'
        - - !Ref RootDomainName
          - 'tracing'

  PublicWebsiteWwwCloudfront:
    Type: AWS::CloudFront::Distribution
    DependsOn:
      - PublicWebsiteWwwBucket
      - InternalTracingBucket
    Properties:
      DistributionConfig:
        Comment: CloudFront to S3 - www
        Origins:
        - DomainName: !Join
            - '.'
            - - 'www'
              - !Ref 'RootDomainName'
              - !FindInMap [RegionMap, !Ref 'AWS::Region', S3WebsiteEndpoint]
          Id: S3WwwOrigin
          CustomOriginConfig:
            HTTPPort: '80'
            HTTPSPort: '443'
            OriginProtocolPolicy: http-only
        Enabled: true
        HttpVersion: 'http2'
        DefaultRootObject: index.html
        Aliases:
          - !Join
            - '.'
            - - 'www'
              - !Ref 'RootDomainName'
        DefaultCacheBehavior:
          AllowedMethods:
          - GET
          - HEAD
          Compress: true
          DefaultTTL: 3600
          TargetOriginId: S3WwwOrigin
          ForwardedValues:
            QueryString: true
            Cookies:
              Forward: none
          ViewerProtocolPolicy: redirect-to-https
          LambdaFunctionAssociations:
            - !If
                - IncludeLambdaEdge
                - EventType: 'origin-response'
                  LambdaFunctionARN: !Join
                    - ':'
                    - - !GetAtt [PoliciesEdgeLambda, Arn]
                      - !GetAtt [PoliciesEdgeLambdaVersion, Version]
                - !Ref 'AWS::NoValue'
        PriceClass: PriceClass_All
        ViewerCertificate:
          AcmCertificateArn: !Ref SubdomainCertificate
          SslSupportMethod: sni-only
        Logging:
          Bucket: !GetAtt [InternalTracingBucket, DomainName]
          IncludeCookies: false
          Prefix: 'www'

  PublicWebsiteRootCloudfront:
    Type: AWS::CloudFront::Distribution
    DependsOn:
      - PublicWebsiteRootBucket
    Properties:
      DistributionConfig:
        Comment: CloudFront to S3 - root
        Origins:
        - DomainName: !Join
            - '.'
            - - !Ref 'RootDomainName'
              - !FindInMap [RegionMap, !Ref 'AWS::Region', S3WebsiteEndpoint]
          Id: S3RootOrigin
          CustomOriginConfig:
            HTTPPort: '80'
            HTTPSPort: '443'
            OriginProtocolPolicy: http-only
        Enabled: true
        HttpVersion: 'http2'
        DefaultRootObject: index.html
        Aliases:
          - !Ref 'RootDomainName'
        DefaultCacheBehavior:
          AllowedMethods:
          - GET
          - HEAD
          Compress: true
          TargetOriginId: S3RootOrigin
          ForwardedValues:
            QueryString: true
            Cookies:
              Forward: none
          ViewerProtocolPolicy: redirect-to-https
          LambdaFunctionAssociations:
            - !If
                - IncludeLambdaEdge
                - EventType: 'origin-response'
                  LambdaFunctionARN: !Join
                    - ':'
                    - - !GetAtt [PoliciesEdgeLambda, Arn]
                      - !GetAtt [PoliciesEdgeLambdaVersion, Version]
                - !Ref 'AWS::NoValue'
        PriceClass: PriceClass_All
        ViewerCertificate:
          AcmCertificateArn: !Ref RootCertificate
          SslSupportMethod: sni-only

  HostedZone:
    Type: 'AWS::Route53::HostedZone'
    Properties:
      Name: !Ref RootDomainName

  DNS:
    Type: AWS::Route53::RecordSetGroup
    Properties:
      HostedZoneName: !Sub 
        - ${Domain}.
        - Domain: !Ref RootDomainName
      RecordSets:
        - Name: !Ref 'RootDomainName'
          Type: A
          AliasTarget:
            HostedZoneId: Z2FDTNDATAQYW2
            DNSName: !GetAtt [PublicWebsiteRootCloudfront, DomainName]
        - Name: !Join
          - '.'
          - - 'www'
            - !Ref 'RootDomainName'
          Type: A
          AliasTarget:
            HostedZoneId: Z2FDTNDATAQYW2
            DNSName: !GetAtt [PublicWebsiteWwwCloudfront, DomainName]

  PoliciesEdgeLambdaRole:
    Type: 'AWS::IAM::Role'
    Condition: IncludeLambdaEdge
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Action: 'sts:AssumeRole'
            Principal:
              Service:
                - lambda.amazonaws.com
                - edgelambda.amazonaws.com
                - replicator.lambda.amazonaws.com
            Effect: Allow
      Policies:
        - PolicyName: EdgePoliciesLambdaPolicy
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Action:
                  - 'xray:PutTraceSegments'
                  - 'xray:PutTelemetryRecords'
                  - 'lambda:GetFunction'
                  - 'lambda:EnableReplication*'
                  - 'lambda:InvokeFunction'
                  - 'logs:CreateLogGroup'
                  - 'logs:CreateLogStream'
                  - 'logs:PutLogEvents'
                Effect: Allow
                Resource: '*'

  PoliciesEdgeLambda: 
    Type: 'AWS::Lambda::Function'
    Condition: IncludeLambdaEdge
    Properties: 
      Handler: 'index.handler'
      Role: 
        Fn::GetAtt: 
          - 'PoliciesEdgeLambdaRole'
          - 'Arn'
      Code: 
        ZipFile: |
          'use strict';
          exports.handler = (event, context, callback) => {
              const response = event.Records[0].cf.response;
              var addHeader = (header, value) => {
                response.headers[header.toLowerCase()] = [{
                  key: header,
                  value: value
                }];
              };
              addHeader('Strict-Transport-Security', 'max-age=31536000; includeSubdomains; preload');
              addHeader('X-Content-Type-Options', 'nosniff');
              addHeader('Content-Security-Policy', "default-src 'self'");
              addHeader('X-Frame-Options', 'DENY');
              addHeader('X-XSS-Protection', '1; mode=block')
              callback(null, response);
          };
      Runtime: 'nodejs6.10'
      Timeout: '25'
      TracingConfig:
        Mode: 'Active'

  PoliciesEdgeLambdaVersion: 
    Type: 'AWS::Lambda::Version'
    Condition: IncludeLambdaEdge
    Properties: 
      FunctionName: 
        Ref: 'PoliciesEdgeLambda'