feat(electrum): Handle spks with expected txids

evanlinjin · LagginTimes · evanlinjin · commit 3c09f752410e · 2025-03-14T10:57:14.000+11:00
Co-authored-by: Wei Chen &lt;wzc110@gmail.com&gt;
diff --git a/crates/electrum/src/bdk_electrum_client.rs b/crates/electrum/src/bdk_electrum_client.rs
@@ -1,14 +1,13 @@
 use bdk_core::{
     bitcoin::{block::Header, BlockHash, OutPoint, ScriptBuf, Transaction, Txid},
-    collections::{BTreeMap, HashMap},
-    spk_client::{FullScanRequest, FullScanResponse, SyncRequest, SyncResponse},
+    collections::{BTreeMap, HashMap, HashSet},
+    spk_client::{
+        FullScanRequest, FullScanResponse, SpkWithExpectedTxids, SyncRequest, SyncResponse,
+    },
     BlockId, CheckPoint, ConfirmationBlockTime, TxUpdate,
 };
 use electrum_client::{ElectrumApi, Error, HeaderNotification};
-use std::{
-    collections::HashSet,
-    sync::{Arc, Mutex},
-};
+use std::sync::{Arc, Mutex};
 
 /// We include a chain suffix of a certain length for the purpose of robustness.
 const CHAIN_SUFFIX_LENGTH: u32 = 8;
