AutistMask/tests/stateMerge.test.js

// saveState() used to write the entire state blob every time
// (src/shared/state.js). Every extension page — the toolbar popup, a dApp
// approval window opened by the background, backgroundRefresh() in
// src/background/index.js — holds its own in-memory `state`, loaded once,
// and src/popup/views/helpers.js showView() saves on EVERY navigation. So
// any second page that saved after a first page had written something new
// overwrote it, with no attacker and no unusual input: a whole wallet, name,
// addresses and encrypted secret included, silently gone
// (https://git.eeqj.de/sneak/AutistMask/issues/304).
//
// Both cases below drive the real state.js module through two independent
// module registries sharing one storage backend, the way two real extension
// pages share one chrome.storage.local. The shared stub structured-clones on
// both get and set — a stub that hands back the object it was given aliases
// the caller's own mutation and would make this entire defect class invisible
// (see https://git.eeqj.de/sneak/AutistMask/issues/324).

const { makeStorageStub } = require("./support/storageStub");

// One extension page: a fresh module registry over the shared storage.
// state.js resolves the storage API at require time, so the stub has to be
// installed before the module is loaded, and `state` is a module-level
// singleton, so each page needs its own registry to hold its own copy.
function loadPage(storage) {
    jest.resetModules();
    globalThis.chrome = { storage: { local: storage.local } };
    return {
        state: require("../src/shared/state"),
        helpers: require("../src/popup/views/helpers"),
    };
}

function wallet(name, secret, address) {
    return {
        type: "hd",
        name,
        xpub: "xpub-" + name,
        encryptedSecret: secret,
        nextIndex: 1,
        addresses: [{ address, balance: "0", tokenBalances: [] }],
    };
}

const W1 = wallet(
    "Wallet 1",
    "secret-one",
    "0x66133E8ea0f5D1d612D2502a968757D1048c214a",
);
const W2 = wallet(
    "Wallet 2",
    "secret-two",
    "0xdAC17F958D2ee523a2206206994597C13D831ec7",
);

// Minimal DOM: showView() toggles view elements, clears the flash line and
// creates/removes the debug banner. Nothing here is asserted; it only has to
// answer without throwing, the way the popup's own index.html would.
function makeElement(id) {
    const classes = new Set();
    return {
        id,
        textContent: "",
        style: {},
        classList: {
            add: (...n) => n.forEach((c) => classes.add(c)),
            remove: (...n) => n.forEach((c) => classes.delete(c)),
            toggle: (c, force) => {
                const on = force === undefined ? !classes.has(c) : force;
                if (on) classes.add(c);
                else classes.delete(c);
                return on;
            },
        },
        remove: () => {},
    };
}

function makeDocument() {
    const els = new Map();
    return {
        getElementById(id) {
            if (id === "debug-banner") return null;
            if (!els.has(id)) els.set(id, makeElement(id));
            return els.get(id);
        },
        createElement: () => makeElement("created"),
        body: { prepend: () => {} },
    };
}

afterEach(() => {
    delete globalThis.chrome;
    delete globalThis.document;
});

describe("a save from a page that never saw a wallet another page added", () => {
    // The first DoD case on the issue: add a wallet in one page, then force
    // a save from a second page loaded before that wallet existed. Both
    // wallets must survive.
    test("both wallets are in storage afterwards", async () => {
        const storage = makeStorageStub();
        await storage.set({ autistmask: { wallets: [W1] } });

        // Loaded while storage held only Wallet 1, and never reloads —
        // the approval window in the reproduction, or a second popup that
        // has been open for a while.
        const stale = loadPage(storage);
        await stale.state.loadState();
        expect(stale.state.state.wallets).toHaveLength(1);

        // A second page, loaded after, adds a wallet — the exact sequence
        // src/popup/views/addWallet.js uses.
        const fresh = loadPage(storage);
        await fresh.state.loadState();
        fresh.state.state.wallets.push(W2);
        fresh.state.state.hasWallet = true;
        await fresh.state.saveState();

        expect(
            (await storage.get("autistmask")).autistmask.wallets,
        ).toHaveLength(2);

        // The stale page saves something that has nothing to do with
        // wallets — exactly what showView() does on every navigation, and
        // what backgroundRefresh() does after a balance poll.
        stale.state.state.currentView = "settings";
        await stale.state.saveState();

        const persisted = (await storage.get("autistmask")).autistmask;
        expect(persisted.wallets.map((w) => w.name)).toEqual([
            "Wallet 1",
            "Wallet 2",
        ]);
        expect(persisted.wallets.map((w) => w.encryptedSecret)).toEqual([
            "secret-one",
            "secret-two",
        ]);
    });
});

