Add browser extension login 2FA state remember logic (#1707)

apps/browser-extension/src/entrypoints/background.ts

@@ -10,10 +10,13 @@ import { setupContextMenus } from '@/entrypoints/background/ContextMenu';
 import { handleGetWebAuthnSettings, handleWebAuthnCreate, handleWebAuthnGet, handlePasskeyPopupResponse, handleGetRequestData } from '@/entrypoints/background/PasskeyHandler';
 import { handleOpenPopup, handlePopupWithItem, handleOpenPopupCreateCredential, handleToggleContextMenu } from '@/entrypoints/background/PopupMessageHandler';
 import { handleStoreSavePromptState, handleGetSavePromptState, handleClearSavePromptState } from '@/entrypoints/background/SavePromptStateHandler';
+import { handleStoreTwoFactorState, handleGetTwoFactorState, handleClearTwoFactorState } from '@/entrypoints/background/TwoFactorStateHandler';
 import { handleCheckAuthStatus, handleClearPersistedFormValues, handleClearSession, handleClearVaultData, handleLockVault, handleCreateItem, handleGetFilteredItems, handleGetSearchItems, handleGetDefaultEmailDomain, handleGetDefaultIdentitySettings, handleGetEncryptionKey, handleGetEncryptionKeyDerivationParams, handleGetPasswordSettings, handleGetPersistedFormValues, handleGetVault, handlePersistFormValues, handleStoreEncryptionKey, handleStoreEncryptionKeyDerivationParams, handleStoreVaultMetadata, handleSyncVault, handleUploadVault, handleGetEncryptedVault, handleStoreEncryptedVault, handleGetSyncState, handleMarkVaultClean, handleGetServerRevision, handleFullVaultSync, handleCheckLoginDuplicate, handleSaveLoginCredential, handleGetLoginSaveSettings, handleSetLoginSaveEnabled } from '@/entrypoints/background/VaultMessageHandler';
 
 import { EncryptionKeyDerivationParams } from "@/utils/dist/core/models/metadata";
+import type { LoginResponse } from "@/utils/dist/core/models/webapi";
 import { LocalPreferencesService } from '@/utils/LocalPreferencesService';
+import type { SavePromptPersistedState } from "@/utils/loginDetector";
 
 import { defineBackground, browser } from '#imports';
 
@@ -68,11 +71,16 @@ export default defineBackground({
     onMessage('GET_LOGIN_SAVE_SETTINGS', () => handleGetLoginSaveSettings());
     onMessage('SET_LOGIN_SAVE_ENABLED', ({ data }) => handleSetLoginSaveEnabled(data as boolean));
 
-    // Save prompt state persistence (for surviving page navigation)
-    onMessage('STORE_SAVE_PROMPT_STATE', ({ data, sender }) => handleStoreSavePromptState({ tabId: sender.tabId!, state: data as unknown as import('@/utils/loginDetector').SavePromptPersistedState }));
+    // Remember login save state (for surviving page navigation)
+    onMessage('STORE_SAVE_PROMPT_STATE', ({ data, sender }) => handleStoreSavePromptState({ tabId: sender.tabId!, state: data as SavePromptPersistedState }));
     onMessage('GET_SAVE_PROMPT_STATE', ({ sender }) => handleGetSavePromptState({ tabId: sender.tabId! }));
     onMessage('CLEAR_SAVE_PROMPT_STATE', ({ sender }) => handleClearSavePromptState({ tabId: sender.tabId! }));
 
+    // Two-factor authentication state persistence
+    onMessage('STORE_TWO_FACTOR_STATE', ({ data }) => handleStoreTwoFactorState(data as { username: string; loginResponse: LoginResponse; passwordHashString: string; passwordHashBase64: string; rememberMe: boolean }));
+    onMessage('GET_TWO_FACTOR_STATE', () => handleGetTwoFactorState());
+    onMessage('CLEAR_TWO_FACTOR_STATE', () => handleClearTwoFactorState());
+
     // Clipboard management messages
     onMessage('CLIPBOARD_COPIED', () => handleClipboardCopied());
     onMessage('CANCEL_CLIPBOARD_CLEAR', () => handleCancelClipboardClear());

apps/browser-extension/src/entrypoints/background/TwoFactorStateHandler.ts

@@ -0,0 +1,70 @@
+/**
+ * In-memory 2FA state handler for persisting login state during popup close/reopen.
+ *
+ * This handler stores 2FA login state ONLY in memory, and the state automatically
+ * expires after a short timeout.
+ *
+ * This allows users to enter username/password and get to 2FA prompt, close popup
+ * to switch to authenticator app, and reopen popup and continue from 2FA prompt
+ * without re-entering credentials.
+ */
+
+import type { LoginResponse } from '@/utils/dist/core/models/webapi';
+
+/**
+ * The 2FA state that is persisted in memory.
+ */
+export type TwoFactorState = {
+  username: string;
+  loginResponse: LoginResponse;
+  passwordHashString: string;
+  passwordHashBase64: string;
+  rememberMe: boolean;
+  timestamp: number;
+};
+
+/**
+ * Timeout for automatic state expiration (60 seconds).
+ */
+const STATE_EXPIRY_MS = 60 * 1000;
+
+/**
+ * In-memory storage for 2FA state.
+ * Intentionally NOT persisted to any storage - lives only in service worker memory.
+ */
+let twoFactorState: TwoFactorState | null = null;
+
+/**
+ * Store 2FA state in memory with current timestamp.
+ */
+export function handleStoreTwoFactorState(state: Omit<TwoFactorState, 'timestamp'>): void {
+  twoFactorState = {
+    ...state,
+    timestamp: Date.now(),
+  };
+}
+
+/**
+ * Retrieve 2FA state from memory.
+ * Returns null if no state exists or if the state has expired.
+ */
+export function handleGetTwoFactorState(): TwoFactorState | null {
+  if (!twoFactorState) {
+    return null;
+  }
+
+  // Check if state has expired
+  if (Date.now() - twoFactorState.timestamp > STATE_EXPIRY_MS) {
+    twoFactorState = null;
+    return null;
+  }
+
+  return twoFactorState;
+}
+
+/**
+ * Clear 2FA state from memory.
+ */
+export function handleClearTwoFactorState(): void {
+  twoFactorState = null;
+}

apps/browser-extension/src/entrypoints/popup/pages/auth/Login.tsx

@@ -3,6 +3,7 @@ import { useTranslation } from 'react-i18next';
 import { useNavigate } from 'react-router-dom';
 import { sendMessage } from 'webext-bridge/popup';
 
+import type { TwoFactorState } from '@/entrypoints/background/TwoFactorStateHandler';
 import Button from '@/entrypoints/popup/components/Button';
 import MobileUnlockModal from '@/entrypoints/popup/components/Dialogs/MobileUnlockModal';
 import HeaderButton from '@/entrypoints/popup/components/HeaderButton';
@@ -29,6 +30,9 @@ import { storage } from '#imports';
 /** Track if username prefill has been attempted (only do it once on mount) */
 let usernamePrefillAttempted = false;
 
+/** Track if 2FA state restoration has been attempted (only do it once on mount) */
+let twoFactorStateRestoreAttempted = false;
+
 /**
  * Login page
  */
@@ -174,7 +178,7 @@ const Login: React.FC = () => {
 
   useEffect(() => {
     /**
-     * Load the client URL and check for saved username (from forced logout).
+     * Load the client URL, check for saved username, and restore 2FA state if available.
      */
     const loadInitialData = async () : Promise<void> => {
       // Load client URL
@@ -185,6 +189,27 @@ const Login: React.FC = () => {
       }
       setClientUrl(clientUrl);
 
+      /*
+       * Check for persisted 2FA state (from popup close during 2FA entry).
+       * This allows users to close the popup to switch to their authenticator app
+       * and continue where they left off when reopening.
+       */
+      if (!twoFactorStateRestoreAttempted) {
+        twoFactorStateRestoreAttempted = true;
+        const savedState = await sendMessage('GET_TWO_FACTOR_STATE', null, 'background') as TwoFactorState | null;
+        if (savedState) {
+          // Restore the 2FA state
+          setCredentials({ username: savedState.username, password: '' });
+          setLoginResponse(savedState.loginResponse);
+          setPasswordHashString(savedState.passwordHashString);
+          setPasswordHashBase64(savedState.passwordHashBase64);
+          setRememberMe(savedState.rememberMe);
+          setTwoFactorRequired(true);
+          setIsInitialLoading(false);
+          return;
+        }
+      }
+
       /*
        * Check for saved username (from forced logout) and prefill once on mount
        * If user clears it, don't repopulate
@@ -263,6 +288,19 @@ const Login: React.FC = () => {
         // Store password hash base64 as we need it for decryption
         setPasswordHashBase64(passwordHashBase64);
         setTwoFactorRequired(true);
+
+        /*
+         * Persist 2FA state to background script so user can
+         * close popup to switch to authenticator app and continue when reopening
+         */
+        await sendMessage('STORE_TWO_FACTOR_STATE', {
+          username: normalizedUsername,
+          loginResponse,
+          passwordHashString,
+          passwordHashBase64,
+          rememberMe,
+        }, 'background');
+
         // Show app.
         hideLoading();
         return;
@@ -329,6 +367,9 @@ const Login: React.FC = () => {
         throw new Error(t('common.errors.unknownError'));
       }
 
+      // Clear any persisted 2FA state since login is successful
+      await sendMessage('CLEAR_TWO_FACTOR_STATE', null, 'background');
+
       // Handle successful authentication
       await handleSuccessfulAuth(
         twoFaUsername,
@@ -449,8 +490,10 @@ const Login: React.FC = () => {
             </Button>
             <Button
               type="button"
-              onClick={() => {
-                // Reset the form.
+              onClick={async () => {
+                // Clear persisted 2FA state
+                await sendMessage('CLEAR_TWO_FACTOR_STATE', null, 'background');
+                // Reset the form
                 setCredentials({
                   username: '',
                   password: ''