Skip to content

Commit

Permalink
feat(chain)!: KeychainTxOutIndex uses a universal lookahead
Browse files Browse the repository at this point in the history
The wallet is currently created without setting any lookahead value for
the keychain. This implicitly makes it a lookahead of 0. As this is a
high-level interface we should avoid footguns and aim for a reasonable
default.

Instead of simply patching it for wallet, we alter `KeychainTxOutIndex`
to have a default lookahead value. Additionally, instead of a
per-keychain lookahead, the constructor asks for a `lookahead` value.
This avoids the footguns of having methods which allows the caller the
decrease the `lookahead` (and therefore panicing). This also simplifies
the API.

Co-authored-by: Antoine Poisot <[email protected]>
Co-authored-by: 志宇 <[email protected]>
  • Loading branch information
darosior and evanlinjin committed Dec 28, 2023
1 parent 01698ae commit 4fd539b
Show file tree
Hide file tree
Showing 5 changed files with 56 additions and 103 deletions.
100 changes: 35 additions & 65 deletions crates/chain/src/keychain/txout_index.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,12 +5,13 @@ use crate::{
spk_iter::BIP32_MAX_INDEX,
SpkIterator, SpkTxOutIndex,
};
use alloc::vec::Vec;
use bitcoin::{OutPoint, Script, TxOut};
use core::{fmt::Debug, ops::Deref};

use crate::Append;

const DEFAULT_LOOKAHEAD: u32 = 1_000;

/// A convenient wrapper around [`SpkTxOutIndex`] that relates script pubkeys to miniscript public
/// [`Descriptor`]s.
///
Expand Down Expand Up @@ -46,7 +47,7 @@ use crate::Append;
/// # let secp = bdk_chain::bitcoin::secp256k1::Secp256k1::signing_only();
/// # let (external_descriptor,_) = Descriptor::<DescriptorPublicKey>::parse_descriptor(&secp, "tr([73c5da0a/86'/0'/0']xprv9xgqHN7yz9MwCkxsBPN5qetuNdQSUttZNKw1dcYTV4mkaAFiBVGQziHs3NRSWMkCzvgjEe3n9xV8oYywvM8at9yRqyaZVz6TYYhX98VjsUk/0/*)").unwrap();
/// # let (internal_descriptor,_) = Descriptor::<DescriptorPublicKey>::parse_descriptor(&secp, "tr([73c5da0a/86'/0'/0']xprv9xgqHN7yz9MwCkxsBPN5qetuNdQSUttZNKw1dcYTV4mkaAFiBVGQziHs3NRSWMkCzvgjEe3n9xV8oYywvM8at9yRqyaZVz6TYYhX98VjsUk/1/*)").unwrap();
/// # let descriptor_for_user_42 = external_descriptor.clone();
/// # let (descriptor_for_user_42, _) = Descriptor::<DescriptorPublicKey>::parse_descriptor(&secp, "tr([73c5da0a/86'/0'/0']xprv9xgqHN7yz9MwCkxsBPN5qetuNdQSUttZNKw1dcYTV4mkaAFiBVGQziHs3NRSWMkCzvgjEe3n9xV8oYywvM8at9yRqyaZVz6TYYhX98VjsUk/2/*)").unwrap();
/// txout_index.add_keychain(MyKeychain::External, external_descriptor);
/// txout_index.add_keychain(MyKeychain::Internal, internal_descriptor);
/// txout_index.add_keychain(MyKeychain::MyAppUser { user_id: 42 }, descriptor_for_user_42);
Expand All @@ -65,17 +66,12 @@ pub struct KeychainTxOutIndex<K> {
// last revealed indexes
last_revealed: BTreeMap<K, u32>,
// lookahead settings for each keychain
lookahead: BTreeMap<K, u32>,
lookahead: u32,
}

impl<K> Default for KeychainTxOutIndex<K> {
fn default() -> Self {
Self {
inner: SpkTxOutIndex::default(),
keychains: BTreeMap::default(),
last_revealed: BTreeMap::default(),
lookahead: BTreeMap::default(),
}
Self::new(DEFAULT_LOOKAHEAD)
}
}

