Skip to content

Commit

Permalink
keychain::ChangeSet includes the descriptor
Browse files Browse the repository at this point in the history 
- The KeychainTxOutIndex's internal SpkIterator now uses DescriptorId
  instead of K. The DescriptorId -> K translation is made at the
  KeychainTxOutIndex level.
- The keychain::Changeset is now a struct, which includes a map for last
  revealed indexes, and one for newly added keychains and their
  descriptor.

API changes in bdk:
- Wallet::keychains returns a `impl Iterator` instead of `BTreeMap`
- Wallet::load doesn't take descriptors anymore, since they're stored in
  the db
- Wallet::new_or_load checks if the loaded descriptor from db is the
  same as the provided one

API changes in bdk_chain:
- `ChangeSet` is now a struct, which includes a map for last revealed
  indexes, and one for keychains and descriptors.
- `KeychainTxOutIndex::inner` returns a `SpkIterator<(DescriptorId, u32)>`
- `KeychainTxOutIndex::outpoints` returns a `impl Iterator` instead of `&BTreeSet`
- `KeychainTxOutIndex::keychains` returns a `impl Iterator` instead of
  `&BTreeMap`
- `KeychainTxOutIndex::txouts` doesn't return a ExactSizeIterator
  anymore
- `KeychainTxOutIndex::last_revealed_indices` returns a `BTreeMap`
  instead of `&BTreeMap`
- `KeychainTxOutIndex::add_keychain` has been renamed to
  `KeychainTxOutIndex::insert_descriptor`, and now it returns a
  ChangeSet
  • Loading branch information
@danielabrozzoni
danielabrozzoni committed Feb 26, 2024
1 parent e3de705 commit 0a5a82a
Show file tree
Hide file tree
Showing 14 changed files with 689 additions and 306 deletions.
162 changes: 99 additions & 63 deletions crates/bdk/src/wallet/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -342,6 +342,8 @@ pub enum LoadError<L> {
MissingNetwork,
/// Data loaded from persistence is missing genesis hash.
MissingGenesis,
/// Data loaded from persistence is missing descriptor.
MissingDescriptor,
}

impl<L> fmt::Display for LoadError<L>
Expand All @@ -357,6 +359,7 @@ where
}
LoadError::MissingNetwork => write!(f, "loaded data is missing network type"),
LoadError::MissingGenesis => write!(f, "loaded data is missing genesis hash"),
LoadError::MissingDescriptor => write!(f, "loaded data is missing descriptor"),
}
}
}
Expand Down Expand Up @@ -394,6 +397,13 @@ pub enum NewOrLoadError<W, L> {
/// The network type loaded from persistence.
got: Option<Network>,
},
/// The loaded desccriptor does not match what was provided.
LoadedDescriptorDoesNotMatch {
/// The descriptor loaded from persistence.
got: Option<ExtendedDescriptor>,
/// The keychain of the descriptor not matching
keychain: KeychainKind,
},
}

impl<W, L> fmt::Display for NewOrLoadError<W, L>
Expand All @@ -415,6 +425,13 @@ where
NewOrLoadError::LoadedNetworkDoesNotMatch { expected, got } => {
write!(f, "loaded network type is not {}, got {:?}", expected, got)
}
NewOrLoadError::LoadedDescriptorDoesNotMatch { got, keychain } => {
write!(
f,
"loaded descriptor is different from what was provided, got {:?} for keychain {:?}",
got, keychain
)
}
}
}
}
Expand Down Expand Up @@ -553,27 +570,18 @@ impl<D> Wallet<D> {
}

/// Load [`Wallet`] from the given persistence backend.
pub fn load<E: IntoWalletDescriptor>(
descriptor: E,
change_descriptor: Option<E>,
mut db: D,
) -> Result<Self, LoadError<D::LoadError>>
pub fn load(mut db: D) -> Result<Self, LoadError<D::LoadError>>
where
D: PersistBackend<ChangeSet>,
{
let changeset = db
.load_from_persistence()
.map_err(LoadError::Load)?
.ok_or(LoadError::NotInitialized)?;
Self::load_from_changeset(descriptor, change_descriptor, db, changeset)
Self::load_from_changeset(db, changeset)
}

