Skip to content

Commit e692768

Browse files
LagginTimesevanlinjin
authored andcommitted
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.
1 parent 20bbb4d commit e692768

File tree

1 file changed

+76
-1
lines changed

1 file changed

+76
-1
lines changed

crates/core/src/spk_client.rs

Lines changed: 76 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
//! Helper types for spk-based blockchain clients.
22
use crate::{
33
alloc::{boxed::Box, collections::VecDeque, vec::Vec},
4-
collections::BTreeMap,
4+
collections::{BTreeMap, HashMap, HashSet},
55
CheckPoint, ConfirmationBlockTime, Indexed,
66
};
77
use bitcoin::{OutPoint, Script, ScriptBuf, Txid};
@@ -86,6 +86,28 @@ impl SyncProgress {
8686
}
8787
}
8888

89+
/// [`Script`] with expected [`Txid`] histories.
90+
#[derive(Debug, Clone)]
91+
pub struct SpkWithExpectedTxids {
92+
/// Script pubkey.
93+
pub spk: ScriptBuf,
94+
95+
/// [`Txid`]s that we expect to appear in the chain source's spk history response.
96+
///
97+
/// Any transaction listed here that is missing form the spk history response shoud be
98+
/// considred evicted from the mempool.
99+
pub expected_txids: HashSet<Txid>,
100+
}
101+
102+
impl From<ScriptBuf> for SpkWithExpectedTxids {
103+
fn from(spk: ScriptBuf) -> Self {
104+
Self {
105+
spk,
106+
expected_txids: HashSet::new(),
107+
}
108+
}
109+
}
110+
89111
/// Builds a [`SyncRequest`].
90112
///
91113
/// Construct with [`SyncRequest::builder`].
@@ -153,6 +175,20 @@ impl<I> SyncRequestBuilder<I> {
153175
self
154176
}
155177

178+
/// Add transactions that are expected to exist under then given spks.
179+
///
180+
/// This is useful for detecting a malicious replacement of an incoming transaction.
181+
pub fn expected_spk_txids(mut self, txs: impl IntoIterator<Item = (ScriptBuf, Txid)>) -> Self {
182+
for (spk, txid) in txs {
183+
self.inner
184+
.spk_expected_txids
185+
.entry(spk)
186+
.or_default()
187+
.insert(txid);
188+
}
189+
self
190+
}
191+
156192
/// Add [`Txid`]s that will be synced against.
157193
pub fn txids(mut self, txids: impl IntoIterator<Item = Txid>) -> Self {
158194
self.inner.txids.extend(txids);
@@ -208,6 +244,7 @@ pub struct SyncRequest<I = ()> {
208244
chain_tip: Option<CheckPoint>,
209245
spks: VecDeque<(I, ScriptBuf)>,
210246
spks_consumed: usize,
247+
spk_expected_txids: HashMap<ScriptBuf, HashSet<Txid>>,
211248
txids: VecDeque<Txid>,
212249
txids_consumed: usize,
213250
outpoints: VecDeque<OutPoint>,
@@ -237,6 +274,7 @@ impl<I> SyncRequest<I> {
237274
chain_tip: None,
238275
spks: VecDeque::new(),
239276
spks_consumed: 0,
277+
spk_expected_txids: HashMap::new(),
240278
txids: VecDeque::new(),
241279
txids_consumed: 0,
242280
outpoints: VecDeque::new(),
@@ -292,6 +330,23 @@ impl<I> SyncRequest<I> {
292330
Some(spk)
293331
}
294332

333+
/// Advances the sync request and returns the next [`ScriptBuf`] with corresponding [`Txid`]
334+
/// history.
335+
///
336+
/// Returns [`None`] when there are no more scripts remaining in the request.
337+
pub fn next_spk_with_expected_txids(&mut self) -> Option<SpkWithExpectedTxids> {
338+
let next_spk = self.next_spk()?;
339+
let spk_history = self
340+
.spk_expected_txids
341+
.get(&next_spk)
342+
.cloned()
343+
.unwrap_or_default();
344+
Some(SpkWithExpectedTxids {
345+
spk: next_spk,
346+
expected_txids: spk_history,
347+
})
348+
}
349+
295350
/// Advances the sync request and returns the next [`Txid`].
296351
///
297352
/// Returns [`None`] when there are no more txids remaining in the request.
@@ -317,6 +372,13 @@ impl<I> SyncRequest<I> {
317372
SyncIter::<I, ScriptBuf>::new(self)
318373
}
319374

375+
/// Iterate over [`ScriptBuf`]s with corresponding [`Txid`] histories contained in this request.
376+
pub fn iter_spks_with_expected_txids(
377+
&mut self,
378+
) -> impl ExactSizeIterator<Item = SpkWithExpectedTxids> + '_ {
379+
SyncIter::<I, SpkWithExpectedTxids>::new(self)
380+
}
381+
320382
/// Iterate over [`Txid`]s contained in this request.
321383
pub fn iter_txids(&mut self) -> impl ExactSizeIterator<Item = Txid> + '_ {
322384
SyncIter::<I, Txid>::new(self)
@@ -556,6 +618,19 @@ impl<I> Iterator for SyncIter<'_, I, ScriptBuf> {
556618
}
557619
}
558620

621+
impl<I> Iterator for SyncIter<'_, I, SpkWithExpectedTxids> {
622+
type Item = SpkWithExpectedTxids;
623+
624+
fn next(&mut self) -> Option<Self::Item> {
625+
self.request.next_spk_with_expected_txids()
626+
}
627+
628+
fn size_hint(&self) -> (usize, Option<usize>) {
629+
let remaining = self.request.spks.len();
630+
(remaining, Some(remaining))
631+
}
632+
}
633+
559634
impl<I> Iterator for SyncIter<'_, I, Txid> {
560635
type Item = Txid;
561636

0 commit comments

Comments
 (0)