From b52768e6fbce0c755cec4aab9d2cb10e5821d504 Mon Sep 17 00:00:00 2001
From: Linus Rath <139418639+rathlinus@users.noreply.github.com>
Date: Sat, 29 Aug 2026 17:57:32 +0200
Subject: [PATCH] fix: keep the live session on a failed add or switch, recover
 from revoked sessions, refresh display names and sign in with a TOTP code
 (#900)

---
 src/components/QrScanModal.tsx          |   2 +-
 src/lib/login-errors.ts                 |  39 +++++++
 src/screens/LoginScreen.tsx             |  29 ++++-
 src/screens/login/PasswordStep.tsx      |  27 ++++-
 src/stores/__tests__/auth-store.test.ts |   6 +
 src/stores/auth-store.ts                | 149 ++++++++++++++++++++----
 6 files changed, 223 insertions(+), 29 deletions(-)

diff --git a/src/components/QrScanModal.tsx b/src/components/QrScanModal.tsx
index d50d08e0..c7f052a7 100644
--- a/src/components/QrScanModal.tsx
+++ b/src/components/QrScanModal.tsx
@@ -63,7 +63,7 @@ export function QrScanModal({ visible, onClose, onScanned }: QrScanModalProps) {
               <View style={styles.hintWrap}>
                 <Text style={styles.hintTitle}>Open Bulwark on your computer</Text>
                 <Text style={styles.hint}>
-                  Settings → Devices → Add phone shows a code. Point the camera at it.
+                  Settings → Security → Link device shows a code. Point the camera at it.
                 </Text>
               </View>
             </View>
diff --git a/src/lib/login-errors.ts b/src/lib/login-errors.ts
index 880b8889..595da995 100644
--- a/src/lib/login-errors.ts
+++ b/src/lib/login-errors.ts
@@ -36,6 +36,38 @@ export function describeLoginError(err: unknown, context: LoginErrorContext = {}
   const lower = message.toLowerCase();
   const host = hostLabel(context.serverUrl);
 
+  if (name === 'TotpRequiredError' || message === 'TOTP_REQUIRED' || lower.includes('two-factor code required')) {
+    return {
+      title: 'Enter your two-factor code',
+      detail: 'This account is protected with two-factor sign-in. Type the 6-digit code from your authenticator app.',
+    };
+  }
+
+  if (name === 'TotpLoginError' && lower.includes('invalid')) {
+    return {
+      title: "That didn't work",
+      detail: 'Check your password and the current code from your authenticator app, then try again.',
+    };
+  }
+
+  if (name === 'RateLimitError' || /\b429\b/.test(message) || lower.includes('rate limited')) {
+    const retryMs = (err as { retryAfterMs?: number } | null)?.retryAfterMs;
+    const seconds = typeof retryMs === 'number' && Number.isFinite(retryMs) ? Math.max(1, Math.round(retryMs / 1000)) : null;
+    return {
+      title: 'Too many attempts',
+      detail: seconds
+        ? `The server asked us to wait. Try again in about ${seconds} seconds.`
+        : 'The server asked us to wait a moment. Try again shortly.',
+    };
+  }
+
+  if (lower.includes('session discovery failed') && /\b402\b/.test(message)) {
+    return {
+      title: 'Enter your two-factor code',
+      detail: 'This account is protected with two-factor sign-in. Type the 6-digit code from your authenticator app.',
+    };
+  }
+
   if (name === 'AuthenticationError' || lower.includes('invalid username or password')) {
     return {
       title: "That didn't work",
@@ -76,6 +108,13 @@ export function describeLoginError(err: unknown, context: LoginErrorContext = {}
     };
   }
 
+  if (lower.includes('session discovery failed') && /\b5\d\d\b/.test(message)) {
+    return {
+      title: `${host} is having trouble`,
+      detail: 'The server answered with an error. Try again in a few minutes.',
+    };
+  }
+
   if (lower.includes('pairing code')) {
     return {
       title: 'That code has expired',
diff --git a/src/screens/LoginScreen.tsx b/src/screens/LoginScreen.tsx
index cf998cc4..e36a0d26 100644
--- a/src/screens/LoginScreen.tsx
+++ b/src/screens/LoginScreen.tsx
@@ -55,6 +55,10 @@ export default function LoginScreen({ onLogin, isAddMode = false, onCancel }: Lo
   // when discovery actually found one.
   const [serverDiscovered, setServerDiscovered] = React.useState(false);
   const [password, setPassword] = React.useState('');
+  // Second factor. Shown only after the server asked for it (402), so the
+  // common no-2FA sign-in stays a two-field form.
+  const [totp, setTotp] = React.useState('');
+  const [totpRequired, setTotpRequired] = React.useState(false);
   const [failedDomain, setFailedDomain] = React.useState<string | null>(null);
   const [notice, setNotice] = React.useState<LoginErrorCopy | null>(null);
   const [searching, setSearching] = React.useState(false);
@@ -198,18 +202,31 @@ export default function LoginScreen({ onLogin, isAddMode = false, onCancel }: Lo
       return;
     }
 
+    if (totpRequired && !/^\d{6,8}$/.test(totp.trim())) {
+      setNotice({ title: 'Enter the 6-digit code from your authenticator app' });
+      return;
+    }
+
     setNotice(null);
     setBusy('connecting');
     const wasAuthenticated = useAuthStore.getState().isAuthenticated;
     try {
-      await login(target, email.trim(), password, { addAccount: isAddMode });
+      await login(target, email.trim(), password, {
+        addAccount: isAddMode,
+        totp: totpRequired ? totp.trim() : undefined,
+      });
       finishIfSignedIn(wasAuthenticated);
     } catch (err) {
+      if (err instanceof Error && err.name === 'TotpRequiredError') {
+        setTotpRequired(true);
+        setNotice(describeLoginError(err, { serverUrl: target }));
+        return;
+      }
       setNotice(describeLoginError(err, { serverUrl: target }));
     } finally {
       setBusy(null);
     }
-  }, [email, finishIfSignedIn, isAddMode, login, password, serverInput, serverUrl]);
+  }, [email, finishIfSignedIn, isAddMode, login, password, serverInput, serverUrl, totp, totpRequired]);
 
   const handleScanned = React.useCallback(
     async (data: string) => {
@@ -218,7 +235,7 @@ export default function LoginScreen({ onLogin, isAddMode = false, onCancel }: Lo
       if (!payload) {
         setNotice({
           title: "That code isn't a Bulwark sign-in code",
-          detail: 'Open Bulwark on the web, then Settings → Devices to show one.',
+          detail: 'Open Bulwark on the web, then Settings → Security → Link device to show one.',
         });
         return;
       }
@@ -364,6 +381,12 @@ export default function LoginScreen({ onLogin, isAddMode = false, onCancel }: Lo
               setNotice(null);
               setPassword(value);
             }}
+            totp={totp}
+            totpRequired={totpRequired}
+            onChangeTotp={(value) => {
+              setNotice(null);
+              setTotp(value.replace(/\s+/g, ''));
+            }}
             onSubmit={() => void handlePasswordSubmit()}
             notice={notice}
           />
diff --git a/src/screens/login/PasswordStep.tsx b/src/screens/login/PasswordStep.tsx
index e334da36..f454b3f2 100644
--- a/src/screens/login/PasswordStep.tsx
+++ b/src/screens/login/PasswordStep.tsx
@@ -12,6 +12,10 @@ interface PasswordStepProps {
   password: string;
   onChangeEmail: (value: string) => void;
   onChangePassword: (value: string) => void;
+  /** Second factor, shown once the server answered 402 "MFA code required". */
+  totp?: string;
+  totpRequired?: boolean;
+  onChangeTotp?: (value: string) => void;
   onSubmit: () => void;
   notice?: { title: string; detail?: string } | null;
 }
@@ -28,6 +32,9 @@ export default function PasswordStep({
   password,
   onChangeEmail,
   onChangePassword,
+  totp = '',
+  totpRequired = false,
+  onChangeTotp,
   onSubmit,
   notice,
 }: PasswordStepProps) {
@@ -65,8 +72,8 @@ export default function PasswordStep({
         secureTextEntry={!showPassword}
         autoComplete="current-password"
         textContentType="password"
-        returnKeyType="go"
-        onSubmitEditing={onSubmit}
+        returnKeyType={totpRequired ? 'next' : 'go'}
+        onSubmitEditing={totpRequired ? undefined : onSubmit}
         rightIcon={
           <Pressable
             onPress={() => setShowPassword((v) => !v)}
@@ -79,6 +86,22 @@ export default function PasswordStep({
         }
       />
 
+      {totpRequired ? (
+        <Input
+          label="Two-factor code"
+          placeholder="123456"
+          value={totp}
+          onChangeText={(value) => onChangeTotp?.(value)}
+          autoFocus
+          keyboardType="number-pad"
+          autoComplete="one-time-code"
+          textContentType="oneTimeCode"
+          maxLength={8}
+          returnKeyType="go"
+          onSubmitEditing={onSubmit}
+        />
+      ) : null}
+
       {notice ? <LoginNotice title={notice.title} detail={notice.detail} /> : null}
 
       <Button variant="default" size="md" onPress={onSubmit}>
diff --git a/src/stores/__tests__/auth-store.test.ts b/src/stores/__tests__/auth-store.test.ts
index 918b200b..2c036059 100644
--- a/src/stores/__tests__/auth-store.test.ts
+++ b/src/stores/__tests__/auth-store.test.ts
@@ -10,6 +10,12 @@ vi.mock('../../api/jmap-client', () => ({
     clearAccountCredentials: vi.fn(async () => undefined),
     clearAllCredentials: vi.fn(async () => undefined),
     reset: vi.fn(),
+    snapshot: vi.fn(() => ({ session: null, credentials: null, accountId: null })),
+    restoreSnapshot: vi.fn(),
+    onAuthFailure: vi.fn(() => () => undefined),
+    onTokenRefresh: vi.fn(() => () => undefined),
+    hasAccountCapability: vi.fn(() => false),
+    request: vi.fn(async () => { throw new Error('not mocked'); }),
     accountId: 'acc-1',
     currentSession: { apiUrl: 'https://mail.example.com/jmap/' },
     username: 'user',
diff --git a/src/stores/auth-store.ts b/src/stores/auth-store.ts
index 11f3830f..dbc7460f 100644
--- a/src/stores/auth-store.ts
+++ b/src/stores/auth-store.ts
@@ -1,6 +1,8 @@
 import { create } from 'zustand';
 import { jmapClient, AuthenticationError, NetworkError } from '../api/jmap-client';
 import type { JMAPSession } from '../api/types';
+import { getIdentities } from '../api/identity';
+import { fetchPrincipal, isStalwartSupported } from '../api/account-security';
 import { useAccountStore } from './account-store';
 import { useEmailStore } from './email-store';
 import { useContactsStore } from './contacts-store';
@@ -39,12 +41,19 @@ export interface AuthState {
   activeAccountId: string | null;
   client: typeof jmapClient | null;
 
-  login: (serverUrl: string, username: string, password: string, opts?: { addAccount?: boolean }) => Promise<void>;
+  login: (
+    serverUrl: string,
+    username: string,
+    password: string,
+    opts?: { addAccount?: boolean; totp?: string },
+  ) => Promise<void>;
   loginViaWebmail: (webmailUrl: string, opts?: { addAccount?: boolean }) => Promise<void>;
   loginViaPairing: (webmailUrl: string, code: string, opts?: { addAccount?: boolean }) => Promise<void>;
   logout: () => Promise<void>;
   logoutAll: () => Promise<void>;
   switchAccount: (accountId: string) => Promise<void>;
+  /** Sign a non-active account out and drop its caches; the active one stays. */
+  removeAccount: (accountId: string) => Promise<void>;
   restoreSession: () => Promise<boolean>;
   retrySession: () => Promise<boolean>;
   clearError: () => void;
@@ -98,6 +107,38 @@ function refetchFeatureStores(): void {
   }
 }
 
+// Refresh the registry's display name / address from the server (#900): the
+// primary identity first, then the Stalwart principal's "Full name" when the
+// account advertises the extension. Fire-and-forget; a failure keeps whatever
+// the registry already had.
+async function syncAccountDisplayName(accountId: string): Promise<void> {
+  try {
+    const accountStore = useAccountStore.getState();
+    const entry = accountStore.getAccountById(accountId);
+    if (!entry) return;
+    const updates: { displayName?: string; email?: string } = {};
+    const identities = await getIdentities().catch(() => []);
+    const primary = identities.find((i) => i.email?.toLowerCase() === entry.email?.toLowerCase())
+      ?? identities.find((i) => i.email?.toLowerCase() === entry.username?.toLowerCase())
+      ?? identities[0];
+    if (primary?.name?.trim()) updates.displayName = primary.name.trim();
+    if (primary?.email && !entry.email.includes('@')) updates.email = primary.email;
+    if (isStalwartSupported()) {
+      const principal = await fetchPrincipal().catch(() => null);
+      if (principal?.displayName?.trim()) updates.displayName = principal.displayName.trim();
+      if (!updates.email && principal?.emails[0] && !entry.email.includes('@')) {
+        updates.email = principal.emails[0];
+      }
+    }
+    if (Object.keys(updates).length === 0) return;
+    if (useAccountStore.getState().getAccountById(accountId)) {
+      useAccountStore.getState().updateAccount(accountId, updates);
+    }
+  } catch {
+    // cosmetic - never block sign-in on it
+  }
+}
+
 // Shared tail of the OAuth sign-in flows (browser handoff and cross-device QR
 // pairing both end here). Bootstraps a JMAP session from the token bundle,
 // registers the account, and flips the store to connected. Throws on failure
@@ -108,17 +149,24 @@ async function completeOAuthHandoff(
   result: Extract<HandoffResult, { flow: 'oauth' }>,
   opts?: { addAccount?: boolean },
 ): Promise<void> {
-  if (opts?.addAccount && get().isAuthenticated) {
-    jmapClient.reset();
+  // Adding an account must not destroy the live one: the singleton keeps the
+  // previous connection until the new sign-in has actually succeeded, and a
+  // failure puts it straight back.
+  const previous = opts?.addAccount && get().isAuthenticated ? jmapClient.snapshot() : null;
+
+  let connected: { session: JMAPSession; username: string; accountId: string };
+  try {
+    connected = await jmapClient.connectWithOAuth(result.serverUrl, result.tokens);
+  } catch (err) {
+    if (previous) jmapClient.restoreSnapshot(previous);
+    throw err;
+  }
+  const { session, username, accountId } = connected;
+  if (previous) {
     useContactsStore.getState().reset();
     useCalendarStore.getState().reset();
   }
 
-  const { session, username, accountId } = await jmapClient.connectWithOAuth(
-    result.serverUrl,
-    result.tokens,
-  );
-
   const accountStore = useAccountStore.getState();
   accountStore.addAccount({
     serverUrl: result.serverUrl.replace(/\/+$/, ''),
@@ -133,6 +181,7 @@ async function completeOAuthHandoff(
   useEmailStore.getState().setActiveAccount(accountId);
 
   applyConnectedState(set, session, result.serverUrl.replace(/\/+$/, ''), username, accountId);
+  void syncAccountDisplayName(accountId);
 }
 
 function applyConnectedState(
@@ -170,17 +219,23 @@ export const useAuthStore = create<AuthState>((set, get) => ({
 
   login: async (serverUrl, username, password, opts) => {
     set({ isLoading: true, error: null });
+    // Adding an additional account: keep the live connection until the new
+    // sign-in succeeded so a typo doesn't kill the current session.
+    const previous = opts?.addAccount && get().isAuthenticated ? jmapClient.snapshot() : null;
     try {
-      // Adding an additional account - snapshot the current account away so
-      // its cache survives, then reset the JMAP client for the new login.
-      // Contacts/calendar are still single-bucket, so wipe those.
-      if (opts?.addAccount && get().isAuthenticated) {
-        jmapClient.reset();
+      let session: JMAPSession;
+      try {
+        session = await jmapClient.connect(serverUrl, username, password, opts?.totp);
+      } catch (err) {
+        if (previous) jmapClient.restoreSnapshot(previous);
+        throw err;
+      }
+      // Contacts/calendar are still single-bucket, so wipe those now that
+      // the new account is the one the client serves.
+      if (previous) {
         useContactsStore.getState().reset();
         useCalendarStore.getState().reset();
       }
-
-      const session = await jmapClient.connect(serverUrl, username, password);
       const accountId = generateAccountId(username, serverUrl.replace(/\/+$/, ''));
 
       const accountStore = useAccountStore.getState();
@@ -199,12 +254,15 @@ export const useAuthStore = create<AuthState>((set, get) => ({
       useEmailStore.getState().setActiveAccount(accountId);
 
       applyConnectedState(set, session, serverUrl.replace(/\/+$/, ''), username, accountId);
+      void syncAccountDisplayName(accountId);
     } catch (err) {
-      const message = err instanceof AuthenticationError
-        ? 'Invalid username or password'
-        : err instanceof Error
-          ? err.message
-          : 'Connection failed';
+      const message = err instanceof Error && err.name === 'TotpRequiredError'
+        ? 'Two-factor code required'
+        : err instanceof AuthenticationError
+          ? 'Invalid username or password'
+          : err instanceof Error
+            ? err.message
+            : 'Connection failed';
       set({ isLoading: false, error: message });
       throw err;
     }
@@ -384,13 +442,21 @@ export const useAuthStore = create<AuthState>((set, get) => ({
     // jmapClient first. If it fails, restore the previous active account
     // so we don't leave the user stranded on a half-switched state.
     const previousActive = get().activeAccountId;
+    // loadAccount overwrites the client's credentials/session; keep the live
+    // connection around so a failed switch can put it back instead of
+    // leaving the previous account dead until relaunch.
+    const previousClient = jmapClient.snapshot();
+    const restorePrevious = () => {
+      jmapClient.restoreSnapshot(previousClient);
+      if (previousActive) useEmailStore.getState().setActiveAccount(previousActive);
+    };
     try {
       const ok = await jmapClient.loadAccount(accountId);
       if (!ok) {
         // Credentials missing - evict stale entry and surface error
         accountStore.removeAccount(accountId);
         useEmailStore.getState().removeAccount(accountId);
-        if (previousActive) useEmailStore.getState().setActiveAccount(previousActive);
+        restorePrevious();
         set({ isLoading: false, error: 'Session expired for this account' });
         return;
       }
@@ -399,13 +465,17 @@ export const useAuthStore = create<AuthState>((set, get) => ({
         await jmapClient.clearAccountCredentials(accountId).catch(() => undefined);
         accountStore.removeAccount(accountId);
         useEmailStore.getState().removeAccount(accountId);
-        if (previousActive) useEmailStore.getState().setActiveAccount(previousActive);
+        restorePrevious();
         set({ isLoading: false, error: 'Session expired for this account' });
         return;
       }
       // NetworkError or anything else - keep the previous active account
       // intact instead of stranding the user on a half-switched state.
-      if (previousActive) useEmailStore.getState().setActiveAccount(previousActive);
+      restorePrevious();
+      accountStore.updateAccount(accountId, {
+        hasError: true,
+        errorMessage: err instanceof Error ? err.message : 'Failed to switch account',
+      });
       set({
         isLoading: false,
         error: err instanceof Error ? err.message : 'Failed to switch account',
@@ -429,6 +499,20 @@ export const useAuthStore = create<AuthState>((set, get) => ({
 
     applyConnectedState(set, session, target.serverUrl, target.username, accountId);
     refetchFeatureStores();
+    void syncAccountDisplayName(accountId);
+  },
+
+  removeAccount: async (accountId) => {
+    if (get().activeAccountId === accountId) {
+      await get().logout();
+      return;
+    }
+    const accountStore = useAccountStore.getState();
+    if (!accountStore.getAccountById(accountId)) return;
+    await teardownPushNotificationsForAccount(accountId).catch(() => undefined);
+    await jmapClient.clearAccountCredentials(accountId).catch(() => undefined);
+    useEmailStore.getState().removeAccount(accountId);
+    accountStore.removeAccount(accountId);
   },
 
   restoreSession: async () => {
@@ -537,6 +621,7 @@ export const useAuthStore = create<AuthState>((set, get) => ({
       // is live. Feature stores show persisted data immediately; this swaps
       // in fresh data once the network round-trip completes.
       refetchFeatureStores();
+      void syncAccountDisplayName(target.id);
       return true;
     } catch {
       set({ isLoading: false, hasRestoredSession: true });
@@ -592,3 +677,21 @@ export const useAuthStore = create<AuthState>((set, get) => ({
 
   clearError: () => set({ error: null }),
 }));
+
+// A 401 on a live session (revoked password/token, expired refresh token)
+// used to leave the user on a dead session until relaunch: nothing outside
+// the login flow handled `AuthenticationError`, and `retrySession` bailed
+// because `session` was still the stale object. Drop the session and retry
+// once; if the credentials really are dead, `retrySession` evicts the
+// account and shows the login screen with "Session expired".
+let authFailureInFlight = false;
+jmapClient.onAuthFailure(() => {
+  if (authFailureInFlight) return;
+  const state = useAuthStore.getState();
+  if (!state.isAuthenticated || !state.session) return;
+  authFailureInFlight = true;
+  useAuthStore.setState({ session: null });
+  void state.retrySession().finally(() => {
+    authFailureInFlight = false;
+  });
+});