Authentication at the edge with Lambda@Edge and Cognito

Michael Wittig – Updated 24 Nov 2022

For many years, we used a hosting partner for serving the Rapid Docker on AWS Video Course. When someone bought the video course, we created a user account with our partner. The hosting partner provided a website to watch the videos and a login form. For many reasons, we migrated the video course to a new home -our home- cloudonaut.io. In this article, I outline the architecture, the user flow, and go through code snippets to implement user authentication at the edge with Lambda@Edge functions and a Cognito user pool.

Authentication at the edge

This post is not about video hosting on AWS. We cover video hosting on AWS in this post.

Architecture

From a high-level perspective, the following components are used:

Cognito user pool stores users, hosts login UI, and issues JWT tokens.
Lambda@Edge functions check if the request contains a cookie with a valid JWT token and implement a tiny backend to implement the OAuth 2.0 Authorization Code Flow.
CloudFront distribution delivers the content to the end-users and triggers Lambda@Edge functions.
S3 bucket stores the content served by CloudFront.

The following figure shows the high-level architecture.

Architecture: Authentication at the edge with Lambda@Edge and Cognito

Let’s continue with the user perspective.

User flow

The user starts the journey by visiting a protected web page. The following figure shows what happens next.

User flow: Authentication at the edge with Lambda@Edge and Cognito

The following steps are executed:

