fix(wallet)!: Change method LoadParams::descriptors

ValuedMammal · ValuedMammal · commit 3951110bb596 · 2024-08-13T19:25:09.000-04:00
to just `descriptor` that takes a `KeychainKind` and optional
`D: IntoWalletDescriptor` representing the expected value of
the descriptor in the changeset.

Add method `LoadParams::extract_keys` that will use any private
keys in the provided descriptors to add a signer to the wallet.
diff --git a/crates/wallet/README.md b/crates/wallet/README.md
@@ -79,8 +79,10 @@ let network = Network::Testnet;
 let descriptor = "wpkh(tprv8ZgxMBicQKsPdcAqYBpzAFwU5yxBUo88ggoBqu1qPcHUfSbKK1sKMLmC7EAk438btHQrSdu3jGGQa6PA71nvH5nkDexhLteJqkM4dQmWF9g/84'/1'/0'/0/*)";
 let change_descriptor = "wpkh(tprv8ZgxMBicQKsPdcAqYBpzAFwU5yxBUo88ggoBqu1qPcHUfSbKK1sKMLmC7EAk438btHQrSdu3jGGQa6PA71nvH5nkDexhLteJqkM4dQmWF9g/84'/1'/0'/1/*)";
 let wallet_opt = Wallet::load()
-    .descriptors(descriptor, change_descriptor)
     .network(network)
+    .descriptor(KeychainKind::External, Some(descriptor))
+    .descriptor(KeychainKind::Internal, Some(change_descriptor))
+    .extract_keys()
     .load_wallet(&mut db)
     .expect("wallet");
 let mut wallet = match wallet_opt {
diff --git a/crates/wallet/src/wallet/mod.rs b/crates/wallet/src/wallet/mod.rs
@@ -246,9 +246,9 @@ pub enum LoadMismatch {
         /// Keychain identifying the descriptor.
         keychain: KeychainKind,
         /// The loaded descriptor.
-        loaded: ExtendedDescriptor,
+        loaded: Option<ExtendedDescriptor>,
         /// The expected descriptor.
-        expected: ExtendedDescriptor,
+        expected: Option<ExtendedDescriptor>,
     },
 }
 
