wasm-worker-host.mjs raw

   1  // wasm-worker-host.mjs — Worker bootstrap for signer.wasm.
   2  // Receives init messages from offscreen-supervisor.mjs, instantiates wasm,
   3  // proxies storage calls to supervisor, routes ext messages.
   4  //
   5  // Signer-specific bridge functions (ext_*) are defined here.
   6  // Common bridge (SAB channels, spawn, subtle_random_bytes) comes from bridge-common.mjs.
   7  
   8  import {
   9    fromSlice, sabForceClose,
  10    makeCoreHelpers, makeWasi, makeCommonBridge,
  11    computeBuildHash, createSpawnedWorker,
  12  } from './bridge-common.mjs';
  13  
  14  // Note: all wasm runs in a Worker context. channel_recv uses Atomics.wait
  15  // which is forbidden on the main page thread.
  16  
  17  // ── Shared state ───────────────────────────────────────────────────────────
  18  
  19  let mem, xp;
  20  
  21  const _sabTable = new Map();           // int32 → SharedArrayBuffer
  22  const _sabSeqRef = { value: 0 };       // next SAB handle
  23  const _spawnedWorkerChans = new Map(); // Worker → SAB[] (dead-Worker cleanup)
  24  const _wasmUrlRef = { value: null };
  25  const _buildHashRef = { value: null };
  26  
  27  // Signer-specific state
  28  const _storageCbs = new Map(); // msgId → callback (storage proxy)
  29  let _storageSeq = 0;
  30  const _session = new Map();    // in-memory session storage
  31  let _extMsgCbID = -1;          // cbID registered by signer during _start
  32  
  33  // ── Core helpers via factory ───────────────────────────────────────────────
  34  
  35  const h = makeCoreHelpers(
  36    () => mem,
  37    () => xp,
  38  );
  39  const { readStr, readBytes, writeStr, writeI32, writeBytes, cb0, cbs, cbdata } = h;
  40  
  41  // ── Bridge ─────────────────────────────────────────────────────────────────
  42  
  43  const bridge = {
  44    // --- signer-specific: storage proxied to supervisor via postMessage ---
  45    ext_storage_get(kPtr, kLen, kCap, cbID) {
  46      const msgId = _storageSeq++;
  47      _storageCbs.set(msgId, val => cbs(cbID, val));
  48      self.postMessage({ type: 'storage-get', msgId, key: readStr(kPtr, kLen) });
  49    },
  50    ext_storage_set(kPtr, kLen, kCap, vPtr, vLen, vCap) {
  51      self.postMessage({ type: 'storage-set', key: readStr(kPtr, kLen), value: readStr(vPtr, vLen) });
  52    },
  53    ext_storage_remove(kPtr, kLen, kCap) {
  54      self.postMessage({ type: 'storage-remove', key: readStr(kPtr, kLen) });
  55    },
  56  
  57    // --- signer-specific: ext message routing ---
  58    ext_on_message(cbID) { _extMsgCbID = cbID; },
  59    ext_on_message_respond(reqID, ptr, len, cap) {
  60      self.postMessage({ id: reqID, result: readStr(ptr, len) });
  61    },
  62  
  63    // --- signer-specific: in-memory session storage ---
  64    ext_session_get(kPtr, kLen, kCap, cbID) {
  65      cbs(cbID, _session.get(readStr(kPtr, kLen)) || '');
  66    },
  67    ext_session_set(kPtr, kLen, kCap, vPtr, vLen, vCap) {
  68      _session.set(readStr(kPtr, kLen), readStr(vPtr, vLen));
  69    },
  70    ext_is_in_page() { return 0; },
  71  
  72    // --- common: SAB channels, spawn, subtle_random_bytes ---
  73    ...makeCommonBridge(h, _sabTable, _sabSeqRef, _wasmUrlRef, _buildHashRef, _spawnedWorkerChans),
  74  };
  75  
  76  const wasi = makeWasi(() => mem);
  77  
  78  
  79  // ── Message handler ────────────────────────────────────────────────────────
  80  
  81  self.onmessage = async function(e) {
  82    const d = e.data;
  83  
  84    // Storage proxy response from supervisor
  85    if (d.type === 'storage-result') {
  86      const fn = _storageCbs.get(d.msgId);
  87      if (fn) { _storageCbs.delete(d.msgId); fn(d.value || ''); }
  88      return;
  89    }
  90  
  91    // Root boot: instantiate wasm and run _start
  92    if (d.type === 'init' && d.mode === 'root') {
  93      try {
  94        const wasmBytes = await fetch(d.wasmUrl).then(r => r.arrayBuffer());
  95        const buildHash = await computeBuildHash(wasmBytes);
  96        _wasmUrlRef.value = d.wasmUrl;
  97        _buildHashRef.value = buildHash;
  98        self.postMessage({ type: 'hello', buildHash });
  99  
 100        const { instance } = await WebAssembly.instantiate(wasmBytes,
 101          { bridge, wasi_snapshot_preview1: wasi });
 102        mem = instance.exports.memory;
 103        xp = instance.exports;
 104        xp._start();
 105        self.postMessage({ type: 'ready' });
 106      } catch (err) {
 107        self.postMessage({ type: 'error', fatal: true, message: String(err) });
 108      }
 109      return;
 110    }
 111  
 112    // Spawn boot: instantiate wasm and call __spawn_entry
 113    if (d.type === 'init' && d.mode === 'spawn') {
 114      try {
 115        const wasmBytes = await fetch(d.wasmUrl).then(r => r.arrayBuffer());
 116        const childHash = await computeBuildHash(wasmBytes);
 117        if (d.buildHash && childHash !== d.buildHash) {
 118          self.postMessage({ type: 'error', fatal: true,
 119            message: 'spawn_domain: build hash mismatch parent=' + d.buildHash + ' child=' + childHash });
 120          self.close();
 121          return;
 122        }
 123        _wasmUrlRef.value = d.wasmUrl;
 124        _buildHashRef.value = childHash;
 125  
 126        const { instance } = await WebAssembly.instantiate(wasmBytes,
 127          { bridge, wasi_snapshot_preview1: wasi });
 128        mem = instance.exports.memory;
 129        xp = instance.exports;
 130  
 131        // Register passed channel SABs, write handles to linear memory,
 132        // then call __spawn_entry. bootstrap allocs are leaked intentionally:
 133        // self.close() reclaims all linear memory when the spawn completes.
 134        const chanSabs = d.chanSabs || [];
 135        const argBytes = d.argBytes || new Uint8Array(0);
 136  
 137        const argPtr = argBytes.length > 0 ? xp.__alloc(argBytes.length) : 0;
 138        if (argBytes.length > 0)
 139          new Uint8Array(mem.buffer, argPtr, argBytes.length).set(argBytes);
 140  
 141        const nChans = chanSabs.length;
 142        const chanHandlesPtr = nChans > 0 ? xp.__alloc(nChans * 4) : 0;
 143        if (nChans > 0) {
 144          const dv = new DataView(mem.buffer);
 145          chanSabs.forEach((sab, idx) => {
 146            const handle = _sabSeqRef.value++;
 147            _sabTable.set(handle, sab);
 148            dv.setInt32(chanHandlesPtr + idx * 4, handle, true);
 149          });
 150        }
 151  
 152        xp.__spawn_entry(d.fnIdx, argPtr, argBytes.length, chanHandlesPtr, nChans);
 153        self.postMessage({ type: 'exit' });
 154      } catch (err) {
 155        self.postMessage({ type: 'error', fatal: true, message: String(err) });
 156      }
 157      self.close();
 158      return;
 159    }
 160  
 161    // Signer request: { id, method, params, senderTabId }
 162    if (d.method !== undefined && _extMsgCbID >= 0 && xp) {
 163      const paramsJSON = typeof d.params === 'string'
 164        ? d.params
 165        : JSON.stringify(d.params || {});
 166      const [mPtr, mLen] = writeStr(d.method);
 167      const [pPtr, pLen] = writeStr(paramsJSON);
 168      xp.__hook_ext_on_message(d.id, mPtr, mLen, pPtr, pLen, d.senderTabId || 0);
 169    }
 170  };
 171