Reduce code duplication to reveal masp nullifiers

anoma · Dec 1, 2023 · 144536f · 144536f
1 parent f1156c1
commit 144536f
Show file tree

Hide file tree

Showing 5 changed files with 37 additions and 42 deletions.
diff --git a/core/src/ledger/masp_utils.rs b/core/src/ledger/masp_utils.rs
@@ -0,0 +1,30 @@
+//! MASP utilities
+
+use masp_primitives::transaction::Transaction;
+
+use super::storage_api::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;
+
+/// Writes the nullifiers of the provided masp transaction to storage
+pub fn reveal_nullifiers(
+    ctx: &mut impl StorageWrite,
+    transaction: &Transaction,
+) -> Result<()> {
+    for description in transaction
+        .sapling_bundle()
+        .map_or(&vec![], |description| &description.shielded_spends)
+    {
+        let nullifier_key = Key::from(MASP.to_db_key())
+            .push(&MASP_NULLIFIERS_KEY_PREFIX.to_owned())
+            .expect("Cannot obtain a storage key")
+            .push(&Hash(description.nullifier.0))
+            .expect("Cannot obtain a storage key");
+        ctx.write(&nullifier_key, ())?;
+    }
+
+    Ok(())
+}
diff --git a/core/src/ledger/mod.rs b/core/src/ledger/mod.rs
@@ -6,6 +6,7 @@ pub mod governance;
 pub mod ibc;
 pub mod inflation;
 pub mod masp_conversions;
+pub mod masp_utils;
 pub mod parameters;
 pub mod pgf;
 pub mod replay_protection;

diff --git a/shared/src/ledger/native_vp/ibc/context.rs b/shared/src/ledger/native_vp/ibc/context.rs
@@ -5,8 +5,7 @@ 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::types::hash::Hash;
-use namada_core::types::token::MASP_NULLIFIERS_KEY_PREFIX;
+use namada_core::ledger::masp_utils;
 
 use crate::ledger::ibc::storage::is_ibc_key;
 use crate::ledger::native_vp::CtxPreStorageRead;
@@ -239,19 +238,7 @@ where
         );
         self.write(&current_tx_key, record.serialize_to_vec())?;
         self.write(&head_tx_key, (current_tx_idx + 1).serialize_to_vec())?;
-        for description in shielded
-            .masp_tx
-            .sapling_bundle()
-            .map_or(&vec![], |description| &description.shielded_spends)
-        {
-            // Reveal the nullifier to prevent double spending
-            let nullifier_key = Key::from(masp_addr.to_db_key())
-                .push(&MASP_NULLIFIERS_KEY_PREFIX.to_owned())
-                .expect("Cannot obtain a storage key")
-                .push(&Hash(description.nullifier.0))
-                .expect("Cannot obtain a storage key");
-            self.write(&nullifier_key, ())?;
-        }
+        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())

diff --git a/shared/src/vm/host_env.rs b/shared/src/vm/host_env.rs
@@ -10,10 +10,10 @@ 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::internal::KeyVal;
 use namada_core::types::storage::TX_INDEX_LENGTH;
-use namada_core::types::token::MASP_NULLIFIERS_KEY_PREFIX;
 use namada_core::types::transaction::TxSentinel;
 use namada_core::types::validity_predicate::VpSentinel;
 use thiserror::Error;
@@ -2542,19 +2542,7 @@ where
         );
         self.write(&current_tx_key, record)?;
         self.write(&head_tx_key, current_tx_idx + 1)?;
-        for description in shielded
-            .masp_tx
-            .sapling_bundle()
-            .map_or(&vec![], |description| &description.shielded_spends)
-        {
-            // Reveal the nullifier to prevent double spending
-            let nullifier_key = Key::from(masp_addr.to_db_key())
-                .push(&MASP_NULLIFIERS_KEY_PREFIX.to_owned())
-                .expect("Cannot obtain a storage key")
-                .push(&Hash(description.nullifier.0))
-                .expect("Cannot obtain a storage key");
-            self.write(&nullifier_key, ())?;
-        }
+        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())

diff --git a/tx_prelude/src/token.rs b/tx_prelude/src/token.rs
@@ -1,6 +1,6 @@
 use masp_primitives::transaction::Transaction;
+use namada_core::ledger::masp_utils;
 use namada_core::types::address::{Address, MASP};
-use namada_core::types::hash::Hash;
 use namada_core::types::storage::KeySeg;
 use namada_core::types::token;
 pub use namada_core::types::token::*;
@@ -61,18 +61,7 @@ pub fn handle_masp_tx(
     );
     ctx.write(&current_tx_key, record)?;
     ctx.write(&head_tx_key, current_tx_idx + 1)?;
-    for description in shielded
-        .sapling_bundle()
-        .map_or(&vec![], |description| &description.shielded_spends)
-    {
-        // Reveal the nullifier to prevent double spending
-        let nullifier_key = storage::Key::from(masp_addr.to_db_key())
-            .push(&MASP_NULLIFIERS_KEY_PREFIX.to_owned())
-            .expect("Cannot obtain a storage key")
-            .push(&Hash(description.nullifier.0))
-            .expect("Cannot obtain a storage key");
-        ctx.write(&nullifier_key, ())?;
-    }
+    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())