theseus/bundled-addons/aegis/lib/wallet.js

// Wallet state machine on top of an electrum client and a WalletKeys tree:
// address discovery (gap limit), balance, history with per-tx deltas, UTXO
// set and send construction. Knows nothing about UI or IPC.
//
// 0.7.0: CashTokens read + coin-selection guard. Every UTXO fetched from
// listunspent is enriched with its scriptPubKey and passed through
// cashtokens.decodePrefixedScript. Token UTXOs are tagged { token: {…} }
// and pooled into state.tokenBalances (category → aggregate); they are
// deliberately EXCLUDED from plain-BCH coin selection so no token UTXO
// gets accidentally spent (and its category burned) on a routine send.
const cashtokens = require("./cashtokens.js");

module.exports = function makeWallet({ client, keys, tx, cashaddr, sha256, storage, log = () => {}, onChange = () => {} }) {
  const GAP = 20;
  const HISTORY_LIMIT = 25;
  const sats = (bch) => Math.round(Number(bch) * 1e8);

  const state = {
    used: new Set(),            // "branch/index" with history
    watched: new Map(),         // scripthash -> entry
    height: 0,
    balance: { confirmed: 0, unconfirmed: 0 },
    utxos: [],                  // { txid, vout, value, height, entry, token? }
    tokenBalances: {},          // { <categoryHex>: { fungible: bigint, nfts: [...], utxoIds: [...] } }
    history: [],                // newest first
    receiveIndex: 0,
    scanning: false,
    error: null,
  };
  // Verbose transactions are public chain data; caching them on disk saves a
  // round of fetches on every launch.
  // Cached transactions are slimmed on the way in, so a schema change to
  // that slim shape has to invalidate them. v2 adds vout.tokenData; entries
  // written by v1 carry no token information at all and an absent field is
  // indistinguishable from "no token", so they are dropped once rather than
  // trusted. Only the pre-1.5 fallback path reads this for classification,
  // but a warm v1 cache there would silently report a token wallet as empty.
  // v3 because v2 was unbounded. getTx() kept every transaction it ever
  // fetched, with full vin/vout arrays, and a busy chipnet test wallet grew
  // this to 7.7 MB. That matters enormously, because the host's addon store
  // is ONE JSON file per add-on that is read, parsed, stringified and written
  // whole and SYNCHRONOUSLY on the Electron main thread — the thread that
  // drives the entire browser. A 7.7 MB cache turned every storage access
  // into a multi-hundred-millisecond freeze of all of Theseus. Bumping the
  // version also discards existing oversized caches on first load.
  // v4: v2 was unbounded and v3 capped the PERSISTED set at 400, which is
  // exactly the wrong number here — one consolidation in this wallet has 400
  // inputs, so a single history row needed 401 entries and the cap thrashed.
  // The persisted set is now just the displayed window; parents live in RAM.
  const TX_CACHE_VERSION = 4;
  // Guard on RAM, not on correctness. Parents are re-fetchable; this only
  // stops a long session from growing without limit.
  const TX_MEM_MAX = 20000;
  const TX_FETCH_CONCURRENCY = 12;
  let txCache = storage.get("txCache", {}) || {};
  if (storage.get("txCacheVersion", 1) !== TX_CACHE_VERSION) {
    txCache = {};
    storage.set("txCache", txCache);
    storage.set("txCacheVersion", TX_CACHE_VERSION);
  }
  // Only write when something actually changed. loadHistory() used to persist
  // the cache on every refresh, so an idle wallet rewrote the whole file on
  // every poll for nothing.
  let txDirty = false;
  // Process-lifetime, never persisted. Seeded from the persisted window so a
  // restart does not refetch what is on screen.
  const memTx = new Map();
  for (const [id, t] of Object.entries(txCache)) memTx.set(id, t);
  function rememberTx(id, slim) {
    if (memTx.size >= TX_MEM_MAX) {
      // Drop the oldest insertion; Map preserves insertion order.
      const oldest = memTx.keys().next();
      if (!oldest.done) memTx.delete(oldest.value);
    }
    memTx.set(id, slim);
  }
  // Fetch many transactions with bounded parallelism. Already-known ids cost
  // nothing, so passing the same list twice is cheap.
  async function getTxMany(ids) {
    const want = [...new Set(ids)].filter((id) => {
      const c = memTx.get(id);
      return !(c && c.confirmations > 0);
    });
    for (let i = 0; i < want.length; i += TX_FETCH_CONCURRENCY) {
      const batch = want.slice(i, i + TX_FETCH_CONCURRENCY);
      await Promise.all(batch.map((id) => getTx(id).catch(() => null)));
    }
  }
  let refreshTimer = null;
  let subscribedHeaders = false;

  function key(e) { return e.branch + "/" + e.index; }
  function watch(e) { if (!state.watched.has(e.scripthash)) state.watched.set(e.scripthash, e); }

  async function historyOf(e) {
    const h = await client.call("blockchain.scripthash.get_history", [e.scripthash]);
    return Array.isArray(h) ? h : [];
  }

  // Walk both branches until GAP consecutive unused indexes, always covering
  // the user's chosen receive cursor so its lookahead stays subscribed.
  async function scan() {
    const cursor = Number(storage.get("receiveCursor", 0)) || 0;
    for (const branch of [0, 1]) {
      let gap = 0, i = 0;
      const minIndex = branch === 0 ? cursor + 1 : 0;
      while (gap < GAP || i < minIndex + GAP) {
        const batch = [];
        for (let k = 0; k < 10; k++) batch.push(keys.entry(branch, i + k));
        const results = await Promise.all(batch.map(historyOf));
        for (let k = 0; k < batch.length; k++) {
          const e = batch[k]; watch(e);
          if (results[k].length) { state.used.add(key(e)); gap = 0; } else gap++;
          i++;
          if (gap >= GAP && i >= minIndex + GAP) break;
        }
      }
    }
    // Current receive address: first unused at or after the cursor.
    let r = cursor;
    while (state.used.has("0/" + r)) r++;
    state.receiveIndex = r;
    watch(keys.entry(0, r));
  }

  async function subscribeAll() {
    if (!subscribedHeaders) {
      subscribedHeaders = true;
      const tip = await client.subscribe("blockchain.headers.subscribe", []);
      if (tip && tip.height) state.height = tip.height;
    }
    await Promise.all([...state.watched.values()].map((e) =>
      client.subscribe("blockchain.scripthash.subscribe", [e.scripthash]).catch(() => {})));
  }

  async function loadUtxos() {
    const lists = await Promise.all([...state.watched.values()].map(async (e) => {
      const u = await client.call("blockchain.scripthash.listunspent", [e.scripthash]);
      return (Array.isArray(u) ? u : []).map((x) => ({
        txid: x.tx_hash, vout: x.tx_pos, value: x.value, height: x.height, entry: e,
        tokenData: x.token_data || null,
      }));
    }));
    const utxos = lists.flat();
    // Classify each UTXO as bare BCH or CashToken.
    //
    // Two routes. When the server negotiated protocol >= 1.5 it reports
    // `token_data` on listunspent itself, and — this is the part that
    // matters — a UTXO WITHOUT token_data at that protocol is definitively
    // not a token UTXO. So the whole set is classified from the one
    // listunspent call, with zero further round-trips.
    //
    // Below 1.5 the server says nothing, so we fall back to fetching each
    // UTXO's parent transaction and decoding the token prefix off its
    // scriptPubKey. That is one request per UTXO: correct, cached to disk,
    // and completely impractical on a faucet-fed chipnet address — a real
    // one here holds 28,289 UTXOs, so a first scan meant ~28k requests and
    // read as a hung wallet rather than as work in progress.
    //
    // Failures on the fallback path are tolerated: an unclassifiable UTXO is
    // treated as bare BCH, which is the conservative choice — the coin
    // selector may spend it as plain value, but it will never be pulled
    // into a token send.
    const tokenBalances = {};
    // Electrum's token shape -> the shape cashtokens.decodePrefixedScript
    // returns, so everything downstream is identical whichever route found
    // it. The two speak different dialects and must be reconciled here or an
    // identical UTXO would describe itself differently depending on which
    // server answered: the decoder yields a NUMERIC capability (0/1/2) with
    // the labels immutable/mutable/minting, while Electrum sends a STRING
    // and calls 0 "none". The decoder's vocabulary wins — it is the one
    // already established here and in the CHIP.
    const CAP_CODE = { none: 0, immutable: 0, mutable: 1, minting: 2 };
    const CAP_LABEL = ["immutable", "mutable", "minting"];
    const tokenFromElectrum = (td) => {
      if (!td || !td.category) return null;
      let amount = 0n;
      try { amount = BigInt(td.amount || 0); } catch (_e) { amount = 0n; }
      const nft = td.nft || null;
      const code = nft ? (CAP_CODE[String(nft.capability || "none").toLowerCase()] ?? 0) : 0;
      return {
        categoryHex: String(td.category),
        hasAmount: amount > 0n,
        amount,
        hasNft: !!nft,
        commitmentHex: nft ? String(nft.commitment || "") : null,
        capability: nft ? code : 0,
        capabilityLabel: nft ? CAP_LABEL[code] : null,
      };
    };
    const addToken = (u, token) => {
      u.token = token;
      const cat = token.categoryHex;
      if (!tokenBalances[cat]) tokenBalances[cat] = { fungible: 0n, nfts: [], utxoIds: [] };
      if (token.hasAmount) tokenBalances[cat].fungible += token.amount;
      if (token.hasNft) {
        tokenBalances[cat].nfts.push({
          utxoId: `${u.txid}:${u.vout}`,
          commitmentHex: token.commitmentHex,
          capability: token.capability,
          capabilityLabel: token.capabilityLabel,
        });
      }
      tokenBalances[cat].utxoIds.push(`${u.txid}:${u.vout}`);
    };

    if (client.hasTokenData) {
      for (const u of utxos) {
        const token = tokenFromElectrum(u.tokenData);
        if (token) addToken(u, token);
      }
    } else {
      await Promise.all(utxos.map(async (u) => {
        try {
          const t = await getTx(u.txid);
          const out = t.vout[u.vout];
          if (!out) return;
          u.scriptHex = out.scriptHex;
          // Prefer the server's own tokenData; fall back to decoding a
          // prefix out of the script for a server that embeds it there.
          const token = tokenFromElectrum(out.tokenData)
            || (out.scriptHex ? cashtokens.decodePrefixedScript(tx.fromHex(out.scriptHex)).token : null);
          if (token) addToken(u, token);
        } catch (e) {
          log("utxo classify failed:", u.txid + ":" + u.vout, e?.message || e);
        }
      }));
    }
    state.utxos = utxos;
    // Serialize BigInt fungible amounts as decimal strings for the snapshot
    // (JSON.stringify chokes on BigInt otherwise).
    const serializedBalances = {};
    for (const [cat, bal] of Object.entries(tokenBalances)) {
      serializedBalances[cat] = {
        fungible: bal.fungible.toString(),
        nfts: bal.nfts,
        utxoCount: bal.utxoIds.length,
      };
    }
    state.tokenBalances = serializedBalances;
    // Balance number is BCH sat only — token UTXOs still carry a small
    // BCH value (dust minimum for the prefix), but treating that as
    // spendable would let a routine send burn the token. Track total
    // separately as bareBalance so the panel can still show "there's
    // BCH sitting in token UTXOs".
    let confirmed = 0, unconfirmed = 0, tokenLocked = 0;
    for (const u of utxos) {
      if (u.token) { tokenLocked += u.value; continue; }
      if (u.height > 0) confirmed += u.value; else unconfirmed += u.value;
    }
    state.balance = { confirmed, unconfirmed, tokenLocked };
  }

  async function getTx(txid) {
    const c = memTx.get(txid) || txCache[txid];
    if (c && c.confirmations > 0) return c;
    const raw = await client.call("blockchain.transaction.get", [txid, true]);
    const slim = {
      txid,
      confirmations: raw.confirmations || 0,
      time: raw.blocktime || raw.time || 0,
      vin: (raw.vin || []).map((i) => ({ txid: i.txid, vout: i.vout })),
      // tokenData was being dropped here, and that was the whole bug: the
      // server reports CashTokens in this field, NOT inside
      // scriptPubKey.hex, which Fulcrum returns with the token prefix
      // already stripped. So the old classify pass fetched a transaction per
      // UTXO, looked for a prefix that was never there, and concluded "no
      // token" every single time. Keep it.
      vout: (raw.vout || []).map((o) => ({
        value: sats(o.value),
        scriptHex: o.scriptPubKey && o.scriptPubKey.hex,
        tokenData: o.tokenData || o.token_data || null,
      })),
      size: raw.size || 0,
    };
    rememberTx(txid, slim);
    txDirty = true;
    return slim;
  }

  async function loadHistory() {
    const entries = [...state.watched.values()].filter((e) => state.used.has(key(e)));
    const merged = new Map();
    const lists = await Promise.all(entries.map(historyOf));
    for (const list of lists) for (const h of list) {
      const prev = merged.get(h.tx_hash);
      if (!prev || (h.height > 0 && prev.height <= 0)) merged.set(h.tx_hash, { txid: h.tx_hash, height: h.height });
    }
    const ordered = [...merged.values()].sort((a, b) => {
      const ha = a.height > 0 ? a.height : Infinity, hb = b.height > 0 ? b.height : Infinity;
      return hb - ha;
    }).slice(0, HISTORY_LIMIT);
    const ours = new Set([...state.watched.values()].map((e) => e.scriptHex));
    // Fetch in two parallel waves instead of one long sequential chain. This
    // loop used to await getTx() per displayed transaction and then again per
    // INPUT of each, one at a time. A consolidation here has 400 inputs, so a
    // single history row cost 400 round trips in series — tens of seconds,
    // during which every response still had to be parsed on the main thread.
    // Median input count is 2, so the common case gets faster too.
    await getTxMany(ordered.map((h) => h.txid));
    const parentIds = [];
    for (const h of ordered) {
      const t = await getTx(h.txid);
      for (const i of t.vin) if (i.txid) parentIds.push(i.txid);
    }
    await getTxMany(parentIds);

    const out = [];
    for (const h of ordered) {
      const t = await getTx(h.txid);
      let received = 0, spent = 0, inputsTotal = 0, outputsTotal = 0, allInputsOurs = true;
      for (const o of t.vout) { outputsTotal += o.value; if (ours.has(o.scriptHex)) received += o.value; }
      for (const i of t.vin) {
        if (!i.txid) continue;   // coinbase
        const p = await getTx(i.txid);   // warm after the wave above
        const po = p.vout[i.vout];
        if (!po) continue;
        inputsTotal += po.value;
        if (ours.has(po.scriptHex)) spent += po.value; else allInputsOurs = false;
      }
      const delta = received - spent;
      let to = null;
      if (delta < 0) {
        const ext = t.vout.find((o) => !ours.has(o.scriptHex));
        if (ext && ext.scriptHex) to = scriptToAddress(ext.scriptHex);
      }
      out.push({
        txid: t.txid, height: h.height, confirmations: t.confirmations, time: t.time,
        delta, fee: allInputsOurs && inputsTotal ? inputsTotal - outputsTotal : null, to,
      });
    }
    state.history = out;
    // Persist ONLY the displayed window. Parent transactions are needed to
    // compute a delta and nothing else, and persisting every one ever seen is
    // what grew this store to 7.4 MB across 6,981 entries — which the host
    // then re-parsed on every storage read, on the main thread. They stay in
    // the in-memory map instead, so a session never refetches them and the
    // file stays a few tens of KB.
    const keep = {};
    for (const h of ordered) {
      const t = memTx.get(h.txid);
      if (t) keep[h.txid] = t;
    }
    const before = JSON.stringify(txCache);
    const after = JSON.stringify(keep);
    txCache = keep;
    if (after !== before) storage.set("txCache", keep);
    txDirty = false;
  }

  function scriptToAddress(scriptHex) {
    try {
      if (/^76a914[0-9a-f]{40}88ac$/.test(scriptHex)) return cashaddr.encode(keys.prefix, 0, tx.fromHex(scriptHex.slice(6, 46)));
      if (/^a914[0-9a-f]{40}87$/.test(scriptHex)) return cashaddr.encode(keys.prefix, 1, tx.fromHex(scriptHex.slice(4, 44)));
    } catch {}
    return null;
  }

  // A manual refresh used to return here the moment a background poll was
  // in flight, and refreshChain reported ok:true for it — so the Refresh
  // button no-opped and claimed success, which is exactly when a user is
  // most likely to press it. A forced refresh now waits for the in-flight
  // pass and then does real work; a background poll still yields.
  let inflight = null;
  async function refresh(full = false) {
    if (state.scanning) {
      if (!full) return;
      try { await inflight; } catch { /* its own error is already on state */ }
      if (state.scanning) return;        // another forced pass won the race
    }
    inflight = doRefresh(full);
    return inflight;
  }
  async function doRefresh(full) {
    state.scanning = true; state.error = null; onChange();
    try {
      if (full || !state.watched.size) await scan();
      else { let r = Number(storage.get("receiveCursor", 0)) || 0; while (state.used.has("0/" + r)) r++; state.receiveIndex = r; watch(keys.entry(0, r)); }
      await loadUtxos();
      await loadHistory();
      await subscribeAll();
      // A tx that just landed can mark the current receive address used.
      for (const u of state.utxos) state.used.add(key(u.entry));
      let r = Number(storage.get("receiveCursor", 0)) || 0;
      while (state.used.has("0/" + r)) r++;
      if (r !== state.receiveIndex) { state.receiveIndex = r; watch(keys.entry(0, r)); }
    } catch (e) {
      state.error = e?.message || String(e);
      log("refresh failed:", state.error);
    } finally {
      state.scanning = false;
      onChange();
    }
  }
  function scheduleRefresh(ms = 800) {
    clearTimeout(refreshTimer);
    refreshTimer = setTimeout(() => refresh(false), ms);
  }

  client.onNotify = (method, params) => {
    if (method === "blockchain.headers.subscribe") {
      const h = params && params[0] && params[0].height;
      if (h) { state.height = h; scheduleRefresh(1500); }
    } else if (method === "blockchain.scripthash.subscribe") {
      scheduleRefresh(800);
    }
  };

  function nextUnusedAddress() {
    let r = state.receiveIndex + 1;
    while (state.used.has("0/" + r)) r++;
    storage.set("receiveCursor", r);
    state.receiveIndex = r;
    watch(keys.entry(0, r));
    client.subscribe("blockchain.scripthash.subscribe", [keys.entry(0, r).scripthash]).catch(() => {});
    onChange();
    return current();
  }
  function current() { return keys.entry(0, state.receiveIndex); }
  function changeEntry() {
    let i = 0;
    while (state.used.has("1/" + i)) i++;
    return keys.entry(1, i);
  }

  // targets: [{ to, value }] (value in sats; ignored for sendMax) -> unsigned plan.
  // memo: optional string (UTF-8, ≤220 bytes) — attached as an OP_RETURN
  //       data output. Zero value, no dust check, fee estimate accounts
  //       for the extra bytes. Passing "" disables the memo.
  function plan({ targets, feeRate = 1, sendMax = false, memo = "" }) {
    const rate = Math.min(10, Math.max(1, Number(feeRate) || 1));
    const outs = targets.map((t) => {
      const a = cashaddr.parseAny(t.to, sha256, keys.prefix);
      const script = a.type === 0
        ? Uint8Array.from([0x76, 0xa9, 0x14, ...a.hash, 0x88, 0xac])
        : Uint8Array.from([0xa9, 0x14, ...a.hash, 0x87]);
      return { value: Math.round(Number(t.value) || 0), script, to: a.cashaddr };
    });
    if (memo) outs.push({ value: 0, script: tx.memoScript(memo), data: true, memo });
    // Spend confirmed coins first; unconfirmed only when needed. Token
    // UTXOs are excluded entirely — burning a category by dropping its
    // prefix is not a mistake we can undo, so a plain BCH send must
    // never pull one. Token sends have their own code path with
    // { includeToken: category } later.
    const spendable = state.utxos
      .filter((u) => !u.token)
      .slice()
      .sort((a, b) => (b.height > 0) - (a.height > 0));
    const sel = tx.select(spendable, outs, rate, changeEntry().script, { sendMax });
    // recipients only lists spendable (non-data) outputs, keeping the
    // panel's summary honest — the memo is surfaced separately as .memo.
    const spendable_outs = sel.outputs.filter((o) => !o.data);
    return {
      ...sel, feeRate: rate,
      recipients: spendable_outs.map((o, i) => ({ to: outs[i]?.to, value: o.value })),
      memo: memo || null,
    };
  }

  async function signAndBroadcast(p) {
    const t = { inputs: p.inputs.map((u) => ({ ...u, script: u.entry.script })), outputs: p.outputs };
    const signed = tx.sign(t, (inp, _i, digest) => ({ sig: keys.sign(inp.entry, digest), publicKey: inp.entry.publicKey }));
    const txid = await client.call("blockchain.transaction.broadcast", [signed.hex]);
    if (typeof txid !== "string" || txid.length !== 64) throw new Error("broadcast rejected: " + JSON.stringify(txid));
    log("broadcast", txid);
    scheduleRefresh(1200);
    return { txid, hex: signed.hex, fee: p.fee };
  }

  function snapshot() {
    const cur = current();
    return {
      address: cur.address,
      addressIndex: state.receiveIndex,
      addressPath: cur.path,
      balance: state.balance,
      height: state.height,
      history: state.history,
      utxoCount: state.utxos.length,
      // CashTokens balances, keyed by category hex. Empty object when the
      // wallet holds no token UTXOs. Serialised BigInts (fungible amounts)
      // come across as decimal strings — panel formats via BigInt again.
      tokenBalances: state.tokenBalances,
      scanning: state.scanning,
      error: state.error,
    };
  }

  function dispose() { clearTimeout(refreshTimer); }

  return { refresh, snapshot, nextUnusedAddress, current, plan, signAndBroadcast, dispose, state };
};