2u lambda edge implementation

experiment.ts

import optimizely from '@optimizely/optimizely-sdk'
import optimizelyLogging from '@optimizely/optimizely-sdk/lib/plugins/logger'
import optimizelyEnums from '@optimizely/optimizely-sdk/lib/utils/enums'
import cookie from 'cookie'
import rp from 'request-promise'
import uuidv4 from 'uuid/v4'
import {
    Callback,
    Handler,
    CloudFrontRequest,
    CloudFrontRequestEvent,
    CloudFrontResponse,
    CloudFrontResponseEvent,
    Context,
} from 'aws-lambda'
import { Experiment, HandlerImplementation } from './types'

const delimeter = '&'
const ev = ':'
const optimizelyLogLevel = optimizelyEnums.LOG_LEVEL.INFO

export function makeHandler(implementation: HandlerImplementation): Handler {
    const handler = async (
        event: CloudFrontRequestEvent | CloudFrontResponseEvent,
        context: Context,
        callback: Callback
    ) => {
        let response
        const cf = event.Records[0].cf
        if ('response' in cf) {
            response = cf.response
        }
        try {
            let { result, client } = await implementation(cf.request, response)
            if (client) {
                const closed = await client.close()
                console.debug('Optimizely flush result', JSON.stringify(closed))
            }
            return result
        } catch (e) {
            console.error(e)
            return response || cf.request
        }
    }
    return handler
}

export function serializeExperiments(experiments: Experiment[]): string {
    return experiments
        .filter(({ variation }: Experiment) => !!variation)
        .map(({ name, variation }: Experiment) => name + ev + variation)
        .join(delimeter)
}

export const fetchUserId = (request: CloudFrontRequest): string | null =>
    cookies(request).experiment_user

export const generateUserId = (): string => uuidv4()

export function storedExperiments(request: CloudFrontRequest): Experiment[] {
    const experiments = (cookies(request).experiments || '')
        .split(delimeter)
        .filter((str: string) => str && str.length > 0 && str !== 'null')
        .map((str: string) => {
            const [name, variation] = str.split(ev)
            return { name, variation }
        })
    console.log('Experiments from cookie', JSON.stringify(experiments))
    return experiments
}

export async function experimentClient(
    datafileUrl: string
): Promise<optimizely.Client | null> {
    const datafile = await rp({ uri: datafileUrl, json: true })
    console.log('Successfully fetched datafile')
    const logger = optimizelyLogging.createLogger({
        logToConsole: true,
        logLevel: optimizelyLogLevel,
    })
    return optimizely.createInstance({ datafile, logger })
}

export function requestMatches(
    request: CloudFrontRequest,
    path: string | undefined,
    method: string | undefined
): boolean {
    return !!(
        request.uri.match(path || '.*') && request.method.match(method || '.*')
    )
}

function cookies(request: CloudFrontRequest): any {
    const cookies = request?.headers?.cookie || []
    let allCookiesObj = {}
    for (const cookieHeader of cookies) {
        allCookiesObj = Object.assign(
            allCookiesObj,
            cookie.parse(cookieHeader.value)
        )
    }
    return allCookiesObj
}

types.ts

import { CloudFrontRequest, CloudFrontResponse } from 'aws-lambda'
import optimizely from '@optimizely/optimizely-sdk'

export interface Experiment {
    name: string
    variation: string
}

export interface HandlerImplementationResponse {
    result: CloudFrontRequest | CloudFrontResponse
    client?: optimizely.Client | null
}

export interface HandlerImplementation {
    (request: CloudFrontRequest, response?: CloudFrontResponse): Promise<
        HandlerImplementationResponse
    >
}

viewerRequest.ts

import { Experiment } from './types'
import {
    fetchUserId,
    generateUserId,
    serializeExperiments,
    storedExperiments,
    experimentClient,
    makeHandler,
    requestMatches,
} from './experiment'
import {
    CloudFrontRequestHandler,
    CloudFrontRequest,
    Context,
} from 'aws-lambda'

const handler: CloudFrontRequestHandler = makeHandler(
    async (request: CloudFrontRequest) => {
        let client

        const experiments = storedExperiments(request)
        const toFetch = newExperiments(experiments)
        const toActivate = (process.env.CURRENT_EXPERIMENTS || '').split(',')

        if (process.env.DATAFILE) {
            client = await experimentClient(process.env.DATAFILE)
            const userId = fetchUserId(request) || generateUserId()
            for (const name of toFetch) {
                const variation = client?.getVariation(name, userId)
                if (variation) {
                    experiments.push({ name, variation })
                }
            }
            if (
                requestMatches(
                    request,
                    process.env.ACTIVATE_PATH,
                    process.env.ACTIVATE_METHOD
                )
            ) {
                for (const name of toActivate) {
                    console.log('Activating experiment', name)
                    client?.activate(name, userId)
                }
            }
            storeHeaders(request, userId, experiments)
        }

        if (process.env.ORIGIN_CHOICE_EXPERIMENT && process.env.ORIGIN_PREFIX) {
            const variation = experiments.find(
                ({ name }) => name === process.env.ORIGIN_CHOICE_EXPERIMENT
            )?.variation
            if (variation && variation !== 'control') {
                console.log('Changing origin', variation)
                request.uri = `${process.env.ORIGIN_PREFIX}/${variation}`
            }
        }
        return { result: request, client }
    }
)

