Lambda@Edge source code for A/B testing

abtesting-lambda-oreq.js

'use strict';

const aws = require('aws-sdk');

const COOKIE_KEY = 'ABTesting-SegmentID';

const s3 = new aws.S3({ region: 'eu-west-1' });
const s3Params = {
    Bucket: 'abtesting-ttblog-map',
    Key: 'map.json',
};
const SEGMENT_MAP_TTL = 300000; // TTL of 5 minutes in ms

const fetchSegmentMapFromS3 = async() => {
    const response = await s3.getObject(s3Params).promise();
    return JSON.parse(response.Body.toString('utf-8'));
}

// Cache the segment map across Lambda invocations
let _segmentMap;
let _lastFetchedSegmentMap = 0;
const fetchSegmentMap = async() => {
    if (!_segmentMap || (Date.now() - _lastFetchedSegmentMap) > SEGMENT_MAP_TTL) {
        _segmentMap = await fetchSegmentMapFromS3();
        _lastFetchedSegmentMap = Date.now();
    }

    return _segmentMap;
}

const getCookie = (headers, cookieKey) => {
    if (headers.cookie) {
        for (let cookieHeader of headers.cookie) {
            const cookies = cookieHeader.value.split(';');
            for (let cookie of cookies) {
                const [key, val] = cookie.split('=');
                if (key === cookieKey) {
                    return val;
                }
            }
        }
    }
    return null;
}

const getSegmentId = (headers) => {
    const cookieVal = getCookie(headers, COOKIE_KEY);
    if (cookieVal != null) {
        return parseInt(cookieVal);
    }
    console.error('No segmentId cookie found!');
    return -1;
}

const getSegmentGroup = async(segmentId) => {
    const segmentMap = await fetchSegmentMap();
    for (let group of segmentMap.segmentGroups) {
        if (segmentId >= group.range[0] && segmentId <= group.range[1]) {
            return group;
        }
    }
    console.error(`No origin for segment id ${segmentId}. Check the segment map.`);
}

exports.handler = async(event, context, callback) => {
    const request = event.Records[0].cf.request;
    const headers = request.headers;

    const segmentId = getSegmentId(headers);

    if (segmentId != -1) {
        const segmentGroup = await getSegmentGroup(segmentId);
        headers['host'] = [{ key: 'host', value: segmentGroup.host }];
        request.origin = segmentGroup.origin;
    }

    callback(null, request);
};

abtesting-lambda-ores.js

'use strict';

const COOKIE_KEY = 'ABTesting-SegmentID';

const getCookie = (headers, cookieKey) => {
    if (headers.cookie) {
        for (let cookieHeader of headers.cookie) {
            const cookies = cookieHeader.value.split(';');
            for (let cookie of cookies) {
                const [key, val] = cookie.split('=');
                if (key === cookieKey) {
                    return val;
                }
            }
        }
    }
    return null;
}

// caveat: this will overwrite the set-cookie header from origin
const setCookie = function(response, cookie) {
    console.log(`Setting cookie ${cookie}`);
    response.headers['set-cookie'] = [{ key: "Set-Cookie", value: cookie }];
}

exports.handler = (event, context, callback) => {
    const request = event.Records[0].cf.request;
    const headers = request.headers;
    const response = event.Records[0].cf.response;

    const cookieVal = getCookie(headers, COOKIE_KEY);
    if (cookieVal != null) {
        setCookie(response, `${COOKIE_KEY}=${cookieVal}`);
        callback(null, response);
        return;
    }

    console.log('no segmentid cookie');
    callback(null, response);
}

abtesting-lambda-vreq.js

'use strict';

const COOKIE_KEY = 'ABTesting-SegmentID';
const NUM_SEGMENTS = 10;

const getRandomId = () => Math.floor(Math.random() * (1 + NUM_SEGMENTS));

const getCookie = (headers, cookieKey) => {
    if (headers.cookie) {
        for (let cookieHeader of headers.cookie) {
            const cookies = cookieHeader.value.split(';');
            for (let cookie of cookies) {
                const [key, val] = cookie.split('=');
                if (key === cookieKey) {
                    return val;
                }
            }
        }
    }
    return null;
}

exports.handler = (event, context, callback) => {
    const request = event.Records[0].cf.request;
    const headers = request.headers;

    const headerCookie = getCookie(headers, COOKIE_KEY);
    if (headerCookie != null) {
        callback(null, request);
        return;
    }

    const segmentId = getRandomId();
    console.log(`segmentId: ${segmentId}`)

    const cookie = `${COOKIE_KEY}=${segmentId}`
    console.log(`setting cookie: ${cookie}`);
    headers.cookie = headers.cookie || [];
    headers.cookie.push({ key: 'Cookie', value: cookie });

    callback(null, request);
};

map.json

{
    "segmentGroups": [
        {
            "range": [
                0,
                6
            ],
            "host": "abtesting-ttblog-a.s3.amazonaws.com",
            "origin": {
                "s3": {
                    "authMethod": "none",
                    "domainName": "abtesting-ttblog-a.s3.amazonaws.com",
                    "path": "",
                    "region": "eu-west-1"
                }
            }
        },
        {
            "range": [
                7,
                9
            ],
            "host": "abtesting-ttblog-b.s3.amazonaws.com",
            "origin": {
                "s3": {
                    "authMethod": "none",
                    "domainName": "abtesting-ttblog-b.s3.amazonaws.com",
                    "path": "",
                    "region": "eu-west-1"
                }
            }
        }
    ]
}