Murisi/sequential payment addresses

.changelog/unreleased/improvements/4417-sequential-payment-addresses.md

@@ -0,0 +1,2 @@
+- Generate shielded payment addresses sequentially instead of randomly.
+  ([\#4417](https://github.com/anoma/namada/pull/4417))

crates/apps_lib/src/cli.rs

@@ -877,7 +877,7 @@ pub mod cmds {
 
     /// Generate a payment address from a viewing key or payment address
     #[derive(Clone, Debug)]
-    pub struct WalletGenPaymentAddress(pub args::PayAddressGen<args::CliTypes>);
+    pub struct WalletGenPaymentAddress(pub args::PayAddressGen);
 
     impl SubCmd for WalletGenPaymentAddress {
         const CMD: &'static str = "gen-payment-addr";
@@ -891,9 +891,9 @@ pub mod cmds {
         fn def() -> App {
             App::new(Self::CMD)
                 .about(wrap!(
-                    "Generates a payment address from the given spending key."
+                    "Generate the next payment address for a viewing key."
                 ))
-                .add_args::<args::PayAddressGen<args::CliTypes>>()
+                .add_args::<args::PayAddressGen>()
         }
     }
 
@@ -3368,7 +3368,7 @@ pub mod args {
 
     use data_encoding::HEXUPPER;
     use either::Either;
-    use namada_core::masp::{MaspEpoch, PaymentAddress};
+    use namada_core::masp::{DiversifierIndex, MaspEpoch, PaymentAddress};
     use namada_sdk::address::{Address, EstablishedAddress};
     pub use namada_sdk::args::*;
     use namada_sdk::chain::{ChainId, ChainIdPrefix};
@@ -3400,7 +3400,6 @@ pub mod args {
     use super::context::*;
     use super::utils::*;
     use super::{ArgGroup, ArgMatches};
-    use crate::client::utils::PRE_GENESIS_DIR;
     use crate::config::genesis::AddrOrPk;
     use crate::config::{self, Action, ActionAtHeight};
     use crate::tendermint::Timeout;
@@ -3483,6 +3482,8 @@ pub mod args {
     pub const DECRYPT: ArgFlag = flag("decrypt");
     pub const DESCRIPTION_OPT: ArgOpt<String> = arg_opt("description");
     pub const DISPOSABLE_SIGNING_KEY: ArgFlag = flag("disposable-gas-payer");
+    pub const DIVERSIFIER_INDEX: ArgOpt<DiversifierIndex> =
+        arg_opt("diversifier-index");
     pub const DESTINATION_VALIDATOR: Arg<WalletAddress> =
         arg("destination-validator");
     pub const DISCORD_OPT: ArgOpt<String> = arg_opt("discord-handle");
@@ -3690,6 +3691,7 @@ pub mod args {
     pub const VALUE: Arg<String> = arg("value");
     pub const VOTER_OPT: ArgOpt<WalletAddress> = arg_opt("voter");
     pub const VIEWING_KEY: Arg<WalletViewingKey> = arg("key");
+    pub const VIEWING_KEY_ALIAS: Arg<String> = arg("key");
     pub const VIEWING_KEYS: ArgMulti<WalletViewingKey, GlobStar> =
         arg_multi("viewing-keys");
     pub const VP: ArgOpt<String> = arg_opt("vp");
@@ -7910,53 +7912,32 @@ pub mod args {
         }
     }
 
-    impl CliToSdk<PayAddressGen<SdkTypes>> for PayAddressGen<CliTypes> {
+    impl CliToSdk<PayAddressGen> for PayAddressGen {
         type Error = std::convert::Infallible;
 
         fn to_sdk(
             self,
-            ctx: &mut Context,
-        ) -> Result<PayAddressGen<SdkTypes>, Self::Error> {
-            use namada_sdk::wallet::Wallet;
-
-            use crate::wallet::CliWalletUtils;
-
-            let find_viewing_key = |w: &mut Wallet<CliWalletUtils>| {
-                w.find_viewing_key(&self.viewing_key.raw)
-                    .copied()
-                    .unwrap_or_else(|_| {
-                        eprintln!(
-                            "Unknown viewing key {}",
-                            self.viewing_key.raw
-                        );
-                        safe_exit(1)
-                    })
-            };
-            let viewing_key = if ctx.global_args.is_pre_genesis {
-                let wallet_path =
-                    ctx.global_args.base_dir.join(PRE_GENESIS_DIR);
-                let mut wallet = crate::wallet::load_or_new(&wallet_path);
-                find_viewing_key(&mut wallet)
-            } else {
-                find_viewing_key(&mut ctx.borrow_mut_chain_or_exit().wallet)
-            };
-
-            Ok(PayAddressGen::<SdkTypes> {
+            _ctx: &mut Context,
+        ) -> Result<PayAddressGen, Self::Error> {
+            Ok(PayAddressGen {
                 alias: self.alias,
                 alias_force: self.alias_force,
-                viewing_key,
+                viewing_key: self.viewing_key,
+                diversifier_index: self.diversifier_index,
             })
         }
     }
 
-    impl Args for PayAddressGen<CliTypes> {
+    impl Args for PayAddressGen {
         fn parse(matches: &ArgMatches) -> Self {
             let alias = ALIAS.parse(matches);
             let alias_force = ALIAS_FORCE.parse(matches);
-            let viewing_key = VIEWING_KEY.parse(matches);
+            let diversifier_index = DIVERSIFIER_INDEX.parse(matches);
+            let viewing_key = VIEWING_KEY_ALIAS.parse(matches);
             Self {
                 alias,
                 alias_force,
+                diversifier_index,
                 viewing_key,
             }
         }
@@ -7968,6 +7949,9 @@ pub mod args {
             .arg(ALIAS_FORCE.def().help(wrap!(
                 "Override the alias without confirmation if it already exists."
             )))
+            .arg(DIVERSIFIER_INDEX.def().help(wrap!(
+                "Set the viewing key's current diversifier index beforehand."
+            )))
             .arg(VIEWING_KEY.def().help(wrap!("The viewing key.")))
         }
     }

crates/apps_lib/src/cli/wallet.rs

@@ -11,12 +11,13 @@ use ledger_transport_hid::hidapi::HidApi;
 use ledger_transport_hid::TransportNativeHID;
 use masp_primitives::zip32::ExtendedFullViewingKey;
 use namada_core::chain::BlockHeight;
-use namada_core::masp::{ExtendedSpendingKey, MaspValue, PaymentAddress};
+use namada_core::masp::{
+    DiversifierIndex, ExtendedSpendingKey, MaspValue, PaymentAddress,
+};
 use namada_sdk::address::{Address, DecodeError};
 use namada_sdk::borsh::{BorshDeserialize, BorshSerializeExt};
 use namada_sdk::io::{display_line, edisplay_line, Io};
 use namada_sdk::key::*;
-use namada_sdk::masp::find_valid_diversifier;
 use namada_sdk::wallet::{
     DecryptionError, DerivationPath, DerivationPathError, FindKeyError, Wallet,
 };
@@ -400,29 +401,53 @@ fn payment_address_gen(
     args::PayAddressGen {
         alias,
         alias_force,
-        viewing_key,
-        ..
+        viewing_key: viewing_key_alias,
+        diversifier_index,
     }: args::PayAddressGen,
 ) {
     let mut wallet = load_wallet(ctx);
     let alias = alias.to_lowercase();
-    let viewing_key = viewing_key.as_viewing_key();
-    let (div, _g_d) = find_valid_diversifier(&mut OsRng);
-    let masp_payment_addr = viewing_key
-        .to_payment_address(div)
-        .expect("a PaymentAddress");
+
+    let viewing_key = wallet
+        .find_viewing_key(&viewing_key_alias)
+        .copied()
+        .unwrap_or_else(|_| {
+            eprintln!("Unknown viewing key {}", viewing_key_alias,);
+            cli::safe_exit(1)
+        });
+    let diversifier_index = diversifier_index.unwrap_or_else(|| {
+        wallet
+            .find_diversifier_index(viewing_key_alias.clone())
+            .copied()
+            .unwrap_or_default()
+    });
+    let (diversifier_index, masp_payment_addr) =
+        ExtendedFullViewingKey::from(viewing_key)
+            .find_address(diversifier_index.into())
+            .expect("exhausted payment addresses");
+    let mut next_div_idx = diversifier_index;
+    next_div_idx
+        .increment()
+        .expect("exhausted payment addresses");
     let payment_addr = PaymentAddress::from(masp_payment_addr);
     let alias = wallet
         .insert_payment_addr(alias, payment_addr, alias_force)
         .unwrap_or_else(|| {
             edisplay_line!(io, "Payment address not added");
             cli::safe_exit(1);
         });
+    wallet
+        .insert_diversifier_index(viewing_key_alias, next_div_idx.into())
+        .expect(
+            "must be able to save next diversifier index under the alias of \
+             the viewing key",
+        );
     wallet.save().unwrap_or_else(|err| eprintln!("{}", err));
     display_line!(
         io,
-        "Successfully generated payment address {} with alias {}",
+        "Successfully generated payment address {} at index {} with alias {}",
         payment_addr,
+        DiversifierIndex::from(diversifier_index),
         alias,
     );
 }

crates/core/src/masp.rs

@@ -288,6 +288,93 @@ pub type TokenMap = BTreeMap<String, Address>;
 // enough capacity to store the payment address
 const PAYMENT_ADDRESS_SIZE: usize = 43;
 
+/// Wrapper for masp_primitive's DiversifierIndex
+#[derive(Clone, Debug, Copy, Eq, PartialEq, Default)]
+pub struct DiversifierIndex(masp_primitives::zip32::DiversifierIndex);
+
+impl From<masp_primitives::zip32::DiversifierIndex> for DiversifierIndex {
+    fn from(idx: masp_primitives::zip32::DiversifierIndex) -> Self {
+        Self(idx)
+    }
+}
+
+impl From<DiversifierIndex> for masp_primitives::zip32::DiversifierIndex {
+    fn from(value: DiversifierIndex) -> Self {
+        value.0
+    }
+}
+
+impl TryFrom<u128> for DiversifierIndex {
+    type Error = std::num::TryFromIntError;
+
+    fn try_from(idx: u128) -> Result<DiversifierIndex, Self::Error> {
+        // Diversifier is supposed to be 11 bytes. So right-shifting it by 3
+        // bytes should yield a 64-bit integer.
+        u64::try_from(idx >> 24)?;
+        let mut result = [0; 11];
+        result[..11].copy_from_slice(&idx.to_le_bytes()[0..11]);
+        Ok(masp_primitives::zip32::DiversifierIndex(result).into())
+    }
+}
+
+impl From<DiversifierIndex> for u128 {
+    fn from(div: DiversifierIndex) -> Self {
+        let mut u128_bytes = [0u8; 16];
+        u128_bytes[0..11].copy_from_slice(&div.0.0[..]);
+        u128::from_le_bytes(u128_bytes)
+    }
+}
+
+/// The describing a failure to parse a diversifier index
+#[derive(Clone, Debug, Copy, Default)]
+pub struct ParseDiversifierError;
+
+impl std::fmt::Display for ParseDiversifierError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        "unable to parse diversifier index".fmt(f)
+    }
+}
+
+impl FromStr for DiversifierIndex {
+    type Err = ParseDiversifierError;
+
+    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
+        u128::from_str(s)
+            .map_err(|_| ParseDiversifierError)?
+            .try_into()
+            .map_err(|_| ParseDiversifierError)
+    }
+}
+
+impl Display for DiversifierIndex {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        u128::from(*self).fmt(f)
+    }
+}
+
+impl serde::Serialize for DiversifierIndex {
+    fn serialize<S>(
+        &self,
+        serializer: S,
+    ) -> std::result::Result<S::Ok, S::Error>
+    where
+        S: serde::Serializer,
+    {
+        serde::Serialize::serialize(&self.to_string(), serializer)
+    }
+}
+
+impl<'de> serde::Deserialize<'de> for DiversifierIndex {
+    fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
+    where
+        D: serde::Deserializer<'de>,
+    {
+        use serde::de::Error;
+        let encoded: String = serde::Deserialize::deserialize(deserializer)?;
+        Self::from_str(&encoded).map_err(D::Error::custom)
+    }
+}
+
 /// Wrapper for masp_primitive's FullViewingKey
 #[derive(
     Clone,

crates/sdk/src/args.rs

@@ -15,7 +15,7 @@ use namada_core::dec::Dec;
 use namada_core::ethereum_events::EthAddress;
 use namada_core::keccak::KeccakHash;
 use namada_core::key::{common, SchemeType};
-use namada_core::masp::{MaspEpoch, PaymentAddress};
+use namada_core::masp::{DiversifierIndex, MaspEpoch, PaymentAddress};
 use namada_core::string_encoding::StringEncoded;
 use namada_core::time::DateTimeUtc;
 use namada_core::token::Amount;
@@ -3006,13 +3006,15 @@ pub struct KeyAddressRemove {
 
 /// Generate payment address arguments
 #[derive(Clone, Debug)]
-pub struct PayAddressGen<C: NamadaTypes = SdkTypes> {
-    /// Key alias
+pub struct PayAddressGen {
+    /// Payment address alias
     pub alias: String,
     /// Whether to force overwrite the alias
     pub alias_force: bool,
     /// Viewing key
-    pub viewing_key: C::ViewingKey,
+    pub viewing_key: String,
+    /// Diversifier index to start search at
+    pub diversifier_index: Option<DiversifierIndex>,
 }
 
 /// Bridge pool batch recommendation.

crates/wallet/src/lib.rs

@@ -20,7 +20,7 @@ use namada_core::chain::BlockHeight;
 use namada_core::collections::{HashMap, HashSet};
 use namada_core::key::*;
 use namada_core::masp::{
-    ExtendedSpendingKey, ExtendedViewingKey, PaymentAddress,
+    DiversifierIndex, ExtendedSpendingKey, ExtendedViewingKey, PaymentAddress,
 };
 use namada_core::time::DateTimeUtc;
 use namada_ibc::trace::is_ibc_denom;
@@ -450,6 +450,14 @@ impl<U> Wallet<U> {
         self.store.find_birthday(alias.as_ref())
     }
 
+    /// Find the diversifier index of the given alias
+    pub fn find_diversifier_index(
+        &self,
+        alias: impl AsRef<str>,
+    ) -> Option<&DiversifierIndex> {
+        self.store.find_diversifier_index(alias.as_ref())
+    }
+
     /// Find the payment address with the given alias in the wallet and return
     /// it
     pub fn find_payment_addr(
@@ -1231,6 +1239,18 @@ impl<U: WalletIo> Wallet<U> {
             .map(Into::into)
     }
 
+    /// Insert a diversifier index into the wallet under the given alias. Useful
+    /// for sequential payment address generation.
+    pub fn insert_diversifier_index(
+        &mut self,
+        alias: String,
+        div_index: DiversifierIndex,
+    ) -> Option<String> {
+        self.store
+            .insert_diversifier_index(alias.into(), div_index)
+            .map(Into::into)
+    }
+
     /// Extend this wallet from another wallet (typically pre-genesis).
     /// Note that this method ignores `store.validator_data` if any.
     pub fn extend(&mut self, wallet: Self) {