fix(chain): sanitize derivation index before apply changeset

f3r10 · f3r10 · commit d3232f53d66e · 2025-01-07T09:01:31.000-05:00
diff --git a/crates/chain/src/indexed_tx_graph.rs b/crates/chain/src/indexed_tx_graph.rs
@@ -6,6 +6,7 @@ use alloc::{sync::Arc, vec::Vec};
 use bitcoin::{Block, OutPoint, Transaction, TxOut, Txid};
 
 use crate::{
+    keychain_txout::InsertDescriptorError,
     tx_graph::{self, TxGraph},
     Anchor, BlockId, Indexer, Merge, TxPosInBlock,
 };
@@ -46,8 +47,11 @@ impl<A, I> IndexedTxGraph<A, I> {
 
 impl<A: Anchor, I: Indexer> IndexedTxGraph<A, I> {
     /// Applies the [`ChangeSet`] to the [`IndexedTxGraph`].
-    pub fn apply_changeset(&mut self, changeset: ChangeSet<A, I::ChangeSet>) {
-        self.index.apply_changeset(changeset.indexer);
+    pub fn apply_changeset(
+        &mut self,
+        changeset: ChangeSet<A, I::ChangeSet>,
+    ) -> Result<(), InsertDescriptorError> {
+        self.index.apply_changeset(changeset.indexer)?;
 
         for tx in &changeset.tx_graph.txs {
             self.index.index_tx(tx);
@@ -57,6 +61,7 @@ impl<A: Anchor, I: Indexer> IndexedTxGraph<A, I> {
         }
 
         self.graph.apply_changeset(changeset.tx_graph);
+        Ok(())
     }
 
     /// Determines the [`ChangeSet`] between `self` and an empty [`IndexedTxGraph`].
diff --git a/crates/chain/src/indexer.rs b/crates/chain/src/indexer.rs
@@ -23,7 +23,10 @@ pub trait Indexer {
     fn index_tx(&mut self, tx: &Transaction) -> Self::ChangeSet;
 
     /// Apply changeset to itself.
-    fn apply_changeset(&mut self, changeset: Self::ChangeSet);
+    fn apply_changeset(
+        &mut self,
+        changeset: Self::ChangeSet,
+    ) -> Result<(), keychain_txout::InsertDescriptorError>;
 
     /// Determines the [`ChangeSet`](Indexer::ChangeSet) between `self` and an empty [`Indexer`].
     fn initial_changeset(&self) -> Self::ChangeSet;
diff --git a/crates/chain/src/indexer/keychain_txout.rs b/crates/chain/src/indexer/keychain_txout.rs
@@ -10,11 +10,17 @@ use crate::{
     DescriptorExt, DescriptorId, Indexed, Indexer, KeychainIndexed, SpkIterator,
 };
 use alloc::{borrow::ToOwned, vec::Vec};
-use bitcoin::{Amount, OutPoint, ScriptBuf, SignedAmount, Transaction, TxOut, Txid};
+use bitcoin::{
+    bip32::ChildNumber, Amount, OutPoint, ScriptBuf, SignedAmount, Transaction, TxOut, Txid,
+};
 use core::{
     fmt::Debug,
     ops::{Bound, RangeBounds},
 };
+use miniscript::{
+    descriptor::{DescriptorXKey, Wildcard},
+    ForEachKey,
+};
 
 use crate::Merge;
 
@@ -170,7 +176,7 @@ impl<K: Clone + Ord + Debug> Indexer for KeychainTxOutIndex<K> {
         }
     }
 
-    fn apply_changeset(&mut self, changeset: Self::ChangeSet) {
+    fn apply_changeset(&mut self, changeset: Self::ChangeSet) -> Result<(), InsertDescriptorError> {
         self.apply_changeset(changeset)
     }
 
@@ -767,18 +773,38 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
     }
 
     /// Applies the `ChangeSet<K>` to the [`KeychainTxOutIndex<K>`]
-    pub fn apply_changeset(&mut self, changeset: ChangeSet) {
+    pub fn apply_changeset(&mut self, changeset: ChangeSet) -> Result<(), InsertDescriptorError> {
         for (&desc_id, &index) in &changeset.last_revealed {
+            let descriptor = self.descriptors.get(&desc_id).expect("invariant");
+            // Ensure the keys don't contain any hardened derivation steps or hardened wildcards
+            let descriptor_contains_hardened_steps = descriptor.for_any_key(|k| {
+                if let DescriptorPublicKey::XPub(DescriptorXKey {
+                    derivation_path,
+                    wildcard,
+                    ..
+                }) = k
+                {
+                    return *wildcard == Wildcard::Hardened
+                        || derivation_path.into_iter().any(ChildNumber::is_hardened);
+                }
+
+                false
+            });
+            if descriptor_contains_hardened_steps {
+                return Err(InsertDescriptorError::HardenedDerivationXpub);
+            }
+            descriptor.sanity_check()?;
             let v = self.last_revealed.entry(desc_id).or_default();
             *v = index.max(*v);
             self.replenish_inner_index_did(desc_id, self.lookahead);
         }
+        Ok(())
     }
 }
 
-#[derive(Clone, Debug, PartialEq)]
+#[derive(Debug, PartialEq)]
 /// Error returned from [`KeychainTxOutIndex::insert_descriptor`]
-pub enum InsertDescriptorError<K> {
+pub enum InsertDescriptorError<K = ()> {
     /// The descriptor has already been assigned to a keychain so you can't assign it to another
     DescriptorAlreadyAssigned {
         /// The descriptor you have attempted to reassign
@@ -793,6 +819,16 @@ pub enum InsertDescriptorError<K> {
         /// The descriptor that the keychain is already assigned to
         existing_assignment: Descriptor<DescriptorPublicKey>,
     },
+    /// Miniscript error
+    Miniscript(miniscript::Error),
+    /// The descriptor contains hardened derivation steps on public extended keys
+    HardenedDerivationXpub,
+}
+
+impl From<miniscript::Error> for InsertDescriptorError {
+    fn from(err: miniscript::Error) -> Self {
+        InsertDescriptorError::Miniscript(err)
+    }
 }
 
 impl<K: core::fmt::Debug> core::fmt::Display for InsertDescriptorError<K> {
@@ -816,6 +852,11 @@ impl<K: core::fmt::Debug> core::fmt::Display for InsertDescriptorError<K> {
                     "attempt to re-assign keychain {keychain:?} already assigned to {existing:?}"
                 )
             }
+            InsertDescriptorError::Miniscript(err) => write!(f, "Miniscript error: {}", err),
+            InsertDescriptorError::HardenedDerivationXpub => write!(
+                f,
+                "The descriptor contains hardened derivation steps on public extended keys"
+            ),
         }
     }
 }
diff --git a/crates/chain/src/indexer/spk_txout.rs b/crates/chain/src/indexer/spk_txout.rs
@@ -8,6 +8,8 @@ use crate::{
 };
 use bitcoin::{Amount, OutPoint, ScriptBuf, SignedAmount, Transaction, TxOut, Txid};
 
+use super::keychain_txout::InsertDescriptorError;
+
 /// An index storing [`TxOut`]s that have a script pubkey that matches those in a list.
 ///
 /// The basic idea is that you insert script pubkeys you care about into the index with
@@ -69,8 +71,12 @@ impl<I: Clone + Ord + core::fmt::Debug> Indexer for SpkTxOutIndex<I> {
 
     fn initial_changeset(&self) -> Self::ChangeSet {}
 
-    fn apply_changeset(&mut self, _changeset: Self::ChangeSet) {
+    fn apply_changeset(
+        &mut self,
+        _changeset: Self::ChangeSet,
+    ) -> Result<(), InsertDescriptorError> {
         // This applies nothing.
+        Ok(())
     }
 
     fn is_tx_relevant(&self, tx: &Transaction) -> bool {
diff --git a/crates/chain/tests/test_keychain_txout_index.rs b/crates/chain/tests/test_keychain_txout_index.rs
@@ -617,7 +617,9 @@ fn applying_changesets_one_by_one_vs_aggregate_must_have_same_result() {
         .insert_descriptor(TestKeychain::External, desc.clone())
         .expect("must insert keychain");
     for changeset in changesets {
-        indexer_a.apply_changeset(changeset.clone());
+        indexer_a
+            .apply_changeset(changeset.clone())
+            .expect("must apply changeset");
     }
 
     let mut indexer_b = KeychainTxOutIndex::<TestKeychain>::new(0);
@@ -632,7 +634,9 @@ fn applying_changesets_one_by_one_vs_aggregate_must_have_same_result() {
             agg
         })
         .expect("must aggregate changesets");
-    indexer_b.apply_changeset(aggregate_changesets);
+    indexer_b
+        .apply_changeset(aggregate_changesets)
+        .expect("must apply changeset");
 
     assert_eq!(
         indexer_a.keychains().collect::<Vec<_>>(),
diff --git a/crates/wallet/src/descriptor/error.rs b/crates/wallet/src/descriptor/error.rs
@@ -43,6 +43,10 @@ pub enum Error {
     Hex(bitcoin::hex::HexToBytesError),
     /// The provided wallet descriptors are identical
     ExternalAndInternalAreTheSame,
+    /// Descriptor has already been assigned to a keychain so you can't assign it to another
+    DescriptorAlreadyAssigned,
+    /// The keychain is already assigned to a descriptor so you can't reassign it
+    KeychainAlreadyAssigned,
 }
 
 impl From<crate::keys::KeyError> for Error {
@@ -83,6 +87,12 @@ impl fmt::Display for Error {
             Self::ExternalAndInternalAreTheSame => {
                 write!(f, "External and internal descriptors are the same")
             }
+            Self::DescriptorAlreadyAssigned => {
+                write!(f, "attempt to re-assign descriptor already assigned")
+            }
+            Self::KeychainAlreadyAssigned => {
+                write!(f, "attempt to re-assign keychain already assigned")
+            }
         }
     }
 }
@@ -125,3 +135,20 @@ impl From<crate::descriptor::policy::PolicyError> for Error {
         Error::Policy(err)
     }
 }
+
+impl From<chain::keychain_txout::InsertDescriptorError> for Error {
+    fn from(err: chain::keychain_txout::InsertDescriptorError) -> Self {
+        match err {
+            chain::keychain_txout::InsertDescriptorError::DescriptorAlreadyAssigned { .. } => {
+                Error::DescriptorAlreadyAssigned
+            }
+            chain::keychain_txout::InsertDescriptorError::KeychainAlreadyAssigned { .. } => {
+                Error::KeychainAlreadyAssigned
+            }
+            chain::keychain_txout::InsertDescriptorError::Miniscript(err) => Error::Miniscript(err),
+            chain::keychain_txout::InsertDescriptorError::HardenedDerivationXpub => {
+                Error::HardenedDerivationXpub
+            }
+        }
+    }
+}
diff --git a/crates/wallet/src/wallet/mod.rs b/crates/wallet/src/wallet/mod.rs
@@ -608,8 +608,12 @@ impl Wallet {
             .map_err(LoadError::Descriptor)?;
 
         let mut indexed_graph = IndexedTxGraph::new(index);
-        indexed_graph.apply_changeset(changeset.indexer.into());
-        indexed_graph.apply_changeset(changeset.tx_graph.into());
+        indexed_graph
+            .apply_changeset(changeset.indexer.into())
+            .map_err(|err| LoadError::Descriptor(err.into()))?;
+        indexed_graph
+            .apply_changeset(changeset.tx_graph.into())
+            .map_err(|err| LoadError::Descriptor(err.into()))?;
 
         let stage = ChangeSet::default();
 
@@ -2530,6 +2534,12 @@ fn create_indexer(
                     InsertDescriptorError::KeychainAlreadyAssigned { .. } => {
                         unreachable!("this is the first time we're assigning internal")
                     }
+                    InsertDescriptorError::Miniscript(err) => {
+                        crate::descriptor::error::Error::Miniscript(err)
+                    }
+                    InsertDescriptorError::HardenedDerivationXpub => {
+                        crate::descriptor::error::Error::HardenedDerivationXpub
+                    }
                 }
             })?);
     }
diff --git a/example-crates/example_cli/src/lib.rs b/example-crates/example_cli/src/lib.rs
@@ -827,7 +827,7 @@ pub fn init_or_load<CS: clap::Subcommand, S: clap::Args>(
                 graph.apply_changeset(indexed_tx_graph::ChangeSet {
                     tx_graph: changeset.tx_graph,
                     indexer: changeset.indexer,
-                });
+                })?;
                 graph
             });
 
