Add useVaultMutate hook compatible with browser extension (#900)

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

@@ -2,7 +2,7 @@ import { onMessage, sendMessage } from "webext-bridge/background";
 
 import { setupContextMenus } from '@/entrypoints/background/ContextMenu';
 import { handleOpenPopup, handlePopupWithCredential, handleToggleContextMenu } from '@/entrypoints/background/PopupMessageHandler';
-import { handleCheckAuthStatus, handleClearVault, handleCreateIdentity, handleGetCredentials, handleGetDefaultEmailDomain, handleGetDefaultIdentityLanguage, handleGetDerivedKey, handleGetPasswordSettings, handleGetVault, handleStoreVault, handleSyncVault } from '@/entrypoints/background/VaultMessageHandler';
+import { handleCheckAuthStatus, handleClearVault, handleCreateIdentity, handleGetCredentials, handleGetDefaultEmailDomain, handleGetDefaultIdentityLanguage, handleGetDerivedKey, handleGetPasswordSettings, handleGetVault, handleStoreVault, handleSyncVault, handleUploadVault } from '@/entrypoints/background/VaultMessageHandler';
 
 import { GLOBAL_CONTEXT_MENU_ENABLED_KEY } from '@/utils/Constants';
 
@@ -16,6 +16,7 @@ export default defineBackground({
     // Listen for messages using webext-bridge
     onMessage('CHECK_AUTH_STATUS', () => handleCheckAuthStatus());
     onMessage('STORE_VAULT', ({ data }) => handleStoreVault(data));
+    onMessage('UPLOAD_VAULT', () => handleUploadVault());
     onMessage('SYNC_VAULT', () => handleSyncVault());
     onMessage('GET_VAULT', () => handleGetVault());
     onMessage('CLEAR_VAULT', () => handleClearVault());

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

@@ -9,6 +9,7 @@ import { CredentialsResponse as messageCredentialsResponse } from '@/utils/types
 import { PasswordSettingsResponse as messagePasswordSettingsResponse } from '@/utils/types/messaging/PasswordSettingsResponse';
 import { StringResponse as stringResponse } from '@/utils/types/messaging/StringResponse';
 import { VaultResponse as messageVaultResponse } from '@/utils/types/messaging/VaultResponse';
+import { VaultUploadResponse as messageVaultUploadResponse } from '@/utils/types/messaging/VaultUploadResponse';
 import { WebApiService } from '@/utils/WebApiService';
 
 /**
@@ -299,10 +300,25 @@ export async function handleGetDerivedKey(
   return derivedKey;
 }
 
+/**
+ * Upload the vault to the server.
+ */
+export async function handleUploadVault(
+) : Promise<messageVaultUploadResponse> {
+  try {
+    const sqliteClient = await createVaultSqliteClient();
+    const response = await uploadNewVaultToServer(sqliteClient);
+    return { success: true, status: response.status, newRevisionNumber: response.newRevisionNumber };
+  } catch (error) {
+    console.error('Failed to upload vault:', error);
+    return { success: false, error: 'Failed to upload vault' };
+  }
+}
+
 /**
  * Upload a new version of the vault to the server using the provided sqlite client.
  */
-async function uploadNewVaultToServer(sqliteClient: SqliteClient) : Promise<void> {
+async function uploadNewVaultToServer(sqliteClient: SqliteClient) : Promise<VaultPostResponse> {
   const updatedVaultData = sqliteClient.exportToBase64();
   const derivedKey = await storage.getItem('session:derivedKey') as string;
 
@@ -346,6 +362,8 @@ async function uploadNewVaultToServer(sqliteClient: SqliteClient) : Promise<void
   } else {
     throw new Error('Failed to upload new vault to server');
   }
+
+  return response;
 }
 
 /**

apps/browser-extension/src/entrypoints/popup/context/DbContext.tsx

@@ -14,6 +14,7 @@ type DbContextType = {
   initializeDatabase: (vaultResponse: VaultResponse, derivedKey: string) => Promise<void>;
   clearDatabase: () => void;
   getVaultMetadata: () => Promise<VaultMetadata | null>;
+  setCurrentVaultRevisionNumber: (revisionNumber: number) => Promise<void>;
 }
 
 const DbContext = createContext<DbContextType | undefined>(undefined);
@@ -104,6 +105,17 @@ export const DbProvider: React.FC<{ children: React.ReactNode }> = ({ children }
     return vaultMetadata;
   }, [vaultMetadata]);
 
+  /**
+   * Set the current vault revision number.
+   */
+  const setCurrentVaultRevisionNumber = useCallback(async (revisionNumber: number) => {
+    setVaultMetadata({
+      publicEmailDomains: vaultMetadata?.publicEmailDomains ?? [],
+      privateEmailDomains: vaultMetadata?.privateEmailDomains ?? [],
+      vaultRevisionNumber: revisionNumber,
+    });
+  }, [vaultMetadata]);
+
   /**
    * Check if database is initialized and try to retrieve vault from background
    */
@@ -129,7 +141,8 @@ export const DbProvider: React.FC<{ children: React.ReactNode }> = ({ children }
     initializeDatabase,
     clearDatabase,
     getVaultMetadata,
-  }), [sqliteClient, dbInitialized, dbAvailable, initializeDatabase, clearDatabase, getVaultMetadata]);
+    setCurrentVaultRevisionNumber,
+  }), [sqliteClient, dbInitialized, dbAvailable, initializeDatabase, clearDatabase, getVaultMetadata, setCurrentVaultRevisionNumber]);
 
   return (
     <DbContext.Provider value={contextValue}>

apps/browser-extension/src/entrypoints/popup/hooks/useVaultMutate.ts

@@ -0,0 +1,136 @@
+import { useCallback, useState } from 'react';
+import { sendMessage } from 'webext-bridge/popup';
+
+import { useDb } from '@/entrypoints/popup/context/DbContext';
+import { useVaultSync } from '@/entrypoints/popup/hooks/useVaultSync';
+
+import { VaultUploadResponse as messageVaultUploadResponse } from '@/utils/types/messaging/VaultUploadResponse';
+
+type VaultMutationOptions = {
+  onSuccess?: () => void;
+  onError?: (error: Error) => void;
+}
+
+/**
+ * Hook to execute a vault mutation.
+ */
+export function useVaultMutate() : {
+  executeVaultMutation: (operation: () => Promise<void>, options?: VaultMutationOptions) => Promise<void>;
+  isLoading: boolean;
+  syncStatus: string;
+  } {
+  const [isLoading, setIsLoading] = useState(false);
+  const [syncStatus, setSyncStatus] = useState('Syncing vault');
+  const dbContext = useDb();
+  const { syncVault } = useVaultSync();
+
+  /**
+   * Execute the provided operation (e.g. create/update/delete credential)
+   */
+  const executeMutateOperation = useCallback(async (
+    operation: () => Promise<void>,
+    options: VaultMutationOptions
+  ) : Promise<void> => {
+    setSyncStatus('Saving changes to vault');
+
+    // Execute the provided operation (e.g. create/update/delete credential)
+    await operation();
+
+    setSyncStatus('Uploading vault to server');
+
+    try {
+      // Trigger the background worker to upload the current vault to the server
+      const response = await sendMessage('UPLOAD_VAULT', {}, 'background') as messageVaultUploadResponse;
+
+      /*
+       * If we get here, it means we have a valid connection to the server.
+       * TODO: offline mode is not implemented for browser extension yet.
+       * authContext.setOfflineMode(false);
+       */
+
+      if (response.status === 0 && response.newRevisionNumber) {
+        await dbContext.setCurrentVaultRevisionNumber(response.newRevisionNumber);
+        options.onSuccess?.();
+      } else if (response.status === 1) {
+        throw new Error('Vault merge required. Please login via the web app to merge the multiple pending updates to your vault.');
+      } else {
+        throw new Error('Failed to upload vault to server');
+      }
+    } catch (error) {
+      // Check if it's a network error
+      if (error instanceof Error && (error.message.includes('network') || error.message.includes('timeout'))) {
+        /*
+         * Network error, mark as offline and track pending changes
+         * TODO: offline mode is not implemented for browser extension yet.
+         * authContext.setOfflineMode(true);
+         */
+        options.onError?.(new Error('Network error'));
+        return;
+      }
+      throw error;
+    }
+  }, [dbContext]);
+
+  /**
+   * Hook to execute a vault mutation which uploads a new encrypted vault to the server
+   */
+  const executeVaultMutation = useCallback(async (
+    operation: () => Promise<void>,
+    options: VaultMutationOptions = {}
+  ) => {
+    try {
+      setIsLoading(true);
+      setSyncStatus('Checking for vault updates');
+
+      await syncVault({
+        /**
+         * Handle the status update.
+         */
+        onStatus: (message) => setSyncStatus(message),
+        /**
+         * Handle successful vault sync and continue with vault mutation.
+         */
+        onSuccess: async (hasNewVault) => {
+          if (hasNewVault) {
+            // Vault was changed, but has now been reloaded so we can continue with the operation.
+          }
+          await executeMutateOperation(operation, options);
+        },
+        /**
+         * Handle error during vault sync.
+         */
+        onError: (error) => {
+          /**
+           *Toast.show({
+           *type: 'error',
+           *text1: 'Failed to sync vault',
+           *text2: error,
+           *position: 'bottom'
+           *});
+           */
+          options.onError?.(new Error(error));
+        }
+      });
+    } catch (error) {
+      console.error('Error during vault mutation:', error);
+      /*
+       * Toast.show({
+       *type: 'error',
+       *text1: 'Operation failed',
+       *text2: error instanceof Error ? error.message : 'Unknown error',
+       *position: 'bottom'
+       *});
+       */
+      options.onError?.(error instanceof Error ? error : new Error('Unknown error'));
+    } finally {
+      setIsLoading(false);
+      setSyncStatus('');
+    }
+  }, [syncVault, executeMutateOperation]);
+
+  return {
+    executeVaultMutation,
+    isLoading,
+    syncStatus,
+  };
+}

apps/browser-extension/src/utils/types/messaging/VaultUploadResponse.ts

@@ -0,0 +1,6 @@
+export type VaultUploadResponse = {
+    success: boolean,
+    error?: string,
+    status?: number,
+    newRevisionNumber?: number
+};