Skip to content

Fix KeychainTxOutIndex recovery logic when spk-cache is enabled #288

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Aug 5, 2025
Merged

Fix KeychainTxOutIndex recovery logic when spk-cache is enabled #288

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
112 changes: 67 additions & 45 deletions wallet/src/wallet/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -675,20 +672,19 @@ 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,
params.use_spk_cache,
)
.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,
Expand Down Expand Up @@ -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<ConfirmationBlockTime>,
indexer_changeset: chain::keychain_txout::ChangeSet,
descriptor: ExtendedDescriptor,
change_descriptor: Option<ExtendedDescriptor>,
lookahead: u32,
use_spk_cache: bool,
) -> Result<KeychainTxOutIndex<KeychainKind>, DescriptorError> {
let mut indexer = KeychainTxOutIndex::<KeychainKind>::new(lookahead, use_spk_cache);

assert!(indexer
.insert_descriptor(KeychainKind::External, descriptor)
.expect("first descriptor introduced must succeed"));
) -> Result<IndexedTxGraph<ConfirmationBlockTime, KeychainTxOutIndex<KeychainKind>>, DescriptorError>
{
let (indexed_graph, changeset) = IndexedTxGraph::from_changeset(
chain::indexed_tx_graph::ChangeSet {
tx_graph: tx_graph_changeset,
indexer: indexer_changeset,
},
|idx_cs| -> Result<KeychainTxOutIndex<KeychainKind>, 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.
Expand Down
114 changes: 88 additions & 26 deletions wallet/tests/persisted_wallet.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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,
};
Expand All @@ -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;
Expand All @@ -24,6 +27,57 @@ const DB_MAGIC: &[u8] = &[0x21, 0x24, 0x48];

#[test]
fn wallet_is_persisted() -> anyhow::Result<()> {
type SpkCacheChangeSet = BTreeMap<DescriptorId, BTreeMap<u32, ScriptBuf>>;

/// Check whether the spk-cache field of the changeset contains the expected spk indices.
fn check_cache_cs(
cache_cs: &SpkCacheChangeSet,
expected: impl IntoIterator<Item = (KeychainKind, impl IntoIterator<Item = u32>)>,
msg: impl AsRef<str>,
) {
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::<BTreeSet<u32>>();
(kind, spk_indices)
})
.filter(|(_, spk_indices)| !spk_indices.is_empty())
.collect::<BTreeMap<KeychainKind, BTreeSet<u32>>>();

let expected_cmp = expected
.into_iter()
.map(|(kind, indices)| (kind, indices.into_iter().collect::<BTreeSet<u32>>()))
.filter(|(_, spk_indices)| !spk_indices.is_empty())
.collect::<BTreeMap<KeychainKind, BTreeSet<u32>>>();

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<Db, CreateDb, OpenDb>(
filename: &str,
create_db: CreateDb,
Expand All @@ -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()
Expand Down Expand Up @@ -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")?;
Expand All @@ -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<u32, ScriptBuf> =
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")?;
Expand Down
Loading