describe("the approval-window reproduction", () => {
    // approval window open, add a wallet in the popup, confirm the approval
    // — the exact sequence from the issue. The approval window and the
    // popup are the same popup code with a different starting view, so
    // showView() is the real save path in both: src/popup/views/approval.js
    // showTxApproval() calls showView("approve-tx") when the window opens,
    // and a successful confirm calls
    // src/popup/views/txStatus.js showWait() -> startWait(), which calls
    // showView("wait-tx") — the save that clobbered the second wallet in
    // the reproduction on the issue.
    test("the wallet added in the popup survives confirming the approval", async () => {
        globalThis.document = makeDocument();

        const storage = makeStorageStub();
        await storage.set({ autistmask: { wallets: [W1] } });

        // The background opens the approval window on the approve-tx
        // screen; nothing else has happened yet.
        const approvalWindow = loadPage(storage);
        await approvalWindow.state.loadState();
        approvalWindow.helpers.showView("approve-tx");
        // showView() does not await its own saveState(); an extra save
        // joins the same queue and only resolves once that one has too,
        // which is the black-box way to know it landed.
        await approvalWindow.state.saveState();

        // The user adds a wallet in the popup — a separate page, loaded
        // after the approval window.
        const popup = loadPage(storage);
        await popup.state.loadState();
        popup.state.state.wallets.push(W2);
        popup.state.state.hasWallet = true;
        await popup.state.saveState();
        expect(
            (await storage.get("autistmask")).autistmask.wallets,
        ).toHaveLength(2);

        // The user confirms the approval. The approval window navigates
        // approve-tx -> wait-tx, saving again from state it loaded before
        // Wallet 2 ever existed.
        approvalWindow.helpers.showView("wait-tx");
        await approvalWindow.state.saveState();

        const persisted = (await storage.get("autistmask")).autistmask;
        expect(persisted.wallets.map((w) => w.name)).toEqual([
            "Wallet 1",
            "Wallet 2",
        ]);
        expect(persisted.wallets.map((w) => w.encryptedSecret)).toEqual([
            "secret-one",
            "secret-two",
        ]);
    });
});

// backgroundRefresh() (src/background/index.js) loads state, spends seconds
// on network I/O in refreshBalances() (src/shared/balances.js) mutating
// addr.balance/ensName/tokenBalances IN PLACE on the wallets it already
// knew about, then saves. Precondition 2 on the issue: that refresh window
// overlapping a membership change (add or delete) on another page must not
// clobber or resurrect a wallet — a whole-field diff on `wallets` failed
// this, because "background changed a balance" and "another page changed
// membership" collided as the same field.
describe("background refresh racing a wallet added on another page", () => {
    test("the wallet added elsewhere survives background's stale balance save", async () => {
        const storage = makeStorageStub();
        await storage.set({ autistmask: { wallets: [W1] } });

        // "background": loads first, and its save is the one that lands
        // last, modeling the multi-second network round trip in between.
        const background = loadPage(storage);
        await background.state.loadState();
        background.state.state.wallets[0].addresses[0].balance = "1.2345";

        // A second page, loaded after, adds a wallet while background's
        // refresh is still in flight.
        const popup = loadPage(storage);
        await popup.state.loadState();
        popup.state.state.wallets.push(W2);
        popup.state.state.hasWallet = true;
        await popup.state.saveState();
        expect(
            (await storage.get("autistmask")).autistmask.wallets,
        ).toHaveLength(2);

        // background's save lands last, carrying only its balance update.
        await background.state.saveState();

        const persisted = (await storage.get("autistmask")).autistmask;
        expect(persisted.wallets.map((w) => w.name)).toEqual([
            "Wallet 1",
            "Wallet 2",
        ]);
        expect(persisted.wallets.map((w) => w.encryptedSecret)).toEqual([
            "secret-one",
            "secret-two",
        ]);
        // The balance update itself must not be lost either — this is a
        // merge, not deletion-always-wins.
        expect(persisted.wallets[0].addresses[0].balance).toBe("1.2345");
    });
});

