AutistMask/tests/backgroundStateIsolation.test.js

// What one background handler's state can do to another's while both are in
// flight.
//
// The background used to read and write the module-level `state` singleton in
// src/shared/state.js — one object, shared by every handler in the worker,
// replaced wholesale by any loadState(). Two consequences, both covered here
// and both from https://git.eeqj.de/sneak/AutistMask/issues/324:
//
//   - A transaction attempt captured the chain id at its loadState() and then
//     read the ENDPOINT off the singleton several awaits later. A chain switch
//     committed in that window moved the endpoint under an artifact already
//     verified against the old chain, so it would have gone to the new chain's
//     node — the very thing the verification exists to prevent.
//
//   - backgroundRefresh() handed the singleton's wallets to refreshBalances(),
//     which mutates address objects in place across a multi-second network
//     round trip. Any concurrent handler that loaded state replaced those
//     objects, so the refreshed balances landed on detached ones and the save
//     that followed persisted the pre-refresh values — while still stamping
//     lastBalanceRefresh, suppressing the redo.
//
// Both use the real persistence path over a cloning storage stub. Nothing here
// asserts the absence of a loadState() call; each asserts the OUTCOME, so it
// holds against any implementation that gets the outcome right.

const { Wallet } = require("ethers");
const { networkById } = require("../src/shared/networks");
const { makeStorageStub } = require("./support/storageStub");

const SIGNER_KEY =
    "0x59c6995e998f97a5a0044966f0945389dc9e86dae88c7a8412f4603b6b78690d";
const signer = new Wallet(SIGNER_KEY);
const RECIPIENT = "0x66133E8ea0f5D1d612D2502a968757D1048c214a";

const CONNECTED_ORIGIN = "https://dapp.example";
const CONNECTED_HOSTNAME = "dapp.example";
const EXT_URL = "chrome-extension://autistmask/";

const MAINNET = networkById("mainnet");
const SEPOLIA = networkById("sepolia");

const NONCE = 7;
const REFRESHED_BALANCE = "1.5";

// The transaction the background populates, and the artifact signed from it.
// Its chain is a parameter because the whole subject here is a chain moving
// under work already committed to one.
function populated(chainId) {
    return {
        type: 2,
        chainId,
        nonce: NONCE,
        gasLimit: 100000n,
        maxFeePerGas: 2000000000n,
        maxPriorityFeePerGas: 1000000000n,
        to: RECIPIENT,
        value: 10000000000000000n,
        data: "0x",
    };
}

function storedProfile(networkId) {
    const net = networkById(networkId);
    return {
        hasWallet: true,
        wallets: [
            {
                name: "Wallet 1",
                type: "hd",
                xpub: "xpub-1",
                addresses: [
                    {
                        address: signer.address,
                        balance: "0.0",
                        tokenBalances: [],
                    },
                ],
            },
        ],
        activeAddress: signer.address,
        networkId,
        rpcUrl: net.defaultRpcUrl,
        blockscoutUrl: net.defaultBlockscoutUrl,
        allowedSites: { [signer.address]: [CONNECTED_HOSTNAME] },
        deniedSites: {},
        trackedTokens: [],
        lastBalanceRefresh: 0,
    };
}

async function settle() {
    for (let i = 0; i < 60; i++) await Promise.resolve();
}

function deferred() {
    let resolve;
    const promise = new Promise((res) => {
        resolve = res;
    });
    return { promise, resolve };
}

afterEach(() => {
    delete global.chrome;
});

