From: Wei Chen <wzc110@gmail.com>
Date: Fri, 24 Jan 2025 09:18:56 +0000 (+0800)
Subject: feat(core): Add expected txids to `SyncRequest` spks
X-Git-Tag: wallet-1.2.0~7^2~5
X-Git-Url: http://internal-gitweb-vhost/script/%22https:/primitives/decode/struct.SegwitCodeLengthError.html?a=commitdiff_plain;h=b38569fd147f8fe30f55e1c30cdba779247ba856;p=bdk

feat(core): Add expected txids to `SyncRequest` spks

The spk history returned from Electrum should have these txs present.
Any missing tx will be considered evicted from the mempool.
---

diff --git a/crates/core/src/spk_client.rs b/crates/core/src/spk_client.rs
index dce3b7ae..a82d2417 100644
--- a/crates/core/src/spk_client.rs
+++ b/crates/core/src/spk_client.rs
@@ -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;