describe("background refresh racing a wallet deleted on another page", () => {
    test("the wallet deleted elsewhere stays deleted after background's stale balance save", async () => {
        const storage = makeStorageStub();
        await storage.set({ autistmask: { wallets: [W1, W2] } });

        const background = loadPage(storage);
        await background.state.loadState();
        background.state.state.wallets[0].addresses[0].balance = "1.2345";

        // A second page deletes Wallet 2 while background's refresh is in
        // flight — the same splice deleteWallet.js's removeWalletFromState()
        // does.
        const popup = loadPage(storage);
        await popup.state.loadState();
        popup.state.state.wallets.splice(1, 1);
        popup.state.state.hasWallet = popup.state.state.wallets.length > 0;
        await popup.state.saveState();
        expect(
            (await storage.get("autistmask")).autistmask.wallets,
        ).toHaveLength(1);

        await background.state.saveState();

        const persisted = (await storage.get("autistmask")).autistmask;
        expect(persisted.wallets.map((w) => w.name)).toEqual(["Wallet 1"]);
        expect(persisted.wallets[0].addresses[0].balance).toBe("1.2345");
    });
});

// allowedSites/deniedSites: { [address]: [origin, ...] }. Mutated in place
// from two different contexts — rememberSiteChoice() in
// src/background/index.js pushes a newly approved origin onto
// state.allowedSites[activeAddress], and the Settings "revoke" button
// (forgetOrigin() in src/popup/views/settings.js) filters an origin out of
// state[key][addr] in place, deleting the address key entirely once its list
// is empty — the exact membership-vs-whole-field pattern that made the
// whole-field `wallets` diff unsafe, on a security-relevant field: a stale
// whole-field save here can resurrect a revoked permission or wipe a freshly
// granted one.
const ADDR1 = "0x66133E8ea0f5D1d612D2502a968757D1048c214a";
const ADDR2 = "0xdAC17F958D2ee523a2206206994597C13D831ec7";

function approveSite(pageState, address, origin) {
    if (!pageState.allowedSites[address]) {
        pageState.allowedSites[address] = [];
    }
    if (!pageState.allowedSites[address].includes(origin)) {
        pageState.allowedSites[address].push(origin);
    }
}

function revokeSite(pageState, origin) {
    for (const addr of Object.keys(pageState.allowedSites)) {
        pageState.allowedSites[addr] = pageState.allowedSites[addr].filter(
            (o) => o !== origin,
        );
        if (pageState.allowedSites[addr].length === 0) {
            delete pageState.allowedSites[addr];
        }
    }
}

describe("a dApp approval racing a stale Settings page's later save", () => {
    test("the fresh approval survives Settings revoking an unrelated site", async () => {
        const storage = makeStorageStub();
        await storage.set({
            autistmask: {
                wallets: [W1],
                allowedSites: { [ADDR2]: ["https://other.example"] },
            },
        });

        // Settings loads first, and its save lands last — before either has
        // any idea a dApp approval happened elsewhere in between.
        const settings = loadPage(storage);
        await settings.state.loadState();

        // A dApp approval window, opened later, approves a new site for a
        // different address and saves — the real sequence in
        // rememberSiteChoice(), src/background/index.js.
        const approval = loadPage(storage);
        await approval.state.loadState();
        approveSite(approval.state.state, ADDR1, "https://dapp.example");
        await approval.state.saveState();
        expect(
            (await storage.get("autistmask")).autistmask.allowedSites[ADDR1],
        ).toEqual(["https://dapp.example"]);

        // Settings revokes its own, unrelated site — the real sequence in
        // forgetOrigin(), src/popup/views/settings.js — and saves from state
        // loaded before the dApp approval ever happened.
        revokeSite(settings.state.state, "https://other.example");
        await settings.state.saveState();

        const persisted = (await storage.get("autistmask")).autistmask;
        expect(persisted.allowedSites[ADDR1]).toEqual(["https://dapp.example"]);
        expect(persisted.allowedSites[ADDR2]).toBeUndefined();
    });
});

describe("a revoked site permission against a stale page's later save", () => {
    test("the revocation holds even when the stale page approves something else", async () => {
        const storage = makeStorageStub();
        await storage.set({
            autistmask: {
                wallets: [W1],
                allowedSites: { [ADDR1]: ["https://evil.example"] },
            },
        });

        // A stale page loads while the permission still stands.
        const stale = loadPage(storage);
        await stale.state.loadState();

        // Settings revokes it — forgetOrigin(), src/popup/views/settings.js —
        // from a second page.
        const settings = loadPage(storage);
        await settings.state.loadState();
        revokeSite(settings.state.state, "https://evil.example");
        await settings.state.saveState();
        expect(
            (await storage.get("autistmask")).autistmask.allowedSites[ADDR1],
        ).toBeUndefined();

        // The stale page, unaware of the revoke, approves an unrelated site
        // for a different address and saves — rememberSiteChoice(),
        // src/background/index.js.
        approveSite(stale.state.state, ADDR2, "https://good.example");
        await stale.state.saveState();

        const persisted = (await storage.get("autistmask")).autistmask;
        expect(persisted.allowedSites[ADDR2]).toEqual(["https://good.example"]);
        expect(persisted.allowedSites[ADDR1]).toBeUndefined();
    });
});