Expand Down Expand Up @@ -118,6 +114,24 @@ impl<K: Clone + Ord + Debug> Indexer for KeychainTxOutIndex<K> {
}
}

impl<K> KeychainTxOutIndex<K> {
/// Construct a [`KeychainTxOutIndex`] with the given `lookahead`.
///
/// The lookahead is the number of scripts to cache ahead of the last revealed script index.
/// This is useful to find outputs you own when processing block data that lie beyond the last
/// revealed index. In certain situations, such as when performing an initial scan of the
/// blockchain during wallet import, it may be uncertain or unknown what the last revealed index
/// is.
pub fn new(lookahead: u32) -> Self {
Self {
inner: SpkTxOutIndex::default(),
keychains: BTreeMap::new(),
last_revealed: BTreeMap::new(),
lookahead,
}
}
}

impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
/// Return a reference to the internal [`SpkTxOutIndex`].
pub fn inner(&self) -> &SpkTxOutIndex<(K, u32)> {
Expand Down Expand Up @@ -145,54 +159,22 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
pub fn add_keychain(&mut self, keychain: K, descriptor: Descriptor<DescriptorPublicKey>) {
let old_descriptor = &*self
.keychains
.entry(keychain)
.entry(keychain.clone())
.or_insert_with(|| descriptor.clone());
assert_eq!(
&descriptor, old_descriptor,
"keychain already contains a different descriptor"
);
self.replenish_lookahead(&keychain, self.lookahead);
}

/// Return the lookahead setting for each keychain.
/// Get the lookahead setting.
///
/// Refer to [`set_lookahead`] for a deeper explanation of the `lookahead`.
/// Refer to [`new`] for more information on the `lookahead`.
///
/// [`set_lookahead`]: Self::set_lookahead
pub fn lookaheads(&self) -> &BTreeMap<K, u32> {
&self.lookahead
}

/// Convenience method to call [`set_lookahead`] for all keychains.
///
/// [`set_lookahead`]: Self::set_lookahead
pub fn set_lookahead_for_all(&mut self, lookahead: u32) {
for keychain in &self.keychains.keys().cloned().collect::<Vec<_>>() {
self.set_lookahead(keychain, lookahead);
}
}

/// Set the lookahead count for `keychain`.
///
/// The lookahead is the number of scripts to cache ahead of the last revealed script index. This
/// is useful to find outputs you own when processing block data that lie beyond the last revealed
/// index. In certain situations, such as when performing an initial scan of the blockchain during
/// wallet import, it may be uncertain or unknown what the last revealed index is.
///
/// # Panics
///
/// This will panic if the `keychain` does not exist.
pub fn set_lookahead(&mut self, keychain: &K, lookahead: u32) {
self.lookahead.insert(keychain.clone(), lookahead);
self.replenish_lookahead(keychain);
}

/// Convenience method to call [`lookahead_to_target`] for multiple keychains.
///
/// [`lookahead_to_target`]: Self::lookahead_to_target
pub fn lookahead_to_target_multi(&mut self, target_indexes: BTreeMap<K, u32>) {
for (keychain, target_index) in target_indexes {
self.lookahead_to_target(&keychain, target_index)
}
/// [`new`]: Self::new
pub fn lookahead(&self) -> u32 {
self.lookahead
}

/// Store lookahead scripts until `target_index`.
Expand All @@ -201,22 +183,14 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
pub fn lookahead_to_target(&mut self, keychain: &K, target_index: u32) {
let next_index = self.next_store_index(keychain);
if let Some(temp_lookahead) = target_index.checked_sub(next_index).filter(|&v| v > 0) {
let old_lookahead = self.lookahead.insert(keychain.clone(), temp_lookahead);
self.replenish_lookahead(keychain);

// revert
match old_lookahead {
Some(lookahead) => self.lookahead.insert(keychain.clone(), lookahead),
None => self.lookahead.remove(keychain),
};
self.replenish_lookahead(keychain, temp_lookahead);
}
}

fn replenish_lookahead(&mut self, keychain: &K) {
fn replenish_lookahead(&mut self, keychain: &K, lookahead: u32) {
let descriptor = self.keychains.get(keychain).expect("keychain must exist");
let next_store_index = self.next_store_index(keychain);
let next_reveal_index = self.last_revealed.get(keychain).map_or(0, |v| *v + 1);
let lookahead = self.lookahead.get(keychain).map_or(0, |v| *v);

for (new_index, new_spk) in
SpkIterator::new_with_range(descriptor, next_store_index..next_reveal_index + lookahead)
Expand Down Expand Up @@ -388,25 +362,21 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {

let target_index = if has_wildcard { target_index } else { 0 };
let next_reveal_index = self.last_revealed.get(keychain).map_or(0, |v| *v + 1);
let lookahead = self.lookahead.get(keychain).map_or(0, |v| *v);

debug_assert_eq!(
next_reveal_index + lookahead,
self.next_store_index(keychain)
);
debug_assert!(next_reveal_index + self.lookahead >= self.next_store_index(keychain));

// if we need to reveal new indices, the latest revealed index goes here
let mut reveal_to_index = None;

// if the target is not yet revealed, but is already stored (due to lookahead), we need to
// set the `reveal_to_index` as target here (as the `for` loop below only updates
// `reveal_to_index` for indexes that are NOT stored)
if next_reveal_index <= target_index && target_index < next_reveal_index + lookahead {
if next_reveal_index <= target_index && target_index < next_reveal_index + self.lookahead {
reveal_to_index = Some(target_index);
}

// we range over indexes that are not stored
let range = next_reveal_index + lookahead..=target_index + lookahead;
let range = next_reveal_index + self.lookahead..=target_index + self.lookahead;
for (new_index, new_spk) in SpkIterator::new_with_range(descriptor, range) {
let _inserted = self
.inner
Expand Down
2 changes: 1 addition & 1 deletion crates/chain/src/spk_iter.rs
Original file line number Diff line number Diff line change
Expand Up @@ -148,7 +148,7 @@ mod test {
Descriptor<DescriptorPublicKey>,
Descriptor<DescriptorPublicKey>,
) {
let mut txout_index = KeychainTxOutIndex::<TestKeychain>::default();
let mut txout_index = KeychainTxOutIndex::<TestKeychain>::new(0);

let secp = Secp256k1::signing_only();
let (external_descriptor,_) = Descriptor::<DescriptorPublicKey>::parse_descriptor(&secp, "tr([73c5da0a/86'/0'/0']xprv9xgqHN7yz9MwCkxsBPN5qetuNdQSUttZNKw1dcYTV4mkaAFiBVGQziHs3NRSWMkCzvgjEe3n9xV8oYywvM8at9yRqyaZVz6TYYhX98VjsUk/0/*)").unwrap();
Expand Down
11 changes: 6 additions & 5 deletions crates/chain/tests/test_indexed_tx_graph.rs
Original file line number Diff line number Diff line change
Expand Up @@ -27,9 +27,10 @@ fn insert_relevant_txs() {
let spk_0 = descriptor.at_derivation_index(0).unwrap().script_pubkey();
let spk_1 = descriptor.at_derivation_index(9).unwrap().script_pubkey();

let mut graph = IndexedTxGraph::<ConfirmationHeightAnchor, KeychainTxOutIndex<()>>::default();
let mut graph = IndexedTxGraph::<ConfirmationHeightAnchor, KeychainTxOutIndex<()>>::new(
KeychainTxOutIndex::new(10),
);
graph.index.add_keychain((), descriptor);
graph.index.set_lookahead(&(), 10);

let tx_a = Transaction {
output: vec![
Expand Down Expand Up @@ -118,12 +119,12 @@ fn test_list_owned_txouts() {
let (desc_1, _) = Descriptor::parse_descriptor(&Secp256k1::signing_only(), "tr(tprv8ZgxMBicQKsPd3krDUsBAmtnRsK3rb8u5yi1zhQgMhF1tR8MW7xfE4rnrbbsrbPR52e7rKapu6ztw1jXveJSCGHEriUGZV7mCe88duLp5pj/86'/1'/0'/0/*)").unwrap();
let (desc_2, _) = Descriptor::parse_descriptor(&Secp256k1::signing_only(), "tr(tprv8ZgxMBicQKsPd3krDUsBAmtnRsK3rb8u5yi1zhQgMhF1tR8MW7xfE4rnrbbsrbPR52e7rKapu6ztw1jXveJSCGHEriUGZV7mCe88duLp5pj/86'/1'/0'/1/*)").unwrap();

let mut graph =
IndexedTxGraph::<ConfirmationHeightAnchor, KeychainTxOutIndex<String>>::default();
let mut graph = IndexedTxGraph::<ConfirmationHeightAnchor, KeychainTxOutIndex<String>>::new(
KeychainTxOutIndex::new(10),
);

graph.index.add_keychain("keychain_1".into(), desc_1);
graph.index.add_keychain("keychain_2".into(), desc_2);
graph.index.set_lookahead_for_all(10);

// Get trusted and untrusted addresses

Expand Down
32 changes: 12 additions & 20 deletions crates/chain/tests/test_keychain_txout_index.rs
Original file line number Diff line number Diff line change
Expand Up @@ -18,12 +18,14 @@ enum TestKeychain {
Internal,
}

fn init_txout_index() -> (
fn init_txout_index(
lookahead: u32,
) -> (
bdk_chain::keychain::KeychainTxOutIndex<TestKeychain>,
Descriptor<DescriptorPublicKey>,
Descriptor<DescriptorPublicKey>,
) {
let mut txout_index = bdk_chain::keychain::KeychainTxOutIndex::<TestKeychain>::default();
let mut txout_index = bdk_chain::keychain::KeychainTxOutIndex::<TestKeychain>::new(lookahead);

let secp = bdk_chain::bitcoin::secp256k1::Secp256k1::signing_only();
let (external_descriptor,_) = Descriptor::<DescriptorPublicKey>::parse_descriptor(&secp, "tr([73c5da0a/86'/0'/0']xprv9xgqHN7yz9MwCkxsBPN5qetuNdQSUttZNKw1dcYTV4mkaAFiBVGQziHs3NRSWMkCzvgjEe3n9xV8oYywvM8at9yRqyaZVz6TYYhX98VjsUk/0/*)").unwrap();
Expand All @@ -46,7 +48,7 @@ fn spk_at_index(descriptor: &Descriptor<DescriptorPublicKey>, index: u32) -> Scr
fn test_set_all_derivation_indices() {
use bdk_chain::indexed_tx_graph::Indexer;

let (mut txout_index, _, _) = init_txout_index();
let (mut txout_index, _, _) = init_txout_index(0);
let derive_to: BTreeMap<_, _> =
[(TestKeychain::External, 12), (TestKeychain::Internal, 24)].into();
assert_eq!(
Expand All @@ -64,19 +66,10 @@ fn test_set_all_derivation_indices() {

#[test]
fn test_lookahead() {
let (mut txout_index, external_desc, internal_desc) = init_txout_index();

// ensure it does not break anything if lookahead is set multiple times
(0..=10).for_each(|lookahead| txout_index.set_lookahead(&TestKeychain::External, lookahead));
(0..=20)
.filter(|v| v % 2 == 0)
.for_each(|lookahead| txout_index.set_lookahead(&TestKeychain::Internal, lookahead));

assert_eq!(txout_index.inner().all_spks().len(), 30);
let (mut txout_index, external_desc, internal_desc) = init_txout_index(10);

// given:
// - external lookahead set to 10
// - internal lookahead set to 20
// when:
// - set external derivation index to value higher than last, but within the lookahead value
// expect:
Expand All @@ -97,7 +90,7 @@ fn test_lookahead() {
assert_eq!(
txout_index.inner().all_spks().len(),
10 /* external lookahead */ +
20 /* internal lookahead */ +
10 /* internal lookahead */ +
index as usize + 1 /* `derived` count */
);
assert_eq!(
Expand Down Expand Up @@ -127,7 +120,7 @@ fn test_lookahead() {
}

// given:
// - internal lookahead is 20
// - internal lookahead is 10
// - internal derivation index is `None`
// when:
// - derivation index is set ahead of current derivation index + lookahead
Expand All @@ -148,7 +141,7 @@ fn test_lookahead() {
assert_eq!(
txout_index.inner().all_spks().len(),
10 /* external lookahead */ +
20 /* internal lookahead */ +
10 /* internal lookahead */ +
20 /* external stored index count */ +
25 /* internal stored index count */
);
Expand Down Expand Up @@ -226,8 +219,7 @@ fn test_lookahead() {
// - last used index should change as expected
#[test]
fn test_scan_with_lookahead() {
let (mut txout_index, external_desc, _) = init_txout_index();
txout_index.set_lookahead_for_all(10);
let (mut txout_index, external_desc, _) = init_txout_index(10);

let spks: BTreeMap<u32, ScriptBuf> = [0, 10, 20, 30]
.into_iter()
Expand Down Expand Up @@ -281,7 +273,7 @@ fn test_scan_with_lookahead() {
#[test]
#[rustfmt::skip]
fn test_wildcard_derivations() {
let (mut txout_index, external_desc, _) = init_txout_index();
let (mut txout_index, external_desc, _) = init_txout_index(0);
let external_spk_0 = external_desc.at_derivation_index(0).unwrap().script_pubkey();
let external_spk_16 = external_desc.at_derivation_index(16).unwrap().script_pubkey();
let external_spk_26 = external_desc.at_derivation_index(26).unwrap().script_pubkey();
Expand Down Expand Up @@ -339,7 +331,7 @@ fn test_wildcard_derivations() {

#[test]
fn test_non_wildcard_derivations() {
let mut txout_index = KeychainTxOutIndex::<TestKeychain>::default();
let mut txout_index = KeychainTxOutIndex::<TestKeychain>::new(0);

let secp = bitcoin::secp256k1::Secp256k1::signing_only();
let (no_wildcard_descriptor, _) = Descriptor::<DescriptorPublicKey>::parse_descriptor(&secp, "wpkh([73c5da0a/86'/0'/0']xprv9xgqHN7yz9MwCkxsBPN5qetuNdQSUttZNKw1dcYTV4mkaAFiBVGQziHs3NRSWMkCzvgjEe3n9xV8oYywvM8at9yRqyaZVz6TYYhX98VjsUk/1/0)").unwrap();
Expand Down
14 changes: 2 additions & 12 deletions example-crates/example_bitcoind_rpc_polling/src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -64,9 +64,6 @@ struct RpcArgs {
/// Starting block height to fallback to if no point of agreement if found
#[clap(env = "FALLBACK_HEIGHT", long, default_value = "0")]
fallback_height: u32,
/// The unused-scripts lookahead will be kept at this size
#[clap(long, default_value = "10")]
lookahead: u32,
}

impl From<RpcArgs> for Auth {
Expand Down Expand Up @@ -161,13 +158,9 @@ fn main() -> anyhow::Result<()> {
match rpc_cmd {
RpcCommands::Sync { rpc_args } => {
let RpcArgs {
fallback_height,
lookahead,
..
fallback_height, ..
} = rpc_args;

graph.lock().unwrap().index.set_lookahead_for_all(lookahead);

let chain_tip = chain.lock().unwrap().tip();
let rpc_client = rpc_args.new_client()?;
let mut emitter = Emitter::new(&rpc_client, chain_tip, fallback_height);
Expand Down Expand Up @@ -233,13 +226,10 @@ fn main() -> anyhow::Result<()> {
}
RpcCommands::Live { rpc_args } => {
let RpcArgs {
fallback_height,
lookahead,
..
fallback_height, ..
} = rpc_args;
let sigterm_flag = start_ctrlc_handler();

graph.lock().unwrap().index.set_lookahead_for_all(lookahead);
let last_cp = chain.lock().unwrap().tip();

println!(
Expand Down

0 comments on commit 4fd539b

Please sign in to comment.