fix: per-account push subscriptions so multi-account notifications work #298

app/api/push/preview/route.ts

@@ -1,10 +1,73 @@
+import { cookies } from 'next/headers';
 import { NextRequest, NextResponse } from 'next/server';
 import { logger } from '@/lib/logger';
-import { getStalwartCredentials } from '@/lib/stalwart/credentials';
+import { MAX_ACCOUNT_SLOTS } from '@/lib/account-utils';
+import { readStalwartAuthContextFromStore } from '@/lib/stalwart/auth-context';
+import {
+  getStalwartCredentials,
+  type StalwartCredentials,
+} from '@/lib/stalwart/credentials';
 
 export const runtime = 'nodejs';
 export const dynamic = 'force-dynamic';
 
+interface ResolvedTarget {
+  authHeader: string;
+  apiUrl: string;
+  accountId: string;
+}
+
+// When the SW passes ?accountId=, we need the slot whose JMAP session owns
+// that account - not just "the first signed-in slot", which is what
+// getStalwartCredentials() defaults to. Probe each candidate's session in
+// parallel and return the first match.
+async function resolveTargetForAccount(accountId: string): Promise<ResolvedTarget | null> {
+  const cookieStore = await cookies();
+  const probes: Promise<ResolvedTarget | null>[] = [];
+  for (let slot = 0; slot < MAX_ACCOUNT_SLOTS; slot++) {
+    const ctx = readStalwartAuthContextFromStore(cookieStore, slot);
+    if (!ctx) continue;
+    const serverUrl = ctx.serverUrl.replace(/\/+$/, '');
+    probes.push(
+      (async () => {
+        try {
+          const res = await fetch(`${serverUrl}/.well-known/jmap`, {
+            headers: { Authorization: ctx.authHeader },
+          });
+          if (!res.ok) return null;
+          const session = (await res.json()) as {
+            apiUrl?: string;
+            primaryAccounts?: Record<string, string>;
+          };
+          const mailAccountId = session.primaryAccounts?.['urn:ietf:params:jmap:mail'];
+          if (!session.apiUrl || !mailAccountId) return null;
+          if (mailAccountId !== accountId) return null;
+          return { authHeader: ctx.authHeader, apiUrl: session.apiUrl, accountId: mailAccountId };
+        } catch {
+          return null;
+        }
+      })(),
+    );
+  }
+  const results = await Promise.all(probes);
+  return results.find((r): r is ResolvedTarget => r !== null) ?? null;
+}
+
+async function resolveDefaultTarget(creds: StalwartCredentials): Promise<ResolvedTarget | null> {
+  const sessionRes = await fetch(`${creds.serverUrl}/.well-known/jmap`, {
+    headers: { Authorization: creds.authHeader },
+  });
+  if (!sessionRes.ok) return null;
+  const session = (await sessionRes.json()) as {
+    apiUrl?: string;
+    primaryAccounts?: Record<string, string>;
+  };
+  const apiUrl = session.apiUrl;
+  const accountId = session.primaryAccounts?.['urn:ietf:params:jmap:mail'];
+  if (!apiUrl || !accountId) return null;
+  return { authHeader: creds.authHeader, apiUrl, accountId };
+}
+
 /**
  * GET /api/push/preview
  *
@@ -19,31 +82,38 @@ export const dynamic = 'force-dynamic';
  */
 export async function GET(request: NextRequest) {
   try {
-    const creds = await getStalwartCredentials(request);
+    // SW passes ?accountId=<jmap-account-id> derived from the push payload's
-    if (!creds) {
+    // StateChange so multi-account browsers fetch from the right slot. Older
-      return NextResponse.json({ error: 'Not authenticated' }, { status: 401 });
+    // clients (and the manual /api/push/preview probe) omit it and fall back
+    // to the first signed-in slot.
+    const requestedAccountId = request.nextUrl.searchParams.get('accountId');
+
+    let target: ResolvedTarget | null = null;
+    let authHeader: string;
+    if (requestedAccountId) {
+      target = await resolveTargetForAccount(requestedAccountId);
+      if (!target) {
+        return NextResponse.json({ error: 'Not authenticated' }, { status: 401 });
+      }
+      authHeader = target.authHeader;
+    } else {
+      const creds = await getStalwartCredentials(request);
+      if (!creds) {
+        return NextResponse.json({ error: 'Not authenticated' }, { status: 401 });
+      }
+      target = await resolveDefaultTarget(creds);
+      if (!target) {
+        return NextResponse.json({ error: 'JMAP session failed' }, { status: 502 });
+      }
+      authHeader = creds.authHeader;
     }
 
-    const sessionRes = await fetch(`${creds.serverUrl}/.well-known/jmap`, {
+    const { apiUrl, accountId } = target;
-      headers: { Authorization: creds.authHeader },
-    });
-    if (!sessionRes.ok) {
-      return NextResponse.json({ error: 'JMAP session failed' }, { status: 502 });
-    }
-    const session = (await sessionRes.json()) as {
-      apiUrl?: string;
-      primaryAccounts?: Record<string, string>;
-    };
-    const apiUrl = session.apiUrl;
-    const accountId = session.primaryAccounts?.['urn:ietf:params:jmap:mail'];
-    if (!apiUrl || !accountId) {
-      return NextResponse.json({ error: 'Incomplete JMAP session' }, { status: 502 });
-    }
 
     const inboxRes = await fetch(apiUrl, {
       method: 'POST',
       headers: {
-        Authorization: creds.authHeader,
+        Authorization: authHeader,
         'Content-Type': 'application/json',
       },
       body: JSON.stringify({
@@ -119,7 +189,7 @@ export async function GET(request: NextRequest) {
     const jmapRes = await fetch(apiUrl, {
       method: 'POST',
       headers: {
-        Authorization: creds.authHeader,
+        Authorization: authHeader,
         'Content-Type': 'application/json',
       },
       body: JSON.stringify(requestBody),

components/settings/notification-settings.tsx

@@ -52,11 +52,14 @@ export function NotificationSettings() {
 
   useEffect(() => {
     if (!supported) return;
+    if (!client) return;
+    const accountId = client.getAccountId();
+    if (!accountId) return;
     void (async () => {
-      const enabled = await isWebPushEnabled();
+      const enabled = await isWebPushEnabled(accountId);
-      if (enabled) setPushStatus({ kind: 'enabled' });
+      setPushStatus(enabled ? { kind: 'enabled' } : { kind: 'idle' });
     })();
-  }, [supported]);
+  }, [supported, client]);
 
   const trimmedRelay = relayUrl.trim().replace(/\/+$/, '');
   const isValidRelay = /^https?:\/\/.+/i.test(trimmedRelay);

lib/web-push.ts

@@ -6,8 +6,21 @@
 
 import type { IJMAPClient } from '@/lib/jmap/client-interface';
 
-const DEVICE_CLIENT_ID_KEY = 'bulwark.push.deviceClientId.v1';
+// Per-account keys: a single browser may be signed in to multiple accounts,
-const SUBSCRIPTION_ID_KEY = 'bulwark.push.subscriptionId.v1';
+// each with its own JMAP PushSubscription and its own relay record. Scoping
+// the deviceClientId per account is what makes per-account notifications work
+// at all - the relay keys subscriptions on subscriptionId (= deviceClientId),
+// so a globally-shared key meant re-registering account B overwrote A.
+const DEVICE_CLIENT_ID_PREFIX = 'bulwark.push.deviceClientId.v1.';
+const SUBSCRIPTION_ID_PREFIX = 'bulwark.push.subscriptionId.v1.';
+
+function deviceClientIdKey(accountId: string): string {
+  return DEVICE_CLIENT_ID_PREFIX + accountId;
+}
+
+function subscriptionIdKey(accountId: string): string {
+  return SUBSCRIPTION_ID_PREFIX + accountId;
+}
 
 const BASE_PATH = (process.env.NEXT_PUBLIC_BASE_PATH ?? '').replace(/\/+$/, '');
 const SW_SCOPE = `${BASE_PATH}/`;
@@ -79,14 +92,24 @@ function randomDeviceClientId(): string {
   return Array.from(bytes, (b) => b.toString(16).padStart(2, '0')).join('');
 }
 
-function getOrCreateDeviceClientId(): string {
+function getOrCreateDeviceClientId(accountId: string): string {
-  const existing = localStorage.getItem(DEVICE_CLIENT_ID_KEY);
+  const key = deviceClientIdKey(accountId);
+  const existing = localStorage.getItem(key);
   if (existing) return existing;
   const next = randomDeviceClientId();
-  localStorage.setItem(DEVICE_CLIENT_ID_KEY, next);
+  localStorage.setItem(key, next);
   return next;
 }
 
+function anyOtherAccountHasSubscription(accountId: string): boolean {
+  const skip = subscriptionIdKey(accountId);
+  for (let i = 0; i < localStorage.length; i++) {
+    const k = localStorage.key(i);
+    if (k && k !== skip && k.startsWith(SUBSCRIPTION_ID_PREFIX)) return true;
+  }
+  return false;
+}
+
 // PushManager.subscribe wants the VAPID public key as a BufferSource.
 // Returning a Uint8Array<ArrayBuffer> (not the wider ArrayBufferLike that
 // includes SharedArrayBuffer) keeps strict TS happy on lib.dom 2024+.
@@ -260,7 +283,8 @@ export async function enableWebPush(
     });
   }
 
-  const deviceClientId = getOrCreateDeviceClientId();
+  const accountId = params.client.getAccountId();
+  const deviceClientId = getOrCreateDeviceClientId(accountId);
 
   await registerWithRelay({
     relayBaseUrl,
@@ -278,7 +302,8 @@ export async function enableWebPush(
   // Reuse the JMAP-side PushSubscription if the server still has it, just
   // refreshing the expiry so it doesn't time out between sessions.
   const existingSubs = await params.client.listPushSubscriptions().catch(() => []);
-  const storedServerId = localStorage.getItem(SUBSCRIPTION_ID_KEY);
+  const subIdKey = subscriptionIdKey(accountId);
+  const storedServerId = localStorage.getItem(subIdKey);
   if (storedServerId) {
     const match = existingSubs.find((s) => s.id === storedServerId);
     if (match) {
@@ -286,7 +311,7 @@ export async function enableWebPush(
       if (refreshed) return { subscriptionId: storedServerId };
       await params.client.destroyPushSubscription(storedServerId).catch(() => undefined);
     }
-    localStorage.removeItem(SUBSCRIPTION_ID_KEY);
+    localStorage.removeItem(subIdKey);
   }
 
   // Reap any leftover subscriptions still bound to this device. These pile
@@ -309,7 +334,7 @@ export async function enableWebPush(
 
   const verificationCode = await pollVerificationCode(relayBaseUrl, deviceClientId);
   await params.client.verifyPushSubscription(serverAssignedId, verificationCode);
-  localStorage.setItem(SUBSCRIPTION_ID_KEY, serverAssignedId);
+  localStorage.setItem(subIdKey, serverAssignedId);
 
   return { subscriptionId: serverAssignedId };
 }
@@ -320,37 +345,50 @@ export interface DisableWebPushParams {
 }
 
 // Best-effort teardown: clear the JMAP subscription, the relay mapping, and
-// the browser PushSubscription. Any single failure is swallowed so the user
+// (only when no other accounts still need it) the browser-wide
-// always ends up in a "disabled" state locally.
+// PushSubscription. Any single failure is swallowed so the user always ends
+// up in a "disabled" state locally.
 export async function disableWebPush(params: DisableWebPushParams): Promise<void> {
   const relayBaseUrl = (params.relayBaseUrl ?? DEFAULT_RELAY_BASE_URL).replace(/\/+$/, '');
+  const accountId = params.client.getAccountId();
 
-  const storedServerId = localStorage.getItem(SUBSCRIPTION_ID_KEY);
+  const subIdKey = subscriptionIdKey(accountId);
+  const devIdKey = deviceClientIdKey(accountId);
+
+  const storedServerId = localStorage.getItem(subIdKey);
   if (storedServerId) {
     await params.client.destroyPushSubscription(storedServerId).catch(() => undefined);
-    localStorage.removeItem(SUBSCRIPTION_ID_KEY);
+    localStorage.removeItem(subIdKey);
   }
 
-  const deviceClientId = localStorage.getItem(DEVICE_CLIENT_ID_KEY);
+  const deviceClientId = localStorage.getItem(devIdKey);
   if (deviceClientId && relayBaseUrl) {
     await fetch(
       buildRelayUrl(relayBaseUrl, `/api/push/register/${encodeURIComponent(deviceClientId)}`),
       { method: 'DELETE' },
     ).catch(() => undefined);
   }
+  // Keep the deviceClientId around so a later re-enable for this account
+  // reuses the same relay subscriptionId rather than scattering orphans.
 
-  if (typeof navigator !== 'undefined' && 'serviceWorker' in navigator) {
+  // The browser-wide PushSubscription is shared by every account on this
+  // origin, so only tear it down if no other account is still using it.
+  if (
+    !anyOtherAccountHasSubscription(accountId)
+    && typeof navigator !== 'undefined'
+    && 'serviceWorker' in navigator
+  ) {
     const registration = await navigator.serviceWorker.getRegistration(SW_SCOPE);
     const sub = await registration?.pushManager.getSubscription();
     if (sub) await sub.unsubscribe().catch(() => undefined);
   }
 }
 
-export async function isWebPushEnabled(): Promise<boolean> {
+export async function isWebPushEnabled(accountId: string): Promise<boolean> {
   if (!isWebPushSupported()) return false;
   if (Notification.permission !== 'granted') return false;
   const registration = await navigator.serviceWorker.getRegistration(SW_SCOPE);
   if (!registration) return false;
   const sub = await registration.pushManager.getSubscription();
-  return sub !== null && localStorage.getItem(SUBSCRIPTION_ID_KEY) !== null;
+  return sub !== null && localStorage.getItem(subscriptionIdKey(accountId)) !== null;
 }

public/sw.js

@@ -98,6 +98,16 @@ async function handlePush(event) {
     ? payload.accountLabel
     : "";
 
+  // JMAP StateChange wraps changes in { changed: { [accountId]: {...} } }.
+  // The relay forwards a single account's StateChange per push, so the first
+  // key is the one this notification is for. Without this the preview API
+  // would just fall back to the first signed-in slot and surface mail from
+  // the wrong account.
+  const changed = payload && payload.changed && typeof payload.changed === "object"
+    ? payload.changed
+    : null;
+  const accountId = changed ? Object.keys(changed)[0] || "" : "";
+
   // Best effort: ask the webmail to look up the latest unread email so we can
   // build a useful notification. If the request fails (offline, session
   // expired, server down) we fall back to a generic "New mail" so the user
@@ -105,7 +115,10 @@ async function handlePush(event) {
   let preview = null;
   let previewOk = false;
   try {
-    const res = await fetch(`${BASE_PATH}/api/push/preview`, {
+    const previewUrl = accountId
+      ? `${BASE_PATH}/api/push/preview?accountId=${encodeURIComponent(accountId)}`
+      : `${BASE_PATH}/api/push/preview`;
+    const res = await fetch(previewUrl, {
       credentials: "include",
       cache: "no-store",
     });