chain: add a 'lookahead' parameter to add_keychain

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 the wallet, make the interface foolproof by requiring any caller of `add_keychain` to set a lookahead value. (h/t evanlinjin for the suggestion.)
bitcoindevkit · Nov 28, 2023 · 6f2cb9c · 6f2cb9c
1 parent 55b680c
commit 6f2cb9c
Show file tree

Hide file tree

Showing 6 changed files with 45 additions and 37 deletions.
diff --git a/crates/bdk/src/wallet/mod.rs b/crates/bdk/src/wallet/mod.rs
@@ -75,6 +75,8 @@ use crate::wallet::error::{BuildFeeBumpError, CreateTxError, MiniscriptPsbtError
 
 const COINBASE_MATURITY: u32 = 100;
 
+const DEFAULT_LOOKAHEAD: u32 = 1_000;
+
 /// A Bitcoin wallet
 ///
 /// The `Wallet` struct acts as a way of coherently interfacing with output descriptors and related transactions.
@@ -2375,13 +2377,13 @@ fn create_signers<E: IntoWalletDescriptor>(
 ) -> Result<(Arc<SignersContainer>, Arc<SignersContainer>), crate::descriptor::error::Error> {
     let (descriptor, keymap) = into_wallet_descriptor_checked(descriptor, secp, network)?;
     let signers = Arc::new(SignersContainer::build(keymap, &descriptor, secp));
-    index.add_keychain(KeychainKind::External, descriptor);
+    index.add_keychain(KeychainKind::External, descriptor, DEFAULT_LOOKAHEAD);
 
     let change_signers = match change_descriptor {
         Some(descriptor) => {
             let (descriptor, keymap) = into_wallet_descriptor_checked(descriptor, secp, network)?;
             let signers = Arc::new(SignersContainer::build(keymap, &descriptor, secp));
-            index.add_keychain(KeychainKind::Internal, descriptor);
+            index.add_keychain(KeychainKind::Internal, descriptor, DEFAULT_LOOKAHEAD);
             signers
         }
         None => Arc::new(SignersContainer::new()),

diff --git a/crates/chain/src/keychain/txout_index.rs b/crates/chain/src/keychain/txout_index.rs
@@ -47,9 +47,9 @@ use crate::Append;
 /// # 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();
-/// 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);
+/// txout_index.add_keychain(MyKeychain::External, external_descriptor, 0);
+/// txout_index.add_keychain(MyKeychain::Internal, internal_descriptor, 0);
+/// txout_index.add_keychain(MyKeychain::MyAppUser { user_id: 42 }, descriptor_for_user_42, 0);
 ///
 /// let new_spk_for_user = txout_index.reveal_next_spk(&MyKeychain::MyAppUser{ user_id: 42 });
 /// ```
@@ -142,15 +142,22 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
     /// # Panics
     ///
     /// This will panic if a different `descriptor` is introduced to the same `keychain`.
-    pub fn add_keychain(&mut self, keychain: K, descriptor: Descriptor<DescriptorPublicKey>) {
+    pub fn add_keychain(
+        &mut self,
+        keychain: K,
+        descriptor: Descriptor<DescriptorPublicKey>,
+        lookahead: u32,
+    ) {
         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.lookahead.insert(keychain.clone(), lookahead);
+        self.replenish_lookahead(&keychain);
     }
 
     /// Return the lookahead setting for each keychain.
@@ -390,10 +397,7 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
         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 + 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;

diff --git a/crates/chain/src/spk_iter.rs b/crates/chain/src/spk_iter.rs
@@ -154,8 +154,8 @@ mod test {
         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();
 
-        txout_index.add_keychain(TestKeychain::External, external_descriptor.clone());
-        txout_index.add_keychain(TestKeychain::Internal, internal_descriptor.clone());
+        txout_index.add_keychain(TestKeychain::External, external_descriptor.clone(), 0);
+        txout_index.add_keychain(TestKeychain::Internal, internal_descriptor.clone(), 0);
 
         (txout_index, external_descriptor, internal_descriptor)
     }

diff --git a/crates/chain/tests/test_indexed_tx_graph.rs b/crates/chain/tests/test_indexed_tx_graph.rs
@@ -28,8 +28,7 @@ fn insert_relevant_txs() {
     let spk_1 = descriptor.at_derivation_index(9).unwrap().script_pubkey();
 
     let mut graph = IndexedTxGraph::<ConfirmationHeightAnchor, KeychainTxOutIndex<()>>::default();
-    graph.index.add_keychain((), descriptor);
-    graph.index.set_lookahead(&(), 10);
+    graph.index.add_keychain((), descriptor, 10);
 
     let tx_a = Transaction {
         output: vec![
@@ -121,9 +120,8 @@ fn test_list_owned_txouts() {
     let mut graph =
         IndexedTxGraph::<ConfirmationHeightAnchor, KeychainTxOutIndex<String>>::default();
 
-    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);
+    graph.index.add_keychain("keychain_1".into(), desc_1, 10);
+    graph.index.add_keychain("keychain_2".into(), desc_2, 10);
 
     // Get trusted and untrusted addresses
 

diff --git a/crates/chain/tests/test_keychain_txout_index.rs b/crates/chain/tests/test_keychain_txout_index.rs
@@ -18,7 +18,10 @@ enum TestKeychain {
     Internal,
 }
 
-fn init_txout_index() -> (
+fn init_txout_index(
+    internal_lookahead: u32,
+    external_lookahead: u32,
+) -> (
     bdk_chain::keychain::KeychainTxOutIndex<TestKeychain>,
     Descriptor<DescriptorPublicKey>,
     Descriptor<DescriptorPublicKey>,
@@ -29,8 +32,16 @@ fn init_txout_index() -> (
     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();
 
-    txout_index.add_keychain(TestKeychain::External, external_descriptor.clone());
-    txout_index.add_keychain(TestKeychain::Internal, internal_descriptor.clone());
+    txout_index.add_keychain(
+        TestKeychain::External,
+        external_descriptor.clone(),
+        internal_lookahead,
+    );
+    txout_index.add_keychain(
+        TestKeychain::Internal,
+        internal_descriptor.clone(),
+        external_lookahead,
+    );
 
     (txout_index, external_descriptor, internal_descriptor)
 }
@@ -46,7 +57,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, 0);
     let derive_to: BTreeMap<_, _> =
         [(TestKeychain::External, 12), (TestKeychain::Internal, 24)].into();
     assert_eq!(
@@ -64,15 +75,7 @@ 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, 20);
 
     // given:
     // - external lookahead set to 10
@@ -226,8 +229,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, 10);
 
     let spks: BTreeMap<u32, ScriptBuf> = [0, 10, 20, 30]
         .into_iter()
@@ -281,7 +283,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, 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();
@@ -348,7 +350,7 @@ fn test_non_wildcard_derivations() {
         .unwrap()
         .script_pubkey();
 
-    txout_index.add_keychain(TestKeychain::External, no_wildcard_descriptor);
+    txout_index.add_keychain(TestKeychain::External, no_wildcard_descriptor, 0);
 
     // given:
     // - `txout_index` with no stored scripts

diff --git a/example-crates/example_cli/src/lib.rs b/example-crates/example_cli/src/lib.rs
@@ -31,6 +31,8 @@ pub type KeychainChangeSet<A> = (
 );
 pub type Database<'m, C> = Persist<Store<'m, C>, C>;
 
+const DEFAULT_LOOKAHEAD: u32 = 1_000;
+
 #[derive(Parser)]
 #[clap(author, version, about, long_about = None)]
 #[clap(propagate_version = true)]
@@ -669,7 +671,7 @@ where
 
     let (descriptor, mut keymap) =
         Descriptor::<DescriptorPublicKey>::parse_descriptor(&secp, &args.descriptor)?;
-    index.add_keychain(Keychain::External, descriptor);
+    index.add_keychain(Keychain::External, descriptor, DEFAULT_LOOKAHEAD);
 
     if let Some((internal_descriptor, internal_keymap)) = args
         .change_descriptor
@@ -678,7 +680,7 @@ where
         .transpose()?
     {
         keymap.extend(internal_keymap);
-        index.add_keychain(Keychain::Internal, internal_descriptor);
+        index.add_keychain(Keychain::Internal, internal_descriptor, DEFAULT_LOOKAHEAD);
     }
 
     let mut db_backend = match Store::<'m, C>::open_or_create_new(db_magic, &args.db_path) {