fix(bitcoin-wallet-snap): reveal own output scripts when signing a PSBT

packages/bitcoin-wallet-snap/CHANGELOG.md

@@ -13,6 +13,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
+- Reveal and persist the wallet's own output scripts when signing a PSBT, so change from partner-supplied templates is always covered by routine sync ([#225](https://github.com/MetaMask/internal-snaps/pull/225))
 - Ensure certain errors are stringified correctly ([#179](https://github.com/MetaMask/internal-snaps/pull/179))
 
 ## [2.0.1]

packages/bitcoin-wallet-snap/src/entities/account.ts

@@ -97,6 +97,15 @@ export type BitcoinAccount = {
    */
   revealNextAddress(): AddressInfo;
 
+  /**
+   * Reveals addresses up to and including the derivation index of `script` if it belongs to
+   * this wallet and lies beyond the revealed set.
+   *
+   * @param script - the script to reveal up to.
+   * @returns true if new addresses were revealed.
+   */
+  revealToScript(script: ScriptBuf): boolean;
+
   /**
    * Start a full scan.
    *

packages/bitcoin-wallet-snap/src/infra/BdkAccountAdapter.ts

@@ -156,6 +156,20 @@ export class BdkAccountAdapter implements BitcoinAccount {
     return this.#wallet.reveal_next_address('external');
   }
 
+  revealToScript(script: ScriptBuf): boolean {
+    const indexed = this.#wallet.derivation_of_spk(script);
+    if (!indexed) {
+      return false;
+    }
+    const keychain = indexed[0];
+    const index = indexed[1];
+    const lastRevealed = this.#wallet.derivation_index(keychain);
+    if (lastRevealed !== undefined && lastRevealed >= index) {
+      return false;
+    }
+    return this.#wallet.reveal_addresses_to(keychain, index).length > 0;
+  }
+
   startFullScan(): FullScanRequest {
     return this.#wallet.start_full_scan();
   }

packages/bitcoin-wallet-snap/src/infra/StoredAccountAdapter.ts

@@ -167,6 +167,10 @@ export class StoredAccountAdapter implements BitcoinAccount {
     return this.#unsupported();
   }
 
+  revealToScript(_script: ScriptBuf): boolean {
+    return this.#unsupported();
+  }
+
   sentAndReceived(_tx: Transaction): [Amount, Amount] {
     return this.#unsupported();
   }

packages/bitcoin-wallet-snap/src/use-cases/AccountUseCases.test.ts

@@ -1279,6 +1279,74 @@ describe('AccountUseCases', () => {
       expect(txid).toBe(mockTxid);
       expect(psbt).toBe('mockSignedPsbt');
     });
+
+    it('reveals and persists own output scripts after signing (broadcast: false)', async () => {
+      mockAccount.isMine.mockReturnValue(true);
+      mockAccount.revealToScript.mockReturnValue(true);
+
+      await useCases.signPsbt('account-id', mockPsbt, 'metamask', {
+        fill: false,
+        broadcast: false,
+      });
+
+      expect(mockAccount.revealToScript).toHaveBeenCalledWith(
+        mockOutput.script_pubkey,
+      );
+      expect(mockRepository.update).toHaveBeenCalledWith(mockAccount);
+    });
+
+    it('reveals and persists own output scripts after signing (broadcast: true)', async () => {
+      mockAccount.getTransaction.mockReturnValue(mockWalletTx);
+      mockAccount.isMine.mockReturnValue(true);
+      mockAccount.revealToScript.mockReturnValue(true);
+
+      await useCases.signPsbt('account-id', mockPsbt, 'metamask', {
+        fill: false,
+        broadcast: true,
+      });
+
+      expect(mockAccount.revealToScript).toHaveBeenCalledWith(
+        mockOutput.script_pubkey,
+      );
+      expect(mockRepository.update).toHaveBeenCalledWith(mockAccount);
+    });
+
+    it('does not persist when no output scripts get revealed (broadcast: false)', async () => {
+      mockAccount.isMine.mockReturnValue(true);
+      mockAccount.revealToScript.mockReturnValue(false);
+
+      await useCases.signPsbt('account-id', mockPsbt, 'metamask', {
+        fill: false,
+        broadcast: false,
+      });
+
+      expect(mockRepository.update).not.toHaveBeenCalled();
+    });
+
+    it('only persists once via #broadcast when no output scripts get revealed (broadcast: true)', async () => {
+      mockAccount.getTransaction.mockReturnValue(mockWalletTx);
+      mockAccount.isMine.mockReturnValue(true);
+      mockAccount.revealToScript.mockReturnValue(false);
+
+      await useCases.signPsbt('account-id', mockPsbt, 'metamask', {
+        fill: false,
+        broadcast: true,
+      });
+
+      expect(mockRepository.update).toHaveBeenCalledTimes(1);
+      expect(mockRepository.update).toHaveBeenCalledWith(mockAccount);
+    });
+
+    it('does not call revealToScript for outputs that are not isMine', async () => {
+      mockAccount.isMine.mockReturnValue(false);
+
+      await useCases.signPsbt('account-id', mockPsbt, 'metamask', {
+        fill: false,
+        broadcast: false,
+      });
+
+      expect(mockAccount.revealToScript).not.toHaveBeenCalled();
+    });
   });
 
   describe('fillPsbt', () => {
@@ -1372,6 +1440,7 @@ describe('AccountUseCases', () => {
       expect(mockTxBuilder.drainToByScript).toHaveBeenCalledWith(
         mockOutput.script_pubkey,
       );
+      expect(mockAccount.revealToScript).not.toHaveBeenCalled();
       expect(psbt).toBe(mockFilledPsbt);
     });
 
@@ -1690,6 +1759,7 @@ describe('AccountUseCases', () => {
         mockOutput.script_pubkey,
       );
       expect(mockTxBuilder.addRecipientByScript).not.toHaveBeenCalled();
+      expect(mockAccount.revealToScript).not.toHaveBeenCalled();
       expect(fee).toBe(mockFee);
     });
 

packages/bitcoin-wallet-snap/src/use-cases/AccountUseCases.ts

@@ -567,6 +567,16 @@ export class AccountUseCases {
       : psbt;
     const signedPsbt = account.sign(psbtToSign);
 
+    let revealed = false;
+    for (const txout of psbtToSign.unsigned_tx.output) {
+      if (account.isMine(txout.script_pubkey)) {
+        revealed = account.revealToScript(txout.script_pubkey) || revealed;
+      }
+    }
+    if (revealed) {
+      await this.#repository.update(account);
+    }
+
     if (options.broadcast) {
       const psbtString = signedPsbt.toString();
       const tx = account.extractTransaction(signedPsbt);