From: valued mammal <valuedmammal@protonmail.com>
Date: Tue, 27 Aug 2024 15:54:25 +0000 (-0400)
Subject: fix(wallet): only mark change address used if `create_tx` succeeds
X-Git-Tag: v1.0.0-beta.3~9^2~2
X-Git-Url: http://internal-gitweb-vhost/script/%22https:/struct.EncoderStringWriter.html?a=commitdiff_plain;h=b60d1d29cb8908c354b43c49237acbea373c3dc7;p=bdk

fix(wallet): only mark change address used if `create_tx` succeeds

If no drain script is specified in tx params then we get it from
the change keychain by looking at the next unused address. We want
to mark the change address used so that other callers won't attempt
to use the same address between the time we create the tx and when
it appears on chain.

Before, we marked the index used regardless of whether a change
output is finally added. Then if creating a PSBT failed, we never
restored the unused status of the change address, so creating the
next tx would have revealed an extra one. Now we only mark the change
address used if we successfully create a PSBT and the drain script
is used in the change output.
---

diff --git a/crates/wallet/src/wallet/mod.rs b/crates/wallet/src/wallet/mod.rs
index 9f12fac4..25803f13 100644
--- a/crates/wallet/src/wallet/mod.rs
+++ b/crates/wallet/src/wallet/mod.rs
@@ -88,7 +88,7 @@ use crate::descriptor::{
 use crate::psbt::PsbtUtils;
 use crate::signer::SignerError;
 use crate::types::*;
-use crate::wallet::coin_selection::Excess::{Change, NoChange};
+use crate::wallet::coin_selection::Excess::{self, Change, NoChange};
 use crate::wallet::error::{BuildFeeBumpError, CreateTxError, MiniscriptPsbtError};
 
 use self::coin_selection::Error;
@@ -1468,17 +1468,28 @@ impl Wallet {
             self.preselect_utxos(&params, Some(current_height.to_consensus_u32()));
 
         // get drain script
+        let mut drain_index = Option::<(KeychainKind, u32)>::None;
         let drain_script = match params.drain_to {
             Some(ref drain_recipient) => drain_recipient.clone(),
             None => {
                 let change_keychain = self.map_keychain(KeychainKind::Internal);
-                let ((index, spk), index_changeset) = self
+                let (index, spk) = self
                     .indexed_graph
                     .index
-                    .next_unused_spk(change_keychain)
-                    .expect("keychain must exist");
-                self.indexed_graph.index.mark_used(change_keychain, index);
-                self.stage.merge(index_changeset.into());
+                    .unused_keychain_spks(change_keychain)
+                    .next()
+                    .unwrap_or_else(|| {
+                        let (next_index, _) = self
+                            .indexed_graph
+                            .index
+                            .next_index(change_keychain)
+                            .expect("keychain must exist");
+                        let spk = self
+                            .peek_address(change_keychain, next_index)
+                            .script_pubkey();
+                        (next_index, spk)
+                    });
+                drain_index = Some((change_keychain, index));
                 spk
             }
         };
@@ -1577,6 +1588,18 @@ impl Wallet {
         params.ordering.sort_tx_with_aux_rand(&mut tx, rng);
 
         let psbt = self.complete_transaction(tx, coin_selection.selected, params)?;
+
+        // recording changes to the change keychain
+        if let (Excess::Change { .. }, Some((keychain, index))) = (excess, drain_index) {
+            let (_, index_changeset) = self
+                .indexed_graph
+                .index
+                .reveal_to_target(keychain, index)
+                .expect("must not be None");
+            self.stage.merge(index_changeset.into());
+            self.mark_used(keychain, index);
+        }
+
         Ok(psbt)
     }