fn load_from_changeset<E: IntoWalletDescriptor>(
descriptor: E,
change_descriptor: Option<E>,
db: D,
changeset: ChangeSet,
) -> Result<Self, LoadError<D::LoadError>>
fn load_from_changeset(db: D, changeset: ChangeSet) -> Result<Self, LoadError<D::LoadError>>
where
D: PersistBackend<ChangeSet>,
{
Expand All @@ -582,6 +590,19 @@ impl<D> Wallet<D> {
let chain =
LocalChain::from_changeset(changeset.chain).map_err(|_| LoadError::MissingGenesis)?;
let mut index = KeychainTxOutIndex::<KeychainKind>::default();
let descriptor = changeset
.indexed_tx_graph
.indexer
.keychains_added
.get(&KeychainKind::External)
.ok_or(LoadError::MissingDescriptor)?
.clone();
let change_descriptor = changeset
.indexed_tx_graph
.indexer
.keychains_added
.get(&KeychainKind::Internal)
.cloned();

let (signers, change_signers) =
create_signers(&mut index, &secp, descriptor, change_descriptor, network)
Expand Down Expand Up @@ -625,8 +646,8 @@ impl<D> Wallet<D> {
)
}

/// Either loads [`Wallet`] from persistence, or initializes it if it does not exist (with a
/// custom genesis hash).
/// Either loads [`Wallet`] from persistence, or initializes it if it does not exist, using the
/// provided descriptor, change descriptor, network, and custom genesis hash.
///
/// This method will fail if the loaded [`Wallet`] has different parameters to those provided.
/// This is like [`Wallet::new_or_load`] with an additional `genesis_hash` parameter. This is
Expand All @@ -644,25 +665,23 @@ impl<D> Wallet<D> {
let changeset = db.load_from_persistence().map_err(NewOrLoadError::Load)?;
match changeset {
Some(changeset) => {
let wallet =
Self::load_from_changeset(descriptor, change_descriptor, db, changeset)
.map_err(|e| match e {
LoadError::Descriptor(e) => NewOrLoadError::Descriptor(e),
LoadError::Load(e) => NewOrLoadError::Load(e),
LoadError::NotInitialized => NewOrLoadError::NotInitialized,
LoadError::MissingNetwork => {
NewOrLoadError::LoadedNetworkDoesNotMatch {
expected: network,
got: None,
}
}
LoadError::MissingGenesis => {
NewOrLoadError::LoadedGenesisDoesNotMatch {
expected: genesis_hash,
got: None,
}
}
})?;
let wallet = Self::load_from_changeset(db, changeset).map_err(|e| match e {
LoadError::Descriptor(e) => NewOrLoadError::Descriptor(e),
LoadError::Load(e) => NewOrLoadError::Load(e),
LoadError::NotInitialized => NewOrLoadError::NotInitialized,
LoadError::MissingNetwork => NewOrLoadError::LoadedNetworkDoesNotMatch {
expected: network,
got: None,
},
LoadError::MissingGenesis => NewOrLoadError::LoadedGenesisDoesNotMatch {
expected: genesis_hash,
got: None,
},
LoadError::MissingDescriptor => NewOrLoadError::LoadedDescriptorDoesNotMatch {
got: None,
keychain: KeychainKind::External,
},
})?;
if wallet.network != network {
return Err(NewOrLoadError::LoadedNetworkDoesNotMatch {
expected: network,
Expand All @@ -675,6 +694,36 @@ impl<D> Wallet<D> {
got: Some(wallet.chain.genesis_hash()),
});
}

let expected_descriptor = descriptor
.into_wallet_descriptor(&wallet.secp, network)
.map_err(NewOrLoadError::Descriptor)?
.0;
let wallet_descriptor = wallet.public_descriptor(KeychainKind::External).cloned();
if wallet_descriptor != Some(expected_descriptor) {
return Err(NewOrLoadError::LoadedDescriptorDoesNotMatch {
got: wallet_descriptor,
keychain: KeychainKind::External,
});
}

let expected_change_descriptor = if let Some(c) = change_descriptor {
Some(
c.into_wallet_descriptor(&wallet.secp, network)
.map_err(NewOrLoadError::Descriptor)?
.0,
)
} else {
None
};
let wallet_change_descriptor =
wallet.public_descriptor(KeychainKind::Internal).cloned();
if wallet_change_descriptor != expected_change_descriptor {
return Err(NewOrLoadError::LoadedDescriptorDoesNotMatch {
got: wallet_change_descriptor,
keychain: KeychainKind::Internal,
});
}
Ok(wallet)
}
None => Self::new_with_genesis_hash(
Expand All @@ -700,7 +749,7 @@ impl<D> Wallet<D> {
}

/// Iterator over all keychains in this wallet
pub fn keychains(&self) -> &BTreeMap<KeychainKind, ExtendedDescriptor> {
pub fn keychains(&self) -> impl Iterator<Item = (&KeychainKind, &ExtendedDescriptor)> {
self.indexed_graph.index.keychains()
}

Expand Down Expand Up @@ -833,7 +882,7 @@ impl<D> Wallet<D> {
.filter_chain_unspents(
&self.chain,
self.chain.tip().block_id(),
self.indexed_graph.index.outpoints().iter().cloned(),
self.indexed_graph.index.outpoints(),
)
.map(|((k, i), full_txo)| new_local_utxo(k, i, full_txo))
}
Expand All @@ -847,7 +896,7 @@ impl<D> Wallet<D> {
.filter_chain_txouts(
&self.chain,
self.chain.tip().block_id(),
self.indexed_graph.index.outpoints().iter().cloned(),
self.indexed_graph.index.outpoints(),
)
.map(|((k, i), full_txo)| new_local_utxo(k, i, full_txo))
}
Expand Down Expand Up @@ -1181,7 +1230,7 @@ impl<D> Wallet<D> {
self.indexed_graph.graph().balance(
&self.chain,
self.chain.tip().block_id(),
self.indexed_graph.index.outpoints().iter().cloned(),
self.indexed_graph.index.outpoints(),
|&(k, _), _| k == KeychainKind::Internal,
)
}
Expand Down Expand Up @@ -1271,17 +1320,9 @@ impl<D> Wallet<D> {
where
D: PersistBackend<ChangeSet>,
{
let external_descriptor = self
.indexed_graph
.index
.keychains()
.get(&KeychainKind::External)
.expect("must exist");
let internal_descriptor = self
.indexed_graph
.index
.keychains()
.get(&KeychainKind::Internal);
let keychains: BTreeMap<_, _> = self.indexed_graph.index.keychains().collect();
let external_descriptor = keychains.get(&KeychainKind::External).expect("must exist");
let internal_descriptor = keychains.get(&KeychainKind::Internal);

let external_policy = external_descriptor
.extract_policy(&self.signers, BuildSatisfaction::None, &self.secp)?
Expand Down Expand Up @@ -1892,7 +1933,11 @@ impl<D> Wallet<D> {
///
/// This can be used to build a watch-only version of a wallet
pub fn public_descriptor(&self, keychain: KeychainKind) -> Option<&ExtendedDescriptor> {
self.indexed_graph.index.keychains().get(&keychain)
self.indexed_graph
.index
.keychains()
.find(|(k, _)| *k == &keychain)
.map(|(_, d)| d)
}

/// Finalize a PSBT, i.e., for each input determine if sufficient data is available to pass
Expand Down Expand Up @@ -1943,17 +1988,9 @@ impl<D> Wallet<D> {
.get_utxo_for(n)
.and_then(|txout| self.get_descriptor_for_txout(&txout))
.or_else(|| {
self.indexed_graph
.index
.keychains()
.iter()
.find_map(|(_, desc)| {
desc.derive_from_psbt_input(
psbt_input,
psbt.get_utxo_for(n),
&self.secp,
)
})
self.indexed_graph.index.keychains().find_map(|(_, desc)| {
desc.derive_from_psbt_input(psbt_input, psbt.get_utxo_for(n), &self.secp)
})
});

match desc {
Expand Down Expand Up @@ -2176,7 +2213,6 @@ impl<D> Wallet<D> {
if params.add_global_xpubs {
let all_xpubs = self
.keychains()
.iter()
.flat_map(|(_, desc)| desc.get_extended_keys())
.collect::<Vec<_>>();

Expand Down Expand Up @@ -2551,13 +2587,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);
let _ = index.insert_descriptor(descriptor, KeychainKind::External);

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);
let _ = index.insert_descriptor(descriptor, KeychainKind::Internal);
signers
}
None => Arc::new(SignersContainer::new()),
Expand Down
Loading

0 comments on commit 0a5a82a

Please sign in to comment.