fix(wallets): make WalletDB.deleteAccount atomic

yarn-project/wallets/src/embedded/wallet_db.test.ts

@@ -189,9 +189,10 @@ describe('WalletDB', () => {
       const wrapMap = (map: AztecAsyncMap<string, Buffer>): AztecAsyncMap<string, Buffer> =>
         new Proxy(map, {
           get(target, prop) {
-            if (prop === 'set') {
-              return (key: string, value: Buffer) =>
-                shouldCrash(key) ? Promise.reject(new Error('simulated write failure')) : target.set(key, value);
+            if (prop === 'set' || prop === 'delete') {
+              const write = (target[prop] as (key: string, value?: Buffer) => Promise<void>).bind(target);
+              return (key: string, value?: Buffer) =>
+                shouldCrash(key) ? Promise.reject(new Error('simulated write failure')) : write(key, value);
             }
             const member = Reflect.get(target, prop);
             return typeof member === 'function' ? member.bind(target) : member;
@@ -226,6 +227,34 @@ describe('WalletDB', () => {
       expect(await dbAfterCrash.listAccounts()).toEqual([]);
       expect(await store.openMap<string, Buffer>('aliases').getAsync('accounts:alice')).toBeUndefined();
     });
+
+    it('deletes no account data when a write fails midway through deleteAccount', async () => {
+      const store = await openTmpStore('wallet-db-atomicity-test');
+      const db = new WalletDB(store, () => {});
+      const address = await AztecAddress.random();
+      const data = makeAccountData('schnorr', 'alice');
+      await db.storeAccount(address, data);
+
+      // Fail the alias delete, which happens after all the account entry deletes
+      const failingDb = new WalletDB(
+        crashingStore(store, key => key.startsWith('accounts:')),
+        () => {},
+      );
+      await expect(failingDb.deleteAccount(address)).rejects.toThrow('simulated write failure');
+
+      // Inspect the same underlying store with a fresh WalletDB: the account must be fully intact
+      const dbAfterCrash = new WalletDB(store, () => {});
+      const retrieved = await dbAfterCrash.retrieveAccount(address);
+      expect(retrieved.secretKey).toEqual(data.secretKey);
+      expect(retrieved.salt).toEqual(data.salt);
+      expect(retrieved.type).toEqual('schnorr');
+      expect(retrieved.signingKey).toEqual(data.signingKey.toBuffer());
+
+      const accounts = await dbAfterCrash.listAccounts();
+      expect(accounts).toHaveLength(1);
+      expect(accounts[0].alias).toEqual('alice');
+      expect(accounts[0].item.toString()).toEqual(address.toString());
+    });
   });
 
   describe('all account types', () => {

yarn-project/wallets/src/embedded/wallet_db.ts

@@ -121,19 +121,25 @@ export class WalletDB {
     return addresses;
   }
 
+  /**
+   * Deletes an account's stored data and its alias atomically. Deletion is local to this store;
+   * any state the PXE holds for the account is unaffected.
+   */
   async deleteAccount(address: AztecAddress) {
-    await Promise.all([
-      this.accounts.delete(accountKey('sk', address)),
-      this.accounts.delete(accountKey('salt', address)),
-      this.accounts.delete(accountKey('type', address)),
-      this.accounts.delete(accountKey('signingKey', address)),
-    ]);
-    // Clean up alias if one exists
-    const aliasesByAddress = await this.#readAccountAliases();
-    const alias = aliasesByAddress.get(address.toString());
-    if (alias) {
-      await this.aliases.delete(`accounts:${alias}`);
-    }
+    await this.store.transactionAsync(async () => {
+      await Promise.all([
+        this.accounts.delete(accountKey('sk', address)),
+        this.accounts.delete(accountKey('salt', address)),
+        this.accounts.delete(accountKey('type', address)),
+        this.accounts.delete(accountKey('signingKey', address)),
+      ]);
+      // Clean up alias if one exists
+      const aliasesByAddress = await this.#readAccountAliases();
+      const alias = aliasesByAddress.get(address.toString());
+      if (alias) {
+        await this.aliases.delete(`accounts:${alias}`);
+      }
+    });
   }
 
   async close() {