const newExperiments = (existingExperiments: Experiment[]): string[] =>
    (process.env.ALL_EXPERIMENTS || '')
        .split(',')
        .filter(
            name =>
                !existingExperiments.some(
                    (experiment: Experiment) => experiment.name === name
                )
        )

const storeHeaders = (
    request: CloudFrontRequest,
    userId: string,
    experiments: Experiment[]
): void => {
    console.log('Putting experiments in headers', experiments)
    request.headers.experiments = [
        {
            key: 'experiments',
            value: serializeExperiments(experiments),
        },
    ]
    request.headers.experiment_user = [
        {
            key: 'experiment_user',
            value: userId,
        },
    ]
}

export { handler } // for typescript & tests
exports.handler = handler // for lambda execution

viewerResponse.ts

import {
    fetchUserId,
    experimentClient,
    makeHandler,
    requestMatches,
} from './experiment'
import {
    CloudFrontRequest,
    CloudFrontResponse,
    CloudFrontResponseHandler,
    Context,
} from 'aws-lambda'

const handler: CloudFrontResponseHandler = makeHandler(
    async (request: CloudFrontRequest, response?: CloudFrontResponse) => {
        let client
        if (!response) {
            return { result: request, client }
        }

        let { experimentsValue, userId } = parseExperimentHeaders(request)
        const cookies = []
        if (experimentsValue) {
            cookies.push(`experiments=${experimentsValue}`)
        }
        if (userId) {
            cookies.push(`experiment_user=${userId}`)
        }
        setCookie(response, cookies)

        if (
            process.env.DATAFILE &&
            process.env.EVENT_NAME &&
            requestMatches(
                request,
                process.env.EVENT_PATH,
                process.env.EVENT_METHOD
            ) &&
            success(response)
        ) {
            client = await experimentClient(process.env.DATAFILE)
            client?.track(
                process.env.EVENT_NAME,
                userId || fetchUserId(request)
            )
        }
        return { result: response, client }
    }
)

const parseExperimentHeaders = (request: CloudFrontRequest): any => {
    const experimentsValue = request.headers?.experiments
        ? request.headers.experiments[0]?.value
        : null
    const userId = request.headers?.experiment_user
        ? request.headers.experiment_user[0]?.value
        : null
    return { experimentsValue, userId }
}

const cookieMaxAge = 60 * 60 * 24 * 30
const setCookie = (response: CloudFrontResponse, values: string[]): void => {
    console.log('Setting cookie', values)
    if (!response.headers['set-cookie']) {
        response.headers['set-cookie'] = []
    }

    response.headers['set-cookie'].push({
        key: 'Set-Cookie',
        value: [
            values.join(';'),
            'Secure',
            `Max-Age=${cookieMaxAge}`,
            `SameSite=Strict`,
        ].join(';'),
    })
}

const success = (response: CloudFrontResponse): boolean => {
    const responseStatus = parseInt(response.status)

    return !!(responseStatus >= 200 && responseStatus < 300)
}

export { handler } // for typescript & tests
exports.handler = handler // for lambda execution

webpack.config.js

const webpack = require('webpack')
const path = require('path')
const FileManagerPlugin = require('filemanager-webpack-plugin')
const { readSync } = require('node-yaml')

const config = (source, destination, env = {}) => ({
    entry: `./src/${source}.ts`,
    mode: 'production',
    target: 'node',
    module: {
        rules: [
            {
                test: /\.ts$/,
                use: 'ts-loader',
                exclude: /node_modules/,
            },
        ],
    },
    resolve: {
        extensions: ['.ts', '.js'],
    },
    output: {
        filename: `${destination}/index.js`,
        path: path.resolve(__dirname, 'dist'),
        library: 'index',
        libraryTarget: 'umd', // required from lambda execution
    },
    plugins: [
        new FileManagerPlugin({
            onEnd: {
                archive: [
                    {
                        source: `./dist/${destination}`,
                        destination: `./deploy/${destination}.zip`,
                    },
                ],
            },
        }),
        new webpack.EnvironmentPlugin(env),
    ],
})

module.exports = [
    config('experiment', 'layer'),
    ...readSync('./outputs').outputs.map(({ source, destination, env }) =>
        config(source, destination, env || {})
    ),
]