From 4fea09abeaae20d88088f3b49e650d9c21e73e55 Mon Sep 17 00:00:00 2001
From: valued mammal <valuedmammal@protonmail.com>
Date: Tue, 27 Aug 2024 11:54:25 -0400
Subject: [PATCH] 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. 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.

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. Now we only mark the change address
used if we successfully create a psbt and the drain script is used
in the change output.
---
 crates/wallet/src/wallet/mod.rs | 37 +++++++++++++++++++++++++++------
 1 file changed, 31 insertions(+), 6 deletions(-)

diff --git a/crates/wallet/src/wallet/mod.rs b/crates/wallet/src/wallet/mod.rs
index 5cad6bb86..d7c7c4254 100644
--- a/crates/wallet/src/wallet/mod.rs
+++ b/crates/wallet/src/wallet/mod.rs
@@ -87,7 +87,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;
@@ -1448,17 +1448,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
             }
         };
@@ -1557,6 +1568,20 @@ 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 { .. } = excess {
+            if let Some((keychain, index)) = 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)
     }