improvement(ARRR): improve shielded transactions change notes handling in swaps

mm2src/coins/lp_coins.rs

@@ -2452,6 +2452,10 @@ impl CoinBalance {
         }
     }
 
+    pub fn new_with_unspendable(spendable: BigDecimal, unspendable: BigDecimal) -> CoinBalance {
+        CoinBalance { spendable, unspendable }
+    }
+
     pub fn into_total(self) -> BigDecimal { self.spendable + self.unspendable }
 
     pub fn get_total(&self) -> BigDecimal { &self.spendable + &self.unspendable }

mm2src/coins/z_coin.rs

@@ -1,3 +1,4 @@
+mod change_tracker;
 pub mod storage;
 mod z_balance_streaming;
 mod z_coin_errors;
@@ -47,6 +48,7 @@ use chain::constants::SEQUENCE_FINAL;
 use chain::{Transaction as UtxoTx, TransactionOutput};
 use common::calc_total_pages;
 use common::executor::{AbortableSystem, AbortedError};
+use common::log::debug;
 use common::{log, one_thousand_u32};
 use crypto::privkey::{key_pair_from_secret, secp_privkey_from_hash};
 use crypto::HDPathToCoin;
@@ -93,11 +95,10 @@ use zcash_primitives::{constants::mainnet as z_mainnet_constants, sapling::Payme
                        zip32::ExtendedFullViewingKey, zip32::ExtendedSpendingKey};
 use zcash_proofs::prover::LocalTxProver;
 
-use self::storage::store_change_output;
+use self::change_tracker::ChangeNoteTracker;
 
 cfg_native!(
     use common::{async_blocking, sha256_digest};
-    use zcash_client_sqlite::error::SqliteClientError as ZcashClientError;
     use zcash_client_sqlite::wallet::get_balance;
     use zcash_proofs::default_params_folder;
     use z_rpc::init_native_client;
@@ -212,6 +213,7 @@ pub struct ZCoinFields {
     consensus_params: ZcoinConsensusParams,
     z_balance_event_handler: Option<ZBalanceEventHandler>,
     sync_state_connector: AsyncMutex<SaplingSyncConnector>,
+    change_tracker: ChangeNoteTracker,
 }
 
 impl Transaction for ZTransaction {
@@ -323,22 +325,21 @@ impl ZCoin {
         })
     }
 
