From 6bdb7cf262d7a4d2f362bdddd9a07c1bc10d539e Mon Sep 17 00:00:00 2001
From: adlee-was-taken <aaron.daniel.lee@gmail.com>
Date: Thu, 25 Jun 2026 21:29:08 -0400
Subject: [PATCH] feat(ext/sw): multi-context session (personal + orgs),
 clearAll zeroes all

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014s527M917W47LDrfQ4t47g
---
 .../__tests__/org-session.test.ts             | 17 +++++
 extension/src/service-worker/session.ts       | 67 ++++++++++++-------
 2 files changed, 61 insertions(+), 23 deletions(-)
 create mode 100644 extension/src/service-worker/__tests__/org-session.test.ts

diff --git a/extension/src/service-worker/__tests__/org-session.test.ts b/extension/src/service-worker/__tests__/org-session.test.ts
new file mode 100644
index 0000000..277f2a8
--- /dev/null
+++ b/extension/src/service-worker/__tests__/org-session.test.ts
@@ -0,0 +1,17 @@
+import { describe, expect, it, vi } from 'vitest';
+import type { SessionHandle } from '../../../wasm/relicario_wasm';
+import * as session from '../session';
+
+describe('multi-context session', () => {
+  it('clearAll frees personal and every org handle', () => {
+    const free = vi.fn();
+    const mk = (id: number) => ({ value: id, free } as unknown as SessionHandle);
+    session.setPersonal(mk(1));
+    session.setOrg('org-a', mk(2));
+    session.setOrg('org-b', mk(3));
+    session.clearAll();
+    expect(free).toHaveBeenCalledTimes(3);
+    expect(session.getPersonal()).toBeNull();
+    expect(session.getOrg('org-a')).toBeNull();
+  });
+});
diff --git a/extension/src/service-worker/session.ts b/extension/src/service-worker/session.ts
index 4859646..f47ad69 100644
--- a/extension/src/service-worker/session.ts
+++ b/extension/src/service-worker/session.ts
@@ -1,34 +1,55 @@
-/// Single module-scope "current" SessionHandle.
+/// Multi-context SW session: one personal handle + a Map of org handles.
 ///
-/// α assumes one vault per extension install. The master key lives only
-/// inside WASM linear memory (wrapped in Zeroizing<[u8;32]>); this module
-/// just holds the opaque handle that names it.
+/// The master key for every context lives only inside WASM linear memory
+/// (wrapped in Zeroizing<[u8;32]>); this module holds the opaque SessionHandle
+/// references that name them.  Calling `.free()` on a SessionHandle removes
+/// the entry from the WASM SESSIONS registry and zeroizes master_key and
+/// image_secret.
 ///
-/// Future multi-vault (β+) would replace `current` with
-/// `Map<vaultId, SessionHandle>` and thread `vaultId` through every
-/// handler. Deliberate α simplicity — not an oversight.
+/// Security invariant: a vault lock OR an inactivity-timer expiry must zero
+/// EVERY handle — personal AND every org handle.  `clearAll()` is the
+/// canonical zero-all path.  `clearCurrent()` is aliased to `clearAll()` so
+/// the existing lock handler (popup-only.ts) and inactivity timer (index.ts)
+/// automatically zero org handles without code changes.
 ///
-/// As of Phase 1 of the security-polish series, `impl Drop for SessionHandle`
-/// on the Rust side makes `.free()` the meaningful cleanup call: it removes
-/// the entry from the SESSIONS registry and zeroizes `master_key` and
-/// `image_secret`. Calling `wasm.lock(handle.value)` before `.free()` would
-/// be redundant belt-and-suspenders; this module intentionally does not.
+/// The org master key MUST NEVER be written to localStorage, IndexedDB, or
+/// any persistent store.  Only in-memory SessionHandle references live here.
 
 import type { SessionHandle } from '../../wasm/relicario_wasm';
 
-let current: SessionHandle | null = null;
+let personal: SessionHandle | null = null;
+const orgs = new Map<string, SessionHandle>();
+let context: 'personal' | string = 'personal';
 
-export function setCurrent(h: SessionHandle): void { current = h; }
+export function setPersonal(h: SessionHandle): void { personal = h; }
+export function getPersonal(): SessionHandle | null { return personal; }
 
-export function getCurrent(): SessionHandle | null { return current; }
+export function setOrg(orgId: string, h: SessionHandle): void { orgs.set(orgId, h); }
+export function getOrg(orgId: string): SessionHandle | null { return orgs.get(orgId) ?? null; }
 
+export function setContext(c: 'personal' | string): void { context = c; }
+export function currentContext(): 'personal' | string { return context; }
+
+export function requireCurrentHandle(): SessionHandle {
+  const h = context === 'personal' ? personal : orgs.get(context) ?? null;
+  if (!h) throw new Error('vault_locked');
+  return h;
+}
+
+export function clearAll(): void {
+  if (personal) { personal.free(); personal = null; }
+  for (const [, h] of orgs) h.free();
+  orgs.clear();
+  context = 'personal';
+}
+
+// Back-compat wrappers so existing personal-vault callers compile unchanged.
+// clearCurrent() aliases clearAll() so the lock handler and inactivity timer
+// in popup-only.ts and index.ts zero EVERY handle (personal + all orgs).
+export function setCurrent(h: SessionHandle): void { setPersonal(h); }
+export function getCurrent(): SessionHandle | null { return getPersonal(); }
 export function requireCurrent(): SessionHandle {
-  if (!current) throw new Error('vault_locked');
-  return current;
-}
-
-export function clearCurrent(): void {
-  if (!current) return;
-  current.free();
-  current = null;
+  if (!personal) throw new Error('vault_locked');
+  return personal;
 }
+export function clearCurrent(): void { clearAll(); }