// mergeListByIdentity()'s identity function is not guaranteed collision-free
// — walletIdentity() falls back to one shared "addr:" value for any wallet
// with neither an xpub nor a populated first address (a legacy or corrupt
// record). Two such records created independently on two different pages
// must not silently collapse into one, dropping the loser's
// encryptedSecret with no error and no log.
function legacyWallet(name, secret) {
    return {
        type: "legacy",
        name,
        encryptedSecret: secret,
        nextIndex: 0,
        addresses: [],
    };
}

describe("two wallets independently created with a colliding identity", () => {
    test("both survive, encryptedSecret included, instead of one silently replacing the other", async () => {
        const storage = makeStorageStub();
        await storage.set({ autistmask: { wallets: [W1] } });

        // Both pages load before either has created their malformed wallet,
        // so neither has baseline knowledge of the other's.
        const pageA = loadPage(storage);
        await pageA.state.loadState();
        const pageB = loadPage(storage);
        await pageB.state.loadState();

        pageA.state.state.wallets.push(legacyWallet("Legacy A", "secret-a"));
        pageA.state.state.hasWallet = true;
        await pageA.state.saveState();
        expect(
            (await storage.get("autistmask")).autistmask.wallets,
        ).toHaveLength(2);

        pageB.state.state.wallets.push(legacyWallet("Legacy B", "secret-b"));
        pageB.state.state.hasWallet = true;
        await pageB.state.saveState();

        const persisted = (await storage.get("autistmask")).autistmask;
        const secrets = persisted.wallets.map((w) => w.encryptedSecret);
        expect(secrets).toContain("secret-one");
        expect(secrets).toContain("secret-a");
        expect(secrets).toContain("secret-b");
    });
});

// showView() saves on every navigation without waiting, so the user can change
// something while that save is still waiting on storage. The change is followed
// by its own saveState(), which runs after the first save; it must be stored
// (https://git.eeqj.de/sneak/AutistMask/issues/448).
describe("a change made while an earlier save from the same page is running", () => {
    // Runs `change` inside the next call to `op` (the stub's get or set),
    // before that call does its work.
    function runInside(op, change) {
        const real = op.getMockImplementation();
        op.mockImplementationOnce(async (arg) => {
            change();
            return real(arg);
        });
    }

    test("a network switched during the earlier save's read is stored", async () => {
        const storage = makeStorageStub({
            autistmask: { wallets: [W1], networkId: "sepolia" },
        });
        const { state, saveState, loadState } = loadPage(storage).state;
        await loadState();

        let queued;
        runInside(storage.get, () => {
            state.networkId = "mainnet";
            queued = saveState();
        });
        state.theme = "dark";
        await saveState();
        await queued;

        const stored = storage.read("autistmask");
        expect(stored.theme).toBe("dark");
        expect(stored.networkId).toBe("mainnet");
    });

    test("a wallet added during the earlier save's read is stored", async () => {
        const storage = makeStorageStub({ autistmask: { wallets: [W1] } });
        const { state, saveState, loadState } = loadPage(storage).state;
        await loadState();

        let queued;
        runInside(storage.get, () => {
            state.wallets.push(W2);
            queued = saveState();
        });
        await saveState();
        await queued;

        const stored = storage.read("autistmask");
        expect(stored.wallets.map((w) => w.encryptedSecret)).toEqual([
            "secret-one",
            "secret-two",
        ]);
    });

    test("a wallet added during the earlier save's write is stored", async () => {
        const storage = makeStorageStub({ autistmask: { wallets: [W1] } });
        const { state, saveState, loadState } = loadPage(storage).state;
        await loadState();

        let queued;
        runInside(storage.set, () => {
            state.wallets.push(W2);
            queued = saveState();
        });
        await saveState();
        await queued;

        const stored = storage.read("autistmask");
        expect(stored.wallets.map((w) => w.encryptedSecret)).toEqual([
            "secret-one",
            "secret-two",
        ]);
    });
});