Cognito Secret Storage — implements token storage in a web worker to prevent inadvertent exposure

cognitoSecretStorage.ts

/* eslint-disable no-restricted-globals */
interface SecretCache {
    refreshToken?: string;
}

export interface SecureStorageMessage {
    method?: string,
    key?: string;
    value?: string;
    clear?: boolean;
}

export type SecureStorageEvent = MessageEvent<SecureStorageMessage>;

/**
 * In-memory storage
 */
const secretCache: SecretCache = {};

/**
 * Web worker OnMessage Hook
 *
 * Send a message to the web worker with the 'key' and an optional 'value'
 *
 * If value is supplied, it will set the value. The system will then read the value
 * from secret storage and post it back.
 *
 * If clear is sent with a key, that key is cleared. If sent without a key, all
 * keys are cleared
 *
 */
self.onmessage = async (message: MessageEvent<SecureStorageMessage>): Promise<null> => {
    const {
        method,
        key = null,
        value = null,
    } = message.data;
    if (method === "clear") {
        Object.keys(secretCache).forEach((cacheKey) => {
            delete secretCache[cacheKey];
        });
    } else if (method === "set" && key) {
        secretCache[key] = value;
    } else if (method === "delete" && key) {
        delete secretCache[key];
    }
    if (typeof key === "string") {
        postMessage({
            method,
            key,
            value: secretCache[key] ?? null,
        } as SecureStorageMessage);
    }
    return null;
};

CognitoSecureStorage.ts

import {KeyValueStorageInterface, CookieStorage} from "aws-amplify/utils";
import {getLogger} from "../../util";
import {IAppInstance} from "../types";
import {SecureStorageEvent, SecureStorageMessage} from "../workers/cognitoSecureStorage";

const logger = getLogger("CognitoSecureStorage");

type SecureStorageCallback = (event: SecureStorageEvent) => void;

/**
 * Provides a more secure way to store Cognito tokens
 *
 * - Most data is stored in CookieStorage, which is not saved when the browser closes.
 * - The refresh token is stored in a web worker in memory
 *
 * @see https://auth0.com/blog/secure-browser-storage-the-facts/
 */
export class CognitoSecureStorage implements KeyValueStorageInterface {
    /**
     * Non-secure storage uses cookies
     *
     * @private
     */
    private cookieStorage: CookieStorage;

    /**
     * Secure storage uses a web worker with in-memory storage
     *
     * @private
     */
    private readonly worker: Worker;

    /**
     * Callbacks to coordinate receiving data back from the web worker
     *
     * @private
     */
    private listeners = {
        set: {} as Record<string, SecureStorageCallback>,
        get: {} as Record<string, SecureStorageCallback>,
        remove: {} as Record<string, SecureStorageCallback>,
    };

    /**
     * The list of secure keys to store in the web worker
     *
     * @private
     */
    private readonly secureKeys = ["refreshToken"];

    /**
     * Constructor
     */
    constructor() {
        logger.log("construct");
        this.cookieStorage = new CookieStorage({
            expires: null, // for session only
        });
        if (typeof Worker === "undefined") {
            logger.warn("worker not supported");
            this.worker = null;
        } else {
            this.worker = new Worker(new URL("../../app/workers/cognitoSecureStorage.ts", import.meta.url));
            this.worker.onmessage = (event: SecureStorageEvent) => {
                const listener = this.listeners[event.data.method][event.data.key] ?? null;
                if (listener) {
                    logger.log(`Calling listener for ${event.data.method} ${event.data.key}`);
                    listener(event);
                    delete this.listeners[event.data.method][event.data.key];
                } else {
                    logger.error(`Missing listener for ${event.data.method} ${event.data.key} callback`);
                }
            };
        }
    }

    /**
     * Returns true if the supplied key should use the web worker
     *
     * @param key
     */
    useWebWorkerStorage(key: string): boolean {
        if (!this.worker) {
            return false;
        }
        let useWebWorker = false;
        this.secureKeys.forEach((secureKey) => {
            if (key.endsWith(secureKey)) {
                useWebWorker = true;
            }
        });
        return useWebWorker;
    }

    /**
     * Set Item storage interface
     *
     * @param key
     * @param value
     */
    setItem(key: string, value: string): Promise<void> {
        logger.log(`set item ${key} to ${value}`);
        return new Promise((resolve, reject) => {
            if (this.useWebWorkerStorage(key)) {
                logger.log(`set item via worker ${key} to ${value}`);
                this.listeners.set[key] = (event: SecureStorageEvent) => resolve();
                this.worker.postMessage({method: "set", key, value});
                logger.log(`stored item via worker ${key} to ${value}`);
            } else {
                logger.log(`stored item ${key} to ${value}`);
                resolve(this.cookieStorage.setItem(key, value));
            }
        });
    }

    /**
     * Get Item storage interface
     *
     * @param key
     */
    getItem(key: string): Promise<string | null> {
        logger.log(`get item ${key}`);
        return new Promise((resolve, reject) => {
            if (this.useWebWorkerStorage(key)) {
                logger.log(`get item via worker ${key}`);
                this.listeners.get[key] = (event: SecureStorageEvent): void => resolve(event.data.value);
                this.worker.postMessage({method: "get", key});
            } else {
                logger.log(`got item ${key}`);
                resolve(this.cookieStorage.getItem(key));
            }
        });
    }

    /**
     * Remove Item storage interface
     *
     * @param key
     */
    removeItem(key: string): Promise<void> {
        return new Promise((resolve, reject) => {
            logger.log(`remove item ${key}`);
            if (this.useWebWorkerStorage(key)) {
                this.listeners.remove[key] = (event: SecureStorageEvent): void => resolve();
                this.worker.postMessage({method: "remove", key});
            } else {
                resolve(this.cookieStorage.removeItem(key));
            }
        });
    }

    /**
     * Clear storage interface
     */
    async clear(): Promise<void> {
        return new Promise((resolve, reject) => {
            logger.log("clear");
            if (this.worker) {
                this.worker.postMessage({method: "clear"});
            }
            resolve(this.cookieStorage.clear());
        });
    }
}

index.ts

/**
 * An example set-up of Cognito
 */

// Set up Cognito - this needs to be set up as early as possible to
// support the OAuth2 signin flow, which uses redirects.
Amplify.configure({
    Auth: {
        Cognito: {
            ...config.publicRuntimeConfig.cognito,
        },
    },
});

// Enable web worker storage via feature flag
if (config.publicRuntimeConfig.cognito.useWebWorkerStorage) {
    logger.log("Secure auth storage activated");
    const secureStorage = new CognitoSecureStorage();
    cognitoUserPoolsTokenProvider.setKeyValueStorage(secureStorage);
}

This gist provides an working example of implementing a web worker to store Cognito secrets, as per Auth0's https://auth0.com/blog/secure-browser-storage-the-facts/ page.

It's almost perfect, but sometimes has a race condition where it fails to return the key requested at login.