// The background worker over a cloning storage stub, with the network and the
// clock stubbed out. `opts.refreshBalances` replaces the balance refresh so a
// test can hold one open across another handler's whole turn.
function loadWorker(networkId, opts) {
    const options = opts || {};
    jest.resetModules();

    // Every provider this worker constructs, in order, with the endpoint and
    // the network id it was given. The subject of the first test is which pair
    // reaches the broadcast.
    const providers = [];
    const broadcastTransaction = jest.fn(async () => ({ hash: "0xfeed" }));

    jest.doMock("../src/shared/balances", () => ({
        getProvider: (rpcUrl, networkId2) => {
            const provider = {
                rpcUrl,
                networkId: networkId2,
                broadcastTransaction,
                getNetwork: async () => ({
                    chainId: BigInt(networkById(networkId2).networkVersion),
                }),
                getTransactionCount: async () => NONCE,
                estimateGas: async () => 100000n,
                getFeeData: async () => ({
                    gasPrice: 2000000000n,
                    maxFeePerGas: 2000000000n,
                    maxPriorityFeePerGas: 1000000000n,
                }),
            };
            providers.push(provider);
            return provider;
        },
        refreshBalances:
            options.refreshBalances || jest.fn(async () => undefined),
    }));
    jest.doMock("../src/shared/phishingDomains", () => ({
        isPhishingDomain: () => false,
    }));
    let alarmHandlers = {};
    jest.doMock("../src/shared/alarms", () => ({
        BALANCE_REFRESH_ALARM: "balance",
        BALANCE_REFRESH_PERIOD_MINUTES: 1,
        ensureRecurringAlarms: jest.fn(async () => {}),
        registerAlarmHandlers: jest.fn((handlers) => {
            alarmHandlers = handlers;
        }),
    }));

    const storage = makeStorageStub({ autistmask: storedProfile(networkId) });
    // A hook the tests use to suspend one handler mid-flight, so the other one
    // runs entirely inside its window.
    let getHook = null;
    const realGet = storage.local.get;
    storage.local.get = jest.fn(async (key) => {
        if (getHook) await getHook();
        return realGet(key);
    });

    let messageListener = null;
    // Every popup URL the background opened. The approval id is in it, and
    // that is how the popup learns which approval it is answering.
    const createdUrls = [];
    global.chrome = {
        storage,
        runtime: {
            getURL: (path) => EXT_URL + path,
            onMessage: {
                addListener: (fn) => {
                    messageListener = fn;
                },
            },
            onConnect: { addListener: () => {} },
            lastError: null,
        },
        windows: {
            getLastFocused: (cb) => cb(null),
            create: (createOpts, cb) => {
                createdUrls.push(createOpts.url);
                cb({ id: createdUrls.length });
            },
            remove: (id, cb) => {
                if (cb) cb();
            },
            onRemoved: { addListener: () => {} },
        },
        tabs: {
            query: (queryInfo, cb) => cb([{ id: 1 }]),
            sendMessage: (tabId, message, cb) => {
                if (cb) cb();
            },
        },
        action: { setPopup: () => {} },
    };

    require("../src/background/index");

    function send(msg, sender) {
        let result = null;
        const kept = messageListener(msg, sender, (r) => {
            result = r;
        });
        return { kept, result: () => result };
    }

    function rpc(method, params, origin) {
        return send(
            { type: "AUTISTMASK_RPC", method, params },
            { origin: origin || CONNECTED_ORIGIN },
        );
    }

    return {
        rpc,
        send,
        providers,
        broadcastTransaction,
        persisted: () => storage.read("autistmask"),
        setGetHook: (hook) => {
            getHook = hook;
        },
        fromPopup: { url: EXT_URL + "src/popup/index.html" },
        fireBalanceAlarm: () => alarmHandlers.balance(),
        lastApprovalId: () => {
            const url = createdUrls[createdUrls.length - 1];
            if (!url) return null;
            return new URL(url, EXT_URL).searchParams.get("approval");
        },
    };
}

