]> Untitled Git - bdk/commitdiff
feat(core): Add expected txids to `SyncRequest` spks
authorWei Chen <wzc110@gmail.com>
Fri, 24 Jan 2025 09:18:56 +0000 (17:18 +0800)
committer志宇 <hello@evanlinjin.me>
Fri, 14 Mar 2025 02:15:29 +0000 (13:15 +1100)
The spk history returned from Electrum should have these txs present.
Any missing tx will be considered evicted from the mempool.

crates/core/src/spk_client.rs

index dce3b7ae1886ae4e747e256958d2a1a11972cfb3..a82d24170dc26b44ad48513f8c486fa526556816 100644 (file)
@@ -1,7 +1,7 @@
 //! Helper types for spk-based blockchain clients.
 use crate::{
     alloc::{boxed::Box, collections::VecDeque, vec::Vec},
-    collections::BTreeMap,
+    collections::{BTreeMap, HashMap, HashSet},
     CheckPoint, ConfirmationBlockTime, Indexed,
 };
 use bitcoin::{OutPoint, Script, ScriptBuf, Txid};
@@ -86,6 +86,28 @@ impl SyncProgress {
     }
 }
 
+/// [`Script`] with expected [`Txid`] histories.
+#[derive(Debug, Clone)]
+pub struct SpkWithExpectedTxids {
+    /// Script pubkey.
+    pub spk: ScriptBuf,
+
+    /// [`Txid`]s that we expect to appear in the chain source's spk history response.
+    ///
+    /// Any transaction listed here that is missing from the spk history response should be
+    /// considered evicted from the mempool.
+    pub expected_txids: HashSet<Txid>,
+}
+
+impl From<ScriptBuf> for SpkWithExpectedTxids {
+    fn from(spk: ScriptBuf) -> Self {
+        Self {
+            spk,
+            expected_txids: HashSet::new(),
+        }
+    }
+}
+
 /// Builds a [`SyncRequest`].
 ///
 /// Construct with [`SyncRequest::builder`].
@@ -153,6 +175,20 @@ impl<I> SyncRequestBuilder<I> {
         self
     }
 
+    /// Add transactions that are expected to exist under the given spks.
+    ///
+    /// This is useful for detecting a malicious replacement of an incoming transaction.
+    pub fn expected_spk_txids(mut self, txs: impl IntoIterator<Item = (ScriptBuf, Txid)>) -> Self {
+        for (spk, txid) in txs {
+            self.inner
+                .spk_expected_txids
+                .entry(spk)
+                .or_default()
+                .insert(txid);
+        }
+        self
+    }
+
     /// Add [`Txid`]s that will be synced against.
     pub fn txids(mut self, txids: impl IntoIterator<Item = Txid>) -> Self {
         self.inner.txids.extend(txids);
@@ -208,6 +244,7 @@ pub struct SyncRequest<I = ()> {
     chain_tip: Option<CheckPoint>,
     spks: VecDeque<(I, ScriptBuf)>,
     spks_consumed: usize,
+    spk_expected_txids: HashMap<ScriptBuf, HashSet<Txid>>,
     txids: VecDeque<Txid>,
     txids_consumed: usize,
     outpoints: VecDeque<OutPoint>,
@@ -237,6 +274,7 @@ impl<I> SyncRequest<I> {
                 chain_tip: None,
                 spks: VecDeque::new(),
                 spks_consumed: 0,
+                spk_expected_txids: HashMap::new(),
                 txids: VecDeque::new(),
                 txids_consumed: 0,
                 outpoints: VecDeque::new(),
@@ -292,6 +330,23 @@ impl<I> SyncRequest<I> {
         Some(spk)
     }
 
+    /// Advances the sync request and returns the next [`ScriptBuf`] with corresponding [`Txid`]
+    /// history.
+    ///
+    /// Returns [`None`] when there are no more scripts remaining in the request.
+    pub fn next_spk_with_expected_txids(&mut self) -> Option<SpkWithExpectedTxids> {
+        let next_spk = self.next_spk()?;
+        let spk_history = self
+            .spk_expected_txids
+            .get(&next_spk)
+            .cloned()
+            .unwrap_or_default();
+        Some(SpkWithExpectedTxids {
+            spk: next_spk,
+            expected_txids: spk_history,
+        })
+    }
+
     /// Advances the sync request and returns the next [`Txid`].
     ///
     /// Returns [`None`] when there are no more txids remaining in the request.
@@ -317,6 +372,13 @@ impl<I> SyncRequest<I> {
         SyncIter::<I, ScriptBuf>::new(self)
     }
 
+    /// Iterate over [`ScriptBuf`]s with corresponding [`Txid`] histories contained in this request.
+    pub fn iter_spks_with_expected_txids(
+        &mut self,
+    ) -> impl ExactSizeIterator<Item = SpkWithExpectedTxids> + '_ {
+        SyncIter::<I, SpkWithExpectedTxids>::new(self)
+    }
+
     /// Iterate over [`Txid`]s contained in this request.
     pub fn iter_txids(&mut self) -> impl ExactSizeIterator<Item = Txid> + '_ {
         SyncIter::<I, Txid>::new(self)
@@ -556,6 +618,19 @@ impl<I> Iterator for SyncIter<'_, I, ScriptBuf> {
     }
 }
 
+impl<I> Iterator for SyncIter<'_, I, SpkWithExpectedTxids> {
+    type Item = SpkWithExpectedTxids;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        self.request.next_spk_with_expected_txids()
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let remaining = self.request.spks.len();
+        (remaining, Some(remaining))
+    }
+}
+
 impl<I> Iterator for SyncIter<'_, I, Txid> {
     type Item = Txid;