@@ -401,11 +401,15 @@ impl Wallet {
         ));
 
         let (change_descriptor, change_signers) = match params.change_descriptor {
-            Some(desc_fn) => {
-                let (descriptor, mut keymap) = desc_fn(&secp, network)?;
-                keymap.extend(params.change_descriptor_keymap);
-                let change_signers = Arc::new(SignersContainer::build(keymap, &descriptor, &secp));
-                (Some(descriptor), change_signers)
+            Some(make_desc) => {
+                let (change_descriptor, mut internal_keymap) = make_desc(&secp, network)?;
+                internal_keymap.extend(params.change_descriptor_keymap);
+                let change_signers = Arc::new(SignersContainer::build(
+                    internal_keymap,
+                    &change_descriptor,
+                    &secp,
+                ));
+                (Some(change_descriptor), change_signers)
             }
             None => (None, Arc::new(SignersContainer::new())),
         };
@@ -442,7 +446,7 @@ impl Wallet {
     /// Note that the descriptor secret keys are not persisted to the db. You can either add
     /// signers after-the-fact with [`Wallet::add_signer`] or [`Wallet::set_keymap`]. Or you can
     /// add keys when building the wallet using [`LoadParams::keymap`] and/or
-    /// [`LoadParams::descriptors`].
+    /// [`LoadParams::descriptor`].
     ///
     /// # Synopsis
     ///
@@ -467,7 +471,9 @@ impl Wallet {
     /// let mut conn = bdk_wallet::rusqlite::Connection::open(file_path)?;
     /// let mut wallet = Wallet::load()
     ///     // check loaded descriptors matches these values and extract private keys
-    ///     .descriptors(EXTERNAL_DESC, INTERNAL_DESC)
+    ///     .descriptor(KeychainKind::External, Some(EXTERNAL_DESC))
+    ///     .descriptor(KeychainKind::Internal, Some(INTERNAL_DESC))
+    ///     .extract_keys()
     ///     // you can also manually add private keys
     ///     .keymap(KeychainKind::External, external_keymap)
     ///     .keymap(KeychainKind::Internal, internal_keymap)
@@ -499,42 +505,6 @@ impl Wallet {
         let chain = LocalChain::from_changeset(changeset.local_chain)
             .map_err(|_| LoadError::MissingGenesis)?;
 
-        let mut descriptor_keymap = params.descriptor_keymap;
-        let descriptor = changeset
-            .descriptor
-            .ok_or(LoadError::MissingDescriptor(KeychainKind::External))?;
-        check_wallet_descriptor(&descriptor).map_err(LoadError::Descriptor)?;
-
-        let (change_descriptor, change_signers) = match changeset.change_descriptor {
-            Some(change_descriptor) => {
-                check_wallet_descriptor(&change_descriptor).map_err(LoadError::Descriptor)?;
-                let mut change_descriptor_keymap = params.change_descriptor_keymap;
-
-                // check params match loaded
-                if let Some(exp_change_descriptor) = params.check_change_descriptor {
-                    let (exp_change_descriptor, keymap) =
-                        (exp_change_descriptor)(&secp, network).map_err(LoadError::Descriptor)?;
-                    change_descriptor_keymap.extend(keymap);
-
-                    if change_descriptor.descriptor_id() != exp_change_descriptor.descriptor_id() {
-                        return Err(LoadError::Mismatch(LoadMismatch::Descriptor {
-                            keychain: KeychainKind::Internal,
-                            loaded: change_descriptor,
-                            expected: exp_change_descriptor,
-                        }));
-                    }
-                }
-                let change_signers = Arc::new(SignersContainer::build(
-                    change_descriptor_keymap,
-                    &change_descriptor,
-                    &secp,
-                ));
-                (Some(change_descriptor), change_signers)
-            }
-            None => (None, Arc::new(SignersContainer::new())),
-        };
-
-        // checks
         if let Some(exp_network) = params.check_network {
             if network != exp_network {
                 return Err(LoadError::Mismatch(LoadMismatch::Network {
@@ -551,25 +521,88 @@ impl Wallet {
                 }));
             }
         }
-        if let Some(exp_descriptor) = params.check_descriptor {
-            let (exp_descriptor, keymap) =
-                (exp_descriptor)(&secp, network).map_err(LoadError::Descriptor)?;
-            descriptor_keymap.extend(keymap);
 
-            if descriptor.descriptor_id() != exp_descriptor.descriptor_id() {
+        let descriptor = changeset
+            .descriptor
+            .ok_or(LoadError::MissingDescriptor(KeychainKind::External))?;
+        check_wallet_descriptor(&descriptor).map_err(LoadError::Descriptor)?;
+        let mut external_keymap = params.descriptor_keymap;
+
+        if let Some(expected) = params.check_descriptor {
+            if let Some(make_desc) = expected {
+                let (exp_desc, keymap) =
+                    make_desc(&secp, network).map_err(LoadError::Descriptor)?;
+                if descriptor.descriptor_id() != exp_desc.descriptor_id() {
+                    return Err(LoadError::Mismatch(LoadMismatch::Descriptor {
+                        keychain: KeychainKind::External,
+                        loaded: Some(descriptor),
+                        expected: Some(exp_desc),
+                    }));
+                }
+                if params.extract_keys {
+                    external_keymap.extend(keymap);
+                }
+            } else {
                 return Err(LoadError::Mismatch(LoadMismatch::Descriptor {
                     keychain: KeychainKind::External,
-                    loaded: descriptor,
-                    expected: exp_descriptor,
+                    loaded: Some(descriptor),
+                    expected: None,
                 }));
             }
         }
+        let signers = Arc::new(SignersContainer::build(external_keymap, &descriptor, &secp));
+
+        let (mut change_descriptor, mut change_signers) = (None, Arc::new(SignersContainer::new()));
+        match (changeset.change_descriptor, params.check_change_descriptor) {
+            // empty signer
+            (None, None) => {}
+            (None, Some(expect)) => {
+                // expected desc but none loaded
+                if let Some(make_desc) = expect {
+                    let (exp_desc, _) = make_desc(&secp, network).map_err(LoadError::Descriptor)?;
+                    return Err(LoadError::Mismatch(LoadMismatch::Descriptor {
+                        keychain: KeychainKind::Internal,
+                        loaded: None,
+                        expected: Some(exp_desc),
+                    }));
+                }
+            }
+            // nothing expected
+            (Some(desc), None) => {
+                check_wallet_descriptor(&desc).map_err(LoadError::Descriptor)?;
+                change_descriptor = Some(desc);
+            }
+            (Some(desc), Some(expect)) => match expect {
+                // expected none for existing
+                None => {
+                    return Err(LoadError::Mismatch(LoadMismatch::Descriptor {
+                        keychain: KeychainKind::Internal,
+                        loaded: Some(desc),
+                        expected: None,
+                    }))
+                }
+                // parameters must match
+                Some(make_desc) => {
+                    let (exp_desc, keymap) =
+                        make_desc(&secp, network).map_err(LoadError::Descriptor)?;
+                    if desc.descriptor_id() != exp_desc.descriptor_id() {
+                        return Err(LoadError::Mismatch(LoadMismatch::Descriptor {
+                            keychain: KeychainKind::Internal,
+                            loaded: Some(desc),
+                            expected: Some(exp_desc),
+                        }));
+                    }
+                    let mut internal_keymap = params.change_descriptor_keymap;
+                    if params.extract_keys {
+                        internal_keymap.extend(keymap);
+                    }
+                    change_signers =
+                        Arc::new(SignersContainer::build(internal_keymap, &desc, &secp));
+                    change_descriptor = Some(desc);
+                }
+            },
+        }
 
-        let signers = Arc::new(SignersContainer::build(
-            descriptor_keymap,
-            &descriptor,
-            &secp,
-        ));
         let index = create_indexer(descriptor, change_descriptor, params.lookahead)
             .map_err(LoadError::Descriptor)?;
 
diff --git a/crates/wallet/src/wallet/params.rs b/crates/wallet/src/wallet/params.rs
@@ -144,8 +144,9 @@ pub struct LoadParams {
     pub(crate) lookahead: u32,
     pub(crate) check_network: Option<Network>,
     pub(crate) check_genesis_hash: Option<BlockHash>,
-    pub(crate) check_descriptor: Option<DescriptorToExtract>,
-    pub(crate) check_change_descriptor: Option<DescriptorToExtract>,
+    pub(crate) check_descriptor: Option<Option<DescriptorToExtract>>,
+    pub(crate) check_change_descriptor: Option<Option<DescriptorToExtract>>,
+    pub(crate) extract_keys: bool,
 }
 
 impl LoadParams {
@@ -161,6 +162,7 @@ impl LoadParams {
             check_genesis_hash: None,
             check_descriptor: None,
             check_change_descriptor: None,
+            extract_keys: false,
         }
     }
 
@@ -174,14 +176,21 @@ impl LoadParams {
         self
     }
 
-    /// Checks that `descriptor` of `keychain` matches this, and extracts private keys (if
-    /// available).
-    pub fn descriptors<D>(mut self, descriptor: D, change_descriptor: D) -> Self
+    /// Checks the `expected_descriptor` matches exactly what is loaded for `keychain`.
+    ///
+    /// # Note
+    ///
+    /// You must also specify [`extract_keys`](Self::extract_keys) if you wish to add a signer
+    /// for an expected descriptor containing secrets.
+    pub fn descriptor<D>(mut self, keychain: KeychainKind, expected_descriptor: Option<D>) -> Self
     where
         D: IntoWalletDescriptor + 'static,
     {
-        self.check_descriptor = Some(make_descriptor_to_extract(descriptor));
-        self.check_change_descriptor = Some(make_descriptor_to_extract(change_descriptor));
+        let expected = expected_descriptor.map(|d| make_descriptor_to_extract(d));
+        match keychain {
+            KeychainKind::External => self.check_descriptor = Some(expected),
+            KeychainKind::Internal => self.check_change_descriptor = Some(expected),
+        }
         self
     }
 
@@ -203,6 +212,13 @@ impl LoadParams {
         self
     }
 
+    /// Whether to try extracting private keys from the *provided descriptors* upon loading.
+    /// See also [`LoadParams::descriptor`].
+    pub fn extract_keys(mut self) -> Self {
+        self.extract_keys = true;
+        self
+    }
+
     /// Load [`PersistedWallet`] with the given `Db`.
     pub fn load_wallet<Db>(
         self,
diff --git a/crates/wallet/tests/wallet.rs b/crates/wallet/tests/wallet.rs
@@ -134,10 +134,18 @@ fn wallet_is_persisted() -> anyhow::Result<()> {
         };
 
         // recover wallet
+        {
+            let mut db = open_db(&file_path).context("failed to recover db")?;
+            let _ = Wallet::load()
+                .network(Network::Testnet)
+                .load_wallet(&mut db)?
+                .expect("wallet must exist");
+        }
         {
             let mut db = open_db(&file_path).context("failed to recover db")?;
             let wallet = Wallet::load()
-                .descriptors(external_desc, internal_desc)
+                .descriptor(KeychainKind::External, Some(external_desc))
+                .descriptor(KeychainKind::Internal, Some(internal_desc))
                 .network(Network::Testnet)
                 .load_wallet(&mut db)?
                 .expect("wallet must exist");
@@ -240,7 +248,16 @@ fn wallet_load_checks() -> anyhow::Result<()> {
         );
         assert_matches!(
             Wallet::load()
-                .descriptors(internal_desc, external_desc)
+                .descriptor(KeychainKind::External, Some(internal_desc))
+                .load_wallet(&mut open_db(&file_path)?),
+            Err(LoadWithPersistError::InvalidChangeSet(LoadError::Mismatch(
+                LoadMismatch::Descriptor { .. }
+            ))),
+            "unexpected descriptors check result",
+        );
+        assert_matches!(
+            Wallet::load()
+                .descriptor(KeychainKind::External, Option::<&str>::None)
                 .load_wallet(&mut open_db(&file_path)?),
             Err(LoadWithPersistError::InvalidChangeSet(LoadError::Mismatch(
                 LoadMismatch::Descriptor { .. }
@@ -279,25 +296,39 @@ fn single_descriptor_wallet_persist_and_recover() {
         .network(Network::Testnet)
         .create_wallet(&mut db)
         .unwrap();
-
     let _ = wallet.reveal_addresses_to(KeychainKind::External, 2);
     assert!(wallet.persist(&mut db).unwrap());
 
     // should recover persisted wallet
     let secp = wallet.secp_ctx();
     let (_, keymap) = <Descriptor<DescriptorPublicKey>>::parse_descriptor(secp, desc).unwrap();
-    let wallet = LoadParams::new()
-        .keymap(KeychainKind::External, keymap.clone())
+    assert!(!keymap.is_empty());
+    let wallet = Wallet::load()
+        .descriptor(KeychainKind::External, Some(desc))
+        .extract_keys()
         .load_wallet(&mut db)
         .unwrap()
         .expect("must have loaded changeset");
-
     assert_eq!(wallet.derivation_index(KeychainKind::External), Some(2));
     // should have private key
     assert_eq!(
         wallet.get_signers(KeychainKind::External).as_key_map(secp),
         keymap,
     );
+
+    // should error on wrong internal params
+    let desc = get_test_wpkh();
+    let (exp_desc, _) = <Descriptor<DescriptorPublicKey>>::parse_descriptor(secp, desc).unwrap();
+    let err = Wallet::load()
+        .descriptor(KeychainKind::Internal, Some(desc))
+        .extract_keys()
+        .load_wallet(&mut db);
+    assert_matches!(
+        err,
+        Err(LoadWithPersistError::InvalidChangeSet(LoadError::Mismatch(LoadMismatch::Descriptor { keychain, loaded, expected })))
+        if keychain == KeychainKind::Internal && loaded.is_none() && expected == Some(exp_desc),
+        "single descriptor wallet should refuse change descriptor param"
+    );
 }
 
 #[test]
diff --git a/example-crates/wallet_electrum/src/main.rs b/example-crates/wallet_electrum/src/main.rs
@@ -26,8 +26,10 @@ fn main() -> Result<(), anyhow::Error> {
     let mut db = Store::<bdk_wallet::ChangeSet>::open_or_create_new(DB_MAGIC.as_bytes(), db_path)?;
 
     let wallet_opt = Wallet::load()
-        .descriptors(EXTERNAL_DESC, INTERNAL_DESC)
         .network(NETWORK)
+        .descriptor(KeychainKind::External, Some(EXTERNAL_DESC))
+        .descriptor(KeychainKind::Internal, Some(INTERNAL_DESC))
+        .extract_keys()
         .load_wallet(&mut db)?;
     let mut wallet = match wallet_opt {
         Some(wallet) => wallet,
diff --git a/example-crates/wallet_esplora_async/src/main.rs b/example-crates/wallet_esplora_async/src/main.rs
@@ -23,8 +23,10 @@ async fn main() -> Result<(), anyhow::Error> {
     let mut conn = Connection::open(DB_PATH)?;
 
     let wallet_opt = Wallet::load()
-        .descriptors(EXTERNAL_DESC, INTERNAL_DESC)
         .network(NETWORK)
+        .descriptor(KeychainKind::External, Some(EXTERNAL_DESC))
+        .descriptor(KeychainKind::Internal, Some(INTERNAL_DESC))
+        .extract_keys()
         .load_wallet(&mut conn)?;
     let mut wallet = match wallet_opt {
         Some(wallet) => wallet,
diff --git a/example-crates/wallet_esplora_blocking/src/main.rs b/example-crates/wallet_esplora_blocking/src/main.rs
@@ -22,8 +22,10 @@ fn main() -> Result<(), anyhow::Error> {
     let mut db = Store::<bdk_wallet::ChangeSet>::open_or_create_new(DB_MAGIC.as_bytes(), DB_PATH)?;
 
     let wallet_opt = Wallet::load()
-        .descriptors(EXTERNAL_DESC, INTERNAL_DESC)
         .network(NETWORK)
+        .descriptor(KeychainKind::External, Some(EXTERNAL_DESC))
+        .descriptor(KeychainKind::Internal, Some(INTERNAL_DESC))
+        .extract_keys()
         .load_wallet(&mut db)?;
     let mut wallet = match wallet_opt {
         Some(wallet) => wallet,
diff --git a/example-crates/wallet_rpc/src/main.rs b/example-crates/wallet_rpc/src/main.rs