describe("a chain switch under a transaction already committed to a chain", () => {
    // Item 4 of https://git.eeqj.de/sneak/AutistMask/issues/324.
    //
    // The artifact is verified against the chain read at the top of the
    // attempt. Whatever endpoint it is then broadcast to has to be that same
    // chain's — otherwise the wallet checks a transaction against Sepolia and
    // sends it to a mainnet node. A connected site can switch the chain at any
    // moment, including this one.
    test("the artifact is broadcast to the endpoint of the chain it was verified against", async () => {
        const bg = loadWorker("sepolia");

        // Raise the approval, then find its id from the popup's own fetch.
        bg.rpc("eth_sendTransaction", [
            {
                from: signer.address,
                to: RECIPIENT,
                value: "0x2386f26fc10000",
                data: "0x",
            },
        ]);
        await settle();

        const id = bg.lastApprovalId();
        expect(id).toBeTruthy();

        // The popup signs what it was shown: Sepolia.
        const rawSignedTx = await signer.signTransaction(
            populated(Number(SEPOLIA.networkVersion)),
        );

        // A connected site switches the chain while the attempt is running,
        // and the switch is committed to storage in full before the attempt
        // goes any further.
        //
        // It is fired from inside the attempt's SECOND state read, because
        // that is where the window used to be: the chain id was captured at
        // the first read and the endpoint was taken off the singleton several
        // awaits later, so a switch landing between them moved the endpoint
        // under an artifact already verified against the old chain. An
        // implementation that takes both from one read has no second read for
        // this to fire on, and the switch below runs after the attempt is
        // done instead — which is the point.
        let reads = 0;
        let switched = null;
        const doSwitch = async () => {
            switched = bg.rpc("wallet_switchEthereumChain", [
                { chainId: MAINNET.chainId },
            ]);
            await settle();
        };
        bg.setGetHook(async () => {
            reads++;
            if (reads !== 2) return;
            bg.setGetHook(null);
            await doSwitch();
        });

        const attempt = bg.send(
            {
                type: "AUTISTMASK_TX_RESPONSE",
                id,
                approved: true,
                rawSignedTx,
            },
            { url: bg.fromPopup.url },
        );
        await settle();

        bg.setGetHook(null);
        if (!switched) await doSwitch();
        expect(switched.result()).toEqual({ result: null });
        expect(bg.persisted().networkId).toBe("mainnet");
        await settle();

        // It went out, and it went out to Sepolia's node — the chain the
        // artifact was verified against. Reading the endpoint separately from
        // the chain id put mainnet's here.
        expect(attempt.result()).toEqual({ txHash: "0xfeed" });
        expect(bg.broadcastTransaction).toHaveBeenCalledTimes(1);
        const used = bg.providers[bg.providers.length - 1];
        expect(used.rpcUrl).toBe(SEPOLIA.defaultRpcUrl);
        expect(used.networkId).toBe("sepolia");
    });
});

describe("a balance refresh under another handler's state read", () => {
    // Item 5 of https://git.eeqj.de/sneak/AutistMask/issues/324.
    //
    // The trigger is a same-chain wallet_switchEthereumChain from a connected
    // site: it answers { result: null } and changes nothing, so the ONLY thing
    // it can do to the refresh is what its state read does. On the singleton
    // that read replaced state.wallets, detaching the objects the refresh was
    // mutating.
    test("a chain read arriving mid-refresh does not discard the refresh", async () => {
        const roundTrip = deferred();
        const reachedNetwork = deferred();

        const bg = loadWorker("sepolia", {
            refreshBalances: async (wallets) => {
                reachedNetwork.resolve();
                await roundTrip.promise;
                // In place, on the objects handed in — as balances.js does.
                wallets[0].addresses[0].balance = REFRESHED_BALANCE;
            },
        });

        const refresh = bg.fireBalanceAlarm();
        await reachedNetwork.promise;

        const answered = bg.rpc("wallet_switchEthereumChain", [
            { chainId: SEPOLIA.chainId },
        ]);
        await settle();
        expect(answered.result()).toEqual({ result: null });

        roundTrip.resolve();
        await refresh;

        expect(bg.persisted().wallets[0].addresses[0].balance).toBe(
            REFRESHED_BALANCE,
        );
        expect(bg.persisted().lastBalanceRefresh).toBeGreaterThan(0);
    });

    // The other half of "does not publish a shared object": a wallet added
    // while the refresh was in flight must survive the refresh's own write.
    test("a wallet added mid-refresh survives the refresh's write", async () => {
        const roundTrip = deferred();
        const reachedNetwork = deferred();

        const bg = loadWorker("sepolia", {
            refreshBalances: async (wallets) => {
                reachedNetwork.resolve();
                await roundTrip.promise;
                wallets[0].addresses[0].balance = REFRESHED_BALANCE;
            },
        });

        const refresh = bg.fireBalanceAlarm();
        await reachedNetwork.promise;

        // Another extension page adds a wallet while the round trip is out.
        const during = bg.persisted();
        during.wallets.push({
            name: "Wallet 2",
            type: "hd",
            xpub: "xpub-2",
            addresses: [
                { address: RECIPIENT, balance: "0.0", tokenBalances: [] },
            ],
        });
        global.chrome.storage.write("autistmask", during);

        roundTrip.resolve();
        await refresh;

        const after = bg.persisted();
        expect(after.wallets).toHaveLength(2);
        expect(after.wallets[0].addresses[0].balance).toBe(REFRESHED_BALANCE);
    });
});