connectionPoisoning.js

// Get the global browser map.
const getGlobalBrowserMap = () => {
    const v = globalThis.__RAINBOWS_INTERNALS_CONN_STACK_IDENTIFIERS__;
    if (!v) {
        globalThis.__RAINBOWS_INTERNALS_CONN_STACK_IDENTIFIERS__ = new Map();
        return globalThis.__RAINBOWS_INTERNALS_CONN_STACK_IDENTIFIERS__;
    }
    return v;
};

// Set the connection used for the async function specified then garbage collects it. This will be the connection used
// unless the function goes and spawns another async function.
async function setStackConnectionBrowser(connection, fn) {
    const stackToken = Math.random().toString(36).substring(2);
    const ma = getGlobalBrowserMap();
    ma.set(stackToken, connection);
    const stackDisp = "$_rainbows_conn_" + stackToken + "_$";
    let res, rej;
    const p = new Promise((resolve, reject) => {
        res = resolve;
        rej = reject;
    });
    setTimeout(async () => {
        try {
            const m = {
                [stackDisp]: fn,
            };
            res(await m[stackDisp]());
        } catch (e) {
            rej(e);
        } finally {
            ma.delete(stackToken);
        }
    }, 0);
    return p;
}

const RE_STACK_TOKEN = /\$_rainbows_conn_([a-z0-9]+)_\$/;

// Gets the nearest stack connection. Returns null if there is no connection.
async function getStackConnectionBrowser() {
    let stack = null;
    try {
        throw new Error();
    } catch (e) {
        stack = e.stack;
    }
    if (!stack) {
        // Unable to use stack poisoning here.
        return null;
    }
    const match = RE_STACK_TOKEN.exec(stack);
    if (!match) {
        // No stack token found.
        return null;
    }
    const token = match[1];
    return getGlobalBrowserMap().get(token) || null;
}

// Basically just doing this to brainfuck any bundlers. We are smarter in this rare case.
let asyncHooksName = "node:async_";
asyncHooksName += "hooks";
const asyncHooksOrNull = import(/* @vite-ignore */ asyncHooksName).catch(() => null);

// Get the async local storage.
async function getAsyncLocalStorageNode() {
    const v = globalThis.__RAINBOWS_INTERNALS_CONN_STACK_IDENTIFIERS__;
    if (!v) {
        globalThis.__RAINBOWS_INTERNALS_CONN_STACK_IDENTIFIERS__ = new (await asyncHooksOrNull).AsyncLocalStorage();
        return globalThis.__RAINBOWS_INTERNALS_CONN_STACK_IDENTIFIERS__;
    }
    return v;
}

// Defines the Node implementation of getting hooks using async_hooks.
async function setStackConnectionNode(connection, fn) {
    /** @type {import("node:async_hooks").AsyncLocalStorage} */
    const localStorage = await getAsyncLocalStorageNode();
    return localStorage.run(connection, fn);
}

// Defines the Node implementation of setting hooks using async_hooks.
async function getStackConnectionNode() {
    /** @type {import("node:async_hooks").AsyncLocalStorage} */
    const localStorage = await getAsyncLocalStorageNode();
    return localStorage.getStore() || null;
}

// Defines the indecisive proxy functions that are exported by default.
export let setStackConnection = async (connection, fn) => {
    const v = await asyncHooksOrNull;
    if (v) {
        setStackConnection = setStackConnectionNode;
        return setStackConnectionNode(connection, fn);
    }
    setStackConnection = setStackConnectionBrowser;
    return setStackConnectionBrowser(connection, fn);
};
export let getStackConnection = async () => {
    const v = await asyncHooksOrNull;
    if (v) {
        getStackConnection = getStackConnectionNode;
        return getStackConnectionNode();
    }
    getStackConnection = getStackConnectionBrowser;
    return getStackConnectionBrowser();
};