From e2814e813feb8239d091c2704eb28a4f31664b4d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=BF=97=E5=AE=87?= Date: Sat, 2 Aug 2025 10:06:43 +0000 Subject: [PATCH] fix: `KeychainTxOutIndex` recovery logic when spk-cache is enabled We should ensure we populate the cache (from the changeset) before doing any other operation on `KeychainTxOutIndex` (i.e. inserting descriptors). Otherwise we will end up deriving spks which have already been cached. The test is also corrected to check that we don't duplicate deriving spks. A helper function `check_cache_cs` is added to make the tests more thorough and easier to read. --- wallet/src/wallet/mod.rs | 112 ++++++++++++++++++------------ wallet/tests/persisted_wallet.rs | 114 ++++++++++++++++++++++++------- 2 files changed, 155 insertions(+), 71 deletions(-) diff --git a/wallet/src/wallet/mod.rs b/wallet/src/wallet/mod.rs index d17204891..4923e8ecd 100644 --- a/wallet/src/wallet/mod.rs +++ b/wallet/src/wallet/mod.rs @@ -471,27 +471,24 @@ impl Wallet { None => (None, Arc::new(SignersContainer::new())), }; - let index = create_indexer( + let mut stage = ChangeSet { + descriptor: Some(descriptor.clone()), + change_descriptor: change_descriptor.clone(), + local_chain: chain_changeset, + network: Some(network), + ..Default::default() + }; + + let indexed_graph = make_indexed_graph( + &mut stage, + Default::default(), + Default::default(), descriptor, change_descriptor, params.lookahead, params.use_spk_cache, )?; - let descriptor = index.get_descriptor(KeychainKind::External).cloned(); - let change_descriptor = index.get_descriptor(KeychainKind::Internal).cloned(); - let indexed_graph = IndexedTxGraph::new(index); - let indexed_graph_changeset = indexed_graph.initial_changeset(); - - let stage = ChangeSet { - descriptor, - change_descriptor, - local_chain: chain_changeset, - tx_graph: indexed_graph_changeset.tx_graph, - indexer: indexed_graph_changeset.indexer, - network: Some(network), - }; - Ok(Wallet { signers, change_signers, @@ -675,7 +672,12 @@ impl Wallet { None => Arc::new(SignersContainer::new()), }; - let index = create_indexer( + let mut stage = ChangeSet::default(); + + let indexed_graph = make_indexed_graph( + &mut stage, + changeset.tx_graph, + changeset.indexer, descriptor, change_descriptor, params.lookahead, @@ -683,12 +685,6 @@ impl Wallet { ) .map_err(LoadError::Descriptor)?; - let mut indexed_graph = IndexedTxGraph::new(index); - indexed_graph.apply_changeset(changeset.indexer.into()); - indexed_graph.apply_changeset(changeset.tx_graph.into()); - - let stage = ChangeSet::default(); - Ok(Some(Wallet { signers, change_signers, @@ -2656,35 +2652,61 @@ fn new_local_utxo( } } -fn create_indexer( +fn make_indexed_graph( + stage: &mut ChangeSet, + tx_graph_changeset: chain::tx_graph::ChangeSet, + indexer_changeset: chain::keychain_txout::ChangeSet, descriptor: ExtendedDescriptor, change_descriptor: Option, lookahead: u32, use_spk_cache: bool, -) -> Result, DescriptorError> { - let mut indexer = KeychainTxOutIndex::::new(lookahead, use_spk_cache); - - assert!(indexer - .insert_descriptor(KeychainKind::External, descriptor) - .expect("first descriptor introduced must succeed")); +) -> Result>, DescriptorError> +{ + let (indexed_graph, changeset) = IndexedTxGraph::from_changeset( + chain::indexed_tx_graph::ChangeSet { + tx_graph: tx_graph_changeset, + indexer: indexer_changeset, + }, + |idx_cs| -> Result, DescriptorError> { + let mut idx = KeychainTxOutIndex::from_changeset(lookahead, use_spk_cache, idx_cs); + + let descriptor_inserted = idx + .insert_descriptor(KeychainKind::External, descriptor) + .expect("already checked to be a unique, wildcard, non-multipath descriptor"); + assert!( + descriptor_inserted, + "this must be the first time we are seeing this descriptor" + ); + + let change_descriptor = match change_descriptor { + Some(change_descriptor) => change_descriptor, + None => return Ok(idx), + }; - if let Some(change_descriptor) = change_descriptor { - assert!(indexer - .insert_descriptor(KeychainKind::Internal, change_descriptor) - .map_err(|e| { - use bdk_chain::indexer::keychain_txout::InsertDescriptorError; - match e { - InsertDescriptorError::DescriptorAlreadyAssigned { .. } => { - crate::descriptor::error::Error::ExternalAndInternalAreTheSame - } - InsertDescriptorError::KeychainAlreadyAssigned { .. } => { - unreachable!("this is the first time we're assigning internal") + let change_descriptor_inserted = idx + .insert_descriptor(KeychainKind::Internal, change_descriptor) + .map_err(|e| { + use bdk_chain::indexer::keychain_txout::InsertDescriptorError; + match e { + InsertDescriptorError::DescriptorAlreadyAssigned { .. } => { + crate::descriptor::error::Error::ExternalAndInternalAreTheSame + } + InsertDescriptorError::KeychainAlreadyAssigned { .. } => { + unreachable!("this is the first time we're assigning internal") + } } - } - })?); - } - - Ok(indexer) + })?; + assert!( + change_descriptor_inserted, + "this must be the first time we are seeing this descriptor" + ); + + Ok(idx) + }, + )?; + stage.tx_graph.merge(changeset.tx_graph); + stage.indexer.merge(changeset.indexer); + Ok(indexed_graph) } /// Transforms a [`FeeRate`] to `f64` with unit as sat/vb. diff --git a/wallet/tests/persisted_wallet.rs b/wallet/tests/persisted_wallet.rs index 7715d7ad8..e28737999 100644 --- a/wallet/tests/persisted_wallet.rs +++ b/wallet/tests/persisted_wallet.rs @@ -1,8 +1,9 @@ -use std::collections::BTreeMap; +use std::collections::{BTreeMap, BTreeSet}; use std::path::Path; use anyhow::Context; use assert_matches::assert_matches; +use bdk_chain::DescriptorId; use bdk_chain::{ keychain_txout::DEFAULT_LOOKAHEAD, ChainPosition, ConfirmationBlockTime, DescriptorExt, }; @@ -14,7 +15,9 @@ use bdk_wallet::{ use bitcoin::constants::ChainHash; use bitcoin::hashes::Hash; use bitcoin::key::Secp256k1; -use bitcoin::{absolute, transaction, Amount, BlockHash, Network, ScriptBuf, Transaction, TxOut}; +use bitcoin::{ + absolute, secp256k1, transaction, Amount, BlockHash, Network, ScriptBuf, Transaction, TxOut, +}; use miniscript::{Descriptor, DescriptorPublicKey}; mod common; @@ -24,6 +27,57 @@ const DB_MAGIC: &[u8] = &[0x21, 0x24, 0x48]; #[test] fn wallet_is_persisted() -> anyhow::Result<()> { + type SpkCacheChangeSet = BTreeMap>; + + /// Check whether the spk-cache field of the changeset contains the expected spk indices. + fn check_cache_cs( + cache_cs: &SpkCacheChangeSet, + expected: impl IntoIterator)>, + msg: impl AsRef, + ) { + let secp = secp256k1::Secp256k1::new(); + let (external, internal) = get_test_tr_single_sig_xprv_and_change_desc(); + let (external_desc, _) = external + .into_wallet_descriptor(&secp, Network::Testnet) + .unwrap(); + let (internal_desc, _) = internal + .into_wallet_descriptor(&secp, Network::Testnet) + .unwrap(); + let external_did = external_desc.descriptor_id(); + let internal_did = internal_desc.descriptor_id(); + + let cache_cmp = cache_cs + .iter() + .map(|(did, spks)| { + let kind: KeychainKind; + if did == &external_did { + kind = KeychainKind::External; + } else if did == &internal_did { + kind = KeychainKind::Internal; + } else { + unreachable!(); + } + let spk_indices = spks.keys().copied().collect::>(); + (kind, spk_indices) + }) + .filter(|(_, spk_indices)| !spk_indices.is_empty()) + .collect::>>(); + + let expected_cmp = expected + .into_iter() + .map(|(kind, indices)| (kind, indices.into_iter().collect::>())) + .filter(|(_, spk_indices)| !spk_indices.is_empty()) + .collect::>>(); + + assert_eq!(cache_cmp, expected_cmp, "{}", msg.as_ref()); + } + + fn staged_cache(wallet: &Wallet) -> SpkCacheChangeSet { + wallet.staged().map_or(SpkCacheChangeSet::default(), |cs| { + cs.indexer.spk_cache.clone() + }) + } + fn run( filename: &str, create_db: CreateDb, @@ -46,8 +100,18 @@ fn wallet_is_persisted() -> anyhow::Result<()> { .network(Network::Testnet) .use_spk_cache(true) .create_wallet(&mut db)?; + wallet.reveal_next_address(KeychainKind::External); + check_cache_cs( + &staged_cache(&wallet), + [ + (KeychainKind::External, 0..DEFAULT_LOOKAHEAD + 1), + (KeychainKind::Internal, 0..DEFAULT_LOOKAHEAD), + ], + "cache cs must return initial set + the external index that was just derived", + ); + // persist new wallet changes assert!(wallet.persist(&mut db)?, "must write"); wallet.spk_index().clone() @@ -81,6 +145,7 @@ fn wallet_is_persisted() -> anyhow::Result<()> { .0 ); } + // Test SPK cache { let mut db = open_db(&file_path).context("failed to recover db")?; @@ -90,35 +155,32 @@ fn wallet_is_persisted() -> anyhow::Result<()> { .load_wallet(&mut db)? .expect("wallet must exist"); - let external_did = wallet - .public_descriptor(KeychainKind::External) - .descriptor_id(); - let internal_did = wallet - .public_descriptor(KeychainKind::Internal) - .descriptor_id(); - assert!(wallet.staged().is_none()); - let _addr = wallet.reveal_next_address(KeychainKind::External); - let cs = wallet.staged().expect("we should have staged a changeset"); - assert!(!cs.indexer.spk_cache.is_empty(), "failed to cache spks"); - assert_eq!(cs.indexer.spk_cache.len(), 2, "we persisted two keychains"); - let spk_cache: &BTreeMap = - cs.indexer.spk_cache.get(&external_did).unwrap(); - assert_eq!(spk_cache.len() as u32, 1 + 1 + DEFAULT_LOOKAHEAD); - assert_eq!(spk_cache.keys().last(), Some(&26)); - let spk_cache = cs.indexer.spk_cache.get(&internal_did).unwrap(); - assert_eq!(spk_cache.len() as u32, DEFAULT_LOOKAHEAD); - assert_eq!(spk_cache.keys().last(), Some(&24)); + let revealed_external_addr = wallet.reveal_next_address(KeychainKind::External); + check_cache_cs( + &staged_cache(&wallet), + [( + KeychainKind::External, + [revealed_external_addr.index + DEFAULT_LOOKAHEAD], + )], + "must only persist the revealed+LOOKAHEAD indexed external spk", + ); + // Clear the stage let _ = wallet.take_staged(); - let _addr = wallet.reveal_next_address(KeychainKind::Internal); - let cs = wallet.staged().unwrap(); - assert_eq!(cs.indexer.spk_cache.len(), 1); - let spk_cache = cs.indexer.spk_cache.get(&internal_did).unwrap(); - assert_eq!(spk_cache.len(), 1); - assert_eq!(spk_cache.keys().next(), Some(&25)); + + let revealed_internal_addr = wallet.reveal_next_address(KeychainKind::Internal); + check_cache_cs( + &staged_cache(&wallet), + [( + KeychainKind::Internal, + [revealed_internal_addr.index + DEFAULT_LOOKAHEAD], + )], + "must only persist the revealed+LOOKAHEAD indexed internal spk", + ); } + // SPK cache requires load params { let mut db = open_db(&file_path).context("failed to recover db")?;