mls-engine.ts raw
1 /* eslint-disable @typescript-eslint/no-explicit-any */
2 // mls-engine.ts — Loads marmot.wasm in the signer extension background.
3 // All MLS crypto happens here with direct access to vault keys.
4 // No bus round-trips for signing or encryption.
5
6 import { signEvent, nip44Encrypt, nip44Decrypt, nip04Encrypt, nip04Decrypt } from './background-common';
7 import browser from 'webextension-polyfill';
8
9 // --- State ---
10 let wasmReady = false;
11 const mlsTabIds = new Set<number>(); // Fix 2g: broadcast to all tabs
12 let currentPrivkey = '';
13 let mlsInitPromise: Promise<string> | null = null;
14 let wasmLoadingPromise: Promise<void> | null = null; // Fix 2d: guard concurrent instantiation
15 let lastEventTSInterval: ReturnType<typeof setInterval> | null = null; // Fix 2e: track interval
16
17 // Fix 2b: queue events arriving before init completes
18 let initDone = false;
19 const pendingDeliverEvents: Array<{ subId: number; eventJSON: string }> = [];
20
21 // Persist init params so the extension can auto-reinitialize after MV3
22 // background termination. browser.storage.session is encrypted, cleared
23 // on browser close — same lifecycle as in-memory state, but survives
24 // background script restarts.
25 const MLS_SESSION_KEY = '_mls_init_params';
26
27 interface MlsInitParams {
28 privkey: string;
29 pubkey: string;
30 relayURLs: string[];
31 lastEventTS: number;
32 }
33
34 async function persistInitParams(params: MlsInitParams): Promise<void> {
35 try { await browser.storage.session.set({ [MLS_SESSION_KEY]: params }); } catch (_) {}
36 }
37
38 async function loadInitParams(): Promise<MlsInitParams | null> {
39 try {
40 const r = await browser.storage.session.get(MLS_SESSION_KEY);
41 return (r[MLS_SESSION_KEY] as MlsInitParams) || null;
42 } catch (_) { return null; }
43 }
44
45 // Auto-reinitialize from session storage if background was restarted.
46 async function ensureInit(): Promise<boolean> {
47 if ((globalThis as any)._marmot) return true;
48 const params = await loadInitParams();
49 if (!params) return false;
50 currentPrivkey = params.privkey;
51 mlsInitPromise = (async () => {
52 await loadWasm();
53 setupMarmotInit(params.lastEventTS);
54 const initFn = (globalThis as any)._marmot_init;
55 if (!initFn) return 'error: wasm bridge not ready';
56 const result = await initFn(params.pubkey, ...params.relayURLs);
57 if (!result || result === 'ok') {
58 const m = (globalThis as any)._marmot;
59 if (m) {
60 m.publishKP();
61 const groups = String(m.listGroups() || '[]');
62 if (groups === '[]') m.restoreGroups();
63 m.subscribe();
64 }
65 }
66 initDone = true;
67 return String(result || 'ok');
68 })().catch(e => { mlsInitPromise = null; throw e; });
69 await mlsInitPromise;
70 return !!(globalThis as any)._marmot;
71 }
72
73 // --- IDB Group Store (same schema as the marmot SW used) ---
74 const GDB_NAME = 'marmot-groups';
75 const GDB_VER = 1;
76 const GDB_STORE = 'groups';
77
78 // Open DB for writes — creates the store if it doesn't exist yet.
79 function gdbOpenWrite(): Promise<IDBDatabase> {
80 return new Promise((resolve, reject) => {
81 const req = indexedDB.open(GDB_NAME, GDB_VER);
82 req.onupgradeneeded = (e: any) => {
83 const db = e.target.result as IDBDatabase;
84 if (!db.objectStoreNames.contains(GDB_STORE)) {
85 db.createObjectStore(GDB_STORE);
86 }
87 };
88 req.onsuccess = () => resolve(req.result);
89 req.onerror = () => reject(req.error);
90 });
91 }
92
93 // Open DB for reads — returns null if the DB/store doesn't exist (no side-effects).
94 function gdbOpenRead(): Promise<IDBDatabase | null> {
95 return new Promise((resolve) => {
96 const req = indexedDB.open(GDB_NAME);
97 req.onupgradeneeded = () => {
98 // DB doesn't exist yet — abort so we don't create an empty one.
99 req.result.close();
100 req.transaction?.abort();
101 };
102 req.onsuccess = () => {
103 const db = req.result;
104 if (!db.objectStoreNames.contains(GDB_STORE)) {
105 db.close();
106 resolve(null);
107 return;
108 }
109 resolve(db);
110 };
111 req.onerror = () => resolve(null);
112 req.onblocked = () => resolve(null);
113 });
114 }
115
116 function storeResult(id: number, data: string, err: string) {
117 const m = (globalThis as any)._marmot;
118 if (m?.storeResult) m.storeResult(id, data || '', err || '');
119 }
120
121 function setupStoreCallbacks() {
122 const g = globalThis as any;
123 g._marmot_store_save = (id: number, groupIDHex: string, stateHex: string) => {
124 gdbOpenWrite().then(db => {
125 const tx = db.transaction(GDB_STORE, 'readwrite');
126 tx.objectStore(GDB_STORE).put(stateHex, groupIDHex);
127 tx.oncomplete = () => storeResult(id, '', '');
128 tx.onerror = () => storeResult(id, '', tx.error?.message || 'save failed');
129 }).catch(e => storeResult(id, '', (e as Error).message));
130 };
131 g._marmot_store_load = (id: number, groupIDHex: string) => {
132 gdbOpenRead().then(db => {
133 if (!db) { storeResult(id, '', ''); return; }
134 const tx = db.transaction(GDB_STORE, 'readonly');
135 const req = tx.objectStore(GDB_STORE).get(groupIDHex);
136 req.onsuccess = () => storeResult(id, req.result || '', '');
137 req.onerror = () => storeResult(id, '', req.error?.message || 'load failed');
138 }).catch(e => storeResult(id, '', (e as Error).message));
139 };
140 g._marmot_store_list = (id: number) => {
141 gdbOpenRead().then(db => {
142 if (!db) { storeResult(id, '', ''); return; }
143 const tx = db.transaction(GDB_STORE, 'readonly');
144 const req = tx.objectStore(GDB_STORE).getAllKeys();
145 req.onsuccess = () => storeResult(id, (req.result || []).join(','), '');
146 req.onerror = () => storeResult(id, '', req.error?.message || 'list failed');
147 }).catch(e => storeResult(id, '', (e as Error).message));
148 };
149 g._marmot_store_delete = (id: number, groupIDHex: string) => {
150 gdbOpenRead().then(db => {
151 if (!db) { storeResult(id, '', ''); return; }
152 const tx = db.transaction(GDB_STORE, 'readwrite');
153 tx.objectStore(GDB_STORE).delete(groupIDHex);
154 tx.oncomplete = () => storeResult(id, '', '');
155 tx.onerror = () => storeResult(id, '', tx.error?.message || 'delete failed');
156 }).catch(e => storeResult(id, '', (e as Error).message));
157 };
158 }
159
160 // --- Epoch Check ---
161
162 async function marmotEpochCheck(wasmVer: string): Promise<void> {
163 try {
164 const db = await gdbOpenWrite();
165 const tx = db.transaction(GDB_STORE, 'readonly');
166 const req = tx.objectStore(GDB_STORE).get('__version__');
167 const stored: string | undefined = await new Promise((resolve) => {
168 req.onsuccess = () => resolve(req.result);
169 req.onerror = () => resolve(undefined);
170 });
171 if (stored === wasmVer) { db.close(); return; }
172 console.warn('[mls-engine] version epoch mismatch: stored=' + (stored || 'none') + ' running=' + wasmVer + ' — flushing marmot groups');
173 const flushTx = db.transaction(GDB_STORE, 'readwrite');
174 const store = flushTx.objectStore(GDB_STORE);
175 store.clear();
176 store.put(wasmVer, '__version__');
177 await new Promise<void>((resolve) => {
178 flushTx.oncomplete = () => resolve();
179 flushTx.onerror = () => resolve();
180 });
181 db.close();
182 } catch (e) {
183 console.warn('[mls-engine] epoch check failed:', e);
184 }
185 }
186
187 // --- WASM Loading ---
188
189 // Fix 2d: guard against concurrent instantiation
190 async function loadWasm(): Promise<void> {
191 if (wasmReady) return;
192 if (wasmLoadingPromise) return wasmLoadingPromise;
193 wasmLoadingPromise = doLoadWasm();
194 try {
195 await wasmLoadingPromise;
196 } catch (e) {
197 wasmLoadingPromise = null; // allow retry on failure
198 throw e;
199 }
200 }
201
202 async function doLoadWasm(): Promise<void> {
203 // wasm_exec.js is loaded as a background script (manifest.json),
204 // so globalThis.Go is already available.
205 const GoClass = (globalThis as any).Go;
206 if (!GoClass) throw new Error('Go WASM runtime not loaded (wasm_exec.js missing from background scripts)');
207
208 const go = new GoClass();
209 const wasmUrl = browser.runtime.getURL('marmot.wasm');
210 const wasmResponse = await fetch(wasmUrl);
211 const wasmBytes = await wasmResponse.arrayBuffer();
212 const result = await WebAssembly.instantiate(wasmBytes, go.importObject);
213
214 // Start the Go program. Don't await — it blocks forever (select {}).
215 // The synchronous part sets up globalThis._marmot before yielding.
216 go.run(result.instance).then(
217 () => { console.error('[mls-engine] Go program exited unexpectedly'); wasmReady = false; },
218 (err: any) => { console.error('[mls-engine] Go program crashed:', err); wasmReady = false; }
219 );
220
221 // Verify the Go side registered its API before we proceed.
222 if (!(globalThis as any)._marmot) {
223 throw new Error('Go WASM started but _marmot global not registered');
224 }
225
226 setupStoreCallbacks();
227
228 // Epoch check: flush stale marmot groups if WASM version changed.
229 const wasmVer = (globalThis as any)._marmot?.version;
230 if (wasmVer) await marmotEpochCheck(wasmVer);
231
232 wasmReady = true;
233 }
234
235 // --- Callback Bridge ---
236
237 function setupMarmotInit(lastEventTS: number) {
238 const g = globalThis as any;
239
240 // Override _marmot_init to wire our local callbacks.
241 // Returns a Promise that resolves when NewClient completes (it runs in a
242 // goroutine because store.ListGroups blocks on async IDB callbacks).
243 g._marmot_init = async (pubkeyHex: string, ...relayURLs: string[]): Promise<string> => {
244 const marmot = g._marmot;
245 if (!marmot) return 'error: wasm not loaded';
246
247 const publishFn = (eventJSON: string) => {
248 pushToTab({ cmd: 'publish', event: eventJSON });
249 };
250
251 const subscribeFn = (subId: number, filterJSON: string) => {
252 pushToTab({ cmd: 'subscribe', subId, filter: filterJSON });
253 };
254
255 const cryptoSendFn = (op: string, peerHex: string, data: string, id: number) => {
256 handleCryptoLocal(op, peerHex, data, id);
257 };
258
259 const onDMFn = (senderHex: string, plaintext: string) => {
260 const ts = Math.floor(Date.now() / 1000);
261 pushToTab({
262 cmd: 'dm',
263 peer: senderHex,
264 sender: senderHex,
265 content: plaintext,
266 ts,
267 source: 'marmot',
268 eventId: '',
269 });
270 };
271
272 const onStatusFn = (msg: string) => {
273 pushToTab({ cmd: 'status', msg });
274 };
275
276 // Fix 2e: clear previous interval before creating new.
277 if (lastEventTSInterval) clearInterval(lastEventTSInterval);
278 lastEventTSInterval = setInterval(() => {
279 try {
280 const ts = marmot.lastEventTS?.();
281 if (ts > 0) pushToTab({ cmd: 'mls_ts', ts });
282 } catch (_) {}
283 }, 30000);
284
285 return new Promise<string>((resolve) => {
286 const onReadyFn = (result: string) => {
287 resolve(result || 'ok');
288 };
289 marmot.init(pubkeyHex, publishFn, subscribeFn, cryptoSendFn, onDMFn, onStatusFn, onReadyFn, lastEventTS, ...relayURLs);
290 });
291 };
292 }
293
294 // --- Local Crypto (no round-trip!) ---
295
296 function handleCryptoLocal(op: string, peerHex: string, data: string, id: number) {
297 const marmot = (globalThis as any)._marmot;
298 if (!marmot) return;
299
300 const resolve = (result: string, err: string) => {
301 marmot.cryptoResult(id, result, err);
302 };
303
304 try {
305 switch (op) {
306 case 'signEvent': {
307 const eventTemplate = JSON.parse(data);
308 const signed = signEvent(eventTemplate, currentPrivkey);
309 resolve(JSON.stringify(signed), '');
310 break;
311 }
312 case 'nip44.encrypt':
313 case 'nip44Encrypt':
314 nip44Encrypt(currentPrivkey, peerHex, data).then(
315 r => resolve(r, ''),
316 e => resolve('', (e as Error).message)
317 );
318 break;
319 case 'nip44.decrypt':
320 case 'nip44Decrypt':
321 nip44Decrypt(currentPrivkey, peerHex, data).then(
322 r => resolve(r, ''),
323 e => resolve('', (e as Error).message)
324 );
325 break;
326 case 'nip04.encrypt':
327 case 'nip04Encrypt':
328 nip04Encrypt(currentPrivkey, peerHex, data).then(
329 r => resolve(r, ''),
330 e => resolve('', (e as Error).message)
331 );
332 break;
333 case 'nip04.decrypt':
334 case 'nip04Decrypt':
335 nip04Decrypt(currentPrivkey, peerHex, data).then(
336 r => resolve(r, ''),
337 e => resolve('', (e as Error).message)
338 );
339 break;
340 default:
341 resolve('', 'unsupported crypto op: ' + op);
342 }
343 } catch (err) {
344 resolve('', (err as Error).message);
345 }
346 }
347
348 // --- Push to originating tab ---
349 // Broadcast to all smesh tabs. Re-discover on every send to avoid stale IDs
350 // after idle (tab refresh assigns new ID, old one stays in Set forever).
351
352 async function discoverTabs(): Promise<void> {
353 try {
354 const tabs = await browser.tabs.query({ url: ['*://smesh.lol/*', '*://127.0.0.1:*/*', '*://localhost:*/*'] });
355 mlsTabIds.clear();
356 for (const t of tabs) {
357 if (t.id) mlsTabIds.add(t.id);
358 }
359 } catch (_) {}
360 }
361
362 async function pushToTab(data: any) {
363 // Only discover if we have no known tabs — otherwise use IDs from mlsInit/mlsSendDM.
364 // discoverTabs() clears the set, so calling it unconditionally destroys manually-added
365 // IDs when browser.tabs.query({url:...}) returns empty (requires tabs permission).
366 if (mlsTabIds.size === 0) await discoverTabs();
367 if (mlsTabIds.size === 0) return;
368 const msg = { ext: 'smesh-signer', type: 'mls-push', data };
369 for (const tabId of mlsTabIds) {
370 browser.tabs.sendMessage(tabId, msg).catch(() => {
371 mlsTabIds.delete(tabId);
372 });
373 }
374 }
375
376 // --- Public API (called from background.ts) ---
377
378 export async function mlsInit(
379 privkey: string,
380 pubkey: string,
381 relayURLs: string[],
382 tabId: number,
383 lastEventTS: number = 0
384 ): Promise<string> {
385 mlsTabIds.add(tabId); // Fix 2g: add, don't overwrite
386 currentPrivkey = privkey;
387 persistInitParams({ privkey, pubkey, relayURLs, lastEventTS });
388 // Fix 2f: reset on failure so retry is possible
389 mlsInitPromise = (async () => {
390 await loadWasm();
391 setupMarmotInit(lastEventTS);
392 const initFn = (globalThis as any)._marmot_init;
393 if (!initFn) return 'error: wasm bridge not ready';
394 const result = await initFn(pubkey, ...relayURLs);
395 // Auto-bootstrap: publish key package, restore groups if empty, subscribe.
396 if (!result || result === 'ok') {
397 const m = (globalThis as any)._marmot;
398 if (m) {
399 m.publishKP();
400 // Restore groups from relay backup if IDB was empty.
401 const groups = String(m.listGroups() || '[]');
402 if (groups === '[]') m.restoreGroups();
403 m.subscribe();
404 }
405 }
406 // Fix 2b: mark init done and drain pending events.
407 initDone = true;
408 const pending = pendingDeliverEvents.splice(0);
409 for (const p of pending) {
410 const m = (globalThis as any)._marmot;
411 if (m) m.deliverEvent(p.subId, p.eventJSON);
412 }
413 return String(result || 'ok');
414 })().catch(e => {
415 mlsInitPromise = null; // Fix 2f: allow retry
416 throw e;
417 });
418 return mlsInitPromise;
419 }
420
421 export function mlsSetTab(tabId: number) {
422 mlsTabIds.add(tabId); // Fix 2g
423 }
424
425 export async function mlsSendDM(recipient: string, content: string): Promise<string> {
426 if (mlsInitPromise) await mlsInitPromise;
427 if (!(globalThis as any)._marmot) await ensureInit();
428 const m = (globalThis as any)._marmot;
429 if (!m) return 'error: not initialized';
430 return String(m.sendDM(recipient, content) || 'ok');
431 }
432
433 export async function mlsSubscribe(): Promise<string> {
434 if (mlsInitPromise) await mlsInitPromise;
435 if (!(globalThis as any)._marmot) await ensureInit();
436 const m = (globalThis as any)._marmot;
437 if (!m) return 'error: not initialized';
438 return String(m.subscribe() || 'ok');
439 }
440
441 export async function mlsPublishKP(): Promise<string> {
442 if (mlsInitPromise) await mlsInitPromise;
443 if (!(globalThis as any)._marmot) await ensureInit();
444 const m = (globalThis as any)._marmot;
445 if (!m) return 'error: not initialized';
446 return String(m.publishKP() || 'ok');
447 }
448
449 export async function mlsListGroups(): Promise<string> {
450 if (mlsInitPromise) await mlsInitPromise;
451 const m = (globalThis as any)._marmot;
452 if (!m) return '[]';
453 return String(m.listGroups() || '[]');
454 }
455
456 // Fix 2b: queue events before init completes
457 export function mlsDeliverEvent(subId: number, eventJSON: string): void {
458 if (!initDone) {
459 pendingDeliverEvents.push({ subId, eventJSON });
460 return;
461 }
462 const m = (globalThis as any)._marmot;
463 if (!m) return;
464 m.deliverEvent(subId, eventJSON);
465 }
466
467 export function mlsHandleEvent(eventJSON: string): void {
468 const m = (globalThis as any)._marmot;
469 if (!m) return;
470 m.handleEvent(eventJSON);
471 }
472
473 export async function mlsBackupGroups(): Promise<void> {
474 if (mlsInitPromise) await mlsInitPromise;
475 const m = (globalThis as any)._marmot;
476 if (m?.backupGroups) m.backupGroups();
477 }
478
479 export async function mlsRestoreGroups(): Promise<void> {
480 if (mlsInitPromise) await mlsInitPromise;
481 const m = (globalThis as any)._marmot;
482 if (m?.restoreGroups) m.restoreGroups();
483 }
484
485 export async function mlsRatchetGroup(peerHex: string): Promise<void> {
486 if (mlsInitPromise) await mlsInitPromise;
487 const m = (globalThis as any)._marmot;
488 if (m?.ratchetGroup) m.ratchetGroup(peerHex);
489 }
490
491 export function isMlsMethod(method: string): boolean {
492 return method.startsWith('mls.');
493 }
494