Skip to content

Commit

Permalink
fix(wallet)!: Change method LoadParams::descriptors
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
ValuedMammal committed Aug 13, 2024
1 parent b802714 commit 3951110
Show file tree
Hide file tree
Showing 8 changed files with 166 additions and 76 deletions.
4 changes: 3 additions & 1 deletion crates/wallet/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
147 changes: 90 additions & 57 deletions crates/wallet/src/wallet/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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>,
},
}

Expand Down Expand Up @@ -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())),
};
Expand Down Expand Up @@ -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
///
Expand All @@ -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)
Expand Down Expand Up @@ -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 {
Expand All @@ -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)?;

Expand Down
30 changes: 23 additions & 7 deletions crates/wallet/src/wallet/params.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -161,6 +162,7 @@ impl LoadParams {
check_genesis_hash: None,
check_descriptor: None,
check_change_descriptor: None,
extract_keys: false,
}
}

Expand All @@ -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
}

Expand All @@ -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,
Expand Down
43 changes: 37 additions & 6 deletions crates/wallet/tests/wallet.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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");
Expand Down Expand Up @@ -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 { .. }
Expand Down Expand Up @@ -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]
Expand Down
4 changes: 3 additions & 1 deletion example-crates/wallet_electrum/src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
4 changes: 3 additions & 1 deletion example-crates/wallet_esplora_async/src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
4 changes: 3 additions & 1 deletion example-crates/wallet_esplora_blocking/src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
Loading

0 comments on commit 3951110

Please sign in to comment.