WebAssembly for AWS Lambda@Edge functions

Introduction

They say premature optimisation is the root of all evil. Well we’re going throw that belief out the window today. As we explore building serverless edge functions with WebAssembly (WASM) and assess if there is any performance advantage over using Javascript.

The Motivation

I think it is important to outline the problem we are trying to solve with Lambda@Edge and the accompanying AWS infrastructure. The problem is related to the authentication of users through Auth0 identity provider (IdP). Auth0 has a feature called actions which I like to describe as type of post-login middleware. Essentially it is a flow that excecutes a sets of scripts before a user is issued an access token. Actions aren’t limited to manipulating just user authentication they can also be used for M2M (Machine to Machine) authentication and various other flows. For user authentication specifically, with actions, we can do things among many other processes like:



Predating actions Auth0 offered rules which uses an older Nodejs runtime version and is limited to user authentication flow.

One feature that is missing from actions which was present in rules was the existence of a datastore. In rules we could store data in a global key-value object called rules config and access it in the rules implementation.


This is ultimately the problem that needs to be solved. Creation of a datastore that only Auth0 actions should be able to read from.


Cloudfront which is AWS’s content delivery network (CDN) offering has a feature called edge functions which allows for code to be run as close to the user and manipulate the request and response without having to worry about servers and the other accompanying infrastructure. Lambda@Edge is one of the two ways to write and manage edge functions.


With the use of Lambda@Edge for authorisation, Cloudfront and s3 as an origin which will hold our data in the form of JSON files. We can be confident that when Auth0 actions has to access the s3 datastore through Cloudfront it will be pretty quick in order to not adversely affect the user experience when logging in.

M2M Token ExchangeAuth0 Login PageAuth0 Action User Auth FlowAuth0 Authorization ServerUserGET /auth0-dataCloudfront DistributionLambda@EdgeAuthorizerAuth0 Actions Datastores3 Bucket

About Lambda@Edge

AWS Lambda@Edge functions are used to run code in response to CloudFront events.

These events include:

OriginViewer RequestViewer ResponseOrigin RequestOrigin ResponseUserCloudfront

Lambda@Edge is able to manipulate only one of these events when setting up path behaviours in Cloudfront. The functions are executed in AWS locations closer to the end user which can reduce latency and improve performance. This is achieved by the Lambda being automatically replicated in each AWS region. Therefore AWS Lambda must be given the relevant IAM roles and permissions to allow for this.

There are some rules/restrictions on Lambda@Edge compared to standard Lambda functions in AWS.

The Plan

Javascript Approach

About Rust wasm-pack, WASM Bindgen and WASM in general

“Hybrid” Approach

Initial usage of the WASM module included some Javascript defined logic in particular the response when verifying the JWT and getting said JWT from the request header.

The Rust WebAssembly Module for “Hybrid” Lambda
#[derive(Debug)]
pub enum Status {
    Ok,
    Unauthorized,
    Forbidden,
}

impl fmt::Display for Status {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Status::Ok => write!(f, "200"),
            Status::Unauthorized => write!(f, "401"),
            Status::Forbidden => write!(f, "403"),
        }
    }
}

const JWT_SECRET: &'static str = env!("JWT_SECRET");

#[wasm_bindgen]
pub fn verify(token: &str) -> String {
    
    let decoded_token = decode::<Claims>(&token, &DecodingKey::from_secret(JWT_SECRET.as_ref()), &Validation::new(Algorithm::HS256));
    let valid_permissions = vec!["view:data"];
    match decoded_token {
        Ok(token_data) => {
            console_log!("{:?}", token_data.claims);
            if token_data.claims.permissions.iter().all(|permission| valid_permissions.contains(&permission.as_str())) {
                return Status::Ok.to_string();
            } else {
                return Status::Forbidden.to_string();
            }
        },
        Err(e) => {
            console_log!("{}",e);
            return Status::Unauthorized.to_string();
        }
    }
}
The “Hybrid” Lambda
import { verify } from 'rust-edge-lambda';

export const handler = async (event, _context, callback) => {
    let authToken = '';
    const request = event.Records[0].cf.request;

    if (request.headers['authorization'])
        authToken = request.headers['authorization'][0].value.replace("Bearer ", "")

    try {
        const verifyTokenResponse = verify(authToken);
        switch (verifyTokenResponse) {
            case '401':
                response.status = '401';
                response.statusDescription = 'Unauthorized';
                response.body = JSON.stringify({ error: 'Invalid token' });
                break;
            case '403':
                response.status = '403';
                response.statusDescription = 'Forbidden';
                break;
            case '200':
                callback(null, request);
            default:
                break;
        }
    }

    catch(error) {
        console.error(error);
        response.status = '401';
        response.statusDescription = 'Unauthorized';
        response.body = JSON.stringify({ error: 'Invalid token' });
    }

    callback(null, response);
}

Rust Approach

Performance Comparison

Conclusion