Skip to content

Commit

Permalink
fix(chain): introduce keychain-variant-ranking to KeychainTxOutIndex
Browse files Browse the repository at this point in the history 
This fixes the bug where the result of applying individual changesets
may result in a different outcome than applying the aggregate changeset.

The nature of the changeset allows different keychain types to be
associated with the same descriptor. The previous design did not take
this into account.

Methods of `KeychainTxOutIndex` only returns one keychain type per
spk/txout/outpoint. This PR makes the keychain type returned to be
consistent by ranking keychain variants by `Ord`. The lowest keychain
variant (according to `Ord`) takes precedence.
  • Loading branch information
@evanlinjin
evanlinjin committed May 8, 2024
1 parent 2e3105c commit e1f3b53
Showing 1 changed file with 79 additions and 81 deletions.
160 changes: 79 additions & 81 deletions crates/chain/src/keychain/txout_index.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -160,6 +160,14 @@ const DEFAULT_LOOKAHEAD: u32 = 25;
/// let new_spk_for_user = txout_index.reveal_next_spk(&MyKeychain::MyAppUser{ user_id: 42 });
/// ```
///
/// # Assigning the same descriptor to multiple keychains
///
/// A keychain can only be associated with one descriptor at one time. However, the same descriptor
/// can be assigned to different keychains. For methods that return the keychain alongside
/// associated data (such as txouts, outpoints, spks), the *highest-ranked* keychain will be
/// returned. Keychain variants are ranked by [`Ord`], where the earliest variant has the highest
/// rank.
///
/// [`Ord`]: core::cmp::Ord
/// [`SpkTxOutIndex`]: crate::spk_txout_index::SpkTxOutIndex
/// [`Descriptor`]: crate::miniscript::Descriptor
Expand All @@ -172,17 +180,19 @@ const DEFAULT_LOOKAHEAD: u32 = 25;
/// [`new`]: KeychainTxOutIndex::new
/// [`unbounded_spk_iter`]: KeychainTxOutIndex::unbounded_spk_iter
/// [`all_unbounded_spk_iters`]: KeychainTxOutIndex::all_unbounded_spk_iters
// Under the hood, the KeychainTxOutIndex uses a SpkTxOutIndex that keeps track of spks, indexed by
// descriptors. Users then assign or unassign keychains to those descriptors. It's important to
// note that descriptors, once added, never get removed from the SpkTxOutIndex; this is useful in
// case a user unassigns a keychain from a descriptor and after some time assigns it again.
#[derive(Clone, Debug)]
pub struct KeychainTxOutIndex<K> {
inner: SpkTxOutIndex<(DescriptorId, u32)>,
// keychain -> (descriptor, descriptor id) map
keychains_to_descriptors: BTreeMap<K, (DescriptorId, Descriptor<DescriptorPublicKey>)>,
// descriptor id -> keychain map
descriptor_ids_to_keychain: BTreeMap<DescriptorId, (K, Descriptor<DescriptorPublicKey>)>,
// descriptor id -> keychain set
// This is a reverse map of `keychains_to_descriptors`. Although there is only one descriptor
// per keychain, different keychains can refer to the same descriptor, therefore we have a set
// of keychains per descriptor. When associated data (such as spks, outpoints) are returned with
// a keychain, we return it with the highest-ranked keychain with it. We rank keychains by
// `Ord`, therefore the keychain set is a `BTreeSet`. The earliest keychain variant (according
// to `Ord`) has precedence.
descriptor_ids_to_keychains: HashMap<DescriptorId, BTreeSet<K>>,
// descriptor_id -> descriptor map
// This is a "monotone" map, meaning that its size keeps growing, i.e., we never delete
// descriptors from it. This is useful for revealing spks for descriptors that don't have
Expand All @@ -208,11 +218,9 @@ impl<K: Clone + Ord + Debug> Indexer for KeychainTxOutIndex<K> {
Some((descriptor_id, index)) => {
// We want to reveal spks for descriptors that aren't tracked by any keychain, and
// so we call reveal with descriptor_id
if let Some((_, changeset)) = self.reveal_to_target_with_id(descriptor_id, index) {
changeset
} else {
super::ChangeSet::default()
}
let (_, changeset) = self.reveal_to_target_with_id(descriptor_id, index)
.expect("descriptors are added in a monotone manner, there cannot be a descriptor id with no corresponding descriptor");
changeset
}
None => super::ChangeSet::default(),
}
Expand Down Expand Up @@ -259,9 +267,9 @@ impl<K> KeychainTxOutIndex<K> {
pub fn new(lookahead: u32) -> Self {
Self {
inner: SpkTxOutIndex::default(),
descriptor_ids_to_keychain: BTreeMap::new(),
descriptor_ids_to_descriptors: BTreeMap::new(),
keychains_to_descriptors: BTreeMap::new(),
descriptor_ids_to_keychains: HashMap::new(),
descriptor_ids_to_descriptors: BTreeMap::new(),
last_revealed: BTreeMap::new(),
lookahead,
}
Expand All @@ -270,6 +278,12 @@ impl<K> KeychainTxOutIndex<K> {

/// Methods that are *re-exposed* from the internal [`SpkTxOutIndex`].
impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
/// Get the highest-ranked keychain that is currently associated with the given `desc_id`.
fn keychain_of_desc_id(&self, desc_id: &DescriptorId) -> Option<&K> {
let keychains = self.descriptor_ids_to_keychains.get(desc_id)?;
keychains.iter().next()
}

/// Return a reference to the internal [`SpkTxOutIndex`].
///
/// **WARNING:** The internal index will contain lookahead spks. Refer to
Expand All @@ -284,18 +298,16 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
.outpoints()
.iter()
.filter_map(|((desc_id, index), op)| {
self.descriptor_ids_to_keychain
.get(desc_id)
.map(|(k, _)| ((k.clone(), *index), *op))
let keychain = self.keychain_of_desc_id(desc_id)?;
Some(((keychain.clone(), *index), *op))
})
}

/// Iterate over known txouts that spend to tracked script pubkeys.
pub fn txouts(&self) -> impl DoubleEndedIterator<Item = (K, u32, OutPoint, &TxOut)> + '_ {
self.inner.txouts().filter_map(|((desc_id, i), op, txo)| {
self.descriptor_ids_to_keychain
.get(desc_id)
.map(|(k, _)| (k.clone(), *i, op, txo))
let keychain = self.keychain_of_desc_id(desc_id)?;
Some((keychain.clone(), *i, op, txo))
})
}

Expand All @@ -307,9 +319,8 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
self.inner
.txouts_in_tx(txid)
.filter_map(|((desc_id, i), op, txo)| {
self.descriptor_ids_to_keychain
.get(desc_id)
.map(|(k, _)| (k.clone(), *i, op, txo))
let keychain = self.keychain_of_desc_id(desc_id)?;
Some((keychain.clone(), *i, op, txo))
})
}

Expand All @@ -321,7 +332,7 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
/// This calls [`SpkTxOutIndex::txout`] internally.
pub fn txout(&self, outpoint: OutPoint) -> Option<(K, u32, &TxOut)> {
let ((descriptor_id, index), txo) = self.inner.txout(outpoint)?;
let (keychain, _) = self.descriptor_ids_to_keychain.get(descriptor_id)?;
let keychain = self.keychain_of_desc_id(descriptor_id)?;
Some((keychain.clone(), *index, txo))
}

Expand All @@ -338,9 +349,8 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
/// This calls [`SpkTxOutIndex::index_of_spk`] internally.
pub fn index_of_spk(&self, script: &Script) -> Option<(K, u32)> {
let (desc_id, last_index) = self.inner.index_of_spk(script)?;
self.descriptor_ids_to_keychain
.get(desc_id)
.map(|(k, _)| (k.clone(), *last_index))
let keychain = self.keychain_of_desc_id(desc_id)?;
Some((keychain.clone(), *last_index))
}

/// Returns whether the spk under the `keychain`'s `index` has been used.
Expand Down Expand Up @@ -448,45 +458,43 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
keychain: K,
descriptor: Descriptor<DescriptorPublicKey>,
) -> super::ChangeSet<K> {
let descriptor_id = descriptor.descriptor_id();
// First, we fill the keychain -> (desc_id, descriptor) map
let old_descriptor_opt = self
let mut changeset = super::ChangeSet::<K>::default();
let desc_id = descriptor.descriptor_id();

let old_desc = self
.keychains_to_descriptors
.insert(keychain.clone(), (descriptor_id, descriptor.clone()));

// Then, we fill the descriptor_id -> (keychain, descriptor) map
let old_keychain_opt = self
.descriptor_ids_to_keychain
.insert(descriptor_id, (keychain.clone(), descriptor.clone()));

// If `keychain` already had a `descriptor` associated, different from the `descriptor`
// passed in, we remove it from the descriptor -> keychain map
if let Some((old_desc_id, _)) = old_descriptor_opt {
if old_desc_id != descriptor_id {
self.descriptor_ids_to_keychain.remove(&old_desc_id);
.insert(keychain.clone(), (desc_id, descriptor.clone()));

if let Some((old_desc_id, _)) = old_desc {
// nothing needs to be done if caller reinsterted the same descriptor under the same
// keychain
if old_desc_id == desc_id {
return changeset;
}
// remove keychain from reverse index
let _is_keychain_removed = self
.descriptor_ids_to_keychains
.get_mut(&old_desc_id)
.expect("descriptor must have existed in reverse index")
.remove(&keychain);
assert!(
_is_keychain_removed,
"keychain must have existed in reverse index"
);
}

// Lastly, we fill the desc_id -> desc map
self.descriptor_ids_to_keychains
.entry(desc_id)
.or_default()
.insert(keychain.clone());
self.descriptor_ids_to_descriptors
.insert(descriptor_id, descriptor.clone());

.insert(desc_id, descriptor.clone());
self.replenish_lookahead(&keychain, self.lookahead);

// If both the keychain and descriptor were already inserted and associated, the
// keychains_added changeset must be empty
let keychains_added = if old_keychain_opt.map(|(k, _)| k) == Some(keychain.clone())
&& old_descriptor_opt.map(|(_, d)| d) == Some(descriptor.clone())
{
[].into()
} else {
[(keychain, descriptor)].into()
};

super::ChangeSet {
keychains_added,
last_revealed: [].into(),
}
changeset
.keychains_added
.insert(keychain.clone(), descriptor.clone());
changeset
}

/// Gets the descriptor associated with the keychain. Returns `None` if the keychain doesn't
Expand Down Expand Up @@ -584,11 +592,8 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
.range((start, end))
.map(|((descriptor_id, i), spk)| {
(
&self
.descriptor_ids_to_keychain
.get(descriptor_id)
.expect("Must be here")
.0,
self.keychain_of_desc_id(descriptor_id)
.expect("must have keychain"),
*i,
spk.as_script(),
)
Expand Down Expand Up @@ -673,10 +678,9 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
pub fn last_revealed_indices(&self) -> BTreeMap<K, u32> {
self.last_revealed
.iter()
.filter_map(|(descriptor_id, index)| {
self.descriptor_ids_to_keychain
.get(descriptor_id)
.map(|(k, _)| (k.clone(), *index))
.filter_map(|(desc_id, index)| {
let keychain = self.keychain_of_desc_id(desc_id)?;
Some((keychain.clone(), *index))
})
.collect()
}
Expand Down Expand Up @@ -870,16 +874,11 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
let bounds = self.map_to_inner_bounds(range);
self.inner
.outputs_in_range(bounds)
.map(move |((descriptor_id, i), op)| {
(
&self
.descriptor_ids_to_keychain
.get(descriptor_id)
.expect("must be here")
.0,
*i,
op,
)
.map(move |((desc_id, i), op)| {
let keychain = self
.keychain_of_desc_id(desc_id)
.expect("keychain must exist");
(keychain, *i, op)
})
}

Expand Down Expand Up @@ -941,10 +940,9 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
}
let last_revealed = last_revealed
.into_iter()
.filter_map(|(descriptor_id, index)| {
self.descriptor_ids_to_keychain
.get(&descriptor_id)
.map(|(k, _)| (k.clone(), index))
.filter_map(|(desc_id, index)| {
let keychain = self.keychain_of_desc_id(&desc_id)?;
Some((keychain.clone(), index))
})
.collect();
let _ = self.reveal_to_target_multi(&last_revealed);
Expand Down

0 comments on commit e1f3b53

Please sign in to comment.