diff --git a/zingolib/src/blaze/block_management_reorg_detection.rs b/zingolib/src/blaze/block_management_reorg_detection.rs index 58b999876..fb0647208 100644 --- a/zingolib/src/blaze/block_management_reorg_detection.rs +++ b/zingolib/src/blaze/block_management_reorg_detection.rs @@ -334,7 +334,7 @@ impl BlockManagementData { transaction_metadata_set .write() .await - .remove_txns_at_height(reorg_height); + .invalidate_all_transactions_after_or_at_height(reorg_height); } } diff --git a/zingolib/src/wallet/transaction_records_by_id.rs b/zingolib/src/wallet/transaction_records_by_id.rs index 6e1d9e2e5..15a401728 100644 --- a/zingolib/src/wallet/transaction_records_by_id.rs +++ b/zingolib/src/wallet/transaction_records_by_id.rs @@ -1,8 +1,18 @@ use std::collections::HashMap; +use orchard::note_encryption::OrchardDomain; +use sapling_crypto::note_encryption::SaplingDomain; + +use zcash_note_encryption::Domain; +use zcash_primitives::consensus::BlockHeight; + use zcash_primitives::transaction::TxId; -use crate::wallet::data::TransactionRecord; +use crate::wallet::{ + data::TransactionRecord, + notes::NoteInterface, + traits::{DomainWalletExt, Recipient}, +}; #[derive(Debug)] pub struct TransactionRecordsById(pub HashMap); @@ -20,9 +30,115 @@ impl std::ops::DerefMut for TransactionRecordsById { &mut self.0 } } + +/// Constructors impl TransactionRecordsById { - // Associated function to create a TransactionRecordMap from a HashMap + /// Constructs a new TransactionRecordsById with an empty map. + pub fn new() -> Self { + TransactionRecordsById(HashMap::new()) + } + // Constructs a TransactionRecordMap from a HashMap pub fn from_map(map: HashMap) -> Self { TransactionRecordsById(map) } } +/// Methods to modify the map. +impl TransactionRecordsById { + /// Invalidates all transactions from a given height including the block with block height `reorg_height` + /// + /// All information above a certain height is invalidated during a reorg. + pub fn invalidate_all_transactions_after_or_at_height(&mut self, reorg_height: BlockHeight) { + // First, collect txids that need to be removed + let txids_to_remove = self + .values() + .filter_map(|transaction_metadata| { + if transaction_metadata + .status + .is_confirmed_after_or_at(&reorg_height) + || transaction_metadata + .status + .is_broadcast_after_or_at(&reorg_height) + // tODo: why dont we only remove confirmed transactions. unconfirmed transactions may still be valid in the mempool and may later confirm or expire. + { + Some(transaction_metadata.txid) + } else { + None + } + }) + .collect::>(); + self.invalidate_transactions(txids_to_remove); + } + /// Invalidiates a vec of txids by removing them and then all references to them. + /// + /// A transaction can be invalidated either by a reorg or if it was never confirmed by a miner. + /// This is required in the case that a note was spent in a invalidated transaction. + /// Takes a slice of txids corresponding to the invalidated transactions, searches all notes for being spent in one of those txids, and resets them to unspent. + pub(crate) fn invalidate_transactions(&mut self, txids_to_remove: Vec) { + for txid in &txids_to_remove { + self.remove(txid); + } + + self.invalidate_transaction_specific_transparent_spends(&txids_to_remove); + // roll back any sapling spends in each invalidated tx + self.invalidate_transaction_specific_domain_spends::(&txids_to_remove); + // roll back any orchard spends in each invalidated tx + self.invalidate_transaction_specific_domain_spends::(&txids_to_remove); + } + /// Reverts any spent transparent notes in the given transactions to unspent. + pub(crate) fn invalidate_transaction_specific_transparent_spends( + &mut self, + invalidated_txids: &[TxId], + ) { + self.values_mut().for_each(|transaction_metadata| { + // Update UTXOs to roll back any spent utxos + transaction_metadata + .transparent_notes + .iter_mut() + .for_each(|utxo| { + if utxo.is_spent() && invalidated_txids.contains(&utxo.spent().unwrap().0) { + *utxo.spent_mut() = None; + } + + if utxo.unconfirmed_spent.is_some() + && invalidated_txids.contains(&utxo.unconfirmed_spent.unwrap().0) + { + utxo.unconfirmed_spent = None; + } + }) + }); + } + /// Reverts any spent shielded notes in the given transactions to unspent. + pub(crate) fn invalidate_transaction_specific_domain_spends( + &mut self, + invalidated_txids: &[TxId], + ) where + ::Recipient: Recipient, + ::Note: PartialEq + Clone, + { + self.values_mut().for_each(|transaction_metadata| { + // Update notes to rollback any spent notes + D::to_notes_vec_mut(transaction_metadata) + .iter_mut() + .for_each(|nd| { + // Mark note as unspent if the txid being removed spent it. + if nd.spent().is_some() && invalidated_txids.contains(&nd.spent().unwrap().0) { + *nd.spent_mut() = None; + } + + // Remove unconfirmed spends too + if nd.pending_spent().is_some() + && invalidated_txids.contains(&nd.pending_spent().unwrap().0) + { + *nd.pending_spent_mut() = None; + } + }); + }); + } +} + +impl Default for TransactionRecordsById { + /// Default constructor + fn default() -> Self { + Self::new() + } +} diff --git a/zingolib/src/wallet/transactions.rs b/zingolib/src/wallet/transactions.rs index 60fc05d4d..c63d926fc 100644 --- a/zingolib/src/wallet/transactions.rs +++ b/zingolib/src/wallet/transactions.rs @@ -1,5 +1,3 @@ -use std::collections::HashMap; - use crate::wallet::{data::WitnessTrees, transaction_records_by_id::TransactionRecordsById}; /// HashMap of all transactions in a wallet, keyed by txid. @@ -17,13 +15,13 @@ pub mod recording; impl TxMapAndMaybeTrees { pub(crate) fn new_with_witness_trees() -> TxMapAndMaybeTrees { Self { - transaction_records_by_id: TransactionRecordsById(HashMap::new()), + transaction_records_by_id: TransactionRecordsById::new(), witness_trees: Some(WitnessTrees::default()), } } pub(crate) fn new_treeless() -> TxMapAndMaybeTrees { Self { - transaction_records_by_id: TransactionRecordsById(HashMap::new()), + transaction_records_by_id: TransactionRecordsById::new(), witness_trees: None, } } diff --git a/zingolib/src/wallet/transactions/recording.rs b/zingolib/src/wallet/transactions/recording.rs index b0eb7eff5..91e8d35b5 100644 --- a/zingolib/src/wallet/transactions/recording.rs +++ b/zingolib/src/wallet/transactions/recording.rs @@ -24,85 +24,13 @@ use crate::{ use super::TxMapAndMaybeTrees; impl TxMapAndMaybeTrees { - pub fn remove_txids(&mut self, txids_to_remove: Vec) { - for txid in &txids_to_remove { - self.transaction_records_by_id.remove(txid); - } - self.transaction_records_by_id - .values_mut() - .for_each(|transaction_metadata| { - // Update UTXOs to rollback any spent utxos - transaction_metadata - .transparent_notes - .iter_mut() - .for_each(|utxo| { - if utxo.is_spent() && txids_to_remove.contains(&utxo.spent().unwrap().0) { - *utxo.spent_mut() = None; - } - - if utxo.unconfirmed_spent.is_some() - && txids_to_remove.contains(&utxo.unconfirmed_spent.unwrap().0) - { - utxo.unconfirmed_spent = None; - } - }) - }); - self.remove_domain_specific_txids::(&txids_to_remove); - self.remove_domain_specific_txids::(&txids_to_remove); - } + /// During reorgs, we need to remove all txns at a given height, and all spends that refer to any removed txns. + pub fn invalidate_all_transactions_after_or_at_height(&mut self, reorg_height: u64) { + let reorg_height = BlockHeight::from_u32(reorg_height as u32); - fn remove_domain_specific_txids(&mut self, txids_to_remove: &[TxId]) - where - ::Recipient: Recipient, - ::Note: PartialEq + Clone, - { self.transaction_records_by_id - .values_mut() - .for_each(|transaction_metadata| { - // Update notes to rollback any spent notes - D::to_notes_vec_mut(transaction_metadata) - .iter_mut() - .for_each(|nd| { - // Mark note as unspent if the txid being removed spent it. - if nd.spent().is_some() && txids_to_remove.contains(&nd.spent().unwrap().0) - { - *nd.spent_mut() = None; - } + .invalidate_all_transactions_after_or_at_height(reorg_height); - // Remove unconfirmed spends too - if nd.pending_spent().is_some() - && txids_to_remove.contains(&nd.pending_spent().unwrap().0) - { - *nd.pending_spent_mut() = None; - } - }); - }); - } - - // During reorgs, we need to remove all txns at a given height, and all spends that refer to any removed txns. - pub fn remove_txns_at_height(&mut self, reorg_height: u64) { - let reorg_height = BlockHeight::from_u32(reorg_height as u32); - - // First, collect txids that need to be removed - let txids_to_remove = self - .transaction_records_by_id - .values() - .filter_map(|transaction_metadata| { - if transaction_metadata - .status - .is_confirmed_after_or_at(&reorg_height) - || transaction_metadata - .status - .is_broadcast_after_or_at(&reorg_height) - // tODo: why dont we only remove confirmed transactions. unconfirmed transactions may still be valid in the mempool and may later confirm or expire. - { - Some(transaction_metadata.txid) - } else { - None - } - }) - .collect::>(); - self.remove_txids(txids_to_remove); if let Some(ref mut t) = self.witness_trees { t.witness_tree_sapling .truncate_removing_checkpoint(&(reorg_height - 1)) @@ -114,6 +42,7 @@ impl TxMapAndMaybeTrees { } } + /// Invalidates all those transactions which were broadcast but never 'confirmed' accepted by a miner. pub(crate) fn clear_expired_mempool(&mut self, latest_height: u64) { let cutoff = BlockHeight::from_u32((latest_height.saturating_sub(MAX_REORG as u64)) as u32); @@ -130,7 +59,8 @@ impl TxMapAndMaybeTrees { .iter() .for_each(|t| println!("Removing expired mempool tx {}", t)); - self.remove_txids(txids_to_remove); + self.transaction_records_by_id + .invalidate_transactions(txids_to_remove); } // Check this transaction to see if it is an outgoing transaction, and if it is, mark all received notes with non-textual memos in this