The user visits a protected page (e.g., https://cloudonaut.io/rapid-docker-on-aws/video-course/ch00-01.html) hosted on CloudFront.
CloudFront invokes the viewer request Lambda@Edge function.
The function inspects the cookie header, extracts the cookie named token, and verifies the value to check if it is a valid JWT token issued by Cognito.
Let’s assume no cookie is present. The Lambda@edge function generates an HTTP 302 response to redirect to the Cognito hosted UI.
The user’s browser follows the redirect and loads the Cognito hosted UI with a login screen.
Once the user enters a valid username and password, Cognito returns an HTTP 302 response to redirect to the cloudonaut.io backend (https://cloudonaut.io/api/cognito/login/).
The user’s browser follows the redirect and loads the backend. In this case, we implement the backend with an origin request Lambda@Edge function.
The function exchanges the received authorization code (query parameter code) with an access token (in JWT format). An HTTP 302 response is generated with the Set-Cookie header.
The user’s browser stores the token cookie and follows the redirect with the Cookie header containing the access token. The user can access the protected page.

For a better understanding, let’s dive into Lambda@Edge.

How Lambda@Edge works

When CloudFront receives a request, it can invoke the so-called viewer request Lambda@Edge function. You can inspect the HTTP headers at this point and generate HTTP responses (such as a 302 redirect). The viewer request function execution is limited to 5 seconds and 128 MB of memory. The function’s code and libraries must fit into a 1 MB zip file.

Warning
The snippets hardcode the region to eu-west-1. If your region is different, replace all occurrences of eu-west-1!

The following snippet shows the implementation to check if a cookie token is part of the HTTP header and if the value is a signed JWT issued from Cognito.

// required libraries
const cookie = require('cookie'); // I use version 0.5.0
const jose = require('jose'); // I use version 4.8.3

// TODO fill config values with outputs from CloudFormation shown later in the article
const config = {
  cognitoUserPoolId: '',
  cognitoClientId: '',
  cognitoDomainName: ''
};

// download jwks.json from https://cognito-idp.eu-west-1.amazonaws.com/${config.cognitoUserPoolId}/.well-known/jwks.json
// according to AWS support, the keys are not rotated so you can do this once and include the file to avoid timeout issues
const jwks = jose.createLocalJWKSet(require('./jwks.json')); 

async function verifyToken(cf) {
  if (cf.request.headers.cookie) {
    const cookies = cookie.parse(cf.request.headers.cookie[0].value);
    try {
      const { payload } = await jose.jwtVerify(cookies.token, jwks, {
        issuer: `https://cognito-idp.eu-west-1.amazonaws.com/${config.cognitoUserPoolId}`
      });
      if (payload.client_id === config.cognitoClientId) {
        return true;
      }
    } catch(err) {
      console.log(`token error: ${err.name} ${err.message}`);
    }
  }
  return false;
}

exports.handler = async function(event) {
  const cf = event.Records[0].cf;
  // check if path is protected and requires the user to be logged in
  if (
    cf.request.uri.startsWith('/rapid-docker-on-aws/video-course/') ||
    cf.request.uri.startsWith('/media/cloudonaut/rapid-docker-on-aws/')
  ) {
    const valid = await verifyToken(cf, 'rapid-docker-on-aws-video-course');
    if (valid === true) {
      return cf.request;
    } else {
      return {
        status: '302',
        statusDescription: 'Found',
        headers: {
          location: [{ // instructs browser to redirect after receiving the response
            key: 'Location',
            value: `https://${config.cognitoDomainName}.auth.eu-west-1.amazoncognito.com/login?client_id=${config.cognitoClientId}&response_type=code&scope=email+openid&redirect_uri=https%3A%2F%2Fcloudonaut.io%2Fapi%2Fcognito%2Flogin%2F`,
          }]
        }
      };
    }
  }
  // do nothing: CloudFront continues as usual
  return cf.request;
};

When the viewer request function does not generate a response, the request is passed to the CloudFront origin (e.g., an S3 bucket).

When CloudFront can not serve the request from the cache, the origin request Lambda@Edge function is invoked just before the request to the origin is made. The main difference is that this function can take up to 30 seconds, use as much memory as Lambda offers, and the uploaded code archive can be up to 50 MB.

In our case, we use a Lambda@Edge function to implement a tiny backend. The backend exchanges an authorization code with an access token and ensures that the response is not cachable by setting the Cache-Control header to no-cache.

Feel free to implement this part using an API Gateway or similar.

// required libraries
const querystring = require('querystring'); // included in Node.js
const cookie = require('cookie'); // I use version 0.5.0
const axios = require('axios'); // I use version 0.27.2

// TODO fill config values with outputs from CloudFormation shown later in the article
const config = {
  cognitoClientId: '',
  cognitoClientSecret: '',
  cognitoDomainName: ''
};

exports.handler = async function(event) {
  const cf = event.Records[0].cf;
  if (cf.request.uri.startsWith('/api/cognito/login/')) {
    const {code} = querystring.parse(cf.request.querystring);
    const res = await axios({
      method: 'POST',
      headers: {
        'content-type': 'application/x-www-form-urlencoded',
        authorization: 'Basic ' + Buffer.from(config.cognitoClientId + ':' + config.cognitoClientSecret).toString('base64')
      },
      data: querystring.stringify({
        grant_type: 'authorization_code',
        redirect_uri: 'https://cloudonaut.io/api/cognito/login/',
        code
      }),
      url: `https://${config.cognitoDomainName}.auth.eu-west-1.amazoncognito.com/oauth2/token`,
    });
    if (res.status === 200) {
      const setCookieValue = cookie.serialize('token', res.data.access_token, {
        maxAge: res.data.expires_in,
        path: '/',
        secure: true
      });
      return {
        status: '302',
        headers: {
          location: [{ // instructs browser to redirect after receiving the response
          	key: 'Location',
          	value: '/rapid-docker-on-aws/video-course/ch00-01.html'
          }],
          'set-cookie': [{ // instructs browser to store a cookie
          	key: 'Set-Cookie',
          	value: setCookieValue
          }],
          'cache-control': [{ // ensures that CloudFront does not cache the response
          	key: 'Cache-Control',
          	value: 'no-cache'
          }]
        }
      };
    } else {
      throw new Error('unexpected status code: ' + res.status);
    }
  }
  // do nothing: CloudFront continues as usual
  return cf.request;
};

Cognito Infrastructure

The following CloudFormation template describes the Cognito user pool, client, and domain infrastructure needed.

---
AWSTemplateFormatVersion: '2010-09-09'
Resources:
  UserPool:
    Type: 'AWS::Cognito::UserPool'
    Properties:
      AccountRecoverySetting:
        RecoveryMechanisms:
        - Name: verified_email
          Priority: 1
      AdminCreateUserConfig:
        AllowAdminCreateUserOnly: true
      AliasAttributes:
      - preferred_username
      AutoVerifiedAttributes:
      - email
      EnabledMfas:
      - SOFTWARE_TOKEN_MFA
      MfaConfiguration: OPTIONAL
      UserPoolName: !Ref 'AWS::StackName'
  UserPoolDomain:
    Type: 'AWS::Cognito::UserPoolDomain'
    Properties:
      Domain: 'cloudonaut-io'
      UserPoolId: !Ref UserPool
  ClientWebsite:
    Type: 'AWS::Cognito::UserPoolClient'
    Properties:
      AccessTokenValidity: 1
      AllowedOAuthFlows:
      - code
      AllowedOAuthFlowsUserPoolClient: true
      AllowedOAuthScopes:
      - phone
      - email
      - openid
      - profile
      CallbackURLs:
      - 'https://cloudonaut.io/api/cognito/login/'
      ClientName: website
      DefaultRedirectURI: 'https://cloudonaut.io/api/cognito/login/'
      ExplicitAuthFlows:
      - ALLOW_USER_SRP_AUTH
      - ALLOW_REFRESH_TOKEN_AUTH # always on
      GenerateSecret: true
      IdTokenValidity: 1
      LogoutURLs:
      - 'https://cloudonaut.io/api/cognito/logout/'
      PreventUserExistenceErrors: ENABLED
      RefreshTokenValidity: 30
      SupportedIdentityProviders:
      - COGNITO
      TokenValidityUnits:
        AccessToken: days
        IdToken: days
        RefreshToken: days
      UserPoolId: !Ref UserPool
Outputs:
  CognitoUserPoolId:
    Value: !Ref UserPool
  CognitoClientId:
    Value: !Ref ClientWebsite
  CognitoDomainName:
  	Value: !Ref UserPoolDomain

To get the client secret, we use the following bash snippet in our deployment pipeline:

COGNITO_STACK_NAME="" # TODO fill with your stack name
USER_POOL_ID="$(aws cloudformation describe-stacks --stack-name $COGNITO_STACK_NAME --query "Stacks[0].Outputs[?OutputKey=='UserPoolId'].OutputValue" --output text)"
CLIENT_ID="$(aws cloudformation describe-stacks --stack-name $COGNITO_STACK_NAME --query "Stacks[0].Outputs[?OutputKey=='Value'].OutputValue" --output text)"
CLIENT_SECRET="$(aws cognito-idp describe-user-pool-client --user-pool-id "${USER_POOL_ID}" --client-id "${CLIENT_ID}" --query UserPoolClient.ClientSecret --output text)"

Lambda@Edge infrastructure

I shared the Lambda@edge code with you already. The missing piece is how to deploy the functions using CloudFormation. To avoid using a JavaScript bundler but still include only the needed libraries, I created two folders (viewer-request-src and origin-request-src) to store the Lambda@Edge function code (lambda.js) together with a package.json.

---
AWSTemplateFormatVersion: '2010-09-09'
Description: 'Static Website: Custom image optimization and routing'
Parameters:
  LogsRetentionInDays:
    Description: 'Specifies the number of days you want to retain log events in the specified log group.'
    Type: Number
    Default: 14
    AllowedValues: [1, 3, 5, 7, 14, 30, 60, 90, 120, 150, 180, 365, 400, 545, 731, 1827, 3653]
Resources:
  ViewerRequestRole:
    Type: 'AWS::IAM::Role'
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
        - Effect: Allow
          Principal:
            Service:
            - 'lambda.amazonaws.com'
            - 'edgelambda.amazonaws.com'
          Action: 'sts:AssumeRole'
  ViewerRequestLambdaPolicy:
    Type: 'AWS::IAM::Policy'
    Properties:
      PolicyDocument:
        Statement:
        - Effect: Allow
          Action:
          - 'logs:CreateLogStream'
          - 'logs:PutLogEvents'
          Resource: !GetAtt 'ViewerRequestLogGroup.Arn'
      PolicyName: lambda
      Roles:
      - !Ref ViewerRequestRole
  ViewerRequestLambdaEdgePolicy:
    Type: 'AWS::IAM::Policy'
    Properties:
      PolicyDocument:
        Statement:
        - Effect: Allow
          Action: 'logs:CreateLogGroup'
          Resource: !Sub 'arn:${AWS::Partition}:logs:*:${AWS::AccountId}:log-group:/aws/lambda/us-east-1.${ViewerRequestFunction}:log-stream:'
        - Effect: Allow
          Action:
          - 'logs:CreateLogStream'
          - 'logs:PutLogEvents'
          Resource: !Sub 'arn:${AWS::Partition}:logs:*:${AWS::AccountId}:log-group:/aws/lambda/us-east-1.${ViewerRequestFunction}:log-stream:*'
      PolicyName: 'lambda-edge'
      Roles:
      - !Ref ViewerRequestRole
  ViewerRequestFunction:
    Type: 'AWS::Lambda::Function'
    Properties:
      Code: './viewer-request-src/' # If you change the code, rename the logical id OriginRequestVersionVX to trigger a new version creation!
      Handler: 'lambda.handler'
      MemorySize: 128
      Role: !GetAtt 'ViewerRequestRole.Arn'
      Runtime: 'nodejs16.x'
      Timeout: 5
  ViewerRequestLogGroup:
    Type: 'AWS::Logs::LogGroup'
    Properties:
      LogGroupName: !Sub '/aws/lambda/${ViewerRequestFunction}'
      RetentionInDays: !Ref LogsRetentionInDays
  ViewerRequestVersionV1:
    Type: 'AWS::Lambda::Version'
    Properties:
      FunctionName: !Ref ViewerRequestFunction
  OriginRequestRole:
    Type: 'AWS::IAM::Role'
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
        - Effect: Allow
          Principal:
            Service:
            - 'lambda.amazonaws.com'
            - 'edgelambda.amazonaws.com'
          Action: 'sts:AssumeRole'
  OriginRequestLambdaPolicy:
    Type: 'AWS::IAM::Policy'
    Properties:
      PolicyDocument:
        Statement:
        - Effect: Allow
          Action:
          - 'logs:CreateLogStream'
          - 'logs:PutLogEvents'
          Resource: !GetAtt 'OriginRequestLogGroup.Arn'
      PolicyName: lambda
      Roles:
      - !Ref OriginRequestRole
  OriginRequestLambdaEdgePolicy:
    Type: 'AWS::IAM::Policy'
    Properties:
      PolicyDocument:
        Statement:
        - Effect: Allow
          Action: 'logs:CreateLogGroup'
          Resource: !Sub 'arn:${AWS::Partition}:logs:*:${AWS::AccountId}:log-group:/aws/lambda/us-east-1.${OriginRequestFunction}:log-stream:'
        - Effect: Allow
          Action:
          - 'logs:CreateLogStream'
          - 'logs:PutLogEvents'
          Resource: !Sub 'arn:${AWS::Partition}:logs:*:${AWS::AccountId}:log-group:/aws/lambda/us-east-1.${OriginRequestFunction}:log-stream:*'
      PolicyName: 'lambda-edge'
      Roles:
      - !Ref OriginRequestRole
  OriginRequestFunction:
    Type: 'AWS::Lambda::Function'
    Properties:
      Code: './origin-request-src/' # If you change the code, rename the logical id OriginRequestVersionVX to trigger a new version creation!
      Handler: 'lambda.handler'
      MemorySize: 1536
      Role: !GetAtt 'OriginRequestRole.Arn'
      Runtime: 'nodejs16.x'
      Timeout: 30
  OriginRequestLogGroup:
    Type: 'AWS::Logs::LogGroup'
    Properties:
      LogGroupName: !Sub '/aws/lambda/${OriginRequestFunction}'
      RetentionInDays: !Ref LogsRetentionInDays
  OriginRequestVersionV1:
    Type: 'AWS::Lambda::Version'
    Properties:
      FunctionName: !Ref OriginRequestFunction
Outputs:
  ViewerRequestLambdaEdgeFunctionVersionARN:
    Description: 'Version ARN of Lambda@Edge viewer request function.'
    Value: !Ref ViewerRequestVersionV1
  OriginRequestLambdaEdgeFunctionVersionARN:
    Description: 'Version ARN of Lambda@Edge origin request function.'
    Value: !Ref OriginRequestVersionV1

To deploy the stack in us-east-1 (region required my Lambda@Edge), run:

aws --region us-east-1 cloudformation package --s3-bucket YOUR_S3_ARTIFACT_BUCKET_NAME --template-file YOUR_TEMPLATE_FILE_NAME.yaml --output-template-file output.yaml
aws --region us-east-1 cloudformation deploy --template-file output.yaml --stack-name YOUR_STACK_NAME --capabilities CAPABILITY_IAM
rm output.yaml

The CloudFront infrastructure can be deployed with our Free Templates for AWS CloudFormation. Set the parameters ViewerRequestLambdaEdgeFunctionVersionARN and OriginRequestLambdaEdgeFunctionVersionARN to the values from the stack you deployed before containing the Lambda@Edge functions.

One modification to the static-website.yaml template is needed. Inside the DefaultCacheBehavior, set:

ForwardedValues:
  Cookies:
    Forward: whitelist
    WhitelistedNames: [token]
  QueryString: true
  QueryStringCacheKeys: [code]

That’s it.

Summary

You can add authentication to any website served by CloudFront by using Lambda@Edge. You can set up a Cognito user pool if you want to use your own identity provider. The described flow works with any other identity provider as long as you receive a JWT access token.

PS: You can even add simple authorization using Cognito user groups. If you add a Cognito user to a group, the group name will show up in the cognito:groups claim in the JWT access token.

Michael Wittig

I’ve been building on AWS since 2012 together with my brother Andreas. We are sharing our insights into all things AWS on cloudonaut and have written the book AWS in Action. Besides that, we’re currently working on bucketAV, HyperEnv for GitHub Actions, and marbot.

Here are the contact options for feedback and questions.