AutistMask/tests/coldWorkerChainSwitch.test.js

// What a chain switch does to a worker that has not loaded state yet.
//
// The MV3 service worker is terminated when idle and revived by the next
// message, and nothing loads state at module scope. The chain-switch handler
// reaches onChainSwitch(), which mutates the module-level `state` singleton
// and then persists EVERY field of it, so a handler that runs before a load
// writes DEFAULT_STATE over the user's stored profile — every wallet, every
// site approval, every tracked token and the custom endpoint
// (https://git.eeqj.de/sneak/AutistMask/issues/316). The same singleton is
// what currentNetwork() answers from, so the same-chain early return also
// compares against the wrong network.
//
// This file therefore uses the REAL state module and never calls loadState()
// itself: the handler has to do it. tests/chainSwitchGate.test.js mocks the
// state module wholesale and tests/networkEndpoints.test.js always loads
// first, so neither can see this.
//
// A site's switch happens only once the user approves it on a prompt
// (https://git.eeqj.de/sneak/AutistMask/issues/408), so every switch here is
// approved the way the popup approves one.

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

const ADDRESS = "0x66133E8ea0f5D1d612D2502a968757D1048c214a";

const CONNECTED_ORIGIN = "https://dapp.example";

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

// The user's own node, and a wallet whose loss is the whole point.
const CUSTOM_RPC = "http://127.0.0.1:8545";
const CUSTOM_BLOCKSCOUT = "http://127.0.0.1:4000/api/v2";
const TOKEN = "0x6B175474E89094C44Da98b954EedeAC495271d0F";

function walletFixture() {
    return [
        {
            name: "Wallet 1",
            type: "hd",
            addresses: [{ address: ADDRESS, balance: "0", tokenBalances: [] }],
        },
    ];
}

// A profile as an installed extension holds it, on `networkId`.
function storedProfile(networkId) {
    return {
        hasWallet: true,
        wallets: walletFixture(),
        activeAddress: ADDRESS,
        networkId,
        rpcUrl: CUSTOM_RPC,
        blockscoutUrl: CUSTOM_BLOCKSCOUT,
        allowedSites: { [ADDRESS]: [CONNECTED_ORIGIN] },
        deniedSites: {},
        trackedTokens: [{ address: TOKEN, symbol: "DAI", decimals: 18 }],
        theme: "dark",
    };
}

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

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

// Load the background worker with the real chain-switch and persistence
// modules behind it, over a storage stub that actually keeps what is written —
// a wipe is only observable against storage that remembers — and that clones
// in both directions, as the real API does. It used to alias, so the record
// the worker held and the "stored" one were a single object; see
// tests/support/storageStub.js.
function loadColdWorker(networkId) {
    jest.resetModules();

    jest.doMock("../src/shared/balances", () => ({
        getProvider: () => ({}),
        refreshBalances: jest.fn(async () => {}),
    }));
    jest.doMock("../src/shared/phishingDomains", () => ({
        isPhishingDomain: () => false,
    }));
    jest.doMock("../src/shared/alarms", () => ({
        BALANCE_REFRESH_ALARM: "balance",
        BALANCE_REFRESH_PERIOD_MINUTES: 1,
        ensureRecurringAlarms: jest.fn(async () => {}),
        registerAlarmHandlers: jest.fn(),
    }));

    const storage = makeStorageStub({ autistmask: storedProfile(networkId) });

    let messageListener = null;
    let connectListener = null;
    // The URL of every approval window opened; the approval id is in it.
    const opened = [];
    const toTabs = [];

    global.chrome = {
        storage,
        runtime: {
            getURL: (path) => "chrome-extension://autistmask/" + path,
            onMessage: {
                addListener: (fn) => {
                    messageListener = fn;
                },
            },
            onConnect: {
                addListener: (fn) => {
                    connectListener = fn;
                },
            },
            lastError: null,
        },
        windows: {
            getLastFocused: (cb) => cb(null),
            create: (options, cb) => {
                opened.push(options.url);
                cb({ id: opened.length });
            },
            remove: (id, cb) => {
                if (cb) cb();
            },
            onRemoved: { addListener: () => {} },
        },
        tabs: {
            query: (queryInfo, cb) => cb([{ id: 1 }]),
            sendMessage: (tabId, message, cb) => {
                toTabs.push(message);
                if (cb) cb();
            },
        },
        action: { setPopup: () => {} },
    };

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

    // The user approves the prompt the request opened, as the popup does: a
    // decision on the port named for the approval, from the extension's own
    // page.
    function approvePrompt() {
        const id = new URL(opened[opened.length - 1]).searchParams.get(
            "approval",
        );
        let onDecision = null;
        const port = {
            name: "approval:" + id,
            sender: { url: "chrome-extension://autistmask/src/popup/" },
            onMessage: {
                addListener: (fn) => {
                    onDecision = fn;
                },
            },
            onDisconnect: { addListener: () => {} },
        };
        connectListener(port);
        onDecision(
            { type: "AUTISTMASK_APPROVAL_DECISION", approved: true },
            port,
        );
    }

    // A connected site asks for `chainId`, and the user approves it.
    async function switchChain(chainId) {
        let result = null;
        messageListener(
            {
                type: "AUTISTMASK_RPC",
                method: "wallet_switchEthereumChain",
                params: [{ chainId }],
            },
            { origin: CONNECTED_ORIGIN },
            (r) => {
                result = r;
            },
        );
        await settle();
        approvePrompt();
        await settle();
        return result;
    }

    return {
        switchChain,
        persisted: () => storage.read("autistmask"),
        chainChangedEvents: () =>
            toTabs.filter((m) => m.eventName === "chainChanged"),
    };
}