-    #[cfg(not(target_arch = "wasm32"))]
-    async fn my_balance_sat(&self) -> Result<u64, MmError<ZcashClientError>> {
+    async fn my_balance_sat(&self) -> Result<u64, MmError<ZCoinBalanceError>> {
         let wallet_db = self.z_fields.light_wallet_db.clone();
-        async_blocking(move || {
+        #[cfg(target_arch = "wasm32")]
+        let balance_sat: u64 = wallet_db.db.get_balance(AccountId::default()).await?.into();
+
+        #[cfg(not(target_arch = "wasm32"))]
+        let balance_sat: u64 = async_blocking(move || {
             let db_guard = wallet_db.db.inner();
             let db_guard = db_guard.lock().unwrap();
-            let balance = get_balance(&db_guard, AccountId::default())?.into();
-            Ok(balance)
+            get_balance(&db_guard, AccountId::default()).map(|v| v.into())
         })
         .await
-    }
+        .map_to_mm(|err| ZCoinBalanceError::BalanceError(err.to_string()))?;
 
-    #[cfg(target_arch = "wasm32")]
-    async fn my_balance_sat(&self) -> Result<u64, MmError<ZCoinBalanceError>> {
-        let wallet_db = self.z_fields.light_wallet_db.clone();
-        Ok(wallet_db.db.get_balance(AccountId::default()).await?.into())
+        Ok(balance_sat)
     }
 
     async fn get_spendable_notes(&self) -> Result<Vec<SpendableNote>, MmError<SpendableNotesError>> {
@@ -376,6 +377,85 @@ impl ZCoin {
         }
     }
 
+    async fn spendable_notes_required_for_tx(
+        &self,
+        total_required: &BigDecimal,
+    ) -> MmResult<Vec<SpendableNote>, GenTxError> {
+        let spendable_notes = self
+            .spendable_notes_ordered()
+            .await
+            .map_err(|err| GenTxError::SpendableNotesError(err.to_string()))?;
+        let changes = self.z_fields.change_tracker.change_notes_iter();
+        let my_z_balance = self
+            .my_balance()
+            .compat()
+            .await
+            .mm_err(|err| GenTxError::Internal(err.to_string()))?;
+
+        if changes.is_empty() {
+            Ok(spendable_notes)
+        } else {
+            if &my_z_balance.spendable < total_required
+                && my_z_balance.unspendable > (total_required - &my_z_balance.spendable)
+            {
+                return MmError::err(GenTxError::InsufficientBalance {
+                    coin: self.ticker().to_string(),
+                    available: my_z_balance.spendable,
+                    required: total_required.clone(),
+                });
+            }
+            if &my_z_balance.spendable >= total_required {
+                // return spendable notes if the total amount from spendable is upto or more
+                // than required amount needed to generate this tx.
+                Ok(spendable_notes)
+            } else {
+                debug!("Calculate pending tx change notes and wait for confirmation");
+                // else, we want to get, calculate the number of pending txs with change notes
+                // needed to be mined to satisfy our required amount.
+                let mut temp_remaining = BigDecimal::from(0);
+                let mut tx_hex = Vec::with_capacity(changes.len());
+                let change_needed_to_gen_tx = total_required - my_z_balance.spendable;
+                for (hex, change) in changes {
+                    temp_remaining += change;
+                    tx_hex.push(hex);
+                    if temp_remaining >= change_needed_to_gen_tx {
+                        break;
+                    }
+                }
+
+                // tx_hex can't be empty at this stage!,
+                for tx in tx_hex {
+                    // If we eventually get pending txs that are needed to be mined for this tx,
+                    // We now need to wait for tx_hex.len() txs confirmations so that this tx can
+                    // be somewhat guaranteed to be generated if it's due to change notes issues.
+                    let confirmations = self.required_confirmations();
+                    let requires_nota = self.requires_notarization();
+                    // 5 minutes(probably will be changed)
+                    let wait_until = 5 * 60;
+                    let input = ConfirmPaymentInput {
+                        payment_tx: tx,
+                        confirmations,
+                        requires_nota,
+                        wait_until,
+                        check_every: 5,
+                    };
+                    self.wait_for_confirmations(input)
+                        .compat()
+                        .await
+                        .map_to_mm(|_| GenTxError::NeededPrevTxConfirmed)?;
+                }
+                // Refetched spendable notes if we are able to get to this stage.
+                debug!("Refetching spendable notes after previous tx has been confirmed");
+                let new_spendable_notes = self
+                    .spendable_notes_ordered()
+                    .await
+                    .map_err(|err| GenTxError::SpendableNotesError(err.to_string()))?;
+
+                Ok(new_spendable_notes)
+            }
+        }
+    }
+
     /// Generates a tx sending outputs from our address
     async fn gen_tx(
         &self,
@@ -390,11 +470,8 @@ impl ZCoin {
         let total_output_sat = t_output_sat + z_output_sat;
         let total_output = big_decimal_from_sat_unsigned(total_output_sat, self.utxo_arc.decimals);
         let total_required = &total_output + &tx_fee;
+        let spendable_notes = self.spendable_notes_required_for_tx(&total_required).await?;
 
-        let spendable_notes = self
-            .spendable_notes_ordered()
-            .await
-            .map_err(|err| GenTxError::SpendableNotesError(err.to_string()))?;
         let mut total_input_amount = BigDecimal::from(0);
         let mut change = BigDecimal::from(0);
 
@@ -442,19 +519,21 @@ impl ZCoin {
             tx_builder.add_sapling_output(z_out.viewing_key, z_out.to_addr, z_out.amount, z_out.memo)?;
         }
 
+        // add change to tx output
+        let change_sat = sat_from_big_decimal(&change, self.utxo_arc.decimals)?;
         if change > BigDecimal::from(0u8) {
-            let change_sat = sat_from_big_decimal(&change, self.utxo_arc.decimals)?;
             received_by_me += change_sat;
+            let change_amount = Amount::from_u64(change_sat).map_to_mm(|_| {
+                GenTxError::NumConversion(NumConversError(format!(
+                    "Failed to get ZCash amount from {}",
+                    change_sat
+                )))
+            })?;
 
             tx_builder.add_sapling_output(
                 Some(self.z_fields.evk.fvk.ovk),
                 self.z_fields.my_z_addr.clone(),
-                Amount::from_u64(change_sat).map_to_mm(|_| {
-                    GenTxError::NumConversion(NumConversError(format!(
-                        "Failed to get ZCash amount from {}",
-                        change_sat
-                    )))
-                })?,
+                change_amount,
                 None,
             )?;
         }
@@ -476,11 +555,11 @@ impl ZCoin {
                 .await?
                 .tx_result?;
 
-        // Store any change outputs we created in this transaction by decrypting them with our keys
-        // and saving them to the wallet database for future spends
-        store_change_output(self.consensus_params_ref(), &self.z_fields.light_wallet_db, &tx)
-            .await
-            .map_to_mm(GenTxError::SaveChangeNotesError)?;
+        if change > BigDecimal::from(0u8) {
+            debug!("found change for txs {change}!");
+            // save change note amount for tracking.
+            self.z_fields.change_tracker.add_change_note(tx.tx_hex(), change);
+        }
 
         let additional_data = AdditionalTxData {
             received_by_me,
@@ -954,6 +1033,7 @@ impl<'a> UtxoCoinBuilder for ZCoinBuilder<'a> {
             consensus_params: self.protocol_info.consensus_params,
             sync_state_connector,
             z_balance_event_handler,
+            change_tracker: ChangeNoteTracker::new(),
         });
 
         let zcoin = ZCoin { utxo_arc, z_fields };
@@ -1129,11 +1209,20 @@ impl MarketCoinOps for ZCoin {
     fn my_balance(&self) -> BalanceFut<CoinBalance> {
         let coin = self.clone();
         let fut = async move {
-            let sat = coin
+            let balance_sat = coin
                 .my_balance_sat()
                 .await
                 .mm_err(|e| BalanceError::WalletStorageError(e.to_string()))?;
-            Ok(CoinBalance::new(big_decimal_from_sat_unsigned(sat, coin.decimals())))
+            let balance = big_decimal_from_sat_unsigned(balance_sat, coin.decimals());
+            let change = coin.z_fields.change_tracker.sum_all_change_notes();
+
+            if balance < change {
+                return MmError::err(BalanceError::WalletStorageError(
+                    "change can't be greater than balance".to_owned(),
+                ));
+            };
+
+            Ok(CoinBalance::new_with_unspendable(&balance - &change, change))
         };
         Box::new(fut.boxed().compat())
     }
@@ -1179,7 +1268,12 @@ impl MarketCoinOps for ZCoin {
     }
 
     fn wait_for_confirmations(&self, input: ConfirmPaymentInput) -> Box<dyn Future<Item = (), Error = String> + Send> {
-        utxo_common::wait_for_confirmations(self.as_ref(), input)
+        let z_fields = self.z_fields.clone();
+        let tx_hex = input.payment_tx.clone();
+        Box::new(utxo_common::wait_for_confirmations(self.as_ref(), input).map(move |_| {
+            common::log::info!("Payment confirm!, removing note from list");
+            z_fields.change_tracker.remove_change_notes(&tx_hex);
+        }))
     }
 
     async fn wait_for_htlc_tx_spend(&self, args: WaitForHTLCTxSpendArgs<'_>) -> TransactionResult {
@@ -1471,12 +1565,24 @@ impl SwapOps for ZCoin {
 
     #[inline]
     async fn validate_maker_payment(&self, input: ValidatePaymentInput) -> ValidatePaymentResult<()> {
-        utxo_common::validate_maker_payment(self, input).await
+        let payment_hex = input.payment_tx.clone();
+        let z_fields = self.z_fields.clone();
+        utxo_common::validate_maker_payment(self, input).await?;
+        z_fields.change_tracker.remove_change_notes(&payment_hex);
+        common::log::info!("Maker Payment confirm!, removing note from list");
+
+        Ok(())
     }
 
     #[inline]
     async fn validate_taker_payment(&self, input: ValidatePaymentInput) -> ValidatePaymentResult<()> {
-        utxo_common::validate_taker_payment(self, input).await
+        let payment_hex = input.payment_tx.clone();
+        let z_fields = self.z_fields.clone();
+        utxo_common::validate_taker_payment(self, input).await?;
+        z_fields.change_tracker.remove_change_notes(&payment_hex);
+        common::log::info!("Taker Payment confirm!, removing note from list");
+
+        Ok(())
     }
 
     #[inline]

mm2src/coins/z_coin/change_tracker.rs

@@ -0,0 +1,40 @@
+use mm2_number::BigDecimal;
+use std::collections::HashMap;
+use std::sync::{Arc, Mutex};
+
+/// Tracks unconfirmed change notes and their associated transaction IDs.
+#[derive(Debug, Clone)]
+pub struct ChangeNoteTracker {
+    tracker: Arc<Mutex<HashMap<Vec<u8>, BigDecimal>>>,
+}
+
+impl ChangeNoteTracker {
+    pub fn new() -> Self {
+        ChangeNoteTracker {
+            tracker: Arc::new(Mutex::new(HashMap::new())),
+        }
+    }
+
+    /// Adds a change note to the tracker.
+    pub fn add_change_note(&self, tx_hex: Vec<u8>, change_note: BigDecimal) {
+        let mut tracker = self.tracker.lock().unwrap();
+        tracker.entry(tx_hex).or_insert_with(|| change_note);
+    }
+
+    /// Removes a change note from the tracker when the transaction is confirmed.
+    pub fn remove_change_notes(&self, tx_hex: &Vec<u8>) -> Option<BigDecimal> {
+        let mut tracker = self.tracker.lock().unwrap();
+        tracker.remove(tx_hex)
+    }
+
+    /// Retrieves all unconfirmed change notes for a given transaction ID.
+    pub fn change_notes_iter(&self) -> Vec<(Vec<u8>, BigDecimal)> {
+        self.tracker.lock().unwrap().clone().into_iter().collect()
+    }
+
+    /// Retrieves all unconfirmed change notes across all transactions.
+    pub fn sum_all_change_notes(&self) -> BigDecimal {
+        let tracker = self.tracker.lock().unwrap();
+        tracker.values().sum()
+    }
+}

mm2src/coins/z_coin/storage.rs

@@ -4,14 +4,12 @@ pub mod blockdb;
 pub use blockdb::*;
 
 pub mod walletdb;
+pub use walletdb::*;
+
 #[cfg(target_arch = "wasm32")] mod z_params;
 #[cfg(target_arch = "wasm32")]
 pub(crate) use z_params::ZcashParamsWasmImpl;
 
-pub use walletdb::*;
-use zcash_extras::wallet::decrypt_and_store_transaction;
-use zcash_primitives::transaction::Transaction;
-
 use crate::z_coin::z_balance_streaming::ZBalanceEventSender;
 use mm2_err_handle::mm_error::MmResult;
 #[cfg(target_arch = "wasm32")]
@@ -192,18 +190,3 @@ pub async fn scan_cached_block(
     // If there are any transactions in the block, return the transaction count
     Ok(txs.len())
 }
-
-/// Processes and stores any change outputs created in the transaction by:
-/// - Decrypting outputs using wallet viewing keys
-/// - Adding decrypted change notes to the wallet database
-/// - Making change notes available for future spends
-pub(crate) async fn store_change_output(
-    params: &ZcoinConsensusParams,
-    shared_db: &WalletDbShared,
-    tx: &Transaction,
-) -> Result<(), String> {
-    let mut data = try_s!(shared_db.db.get_update_ops());
-    try_s!(decrypt_and_store_transaction(params, &mut data, tx).await);
-
-    Ok(())
-}