Refactor to prevent accidental double spends

src/faucet.rs

@@ -1,5 +1,6 @@
 use std::{collections::HashMap, str::FromStr};
 
+use bitcoincore_rpc::bitcoin::secp256k1::PublicKey;
 use bitcoincore_rpc::json::{self as bitcoin_json};
 use sdk_common::sp_client::bitcoin::secp256k1::{
     rand::thread_rng, Keypair, Message as Secp256k1Message, Secp256k1, ThirtyTwoByteHash,
@@ -29,6 +30,7 @@ use sdk_common::sp_client::spclient::Recipient;
 use anyhow::{Error, Result};
 
 use crate::lock_freezed_utxos;
+use crate::scan::check_transaction_alone;
 use crate::{scan::compute_partial_tweak_to_transaction, MutexExt, DAEMON, FAUCET_AMT, WALLET};
 
 fn spend_from_core(dest: XOnlyPublicKey) -> Result<(Transaction, Amount)> {
@@ -61,13 +63,10 @@ fn spend_from_core(dest: XOnlyPublicKey) -> Result<(Transaction, Amount)> {
     }
 }
 
-fn faucet_send(sp_address: SilentPaymentAddress, commitment: &str) -> Result<Transaction> {
+fn faucet_send(sp_address: SilentPaymentAddress, commitment: &str) -> Result<(Transaction, PublicKey)> {
-    let mut first_tx: Option<Transaction> = None;
+    let sp_wallet = WALLET.get().ok_or(Error::msg("Wallet not initialized"))?.get_wallet()?;
-    let final_tx: Transaction;
-
-    let sp_wallet = WALLET.get().ok_or(Error::msg("Wallet not initialized"))?;
     // do we have a sp output available ?
-    let available_outpoints = sp_wallet.get_wallet()?.get_outputs().to_spendable_list();
+    let available_outpoints = sp_wallet.get_outputs().to_spendable_list();
 
     let available_amt = available_outpoints
         .iter()
@@ -102,21 +101,33 @@ fn faucet_send(sp_address: SilentPaymentAddress, commitment: &str) -> Result<Tra
 
         log::debug!("fee estimate for 6 blocks: {}", fee_estimate);
 
-        let wallet = sp_wallet.get_wallet()?;
-
         let freezed_utxos = lock_freezed_utxos()?;
 
         let signed_psbt = create_transaction(
             vec![],
             &freezed_utxos,
-            &wallet,
+            &sp_wallet,
             vec![recipient],
             Some(Vec::from_hex(commitment).unwrap()),
             fee_estimate,
             None
         )?;
 
-        final_tx = signed_psbt.extract_tx()?;
+        let final_tx = signed_psbt.extract_tx()?;
+
+        let partial_tweak = compute_partial_tweak_to_transaction(&final_tx)?;
+
+        let daemon = DAEMON
+            .get()
+            .ok_or(Error::msg("DAEMON not initialized"))?
+            .lock_anyhow()?;
+        let txid = daemon.broadcast(&final_tx)?;
+        log::debug!("Sent tx {}", txid);
+
+        // We immediately add the new tx to our wallet to prevent accidental double spend
+        check_transaction_alone(sp_wallet, &final_tx, &partial_tweak)?;
+
+        Ok((final_tx, partial_tweak))
     } else {
         // let's try to spend directly from the mining address
         let secp = Secp256k1::signing_only();
@@ -159,7 +170,7 @@ fn faucet_send(sp_address: SilentPaymentAddress, commitment: &str) -> Result<Tra
                 .get(0)
                 .expect("Failed to generate keys")
                 .to_owned();
-        let change_sp_address = sp_wallet.get_wallet()?.get_client().get_receiving_address();
+        let change_sp_address = sp_wallet.get_client().get_receiving_address();
         let change_output_key: XOnlyPublicKey =
             generate_recipient_pubkeys(vec![change_sp_address], partial_secret)?
                 .into_values()
@@ -224,56 +235,28 @@ fn faucet_send(sp_address: SilentPaymentAddress, commitment: &str) -> Result<Tra
 
         faucet_tx.input[0].witness.push(final_sig.to_vec());
 
-        first_tx = Some(core_tx);
-
-        final_tx = faucet_tx;
-    }
-
-    {
         let daemon = DAEMON
             .get()
             .ok_or(Error::msg("DAEMON not initialized"))?
             .lock_anyhow()?;
-        // broadcast one or two transactions
+        daemon.broadcast(&core_tx)?;
-        if first_tx.is_some() {
+        let txid = daemon.broadcast(&faucet_tx)?;
-            daemon.broadcast(&first_tx.unwrap())?;
-        }
-        let txid = daemon.broadcast(&final_tx)?;
         log::debug!("Sent tx {}", txid);
-    }
 
-    Ok(final_tx)
+        let partial_tweak = compute_partial_tweak_to_transaction(&faucet_tx)?;
+
+        check_transaction_alone(sp_wallet, &faucet_tx, &partial_tweak)?;
+
+        Ok((faucet_tx, partial_tweak))
+    }
 }
 
 pub fn handle_faucet_request(msg: &FaucetMessage) -> Result<NewTxMessage> {
     let sp_address = SilentPaymentAddress::try_from(msg.sp_address.as_str())?;
     log::debug!("Sending bootstrap coins to {}", sp_address);
     // send bootstrap coins to this sp_address
-    let tx = faucet_send(sp_address, &msg.commitment)?;
+    let (tx, partial_tweak) = faucet_send(sp_address, &msg.commitment)?;
 
-    // get the tweak
-    let partial_tweak = compute_partial_tweak_to_transaction(&tx)?;
-
-    // get current blockheight
-    let blkheight: u32 = DAEMON
-        .get()
-        .unwrap()
-        .lock_anyhow()?
-        .get_current_height()?
-        .try_into()?;
-
-    let sp_wallet = WALLET.get().ok_or(Error::msg("Wallet not initialized"))?;
-
-    // update our sp_client with the change output(s)
-    sp_wallet
-        .get_wallet()?
-        .update_wallet_with_transaction(&tx, blkheight, partial_tweak)?;
-
-    log::debug!("updated the wallet");
-    // save to disk
-    sp_wallet.save()?;
-
-    log::debug!("saved the wallet");
     Ok(NewTxMessage::new(
         serialize(&tx).to_lower_hex_string(),
         Some(partial_tweak.to_string()),

src/main.rs

@@ -250,11 +250,8 @@ fn create_new_tx_message(transaction: Vec<u8>) -> Result<NewTxMessage> {
 
     let partial_tweak = compute_partial_tweak_to_transaction(&tx)?;
 
-    let found = check_transaction_alone(&tx, &partial_tweak)?;
+    let sp_wallet = WALLET.get().ok_or(Error::msg("Wallet not initialized"))?;
-
+    check_transaction_alone(sp_wallet.get_wallet()?, &tx, &partial_tweak)?;
-    if found.len() > 0 {
-        debug!("Found {} modified outputs in {}", found.len(), tx.txid());
-    }
 
     Ok(NewTxMessage::new(
         transaction.to_lower_hex_string(),

src/scan.rs

@@ -1,5 +1,6 @@
 use std::collections::HashMap;
 use std::str::FromStr;
+use std::sync::MutexGuard;
 
 use anyhow::{Error, Result};
 use electrum_client::ElectrumApi;
@@ -12,7 +13,7 @@ use sdk_common::sp_client::silentpayments::receiving::Receiver;
 use sdk_common::sp_client::silentpayments::utils::receiving::{
     calculate_tweak_data, get_pubkey_from_input,
 };
-use sdk_common::sp_client::spclient::{OutputSpendStatus, OwnedOutput};
+use sdk_common::sp_client::spclient::{OutputSpendStatus, OwnedOutput, SpWallet};
 use tokio::time::Instant;
 
 use crate::{electrumclient, MutexExt, DAEMON, WALLET};
@@ -21,6 +22,7 @@ pub fn compute_partial_tweak_to_transaction(tx: &Transaction) -> Result<PublicKe
     let daemon = DAEMON.get().ok_or(Error::msg("DAEMON not initialized"))?;
     let mut outpoints: Vec<(String, u32)> = Vec::with_capacity(tx.input.len());
     let mut pubkeys: Vec<PublicKey> = Vec::with_capacity(tx.input.len());
+    // TODO we should cache transactions to prevent multiple rpc request when transaction spends multiple outputs from the same tx
     for input in tx.input.iter() {
         outpoints.push((
             input.previous_output.txid.to_string(),
@@ -82,9 +84,8 @@ fn get_script_to_secret_map(
     Ok(res)
 }
 
-pub fn check_transaction_alone(tx: &Transaction, tweak_data: &PublicKey) -> Result<HashMap<OutPoint, OwnedOutput>> {
+pub fn check_transaction_alone(mut wallet: MutexGuard<SpWallet>, tx: &Transaction, tweak_data: &PublicKey) -> Result<HashMap<OutPoint, OwnedOutput>> {
-    let sp_wallet = WALLET.get().ok_or(Error::msg("Wallet not initialized"))?;
+    let updates = match wallet.update_wallet_with_transaction(tx, 0, *tweak_data) {
-    let updates = match sp_wallet.get_wallet()?.update_wallet_with_transaction(tx, 0, *tweak_data) {
         Ok(updates) => updates,
         Err(e) => {
             log::debug!("Error while checking transaction: {}", e);
@@ -234,6 +235,7 @@ pub fn scan_blocks(mut n_blocks_to_scan: u32, electrum_url: &str) -> anyhow::Res
     let electrum_client = electrumclient::create_electrum_client(electrum_url)?;
 
     let sp_wallet = WALLET.get().ok_or(Error::msg("Wallet not initialized"))?;
+    let mut wallet = sp_wallet.get_wallet()?;
 
     let core = DAEMON
         .get()
@@ -241,7 +243,7 @@ pub fn scan_blocks(mut n_blocks_to_scan: u32, electrum_url: &str) -> anyhow::Res
         .lock_anyhow()?;
 
     let secp = Secp256k1::new();
-    let scan_height = sp_wallet.get_wallet()?.get_outputs().get_last_scan();
+    let scan_height = wallet.get_outputs().get_last_scan();
     let tip_height: u32 = core.get_current_height()?.try_into()?;
 
     // 0 means scan to tip
@@ -268,9 +270,9 @@ pub fn scan_blocks(mut n_blocks_to_scan: u32, electrum_url: &str) -> anyhow::Res
 
     let mut tweak_data_map = electrum_client.sp_tweaks(start as usize)?;
 
-    let scan_sk = sp_wallet.get_wallet()?.get_client().get_scan_key();
+    let scan_sk = wallet.get_client().get_scan_key();
 
-    let sp_receiver = sp_wallet.get_wallet()?.get_client().sp_receiver.clone();
+    let sp_receiver = wallet.get_client().sp_receiver.clone();
     let start_time = Instant::now();
 
     for (blkheight, blkhash, blkfilter) in filters {
@@ -286,7 +288,7 @@ pub fn scan_blocks(mut n_blocks_to_scan: u32, electrum_url: &str) -> anyhow::Res
 
         // check if owned inputs are spent
         let our_outputs: HashMap<OutPoint, OwnedOutput> =
-            sp_wallet.get_wallet()?.get_outputs().to_outpoints_list();
+            wallet.get_outputs().to_outpoints_list();
 
         let owned_spks: Result<Vec<Vec<u8>>> = our_outputs
             .iter()
@@ -305,8 +307,7 @@ pub fn scan_blocks(mut n_blocks_to_scan: u32, electrum_url: &str) -> anyhow::Res
             let utxo_created_in_block =
                 scan_block_outputs(&sp_receiver, &blk.txdata, blkheight.into(), spk2secret)?;
             if !utxo_created_in_block.is_empty() {
-                sp_wallet
+                wallet
-                    .get_wallet()?
                     .get_mut_outputs()
                     .extend_from(utxo_created_in_block);
             }
@@ -314,12 +315,11 @@ pub fn scan_blocks(mut n_blocks_to_scan: u32, electrum_url: &str) -> anyhow::Res
             // update the list of outputs just in case
             // utxos may be created and destroyed in the same block
             let updated_outputs: HashMap<OutPoint, OwnedOutput> =
-                sp_wallet.get_wallet()?.get_outputs().to_outpoints_list();
+                wallet.get_outputs().to_outpoints_list();
 
             // search inputs and mark as mined
             let utxo_destroyed_in_block = scan_block_inputs(updated_outputs, blk.txdata)?;
             if !utxo_destroyed_in_block.is_empty() {
-                let mut wallet = sp_wallet.get_wallet()?;
                 let outputs = wallet.get_mut_outputs();
                 for outpoint in utxo_destroyed_in_block {
                     outputs.mark_mined(outpoint, blkhash)?;
@@ -335,8 +335,7 @@ pub fn scan_blocks(mut n_blocks_to_scan: u32, electrum_url: &str) -> anyhow::Res
     );
 
     // update last_scan height
-    sp_wallet
+    wallet
-        .get_wallet()?
         .get_mut_outputs()
         .update_last_scan(end);
     WALLET.get().unwrap().save(wallet)?;