Refactor get fee unshield

.changelog/unreleased/improvements/1877-refactor-get-fee-unshield.md

@@ -0,0 +1,2 @@
+- Refactored retrieval of `Transaction` object for fee unshielding.
+  ([\#1877](https://github.com/anoma/namada/pull/1877))

apps/src/lib/node/ledger/shell/mod.rs

@@ -38,7 +38,8 @@ use namada::ledger::pos::namada_proof_of_stake::types::{
     ConsensusValidator, ValidatorSetUpdate,
 };
 use namada::ledger::protocol::{
-    apply_wasm_tx, get_transfer_hash_from_storage, ShellParams,
+    apply_wasm_tx, get_fee_unshielding_transaction,
+    get_transfer_hash_from_storage, ShellParams,
 };
 use namada::ledger::storage::write_log::WriteLog;
 use namada::ledger::storage::{
@@ -1262,20 +1263,10 @@ where
                     return response;
                 }
 
-                let fee_unshield = wrapper
-                    .unshield_section_hash
-                    .and_then(|ref hash| tx.get_section(hash))
-                    .and_then(|section| {
-                        if let Section::MaspTx(transaction) = section.as_ref() {
-                            Some(transaction.to_owned())
-                        } else {
-                            None
-                        }
-                    });
                 // Validate wrapper fees
                 if let Err(e) = self.wrapper_fee_check(
                     &wrapper,
-                    fee_unshield,
+                    get_fee_unshielding_transaction(&tx, &wrapper),
                     &mut TempWlStorage::new(&self.wl_storage.storage),
                     &mut self.vp_wasm_cache.clone(),
                     &mut self.tx_wasm_cache.clone(),

apps/src/lib/node/ledger/shell/prepare_proposal.rs

@@ -5,9 +5,10 @@ use namada::core::ledger::gas::TxGasMeter;
 #[cfg(feature = "abcipp")]
 use namada::ledger::eth_bridge::{EthBridgeQueries, SendValsetUpd};
 use namada::ledger::pos::PosQueries;
+use namada::ledger::protocol::get_fee_unshielding_transaction;
 use namada::ledger::storage::{DBIter, StorageHasher, TempWlStorage, DB};
 use namada::proof_of_stake::find_validator_by_raw_hash;
-use namada::proto::{Section, Tx};
+use namada::proto::Tx;
 use namada::types::address::Address;
 use namada::types::internal::TxInQueue;
 use namada::types::key::tm_raw_hash_to_string;
@@ -249,19 +250,9 @@ where
                 .map_err(|_| ())?;
 
             // Check fees
-            let fee_unshield =
-                wrapper.unshield_section_hash.and_then(|ref hash| {
-                    tx.get_section(hash).and_then(|section| {
-                        if let Section::MaspTx(transaction) = section.as_ref() {
-                            Some(transaction.to_owned())
-                        } else {
-                            None
-                        }
-                    })
-                });
             match self.wrapper_fee_check(
                 &wrapper,
-                fee_unshield,
+                get_fee_unshielding_transaction(&tx, &wrapper),
                 temp_wl_storage,
                 vp_wasm_cache,
                 tx_wasm_cache,

apps/src/lib/node/ledger/shell/process_proposal.rs

@@ -6,6 +6,7 @@ use namada::core::hints;
 use namada::core::ledger::storage::WlStorage;
 use namada::ledger::eth_bridge::{EthBridgeQueries, SendValsetUpd};
 use namada::ledger::pos::PosQueries;
+use namada::ledger::protocol::get_fee_unshielding_transaction;
 use namada::ledger::storage::TempWlStorage;
 use namada::proof_of_stake::find_validator_by_raw_hash;
 use namada::types::internal::TxInQueue;
@@ -889,22 +890,9 @@ where
                     }
 
                     // Check that the fee payer has sufficient balance.
-                    let fee_unshield =
-                        wrapper.unshield_section_hash.and_then(|ref hash| {
-                            tx.get_section(hash).and_then(|section| {
-                                if let Section::MaspTx(transaction) =
-                                    section.as_ref()
-                                {
-                                    Some(transaction.to_owned())
-                                } else {
-                                    None
-                                }
-                            })
-                        });
-
                     match self.wrapper_fee_check(
                         &wrapper,
-                        fee_unshield,
+                        get_fee_unshielding_transaction(&tx, &wrapper),
                         temp_wl_storage,
                         vp_wasm_cache,
                         tx_wasm_cache,

core/src/proto/types.rs

@@ -1500,30 +1500,22 @@ impl Tx {
     /// signatures over it
     #[cfg(feature = "ferveo-tpke")]
     pub fn encrypt(&mut self, pubkey: &EncryptionKey) -> &mut Self {
-        use crate::types::hash::Hash;
         let header_hash = self.header_hash();
         let mut plaintexts = vec![];
         // Iterate backwrds to sidestep the effects of deletion on indexing
         for i in (0..self.sections.len()).rev() {
             match &self.sections[i] {
                 Section::Signature(sig)
                     if sig.targets.contains(&header_hash) => {}
-                Section::MaspTx(_) => {
+                masp_section @ Section::MaspTx(_) => {
                     // Do NOT encrypt the fee unshielding transaction
                     if let Some(unshield_section_hash) = self
                         .header()
                         .wrapper()
                         .expect("Tried to encrypt a non-wrapper tx")
                         .unshield_section_hash
                     {
-                        if unshield_section_hash
-                            == Hash(
-                                self.sections[i]
-                                    .hash(&mut Sha256::new())
-                                    .finalize_reset()
-                                    .into(),
-                            )
-                        {
+                        if unshield_section_hash == masp_section.get_hash() {
                             continue;
                         }
                     }

shared/src/ledger/protocol/mod.rs

@@ -175,20 +175,9 @@ where
             apply_protocol_tx(protocol_tx.tx, tx.data(), wl_storage)
         }
         TxType::Wrapper(ref wrapper) => {
-            let masp_transaction =
-                wrapper.unshield_section_hash.and_then(|ref hash| {
-                    tx.get_section(hash).and_then(|section| {
-                        if let Section::MaspTx(transaction) = section.as_ref() {
-                            Some(transaction.to_owned())
-                        } else {
-                            None
-                        }
-                    })
-                });
-
             let changed_keys = apply_wrapper_tx(
                 wrapper,
-                masp_transaction,
+                get_fee_unshielding_transaction(&tx, wrapper),
                 tx_bytes,
                 ShellParams {
                     tx_gas_meter,
@@ -291,6 +280,24 @@ where
     Ok(changed_keys)
 }
 
+/// Retrieve the Masp `Transaction` for fee unshielding from the provided
+/// transaction, if present
+pub fn get_fee_unshielding_transaction(
+    tx: &Tx,
+    wrapper: &WrapperTx,
+) -> Option<Transaction> {
+    wrapper
+        .unshield_section_hash
+        .and_then(|ref hash| tx.get_section(hash))
+        .and_then(|section| {
+            if let Section::MaspTx(transaction) = section.as_ref() {
+                Some(transaction.to_owned())
+            } else {
+                None
+            }
+        })
+}
+
 /// Charge fee for the provided wrapper transaction. In ABCI returns an error if
 /// the balance of the block proposer overflows. In ABCI plus returns error if:
 /// - The unshielding fails