@@ -138,7 +137,9 @@ impl<E: ElectrumApi> BdkElectrumClient<E> {
         let mut tx_update = TxUpdate::<ConfirmationBlockTime>::default();
         let mut last_active_indices = BTreeMap::<K, u32>::default();
         for keychain in request.keychains() {
-            let spks = request.iter_spks(keychain.clone());
+            let spks = request
+                .iter_spks(keychain.clone())
+                .map(|(spk_i, spk)| (spk_i, SpkWithExpectedTxids::from(spk)));
             if let Some(last_active_index) =
                 self.populate_with_spks(start_time, &mut tx_update, spks, stop_gap, batch_size)?
             {
@@ -209,7 +210,7 @@ impl<E: ElectrumApi> BdkElectrumClient<E> {
             start_time,
             &mut tx_update,
             request
-                .iter_spks()
+                .iter_spks_with_expected_txids()
                 .enumerate()
                 .map(|(i, spk)| (i as u32, spk)),
             usize::MAX,
@@ -247,7 +248,7 @@ impl<E: ElectrumApi> BdkElectrumClient<E> {
         &self,
         start_time: u64,
         tx_update: &mut TxUpdate<ConfirmationBlockTime>,
-        mut spks: impl Iterator<Item = (u32, ScriptBuf)>,
+        mut spks_with_expected_txids: impl Iterator<Item = (u32, SpkWithExpectedTxids)>,
         stop_gap: usize,
         batch_size: usize,
     ) -> Result<Option<u32>, Error> {
@@ -256,17 +257,30 @@ impl<E: ElectrumApi> BdkElectrumClient<E> {
 
         loop {
             let spks = (0..batch_size)
-                .map_while(|_| spks.next())
+                .map_while(|_| spks_with_expected_txids.next())
                 .collect::<Vec<_>>();
             if spks.is_empty() {
                 return Ok(last_active_index);
             }
 
+            let (spks, expected_spk_txids): (Vec<(u32, ScriptBuf)>, HashMap<ScriptBuf, _>) = spks
+                .iter()
+                .map(|(spk_i, spk_with_expected_txids)| {
+                    (
+                        (*spk_i, spk_with_expected_txids.spk.clone()),
+                        (
+                            spk_with_expected_txids.spk.clone(),
+                            spk_with_expected_txids.expected_txids.clone(),
+                        ),
+                    )
+                })
+                .unzip();
+
             let spk_histories = self
                 .inner
                 .batch_script_get_history(spks.iter().map(|(_, s)| s.as_script()))?;
 
-            for ((spk_index, _spk), spk_history) in spks.into_iter().zip(spk_histories) {
+            for ((spk_index, spk), spk_history) in spks.into_iter().zip(spk_histories) {
                 if spk_history.is_empty() {
                     unused_spk_count = unused_spk_count.saturating_add(1);
                     if unused_spk_count >= stop_gap {
@@ -277,6 +291,19 @@ impl<E: ElectrumApi> BdkElectrumClient<E> {
                     unused_spk_count = 0;
                 }
 
+                let spk_history_set = spk_history
+                    .iter()
+                    .map(|res| res.tx_hash)
+                    .collect::<HashSet<_>>();
+
+                if let Some(expected_txids) = expected_spk_txids.get(&spk) {
+                    tx_update.evicted_ats.extend(
+                        expected_txids
+                            .difference(&spk_history_set)
+                            .map(|&txid| (txid, start_time)),
+                    );
+                }
+
                 for tx_res in spk_history {
                     tx_update.txs.push(self.fetch_tx(tx_res.tx_hash)?);
                     match tx_res.height.try_into() {
diff --git a/crates/electrum/tests/test_electrum.rs b/crates/electrum/tests/test_electrum.rs
@@ -5,7 +5,10 @@ use bdk_chain::{
     spk_txout::SpkTxOutIndex,
     Balance, ConfirmationBlockTime, IndexedTxGraph, Indexer, Merge, TxGraph,
 };
-use bdk_core::bitcoin::Network;
+use bdk_core::bitcoin::{
+    key::{Secp256k1, UntweakedPublicKey},
+    Network,
+};
 use bdk_electrum::BdkElectrumClient;
 use bdk_testenv::{
     anyhow,
@@ -14,12 +17,22 @@ use bdk_testenv::{
 };
 use core::time::Duration;
 use electrum_client::ElectrumApi;
-use std::collections::{BTreeSet, HashSet};
+use std::collections::{BTreeSet, HashMap, HashSet};
 use std::str::FromStr;
 
 // Batch size for `sync_with_electrum`.
 const BATCH_SIZE: usize = 5;
 
+pub fn get_test_spk() -> ScriptBuf {
+    const PK_BYTES: &[u8] = &[
+        12, 244, 72, 4, 163, 4, 211, 81, 159, 82, 153, 123, 125, 74, 142, 40, 55, 237, 191, 231,
+        31, 114, 89, 165, 83, 141, 8, 203, 93, 240, 53, 101,
+    ];
+    let secp = Secp256k1::new();
+    let pk = UntweakedPublicKey::from_slice(PK_BYTES).expect("Must be valid PK");
+    ScriptBuf::new_p2tr(&secp, pk, None)
+}
+
 fn get_balance(
     recv_chain: &LocalChain,
     recv_graph: &IndexedTxGraph<ConfirmationBlockTime, SpkTxOutIndex<()>>,
@@ -60,6 +73,122 @@ where
     Ok(update)
 }
 
+// Ensure that a wallet can detect a malicious replacement of an incoming transaction.
+//
+// This checks that both the Electrum chain source and the receiving structures properly track the
+// replaced transaction as missing.
+#[test]
+pub fn detect_receive_tx_cancel() -> anyhow::Result<()> {
+    const SEND_TX_FEE: Amount = Amount::from_sat(1000);
+    const UNDO_SEND_TX_FEE: Amount = Amount::from_sat(2000);
+
+    let env = TestEnv::new()?;
+    let rpc_client = env.rpc_client();
+    let electrum_client = electrum_client::Client::new(env.electrsd.electrum_url.as_str())?;
+    let client = BdkElectrumClient::new(electrum_client);
+
+    let mut graph = IndexedTxGraph::<ConfirmationBlockTime, _>::new(SpkTxOutIndex::<()>::default());
+    let (chain, _) = LocalChain::from_genesis_hash(env.bitcoind.client.get_block_hash(0)?);
+
+    // Get receiving address.
+    let receiver_spk = get_test_spk();
+    let receiver_addr = Address::from_script(&receiver_spk, bdk_chain::bitcoin::Network::Regtest)?;
+    graph.index.insert_spk((), receiver_spk);
+
+    env.mine_blocks(101, None)?;
+
+    // Select a UTXO to use as an input for constructing our test transactions.
+    let selected_utxo = rpc_client
+        .list_unspent(None, None, None, Some(false), None)?
+        .into_iter()
+        // Find a block reward tx.
+        .find(|utxo| utxo.amount == Amount::from_int_btc(50))
+        .expect("Must find a block reward UTXO");
+
+    // Derive the sender's address from the selected UTXO.
+    let sender_spk = selected_utxo.script_pub_key.clone();
+    let sender_addr = Address::from_script(&sender_spk, bdk_chain::bitcoin::Network::Regtest)
+        .expect("Failed to derive address from UTXO");
+
+    // Setup the common inputs used by both `send_tx` and `undo_send_tx`.
+    let inputs = [CreateRawTransactionInput {
+        txid: selected_utxo.txid,
+        vout: selected_utxo.vout,
+        sequence: None,
+    }];
+
+    // Create and sign the `send_tx` that sends funds to the receiver address.
+    let send_tx_outputs = HashMap::from([(
+        receiver_addr.to_string(),
+        selected_utxo.amount - SEND_TX_FEE,
+    )]);
+    let send_tx = rpc_client.create_raw_transaction(&inputs, &send_tx_outputs, None, Some(true))?;
+    let send_tx = rpc_client
+        .sign_raw_transaction_with_wallet(send_tx.raw_hex(), None, None)?
+        .transaction()?;
+
+    // Create and sign the `undo_send_tx` transaction. This redirects funds back to the sender
+    // address.
+    let undo_send_outputs = HashMap::from([(
+        sender_addr.to_string(),
+        selected_utxo.amount - UNDO_SEND_TX_FEE,
+    )]);
+    let undo_send_tx =
+        rpc_client.create_raw_transaction(&inputs, &undo_send_outputs, None, Some(true))?;
+    let undo_send_tx = rpc_client
+        .sign_raw_transaction_with_wallet(undo_send_tx.raw_hex(), None, None)?
+        .transaction()?;
+
+    // Sync after broadcasting the `send_tx`. Ensure that we detect and receive the `send_tx`.
+    let send_txid = env.rpc_client().send_raw_transaction(send_tx.raw_hex())?;
+    env.wait_until_electrum_sees_txid(send_txid, Duration::from_secs(6))?;
+    let sync_request = SyncRequest::builder()
+        .chain_tip(chain.tip())
+        .spks_with_indexes(graph.index.all_spks().clone())
+        .expected_spk_txids(graph.list_expected_spk_txids(&chain, chain.tip().block_id(), ..));
+    let sync_response = client.sync(sync_request, BATCH_SIZE, true)?;
+    assert!(
+        sync_response
+            .tx_update
+            .txs
+            .iter()
+            .any(|tx| tx.compute_txid() == send_txid),
+        "sync response must include the send_tx"
+    );
+    let changeset = graph.apply_update(sync_response.tx_update.clone());
+    assert!(
+        changeset.tx_graph.txs.contains(&send_tx),
+        "tx graph must deem send_tx relevant and include it"
+    );
+
+    // Sync after broadcasting the `undo_send_tx`. Verify that `send_tx` is now missing from the
+    // mempool.
+    let undo_send_txid = env
+        .rpc_client()
+        .send_raw_transaction(undo_send_tx.raw_hex())?;
+    env.wait_until_electrum_sees_txid(undo_send_txid, Duration::from_secs(6))?;
+    let sync_request = SyncRequest::builder()
+        .chain_tip(chain.tip())
+        .spks_with_indexes(graph.index.all_spks().clone())
+        .expected_spk_txids(graph.list_expected_spk_txids(&chain, chain.tip().block_id(), ..));
+    let sync_response = client.sync(sync_request, BATCH_SIZE, true)?;
+    assert!(
+        sync_response
+            .tx_update
+            .evicted_ats
+            .iter()
+            .any(|(txid, _)| *txid == send_txid),
+        "sync response must track send_tx as missing from mempool"
+    );
+    let changeset = graph.apply_update(sync_response.tx_update.clone());
+    assert!(
+        changeset.tx_graph.last_evicted.contains_key(&send_txid),
+        "tx graph must track send_tx as missing"
+    );
+
+    Ok(())
+}
+
 /// If an spk history contains a tx that spends another unconfirmed tx (chained mempool history),
 /// the Electrum API will return the tx with a negative height. This should succeed and not panic.
 #[test]
