From d7eac3ef8c495fe1f1d49d5fab85808ef49695d9 Mon Sep 17 00:00:00 2001
From: Wei Chen <wzc110@gmail.com>
Date: Fri, 8 Nov 2024 03:56:17 +0800
Subject: [PATCH] test(electrum): tests for exploit fixes

---
 crates/electrum/src/bdk_electrum_client.rs | 84 ++++++++++++++++++++++
 1 file changed, 84 insertions(+)

diff --git a/crates/electrum/src/bdk_electrum_client.rs b/crates/electrum/src/bdk_electrum_client.rs
index 1639b488b..96678d965 100644
--- a/crates/electrum/src/bdk_electrum_client.rs
+++ b/crates/electrum/src/bdk_electrum_client.rs
@@ -569,3 +569,87 @@ fn chain_update(
     }
     Ok(tip)
 }
+
+#[cfg(test)]
+mod test {
+    use crate::{bdk_electrum_client::TxUpdate, BdkElectrumClient};
+    use bdk_chain::bitcoin::{Amount, OutPoint, ScriptBuf, Transaction, TxIn, TxOut};
+    use bdk_testenv::{utils::new_tx, TestEnv};
+    use std::{collections::BTreeMap, sync::Arc};
+
+    #[test]
+    fn test_populate_with_txids_without_output() {
+        let env = TestEnv::new().unwrap();
+        let electrum_client =
+            electrum_client::Client::new(env.electrsd.electrum_url.as_str()).unwrap();
+        let client = BdkElectrumClient::new(electrum_client);
+
+        // Setup transaction with no outputs.
+        let tx = new_tx(0);
+
+        // Populate tx_cache with `tx` to make it fetchable.
+        client.populate_tx_cache(vec![tx.clone()]);
+
+        // Test that populate_with_txids does not panic or process a tx with no output.
+        let mut tx_update = TxUpdate::default();
+        let _ = client.populate_with_txids(&mut tx_update, vec![tx.compute_txid()]);
+
+        assert_eq!(tx_update.txs, Vec::new());
+    }
+
+    #[test]
+    fn test_fetch_prev_txout_with_coinbase() {
+        let env = TestEnv::new().unwrap();
+        let electrum_client =
+            electrum_client::Client::new(env.electrsd.electrum_url.as_str()).unwrap();
+        let client = BdkElectrumClient::new(electrum_client);
+
+        // Setup the transactions necessary for fetch_prev_txout to fetch `txouts` for `tx`.
+        let expected_txout = TxOut {
+            value: Amount::from_sat(10_000),
+            script_pubkey: ScriptBuf::default(),
+        };
+
+        let prev_tx = Transaction {
+            output: vec![expected_txout.clone()],
+            ..new_tx(0)
+        };
+
+        let expected_outpoint = OutPoint::new(prev_tx.compute_txid(), 0);
+
+        let tx = Transaction {
+            input: vec![TxIn {
+                previous_output: expected_outpoint,
+                ..Default::default()
+            }],
+            ..new_tx(0)
+        };
+
+        // Populate tx_cache with `prev_tx` to make it fetchable.
+        client.populate_tx_cache(vec![prev_tx.clone()]);
+
+        // Test fetch_prev_txout to see if we get our expected `txouts`.
+        let mut tx_update = TxUpdate {
+            txs: vec![Arc::new(tx.clone())],
+            ..Default::default()
+        };
+        let _ = client.fetch_prev_txout(&mut tx_update);
+
+        assert_eq!(
+            tx_update.txouts,
+            BTreeMap::from([(expected_outpoint, expected_txout)])
+        );
+
+        // Assert that `tx` is now a coinbase transaction.
+        let _ = client.update_coinbase_txid(tx.compute_txid());
+
+        // Test that fetch_prev_txout does not fetch `txouts` for coinbase tx.
+        let mut tx_update = TxUpdate {
+            txs: vec![Arc::new(tx)],
+            ..Default::default()
+        };
+        let _ = client.fetch_prev_txout(&mut tx_update);
+
+        assert_eq!(tx_update.txouts, BTreeMap::default());
+    }
+}