describe("a chain switch on a worker that never loaded state", () => {
    test("keeps the wallets, approvals, tokens and custom endpoint", async () => {
        const bg = loadColdWorker("mainnet");

        const result = await bg.switchChain(SEPOLIA.chainId);
        expect(result).toEqual({ result: null });

        const after = bg.persisted();
        // The switch itself happened.
        expect(after.networkId).toBe("sepolia");
        expect(after.rpcUrl).toBe(SEPOLIA.defaultRpcUrl);

        // And it took nothing else with it. Without the load these come back
        // as [], {}, [] and "system" from DEFAULT_STATE — every wallet in the
        // extension gone, encrypted secrets included.
        expect(after.wallets).toEqual(walletFixture());
        expect(after.hasWallet).toBe(true);
        expect(after.activeAddress).toBe(ADDRESS);
        expect(after.allowedSites).toEqual({ [ADDRESS]: [CONNECTED_ORIGIN] });
        expect(after.trackedTokens).toEqual([
            { address: TOKEN, symbol: "DAI", decimals: 18 },
        ]);
        expect(after.theme).toBe("dark");

        // The user's mainnet endpoint is remembered rather than replaced by
        // the public default, so switching back returns it.
        expect(after.networkEndpoints.mainnet).toEqual({
            rpcUrl: CUSTOM_RPC,
            blockscoutUrl: CUSTOM_BLOCKSCOUT,
        });

        await bg.switchChain(MAINNET.chainId);
        expect(bg.persisted().rpcUrl).toBe(CUSTOM_RPC);
        expect(bg.persisted().blockscoutUrl).toBe(CUSTOM_BLOCKSCOUT);
        expect(bg.persisted().wallets).toEqual(walletFixture());
    });

    test("compares the requested chain against the stored one, not the default", async () => {
        // Stored on Sepolia, asked for mainnet. Reading the unloaded
        // singleton makes this look like the chain already active, so the
        // page is told the switch succeeded while the wallet stays on the
        // testnet it was on.
        const bg = loadColdWorker("sepolia");

        const result = await bg.switchChain(MAINNET.chainId);
        expect(result).toEqual({ result: null });

        expect(bg.persisted().networkId).toBe("mainnet");
        expect(bg.chainChangedEvents()).toEqual([
            {
                type: "AUTISTMASK_EVENT",
                eventName: "chainChanged",
                data: MAINNET.chainId,
            },
        ]);
    });
});