Using Lambda@Edge with CloudFront

In my previous post, I showed how to create a static website in AWS using CloudFormation to create a CloudFront distribution. In this post, I show how to add custom headers to the site using Lambda@Edge.

I expected this to be a very quick post: AWS provides a facility called Lambda@Edge that allows code (written in JavaScript/NodeJS) to be executed by CloudFront, and to modify what is returned by CloudFront. Unfortunately however, it seems that Lambda@Edge isn’t quite up to the usual standard of AWS: it has a number of severe limitations, and the documentation is very poor. This took me several hours of experimenting before I was able to put together the code shown below.

The main problems I found were as follows:

  • Almost all the documentation is in terms of creating edge functions via the console. However using the console for anything other than read-only actions is an anti-pattern: I want everything to be automated, and I want my infrastructure described as code, using CloudFormation if possible.
  • Edge functions can only be created from Lambda versions - you cannot create an edge function from $LATEST.
  • Making this worse is the fact that CloudFormation support for Lambda versions is hopeless - if you look at the documentation, you’ll see that managing changes to versions simply isn’t supported.#
  • Lambda functions can get “stuck” and be unable to be deleted. Relevant StackOverflow question and link to the docs.
  • Updates don’t always seem to “stick” - it seems that sometimes updates have to be re-applied before the changes start appearing. (UPDATE: It seems updated lambda functions only apply if the content at the origin changes - TTL timeouts do not automatically apply lambda changes. So you’ll need to update your source files - I added a timestamp to my relevant files.)
  • Finally, updating CloudFront distributions takes a very long time (tens of minutes), which makes iterative development harder.

In short, I can’t really recommend using Lambda@Edge right now - it all feel too bleeding edge.

Putting all that behind me however, here’s my solution.

Managing Lambda Versions

Since Lambda@Edge only works with versions of lambdas, I needed a way of being able to create and “update” those versions. CloudFormation doesn’t really support updating or managing versions, so I had to fudge it a little.

In the end, I found a workaround: put conditions around the lambda version, and update the CloudFormation stack twice: once with the condition set to false, to delete the version, and again with the condition set to true, to recreate it.

So, my CloudFormation template now includes a condition and a parameter to control it:

Parameters:
  RootDomainName:
    Type: String
  IncludeLambdaEdge:
    Type: String
    AllowedValues: ['true', 'false']
Conditions:
  IncludeLambdaEdge:
    !Equals ['true', !Ref IncludeLambdaEdge]

Debugging

There are various bits of my template that set up debugging facilities for lambdas. I could probably get rid of them now I have it all debugged, but I’ll show those bits here for completeness:

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

I modified my www CloudFront distribution to turn on logging as well:

  PublicWebsiteWwwCloudfront:
    Type: AWS::CloudFront::Distribution
    DependsOn:
      - PublicWebsiteWwwBucket
      - InternalTracingBucket
    Properties:
      DistributionConfig:
        Logging:
          Bucket: !GetAtt [InternalTracingBucket, DomainName]
          IncludeCookies: false
          Prefix: 'www'

(Existing properties snipped for clarity).

Lambda@Edge

Next up is the actual lambda itself:

  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'

You can specify the code either inline or in a S3 bucket - my code is so small I figured inline was OK.

Again, I could remove the TracingConfig now I have it all debugged.

The role for this lambda is as follows:

  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: '*'

Lastly, the lambda version looks like this:

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

It will be destroyed and recreated in the two-pass stack update mechanism.

Putting it all together

At this stage, the lambda is just sitting there: not associated with CloudFront. I can tie the two together as follows:

  PublicWebsiteWwwCloudfront:
    Type: AWS::CloudFront::Distribution
    DependsOn:
      - PublicWebsiteWwwBucket
      - InternalTracingBucket
    Properties:
      DistributionConfig:
        DefaultCacheBehavior:
          LambdaFunctionAssociations:
            - !If
                - IncludeLambdaEdge
                - EventType: 'origin-response'
                  LambdaFunctionARN: !Join
                    - ':'
                    - - !GetAtt [PoliciesEdgeLambda, Arn]
                      - !GetAtt [PoliciesEdgeLambdaVersion, Version]
                - !Ref 'AWS::NoValue'

(Again, existing properties omitted for clarity).

If you prefer, I have put the whole lot in a Gist.

Written on February 5, 2018.