Helper functions for masp keys

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

@@ -7,7 +7,6 @@ use masp_proofs::bls12_381;
 use namada::core::ledger::inflation;
 use namada::core::ledger::masp_conversions::update_allowed_conversions;
 use namada::core::ledger::pgf::ADDRESS as pgf_address;
-use namada::core::types::storage::KeySeg;
 use namada::ledger::events::EventType;
 use namada::ledger::gas::{GasMetering, TxGasMeter};
 use namada::ledger::parameters::storage as params_storage;
@@ -22,13 +21,9 @@ use namada::proof_of_stake::{
     find_validator_by_raw_hash, read_last_block_proposer_address,
     read_pos_params, read_total_stake, write_last_block_proposer_address,
 };
-use namada::types::address::MASP;
 use namada::types::dec::Dec;
 use namada::types::key::tm_raw_hash_to_string;
 use namada::types::storage::{BlockHash, BlockResults, Epoch, Header};
-use namada::types::token::{
-    MASP_NOTE_COMMITMENT_ANCHOR_PREFIX, MASP_NOTE_COMMITMENT_TREE_KEY,
-};
 use namada::types::transaction::protocol::{
     ethereum_tx_data_variants, ProtocolTxType,
 };
@@ -567,21 +562,16 @@ where
         tracing::info!("{}", stats.format_tx_executed());
 
         // Update the MASP commitment tree anchor if the tree was updated
-        let tree_key = Key::from(MASP.to_db_key())
-            .push(&MASP_NOTE_COMMITMENT_TREE_KEY.to_owned())
-            .expect("Cannot obtain a storage key");
+        let tree_key = namada::core::types::token::masp_commitment_tree_key();
         if let Some(StorageModification::Write { value }) =
             self.wl_storage.write_log.read(&tree_key).0
         {
             let updated_tree = CommitmentTree::<Node>::try_from_slice(value)
                 .into_storage_result()?;
-            let anchor_key = Key::from(MASP.to_db_key())
-                .push(&MASP_NOTE_COMMITMENT_ANCHOR_PREFIX.to_owned())
-                .expect("Cannot obtain a storage key")
-                .push(&namada::core::types::hash::Hash(
-                    bls12_381::Scalar::from(updated_tree.root()).to_bytes(),
-                ))
-                .expect("Cannot obtain a storage key");
+            let anchor_key =
+                namada::core::types::token::masp_commitment_anchor_key(
+                    &bls12_381::Scalar::from(updated_tree.root()),
+                );
             self.wl_storage.write(&anchor_key, ())?;
         }
 

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

@@ -12,15 +12,10 @@ use namada::ledger::storage_api::token::{credit_tokens, write_denom};
 use namada::ledger::storage_api::StorageWrite;
 use namada::ledger::{ibc, pos};
 use namada::proof_of_stake::BecomeValidator;
-use namada::types::address::MASP;
 use namada::types::hash::Hash as CodeHash;
 use namada::types::key::*;
 use namada::types::storage::KeySeg;
 use namada::types::time::{DateTimeUtc, TimeZone, Utc};
-use namada::types::token::{
-    MASP_CONVERT_ANCHOR_KEY, MASP_NOTE_COMMITMENT_ANCHOR_PREFIX,
-    MASP_NOTE_COMMITMENT_TREE_KEY,
-};
 use namada::vm::validate_untrusted_wasm;
 use namada_sdk::eth_bridge::EthBridgeStatus;
 use namada_sdk::proof_of_stake::types::ValidatorMetaData;
@@ -184,27 +179,22 @@ where
         let empty_commitment_tree: CommitmentTree<Node> =
             CommitmentTree::empty();
         let anchor = empty_commitment_tree.root();
-        let note_commitment_tree_key = Key::from(MASP.to_db_key())
-            .push(&MASP_NOTE_COMMITMENT_TREE_KEY.to_owned())
-            .expect("Cannot obtain a storage key");
+        let note_commitment_tree_key =
+            namada::core::types::token::masp_commitment_tree_key();
         self.wl_storage
             .write(&note_commitment_tree_key, empty_commitment_tree)
             .unwrap();
-        let commitment_tree_anchor_key = Key::from(MASP.to_db_key())
-            .push(&MASP_NOTE_COMMITMENT_ANCHOR_PREFIX.to_owned())
-            .expect("Cannot obtain a storage key")
-            .push(&namada::core::types::hash::Hash(
-                bls12_381::Scalar::from(anchor).to_bytes(),
-            ))
-            .expect("Cannot obtain a storage key");
+        let commitment_tree_anchor_key =
+            namada::core::types::token::masp_commitment_anchor_key(
+                &bls12_381::Scalar::from(anchor),
+            );
         self.wl_storage
             .write(&commitment_tree_anchor_key, ())
             .unwrap();
 
         // Init masp convert anchor
