Refactors handle_masp_tx

anoma · Dec 4, 2023 · 65d031a · 65d031a
1 parent 144536f
commit 65d031a
Show file tree

Hide file tree

Showing 5 changed files with 60 additions and 121 deletions.
diff --git a/core/src/ledger/masp_utils.rs b/core/src/ledger/masp_utils.rs
@@ -2,15 +2,18 @@
 
 use masp_primitives::transaction::Transaction;
 
-use super::storage_api::StorageWrite;
+use super::storage_api::{StorageRead, StorageWrite};
 use crate::ledger::storage_api::Result;
 use crate::types::address::MASP;
 use crate::types::hash::Hash;
-use crate::types::storage::{Key, KeySeg};
-use crate::types::token::MASP_NULLIFIERS_KEY_PREFIX;
+use crate::types::storage::{BlockHeight, Epoch, Key, KeySeg, TxIndex};
+use crate::types::token::{
+    Transfer, HEAD_TX_KEY, MASP_NULLIFIERS_KEY_PREFIX, PIN_KEY_PREFIX,
+    TX_KEY_PREFIX,
+};
 
-/// Writes the nullifiers of the provided masp transaction to storage
-pub fn reveal_nullifiers(
+// Writes the nullifiers of the provided masp transaction to storage
+fn reveal_nullifiers(
     ctx: &mut impl StorageWrite,
     transaction: &Transaction,
 ) -> Result<()> {
@@ -28,3 +31,43 @@ pub fn reveal_nullifiers(
 
     Ok(())
 }
+
+/// Handle a MASP transaction.
+pub fn handle_masp_tx(
+    ctx: &mut (impl StorageRead + StorageWrite),
+    transfer: &Transfer,
+    shielded: &Transaction,
+) -> Result<()> {
+    let masp_addr = MASP;
+    let head_tx_key = Key::from(masp_addr.to_db_key())
+        .push(&HEAD_TX_KEY.to_owned())
+        .expect("Cannot obtain a storage key");
+    let current_tx_idx: u64 =
+        ctx.read(&head_tx_key).unwrap_or(None).unwrap_or(0);
+    let current_tx_key = Key::from(masp_addr.to_db_key())
+        .push(&(TX_KEY_PREFIX.to_owned() + &current_tx_idx.to_string()))
+        .expect("Cannot obtain a storage key");
+    // Save the Transfer object and its location within the blockchain
+    // so that clients do not have to separately look these
+    // up
+    let record: (Epoch, BlockHeight, TxIndex, Transfer, Transaction) = (
+        ctx.get_block_epoch()?,
+        ctx.get_block_height()?,
+        ctx.get_tx_index()?,
+        transfer.clone(),
+        shielded.clone(),
+    );
+    ctx.write(&current_tx_key, record)?;
+    ctx.write(&head_tx_key, current_tx_idx + 1)?;
+    reveal_nullifiers(ctx, shielded)?;
+
+    // If storage key has been supplied, then pin this transaction to it
+    if let Some(key) = &transfer.key {
+        let pin_key = Key::from(masp_addr.to_db_key())
+            .push(&(PIN_KEY_PREFIX.to_owned() + key))
+            .expect("Cannot obtain a storage key");
+        ctx.write(&pin_key, current_tx_idx)?;
+    }
+
+    Ok(())
+}
diff --git a/shared/src/ledger/native_vp/ibc/context.rs b/shared/src/ledger/native_vp/ibc/context.rs
@@ -3,7 +3,6 @@
 use std::collections::{BTreeSet, HashMap, HashSet};
 
 use borsh_ext::BorshSerializeExt;
-use masp_primitives::transaction::Transaction;
 use namada_core::ledger::ibc::{IbcCommonContext, IbcStorageContext};
 use namada_core::ledger::masp_utils;
 
@@ -12,15 +11,12 @@ use crate::ledger::native_vp::CtxPreStorageRead;
 use crate::ledger::storage::write_log::StorageModification;
 use crate::ledger::storage::{self as ledger_storage, StorageHasher};
 use crate::ledger::storage_api::{self, StorageRead, StorageWrite};
-use crate::types::address::{Address, InternalAddress, MASP};
+use crate::types::address::{Address, InternalAddress};
 use crate::types::ibc::{IbcEvent, IbcShieldedTransfer};
 use crate::types::storage::{
-    BlockHash, BlockHeight, Epoch, Header, Key, KeySeg, TxIndex,
-};
-use crate::types::token::{
-    self, Amount, DenominatedAmount, Transfer, HEAD_TX_KEY, PIN_KEY_PREFIX,
-    TX_KEY_PREFIX,
+    BlockHash, BlockHeight, Epoch, Header, Key, TxIndex,
 };
+use crate::types::token::{self, Amount, DenominatedAmount};
 use crate::vm::WasmCacheAccess;
 
 /// Result of a storage API call.
@@ -217,36 +213,7 @@ where
     }
 
     fn handle_masp_tx(&mut self, shielded: &IbcShieldedTransfer) -> Result<()> {
-        let masp_addr = MASP;
-        let head_tx_key = Key::from(masp_addr.to_db_key())
-            .push(&HEAD_TX_KEY.to_owned())
-            .expect("Cannot obtain a storage key");
-        let current_tx_idx: u64 =
-            self.ctx.read(&head_tx_key).unwrap_or(None).unwrap_or(0);
-        let current_tx_key = Key::from(masp_addr.to_db_key())
-            .push(&(TX_KEY_PREFIX.to_owned() + &current_tx_idx.to_string()))
-            .expect("Cannot obtain a storage key");
-        // Save the Transfer object and its location within the blockchain
-        // so that clients do not have to separately look these
-        // up
-        let record: (Epoch, BlockHeight, TxIndex, Transfer, Transaction) = (
-            self.ctx.get_block_epoch()?,
-            self.ctx.get_block_height()?,
-            self.ctx.get_tx_index()?,
-            shielded.transfer.clone(),
-            shielded.masp_tx.clone(),
-        );
-        self.write(&current_tx_key, record.serialize_to_vec())?;
-        self.write(&head_tx_key, (current_tx_idx + 1).serialize_to_vec())?;
-        masp_utils::reveal_nullifiers(self, &shielded.masp_tx)?;
-        // If storage key has been supplied, then pin this transaction to it
-        if let Some(key) = &shielded.transfer.key {
-            let pin_key = Key::from(masp_addr.to_db_key())
-                .push(&(PIN_KEY_PREFIX.to_owned() + key))
-                .expect("Cannot obtain a storage key");
-            self.write(&pin_key, current_tx_idx.serialize_to_vec())?;
-        }
-        Ok(())
+        masp_utils::handle_masp_tx(self, &shielded.transfer, &shielded.masp_tx)
     }
 
     fn mint_token(

diff --git a/shared/src/vm/host_env.rs b/shared/src/vm/host_env.rs
@@ -6,12 +6,11 @@ use std::num::TryFromIntError;
 
 use borsh::BorshDeserialize;
 use borsh_ext::BorshSerializeExt;
-use masp_primitives::transaction::Transaction;
 use namada_core::ledger::gas::{
     GasMetering, TxGasMeter, MEMORY_ACCESS_GAS_PER_BYTE,
 };
 use namada_core::ledger::masp_utils;
-use namada_core::types::address::{ESTABLISHED_ADDRESS_BYTES_LEN, MASP};
+use namada_core::types::address::ESTABLISHED_ADDRESS_BYTES_LEN;
 use namada_core::types::internal::KeyVal;
 use namada_core::types::storage::TX_INDEX_LENGTH;
 use namada_core::types::transaction::TxSentinel;
@@ -33,10 +32,9 @@ use crate::types::address::{self, Address};
 use crate::types::hash::Hash;
 use crate::types::ibc::{IbcEvent, IbcShieldedTransfer};
 use crate::types::internal::HostEnvResult;
-use crate::types::storage::{BlockHeight, Epoch, Key, KeySeg, TxIndex};
+use crate::types::storage::{BlockHeight, Epoch, Key, TxIndex};
 use crate::types::token::{
     is_any_minted_balance_key, is_any_minter_key, is_any_token_balance_key,
-    Transfer, HEAD_TX_KEY, PIN_KEY_PREFIX, TX_KEY_PREFIX,
 };
 use crate::vm::memory::VmMemory;
 use crate::vm::prefix_iter::{PrefixIteratorId, PrefixIterators};
@@ -2521,36 +2519,7 @@ where
         &mut self,
         shielded: &IbcShieldedTransfer,
     ) -> Result<(), storage_api::Error> {
-        let masp_addr = MASP;
-        let head_tx_key = Key::from(masp_addr.to_db_key())
-            .push(&HEAD_TX_KEY.to_owned())
-            .expect("Cannot obtain a storage key");
-        let current_tx_idx =
-            self.read::<u64>(&head_tx_key).unwrap_or(None).unwrap_or(0);
-        let current_tx_key = Key::from(masp_addr.to_db_key())
-            .push(&(TX_KEY_PREFIX.to_owned() + &current_tx_idx.to_string()))
-            .expect("Cannot obtain a storage key");
-        // Save the Transfer object and its location within the blockchain
-        // so that clients do not have to separately look these
-        // up
-        let record: (Epoch, BlockHeight, TxIndex, Transfer, Transaction) = (
-            self.get_block_epoch()?,
-            self.get_block_height()?,
-            self.get_tx_index()?,
-            shielded.transfer.clone(),
-            shielded.masp_tx.clone(),
-        );
-        self.write(&current_tx_key, record)?;
-        self.write(&head_tx_key, current_tx_idx + 1)?;
-        masp_utils::reveal_nullifiers(self, &shielded.masp_tx)?;
-        // If storage key has been supplied, then pin this transaction to it
-        if let Some(key) = &shielded.transfer.key {
-            let pin_key = Key::from(masp_addr.to_db_key())
-                .push(&(PIN_KEY_PREFIX.to_owned() + key))
-                .expect("Cannot obtain a storage key");
-            self.write(&pin_key, current_tx_idx)?;
-        }
-        Ok(())
+        masp_utils::handle_masp_tx(self, &shielded.transfer, &shielded.masp_tx)
     }
 
     fn mint_token(

diff --git a/tx_prelude/src/ibc.rs b/tx_prelude/src/ibc.rs
@@ -6,12 +6,13 @@ use std::rc::Rc;
 pub use namada_core::ledger::ibc::{
     IbcActions, IbcCommonContext, IbcStorageContext, ProofSpec, TransferModule,
 };
+use namada_core::ledger::masp_utils;
 use namada_core::ledger::tx_env::TxEnv;
 use namada_core::types::address::{Address, InternalAddress};
 pub use namada_core::types::ibc::{IbcEvent, IbcShieldedTransfer};
 use namada_core::types::token::DenominatedAmount;
 
-use crate::token::{burn, handle_masp_tx, mint, transfer};
+use crate::token::{burn, mint, transfer};
 use crate::{Ctx, Error};
 
 /// IBC actions to handle an IBC message
@@ -52,7 +53,7 @@ impl IbcStorageContext for Ctx {
         &mut self,
         shielded: &IbcShieldedTransfer,
     ) -> Result<(), Error> {
-        handle_masp_tx(self, &shielded.transfer, &shielded.masp_tx)
+        masp_utils::handle_masp_tx(self, &shielded.transfer, &shielded.masp_tx)
     }
 
     fn mint_token(

diff --git a/tx_prelude/src/token.rs b/tx_prelude/src/token.rs
@@ -1,7 +1,5 @@
-use masp_primitives::transaction::Transaction;
-use namada_core::ledger::masp_utils;
-use namada_core::types::address::{Address, MASP};
-use namada_core::types::storage::KeySeg;
+pub use namada_core::ledger::masp_utils;
+use namada_core::types::address::Address;
 use namada_core::types::token;
 pub use namada_core::types::token::*;
 
@@ -33,45 +31,6 @@ pub fn transfer(
     Ok(())
 }
 
-/// Handle a MASP transaction.
-pub fn handle_masp_tx(
-    ctx: &mut Ctx,
-    transfer: &Transfer,
-    shielded: &Transaction,
-) -> TxResult {
-    let masp_addr = MASP;
-    ctx.insert_verifier(&masp_addr)?;
-    let head_tx_key = storage::Key::from(masp_addr.to_db_key())
-        .push(&HEAD_TX_KEY.to_owned())
-        .expect("Cannot obtain a storage key");
-    let current_tx_idx: u64 =
-        ctx.read(&head_tx_key).unwrap_or(None).unwrap_or(0);
-    let current_tx_key = storage::Key::from(masp_addr.to_db_key())
-        .push(&(TX_KEY_PREFIX.to_owned() + &current_tx_idx.to_string()))
-        .expect("Cannot obtain a storage key");
-    // Save the Transfer object and its location within the blockchain
-    // so that clients do not have to separately look these
-    // up
-    let record: (Epoch, BlockHeight, TxIndex, Transfer, Transaction) = (
-        ctx.get_block_epoch()?,
-        ctx.get_block_height()?,
-        ctx.get_tx_index()?,
-        transfer.clone(),
-        shielded.clone(),
-    );
-    ctx.write(&current_tx_key, record)?;
-    ctx.write(&head_tx_key, current_tx_idx + 1)?;
-    masp_utils::reveal_nullifiers(ctx, shielded)?;
-    // If storage key has been supplied, then pin this transaction to it
-    if let Some(key) = &transfer.key {
-        let pin_key = storage::Key::from(masp_addr.to_db_key())
-            .push(&(PIN_KEY_PREFIX.to_owned() + key))
-            .expect("Cannot obtain a storage key");
-        ctx.write(&pin_key, current_tx_idx)?;
-    }
-    Ok(())
-}
-
 /// Mint that can be used in a transaction.
 pub fn mint(
     ctx: &mut Ctx,