fix: stop a refresh after sign-out from writing the removed account's tokens back

src/api/unified-inbox.ts

@@ -6,7 +6,7 @@ import {
 } from './jmap-client';
 import { keywordPointer, mailboxPointer } from './patch-pointer';
 import { secureFetch } from '../lib/client-cert';
-import { refreshOAuthAccessToken, type OAuthTokens } from '../lib/oauth';
+import { freshStoredCredentials } from '../lib/stored-token-refresh';
 import { textConditions } from '../lib/search-utils';
 
 // Aggregated views across accounts ("All inboxes", "All Sent", All mail /
@@ -23,7 +23,6 @@ const EMAIL_LIST_PROPERTIES = [
   'receivedAt', 'from', 'to', 'cc', 'subject', 'preview', 'hasAttachment',
 ];
 
-const TOKEN_REFRESH_LEEWAY_MS = 60_000;
 const SESSION_CACHE_TTL_MS = 10 * 60 * 1000;
 const MAILBOX_CACHE_TTL_MS = 60 * 1000;
 
@@ -166,51 +165,6 @@ function sharedMailAccountsOf(session: JMAPSession, primaryId: string): { id: st
   return out;
 }
 
-// Refresh an about-to-expire OAuth token and persist it so the next fetch (and
-// a later account switch) reuse the fresh one. Returns possibly-updated creds.
-async function ensureFreshCredentials(
-  accountId: string,
-  creds: StoredCredentials,
-): Promise<StoredCredentials> {
-  if (
-    !creds.accessToken ||
-    !creds.refreshToken ||
-    !creds.tokenEndpoint ||
-    !creds.clientId ||
-    creds.expiresAt == null
-  ) {
-    return creds;
-  }
-  if (creds.expiresAt - Date.now() > TOKEN_REFRESH_LEEWAY_MS) return creds;
-  try {
-    const tokens: OAuthTokens = {
-      accessToken: creds.accessToken,
-      refreshToken: creds.refreshToken,
-      expiresAt: creds.expiresAt,
-      tokenEndpoint: creds.tokenEndpoint,
-      clientId: creds.clientId,
-    };
-    const next = await refreshOAuthAccessToken(tokens);
-    const updated: StoredCredentials = {
-      ...creds,
-      accessToken: next.accessToken,
-      refreshToken: next.refreshToken ?? creds.refreshToken,
-      expiresAt: next.expiresAt,
-      tokenEndpoint: next.tokenEndpoint,
-      clientId: next.clientId,
-    };
-    const current = await jmapClient.getStoredCredentials(accountId);
-    if (!current || current.accessToken !== updated.accessToken) {
-      await jmapClient.setStoredCredentials(accountId, updated);
-    }
-    return updated;
-  } catch {
-    // Fall back to the existing (possibly expired) token; the request will
-    // surface a 401 which we report as an error for this account.
-    return creds;
-  }
-}
-
 // ── Per-account session + mailbox cache ────────────────────────────────
 
 interface AccountEntry {
@@ -266,7 +220,7 @@ async function entryFor(accountId: string): Promise<AccountEntry> {
 
   let creds = await jmapClient.getStoredCredentials(accountId);
   if (!creds) throw new Error('No stored credentials');
-  creds = await ensureFreshCredentials(accountId, creds);
+  creds = await freshStoredCredentials(accountId, creds);
   const baseUrl = creds.serverUrl.replace(/\/+$/, '');
   const authHeader = authHeaderFor(creds);
 

src/lib/__tests__/oauth.test.ts

@@ -8,7 +8,7 @@ vi.mock('../client-cert', () => ({
   secureFetch: mockSecureFetch,
 }));
 
-import { refreshOAuthAccessToken, resetRefreshRotations, type OAuthTokens } from '../oauth';
+import { forgetRefreshRotations, refreshOAuthAccessToken, resetRefreshRotations, type OAuthTokens } from '../oauth';
 
 describe('oauth token refresh deduplication', () => {
   beforeEach(() => {
@@ -89,6 +89,42 @@ describe('oauth token refresh deduplication', () => {
     expect(res.refreshToken).toBe('r3');
   });
 
+  it('forgets the whole chain of a signed-out account, from any of its tokens', async () => {
+    const tokens: OAuthTokens = {
+      accessToken: 'a1',
+      refreshToken: 'r1',
+      tokenEndpoint: 'https://auth.example.com/token',
+      clientId: 'swc1.x',
+    };
+    let generation = 1;
+    mockSecureFetch.mockImplementation(async () => {
+      generation += 1;
+      return {
+        ok: true,
+        status: 200,
+        json: async () => ({ access_token: `a${generation}`, refresh_token: `r${generation}`, expires_in: 3600 }),
+      } as Response;
+    });
+    // Another account's chain: q1 -> r2.
+    const other = await refreshOAuthAccessToken({ ...tokens, accessToken: 'b1', refreshToken: 'q1' });
+    // This account's: r1 -> r3 -> r4.
+    await refreshOAuthAccessToken(tokens);
+    await refreshOAuthAccessToken({ ...tokens, accessToken: 'a3', refreshToken: 'r3', expiresAt: Date.now() });
+
+    // Sign-out reads the stored bundle, which may sit anywhere in the chain.
+    forgetRefreshRotations('r3');
+    mockSecureFetch.mockClear();
+    mockSecureFetch.mockResolvedValue({ ok: false, status: 400 } as Response);
+
+    // A stale copy no longer gets a successor out of memory.
+    await expect(refreshOAuthAccessToken(tokens)).rejects.toThrow('400');
+    expect(mockSecureFetch).toHaveBeenCalledTimes(1);
+    // The other account's chain stays.
+    mockSecureFetch.mockClear();
+    expect((await refreshOAuthAccessToken({ ...tokens, accessToken: 'b1', refreshToken: 'q1' })).accessToken).toBe(other.accessToken);
+    expect(mockSecureFetch).not.toHaveBeenCalled();
+  });
+
   it('keeps refreshing with the same token on a server that does not rotate', async () => {
     const tokens: OAuthTokens = {
       accessToken: 'a1',

src/lib/__tests__/stored-token-refresh.test.ts

@@ -0,0 +1,126 @@
+import { describe, it, expect, vi, beforeEach } from 'vitest';
+
+const { mockSecureFetch, store } = vi.hoisted(() => ({
+  mockSecureFetch: vi.fn(),
+  store: new Map<string, unknown>(),
+}));
+
+vi.mock('../client-cert', () => ({ secureFetch: mockSecureFetch }));
+vi.mock('../../api/jmap-client', () => ({
+  jmapClient: {
+    getStoredCredentials: vi.fn(async (id: string) => (store.has(id) ? structuredClone(store.get(id)) : null)),
+    setStoredCredentials: vi.fn(async (id: string, creds: unknown) => { store.set(id, structuredClone(creds)); }),
+  },
+}));
+
+import { jmapClient, type StoredCredentials } from '../../api/jmap-client';
+import { resetRefreshRotations } from '../oauth';
+import { freshStoredCredentials } from '../stored-token-refresh';
+
+const ID = 'u@example.org@https://mail.example.com';
+
+function bundle(generation: number, expiresAt = Date.now() + 5_000): StoredCredentials {
+  return {
+    serverUrl: 'https://mail.example.com',
+    username: 'u@example.org',
+    password: '',
+    accessToken: `a${generation}`,
+    refreshToken: `r${generation}`,
+    expiresAt,
+    tokenEndpoint: 'https://mail.example.com/auth/token',
+    clientId: 'swc1.x',
+    tokenSource: 'totp',
+  };
+}
+
+// A token endpoint that rotates: r<n> buys a<n+1>/r<n+1>, once.
+function rotatingServer() {
+  const spent = new Set<string>();
+  mockSecureFetch.mockImplementation(async (_url: string, init: RequestInit) => {
+    const presented = new URLSearchParams(String(init.body)).get('refresh_token')!;
+    if (spent.has(presented)) return { ok: false, status: 400, json: async () => ({ error: 'invalid_grant' }) } as Response;
+    spent.add(presented);
+    const n = Number(presented.slice(1)) + 1;
+    return {
+      ok: true,
+      status: 200,
+      json: async () => ({ access_token: `a${n}`, refresh_token: `r${n}`, expires_in: 3600 }),
+    } as Response;
+  });
+}
+
+const presented = () => mockSecureFetch.mock.calls.map(([, init]) => new URLSearchParams(String(init.body)).get('refresh_token'));
+
+describe('freshStoredCredentials', () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+    store.clear();
+    resetRefreshRotations();
+    rotatingServer();
+  });
+
+  it('refreshes an expiring bundle and stores the successor', async () => {
+    store.set(ID, bundle(1));
+    const next = await freshStoredCredentials(ID, bundle(1));
+
+    expect(next.accessToken).toBe('a2');
+    expect((store.get(ID) as StoredCredentials).refreshToken).toBe('r2');
+  });
+
+  it('takes the newer stored bundle instead of spending a replaced refresh token', async () => {
+    // The live client refreshed meanwhile and stored r2.
+    store.set(ID, bundle(2, Date.now() + 3_600_000));
+    const next = await freshStoredCredentials(ID, bundle(1));
+
+    expect(next.accessToken).toBe('a2');
+    expect(mockSecureFetch).not.toHaveBeenCalled();
+  });
+
+  it('shares one refresh with a holder refreshing at the same moment', async () => {
+    store.set(ID, bundle(1));
+    const [a, b] = await Promise.all([freshStoredCredentials(ID, bundle(1)), freshStoredCredentials(ID, bundle(1))]);
+
+    expect(presented()).toEqual(['r1']);
+    expect(a.accessToken).toBe('a2');
+    expect(b.accessToken).toBe('a2');
+  });
+
+  it('writes nothing back for an account that signed out', async () => {
+    // A copy taken before the sign-out; the live client had rotated r1 -> r2,
+    // which the rotation memory still knows.
+    store.set(ID, bundle(1));
+    await freshStoredCredentials(ID, bundle(1));
+    store.delete(ID);
+    vi.mocked(jmapClient.setStoredCredentials).mockClear();
+
+    await freshStoredCredentials(ID, bundle(1));
+
+    expect(store.has(ID)).toBe(false);
+    expect(jmapClient.setStoredCredentials).not.toHaveBeenCalled();
+  });
+
+  it('does not overwrite a newer bundle stored while the refresh ran', async () => {
+    store.set(ID, bundle(1));
+    // The iOS notification extension rotates the grant meanwhile (it can't
+    // share this process's memory) and stores its own bundle.
+    const newer = { ...bundle(7, Date.now() + 3_600_000) };
+    mockSecureFetch.mockImplementationOnce(async () => {
+      store.set(ID, newer);
+      return {
+        ok: true,
+        status: 200,
+        json: async () => ({ access_token: 'a2', refresh_token: 'r2', expires_in: 3600 }),
+      } as Response;
+    });
+
+    await freshStoredCredentials(ID, bundle(1));
+
+    expect((store.get(ID) as StoredCredentials).refreshToken).toBe('r7');
+  });
+
+  it('leaves password credentials alone', async () => {
+    const basic: StoredCredentials = { serverUrl: 'https://mail.example.com', username: 'u', password: 'pw' };
+    expect(await freshStoredCredentials(ID, basic)).toBe(basic);
+    expect(jmapClient.getStoredCredentials).not.toHaveBeenCalled();
+  });
+});

src/lib/oauth.ts

@@ -574,6 +574,28 @@ export function resetRefreshRotations(): void {
   rotatedTo.clear();
 }
 
+/**
+ * Drop the rotation chain `refreshToken` is part of, its predecessors and
+ * successors both, when its account signs out. A stale copy of the bundle
+ * refreshed later would otherwise get the successor from here and could
+ * write the removed account's tokens back.
+ */
+export function forgetRefreshRotations(refreshToken: string | undefined): void {
+  if (!refreshToken) return;
+  const chain = new Set<string>([refreshToken]);
+  let grew = true;
+  while (grew) {
+    grew = false;
+    for (const [previous, next] of rotatedTo) {
+      const linked = chain.has(previous) || (!!next.refreshToken && chain.has(next.refreshToken));
+      if (!linked) continue;
+      if (!chain.has(previous)) { chain.add(previous); grew = true; }
+      if (next.refreshToken && !chain.has(next.refreshToken)) { chain.add(next.refreshToken); grew = true; }
+    }
+  }
+  for (const token of chain) rotatedTo.delete(token);
+}
+
 // OAuth refresh — exchanges the refresh token at the original token endpoint
 // for a new access token. Returns the updated bundle so the caller can
 // persist it.

src/lib/push-background-task.ts

@@ -5,7 +5,7 @@ import { CAPABILITIES } from '../api/types';
 import type { Email, JMAPMethodCall, JMAPSession, Mailbox } from '../api/types';
 import { detectDeviceLocale, isSupportedLocale, translate, type LocaleCode } from '../i18n';
 import { secureFetch } from './client-cert';
-import { refreshOAuthAccessToken, type OAuthTokens } from './oauth';
+import { freshStoredCredentials } from './stored-token-refresh';
 import {
   generateEmailAvatarColor,
   getEmailInitials,
@@ -38,7 +38,6 @@ const NOTIFIED_IDS_LIMIT = 200;
 // Legacy state-change payloads carry no ids; look at this many newest unread
 // inbox messages and notify the ones we have not shown yet.
 const LEGACY_QUERY_LIMIT = 5;
-const TOKEN_REFRESH_LEEWAY_MS = 60_000;
 
 const EMAIL_PROPERTIES = [
   'id', 'threadId', 'mailboxIds', 'keywords', 'size', 'receivedAt', 'from', 'subject', 'preview',
@@ -269,48 +268,10 @@ function rewriteApiUrl(session: JMAPSession, serverUrl: string): string {
   return serverOrigin + session.apiUrl.slice(apiOrigin.length);
 }
 
-async function ensureFreshCredentials(
-  accountId: string,
-  creds: StoredCredentials,
-): Promise<StoredCredentials> {
-  if (
-    !creds.accessToken || !creds.refreshToken || !creds.tokenEndpoint
-    || !creds.clientId || creds.expiresAt == null
-  ) {
-    return creds;
-  }
-  if (creds.expiresAt - Date.now() > TOKEN_REFRESH_LEEWAY_MS) return creds;
-  try {
-    const tokens: OAuthTokens = {
-      accessToken: creds.accessToken,
-      refreshToken: creds.refreshToken,
-      expiresAt: creds.expiresAt,
-      tokenEndpoint: creds.tokenEndpoint,
-      clientId: creds.clientId,
-    };
-    const next = await refreshOAuthAccessToken(tokens);
-    const updated: StoredCredentials = {
-      ...creds,
-      accessToken: next.accessToken,
-      refreshToken: next.refreshToken ?? creds.refreshToken,
-      expiresAt: next.expiresAt,
-      tokenEndpoint: next.tokenEndpoint,
-      clientId: next.clientId,
-    };
-    const current = await jmapClient.getStoredCredentials(accountId);
-    if (!current || current.accessToken !== updated.accessToken) {
-      await jmapClient.setStoredCredentials(accountId, updated);
-    }
-    return updated;
-  } catch {
-    return creds;
-  }
-}
-
 async function openDetachedSession(accountId: string): Promise<DetachedSession | null> {
   let creds = await jmapClient.getStoredCredentials(accountId);
   if (!creds) return null;
-  creds = await ensureFreshCredentials(accountId, creds);
+  creds = await freshStoredCredentials(accountId, creds);
   const baseUrl = creds.serverUrl.replace(/\/+$/, '');
   const authHeader = authHeaderFor(creds);
   const res = await secureFetch(`${baseUrl}/.well-known/jmap`, {

src/lib/stored-token-refresh.ts

@@ -0,0 +1,85 @@
+// Token refresh for the holders that work from their own copy of an account's
+// stored credentials instead of the live client: the unified inbox and the
+// background push task. The live client keeps its own (jmap-client
+// ensureFreshToken / persistRefreshedTokens) and follows the same rules.
+//
+// Stalwart 1.0 rotates refresh tokens and revokes the whole grant when a
+// replaced one comes back, so a copy that is behind must not be spent:
+// - before refreshing, the stored bundle is read again and taken when it is
+//   newer (the live client or the iOS notification extension refreshed
+//   meanwhile);
+// - the refresh goes through refreshOAuthAccessToken, whose in-process
+//   memory shares one request between holders refreshing at once and hands
+//   the successor to one that comes late;
+// - the result is written back only over the bundle that was refreshed: not
+//   when the account signed out meanwhile, and not over a newer bundle
+//   someone else wrote.
+
+import { jmapClient, type StoredCredentials } from '../api/jmap-client';
+import { refreshOAuthAccessToken } from './oauth';
+
+const TOKEN_REFRESH_LEEWAY_MS = 60_000;
+
+function expiresSoon(expiresAt: number | undefined): boolean {
+  return expiresAt != null && expiresAt - Date.now() <= TOKEN_REFRESH_LEEWAY_MS;
+}
+
+/**
+ * `creds` with an access token that is not about to expire, refreshed and
+ * stored when needed. Password credentials come back as they are; so does
+ * an account whose credentials are gone from storage (signed out). A failed
+ * refresh returns what there is, and the request that uses it fails on its
+ * own.
+ */
+export async function freshStoredCredentials(
+  accountId: string,
+  creds: StoredCredentials,
+): Promise<StoredCredentials> {
+  if (
+    !creds.accessToken || !creds.refreshToken || !creds.tokenEndpoint
+    || !creds.clientId || creds.expiresAt == null
+  ) {
+    return creds;
+  }
+  if (!expiresSoon(creds.expiresAt)) return creds;
+
+  const stored = await jmapClient.getStoredCredentials(accountId).catch(() => null);
+  if (!stored?.accessToken || !stored.refreshToken) return creds;
+  let base = creds;
+  if (stored.accessToken !== creds.accessToken && (stored.expiresAt ?? 0) > creds.expiresAt) {
+    base = {
+      ...creds,
+      accessToken: stored.accessToken,
+      refreshToken: stored.refreshToken,
+      expiresAt: stored.expiresAt,
+      tokenEndpoint: stored.tokenEndpoint ?? creds.tokenEndpoint,
+      clientId: stored.clientId ?? creds.clientId,
+    };
+    if (!expiresSoon(base.expiresAt)) return base;
+  }
+
+  try {
+    const next = await refreshOAuthAccessToken({
+      accessToken: base.accessToken!,
+      refreshToken: base.refreshToken!,
+      expiresAt: base.expiresAt,
+      tokenEndpoint: base.tokenEndpoint!,
+      clientId: base.clientId!,
+      source: base.tokenSource,
+    });
+    const tokens = {
+      accessToken: next.accessToken,
+      refreshToken: next.refreshToken ?? base.refreshToken,
+      expiresAt: next.expiresAt,
+      tokenEndpoint: next.tokenEndpoint,
+      clientId: next.clientId,
+    };
+    const current = await jmapClient.getStoredCredentials(accountId);
+    if (current && current.refreshToken === base.refreshToken && current.accessToken !== tokens.accessToken) {
+      await jmapClient.setStoredCredentials(accountId, { ...current, ...tokens });
+    }
+    return { ...base, ...tokens };
+  } catch {
+    return base;
+  }
+}

src/stores/__tests__/auth-store-logout.test.ts

@@ -78,6 +78,7 @@ vi.mock('../../api/jmap-client', () => {
       logout: vi.fn(async () => undefined),
       consumeLegacyCredentials: vi.fn(async () => null),
       clearAccountCredentials: vi.fn(async () => undefined),
+      clearAllCredentials: vi.fn(async () => undefined),
       getStoredOAuthTokens: vi.fn(async () => null),
       onAuthFailure: vi.fn(() => () => undefined),
       onTokenRefresh: vi.fn(() => () => undefined),
@@ -110,9 +111,15 @@ vi.mock('../../lib/push-notifications', () => ({
   forgetPushSubscriptionExpiry: vi.fn(async () => undefined),
 }));
 
+vi.mock('../../lib/oauth', async (importOriginal) => {
+  const actual = await importOriginal<typeof import('../../lib/oauth')>();
+  return { ...actual, forgetRefreshRotations: vi.fn(actual.forgetRefreshRotations) };
+});
+
 import AsyncStorage from '@react-native-async-storage/async-storage';
 import type { JMAPSession } from '../../api/types';
-import { AuthenticationError, NetworkError } from '../../api/jmap-client';
+import { AuthenticationError, NetworkError, jmapClient } from '../../api/jmap-client';
+import { forgetRefreshRotations } from '../../lib/oauth';
 import { sessionFingerprint, writeServerFingerprint } from '../../lib/server-fingerprint';
 import { useAccountStore } from '../account-store';
 import { useAuthStore } from '../auth-store';
@@ -208,3 +215,33 @@ describe('signing out of the active account (Stalwart 1.0)', () => {
     expect(useAccountStore.getState().accounts.map((a) => a.id)).toEqual([C]);
   });
 });
+
+describe('what a sign-out leaves behind', () => {
+  // Pairing bundles aren't revoked, so no request goes out.
+  const tokensOf = (id: string) => ({
+    accessToken: `at-${id}`,
+    refreshToken: `rt-${id}`,
+    tokenEndpoint: 'https://mail.example.org/auth/token',
+    clientId: 'swc1.x',
+    source: 'pairing' as const,
+  });
+
+  beforeEach(() => {
+    vi.mocked(forgetRefreshRotations).mockClear();
+    vi.mocked(jmapClient.getStoredOAuthTokens).mockImplementation(async (id: string) => tokensOf(id));
+  });
+
+  it("drops the account's token rotations on each kind of sign-out", async () => {
+    await signedIn(A, B, C);
+    await useAuthStore.getState().removeAccount(C);
+    await settle();
+    expect(vi.mocked(forgetRefreshRotations).mock.calls).toEqual([[`rt-${C}`]]);
+
+    await useAuthStore.getState().logout();
+    await settle();
+    expect(vi.mocked(forgetRefreshRotations).mock.calls.at(-1)).toEqual([`rt-${A}`]);
+
+    await useAuthStore.getState().logoutAll();
+    expect(vi.mocked(forgetRefreshRotations).mock.calls.at(-1)).toEqual([`rt-${B}`]);
+  });
+});

src/stores/auth-store.ts

@@ -27,6 +27,7 @@ import {
   HandoffCancelledError,
   HandoffError,
   PairingError,
+  forgetRefreshRotations,
   type HandoffResult,
 } from '../lib/oauth';
 import { discoverOAuthMetadata, loginWithPkce, probeWebmail, revokeRefreshToken } from '../lib/oauth-native';
@@ -371,11 +372,14 @@ function refuseAddWhenFull(set: (partial: Partial<AuthState>) => void, opts?: {
 // Best-effort RFC 7009 revocation of an account's refresh token on sign-out.
 // QR-paired bundles are left alone: webmail up to 1.11 handed the phone the
 // desktop's own refresh token, and newer ones hand out a separate grant whose
-// refresh token only the webmail's token proxy understands.
+// refresh token only the webmail's token proxy understands. The account's
+// chain also leaves the in-process rotation memory either way, so a stale copy
+// refreshed after the sign-out can't pick up its tokens from there.
 async function revokeStoredRefreshToken(accountId: string): Promise<void> {
   try {
     const entry = useAccountStore.getState().getAccountById(accountId);
     const tokens = await jmapClient.getStoredOAuthTokens(accountId);
+    forgetRefreshRotations(tokens?.refreshToken);
     if (!entry || !tokens || tokens.source === 'pairing') return;
     await revokeRefreshToken(entry.serverUrl, tokens);
   } catch {