-        let convert_anchor_key = Key::from(MASP.to_db_key())
-            .push(&MASP_CONVERT_ANCHOR_KEY.to_owned())
-            .expect("Cannot obtain a storage key");
+        let convert_anchor_key =
+            namada::core::types::token::masp_convert_anchor_key();
         self.wl_storage.write(
             &convert_anchor_key,
             namada::core::types::hash::Hash(

benches/native_vps.rs

@@ -47,13 +47,10 @@ use namada::namada_sdk::masp_primitives::transaction::Transaction;
 use namada::proof_of_stake;
 use namada::proof_of_stake::KeySeg;
 use namada::proto::{Code, Section, Tx};
-use namada::types::address::{InternalAddress, MASP};
+use namada::types::address::InternalAddress;
 use namada::types::eth_bridge_pool::{GasFee, PendingTransfer};
 use namada::types::masp::{TransferSource, TransferTarget};
 use namada::types::storage::{Epoch, TxIndex};
-use namada::types::token::{
-    MASP_NOTE_COMMITMENT_ANCHOR_PREFIX, MASP_NOTE_COMMITMENT_TREE_KEY,
-};
 use namada::types::transaction::governance::{
     InitProposalData, VoteProposalData,
 };
@@ -510,22 +507,16 @@ fn setup_storage_for_masp_verification(
     shielded_ctx.shell.wl_storage.commit_tx();
 
     // Update the anchor in storage
-    let tree_key = namada::core::types::storage::Key::from(MASP.to_db_key())
-        .push(&MASP_NOTE_COMMITMENT_TREE_KEY.to_owned())
-        .expect("Cannot obtain a storage key");
+    let tree_key = namada::core::types::token::masp_commitment_tree_key();
     let updated_tree: CommitmentTree<Node> = shielded_ctx
         .shell
         .wl_storage
         .read(&tree_key)
         .unwrap()
         .unwrap();
-    let anchor_key = namada::core::types::storage::Key::from(MASP.to_db_key())
-        .push(&MASP_NOTE_COMMITMENT_ANCHOR_PREFIX.to_owned())
-        .expect("Cannot obtain a storage key")
-        .push(&namada::core::types::hash::Hash(
-            bls12_381::Scalar::from(updated_tree.root()).to_bytes(),
-        ))
-        .expect("Cannot obtain a storage key");
+    let anchor_key = namada::core::types::token::masp_commitment_anchor_key(
+        &bls12_381::Scalar::from(updated_tree.root()),
+    );
     shielded_ctx
         .shell
         .wl_storage

core/src/ledger/masp_conversions.rs

@@ -19,8 +19,6 @@ use crate::ledger::storage_api::{StorageRead, StorageWrite};
 use crate::types::address::{Address, MASP};
 use crate::types::dec::Dec;
 use crate::types::storage::Epoch;
-#[cfg(feature = "wasm-runtime")]
-use crate::types::storage::{Key, KeySeg};
 use crate::types::token;
 use crate::types::token::MaspDenom;
 use crate::types::uint::Uint;
@@ -215,7 +213,6 @@ where
     use rayon::prelude::ParallelSlice;
 
     use crate::types::address;
-    use crate::types::token::MASP_CONVERT_ANCHOR_KEY;
 
     // The derived conversions will be placed in MASP address space
     let masp_addr = MASP;
@@ -460,11 +457,8 @@ where
     wl_storage.storage.conversion_state.tree =
         FrozenCommitmentTree::merge(&tree_parts);
     // Update the anchor in storage
-    let anchor_key = Key::from(MASP.to_db_key())
-        .push(&MASP_CONVERT_ANCHOR_KEY.to_owned())
-        .expect("Cannot obtain a storage key");
     wl_storage.write(
-        &anchor_key,
+        &crate::types::token::masp_convert_anchor_key(),
         crate::types::hash::Hash(
             bls12_381::Scalar::from(
                 wl_storage.storage.conversion_state.tree.root(),

core/src/ledger/masp_utils.rs

@@ -6,12 +6,10 @@ use masp_primitives::transaction::Transaction;
 
 use super::storage_api::{StorageRead, StorageWrite};
 use crate::ledger::storage_api::{Error, Result};
-use crate::types::address::MASP;
-use crate::types::hash::Hash;
-use crate::types::storage::{BlockHeight, Epoch, Key, KeySeg, TxIndex};
+use crate::types::storage::{BlockHeight, Epoch, TxIndex};
 use crate::types::token::{
-    Transfer, HEAD_TX_KEY, MASP_NOTE_COMMITMENT_TREE_KEY, MASP_NULLIFIERS_KEY,
-    PIN_KEY_PREFIX, TX_KEY_PREFIX,
+    masp_commitment_tree_key, masp_head_tx_key, masp_nullifier_key,
+    masp_pin_tx_key, masp_tx_key, Transfer,
 };
 
 // Writes the nullifiers of the provided masp transaction to storage
@@ -23,12 +21,7 @@ fn reveal_nullifiers(
         .sapling_bundle()
         .map_or(&vec![], |description| &description.shielded_spends)
     {
-        let nullifier_key = Key::from(MASP.to_db_key())
-            .push(&MASP_NULLIFIERS_KEY.to_owned())
-            .expect("Cannot obtain a storage key")
-            .push(&Hash(description.nullifier.0))
-            .expect("Cannot obtain a storage key");
-        ctx.write(&nullifier_key, ())?;
+        ctx.write(&masp_nullifier_key(&description.nullifier), ())?;
     }
 
     Ok(())
@@ -44,9 +37,7 @@ pub fn update_note_commitment_tree(
 ) -> Result<()> {
     if let Some(bundle) = transaction.sapling_bundle() {
         if !bundle.shielded_outputs.is_empty() {
-            let tree_key = Key::from(MASP.to_db_key())
-                .push(&MASP_NOTE_COMMITMENT_TREE_KEY.to_owned())
-                .expect("Cannot obtain a storage key");
+            let tree_key = masp_commitment_tree_key();
             let mut commitment_tree: CommitmentTree<Node> =
                 ctx.read(&tree_key)?.ok_or(Error::SimpleMessage(
                     "Missing note commitment tree in storage",
@@ -74,15 +65,10 @@ pub fn handle_masp_tx(
     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 head_tx_key = masp_head_tx_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");
+    let current_tx_key = masp_tx_key(current_tx_idx);
     // Save the Transfer object and its location within the blockchain
     // so that clients do not have to separately look these
     // up
@@ -103,9 +89,7 @@ pub fn handle_masp_tx(
 
     // 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");
+        let pin_key = masp_pin_tx_key(key);
         ctx.write(&pin_key, current_tx_idx)?;
     }
 

core/src/types/token.rs

@@ -9,6 +9,8 @@ use std::str::FromStr;
 use borsh::{BorshDeserialize, BorshSchema, BorshSerialize};
 use data_encoding::BASE32HEX_NOPAD;
 use ethabi::ethereum_types::U256;
+use masp_primitives::bls12_381::Scalar;
+use masp_primitives::sapling::Nullifier;
 use serde::{Deserialize, Serialize};
 use thiserror::Error;
 
@@ -1197,6 +1199,59 @@ pub fn masp_last_inflation_key(token_address: &Address) -> Key {
     )
 }
 
+/// Get a key for a masp nullifier
+pub fn masp_nullifier_key(nullifier: &Nullifier) -> Key {
+    Key::from(MASP.to_db_key())
+        .push(&MASP_NULLIFIERS_KEY.to_owned())
+        .expect("Cannot obtain a storage key")
+        .push(&Hash(nullifier.0))
+        .expect("Cannot obtain a storage key")
+}
+
+/// Get the key for the masp commitment tree
+pub fn masp_commitment_tree_key() -> Key {
+    Key::from(MASP.to_db_key())
+        .push(&MASP_NOTE_COMMITMENT_TREE_KEY.to_owned())
+        .expect("Cannot obtain a storage key")
+}
+
+/// Get a key for a masp commitment tree anchor
+pub fn masp_commitment_anchor_key(anchor: &Scalar) -> Key {
+    Key::from(MASP.to_db_key())
+        .push(&MASP_NOTE_COMMITMENT_ANCHOR_PREFIX.to_owned())
+        .expect("Cannot obtain a storage key")
+        .push(&Hash(anchor.to_bytes()))
+        .expect("Cannot obtain a storage key")
+}
+
+/// Get the key for the masp convert tree anchor
+pub fn masp_convert_anchor_key() -> Key {
+    Key::from(MASP.to_db_key())
+        .push(&MASP_CONVERT_ANCHOR_KEY.to_owned())
+        .expect("Cannot obtain a storage key")
+}
+
+/// Get the masp head-tx key
+pub fn masp_head_tx_key() -> Key {
+    Key::from(MASP.to_db_key())
+        .push(&HEAD_TX_KEY.to_owned())
+        .expect("Cannot obtain a storage key")
+}
+
+/// Get the masp tx key for the provided tx id
+pub fn masp_tx_key(tx_id: u64) -> Key {
+    Key::from(MASP.to_db_key())
+        .push(&(TX_KEY_PREFIX.to_owned() + &tx_id.to_string()))
+        .expect("Cannot obtain a storage key")
+}
+
+/// Get the masp pin tx key for the provided pin key
+pub fn masp_pin_tx_key(pin_key: &str) -> Key {
+    Key::from(MASP.to_db_key())
+        .push(&(PIN_KEY_PREFIX.to_owned() + pin_key))
+        .expect("Cannot obtain a storage key")
+}
+
 /// Check if the given storage key is for a minter of a unspecified token.
 /// If it is, returns the token.
 pub fn is_any_minter_key(key: &Key) -> Option<&Address> {